Merge pull request #19 from AllenInstitute/feature/update-batch-launc…

…h-templates
AllenInstitute · Jul 31, 2024 · 1dd8f3f · 1dd8f3f
2 parents b768105 + 4d5a13a
commit 1dd8f3f
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/src/aibs_informatics_cdk_lib/constructs_/batch/launch_template.py b/src/aibs_informatics_cdk_lib/constructs_/batch/launch_template.py
@@ -33,6 +33,7 @@ def create_launch_template(
         self,
         descriptor: IBatchEnvironmentDescriptor,
         security_group: Optional[ec2.SecurityGroup] = None,
+        **kwargs,
     ) -> ec2.LaunchTemplate:
         raise NotImplementedError()
 
@@ -46,10 +47,14 @@ def create_launch_template(
         self,
         descriptor: IBatchEnvironmentDescriptor,
         security_group: Optional[ec2.SecurityGroup] = None,
+        python_version: str = "python3.11",
+        **kwargs,
     ) -> ec2.LaunchTemplate:
         user_data = ec2.UserData.custom(
             BatchLaunchTemplateUserData(
-                env_base=self.env_base, batch_env_name=descriptor.get_name()
+                env_base=self.env_base,
+                batch_env_name=descriptor.get_name(),
+                python_version=python_version,
             ).user_data_text
         )
 
@@ -88,12 +93,15 @@ def create_launch_template(
         self,
         descriptor: IBatchEnvironmentDescriptor,
         security_group: Optional[ec2.SecurityGroup] = None,
+        python_version: str = "python3.11",
+        **kwargs,
     ) -> ec2.LaunchTemplate:
         user_data = ec2.UserData.custom(
             EbsBatchLaunchTemplateUserData(
                 env_base=self.env_base,
                 batch_env_name=descriptor.get_name(),
                 docker_volume_device_name=self.docker_volume_device_name,
+                python_version=python_version,
             ).user_data_text
         )
 
@@ -124,11 +132,12 @@ def create_launch_template(
 class BatchLaunchTemplateUserData:
     env_base: EnvBase
     batch_env_name: str
+    python_version: str
     user_data_text: str = field(init=False)
 
     def __post_init__(self):
         self.user_data_text = DEFAULT_LAUNCH_TEMPLATE_USER_DATA.format(
-            python_version="python3.9",
+            python_version=self.python_version,
             config_json=self.config_builder.to_string(),
         )
 
@@ -143,12 +152,13 @@ def get_base64_user_data(self) -> str:
 @dataclass
 class EbsBatchLaunchTemplateUserData(BatchLaunchTemplateUserData):
     docker_volume_device_name: str
+    python_version: str
 
     def __post_init__(self):
         self.user_data_text = EBS_AUTOSCALE_LAUNCH_TEMPLATE_USER_DATA.format(
             device_name=self.docker_volume_device_name,
             config_json=self.config_builder.to_string(),
-            python_version="python3.9",
+            python_version=self.python_version,
         )