Hotfix: Invalid instance type check for HPO/Training (#205)

Training/HPO jobs cannot use KMS keys for volume encryption when launched on instance types that provide their own volume encryption. The previous check for this condition was incorrect because it did not account for the ml.* prefix used by SageMaker instances.
awslabs · Jun 28, 2024 · e23eaad · e23eaad
1 parent 7e2e531
commit e23eaad
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/backend/src/ml_space_lambda/hpo_job/lambda_functions.py b/backend/src/ml_space_lambda/hpo_job/lambda_functions.py
@@ -39,8 +39,9 @@
 def _normalize_job_definition(definition, iam_role, param_file):
     definition["RoleArn"] = iam_role
     definition["OutputDataConfig"]["KmsKeyId"] = param_file["pSMSKMSKeyId"]
+    instance_type = definition["ResourceConfig"]["InstanceType"].removeprefix("ml.")
     definition["ResourceConfig"]["VolumeKmsKeyId"] = (
-        "" if definition["ResourceConfig"]["InstanceType"] in kms_unsupported_instances() else param_file["pSMSKMSKeyId"]
+        "" if instance_type in kms_unsupported_instances() else param_file["pSMSKMSKeyId"]
     )
     definition["VpcConfig"] = {
         "SecurityGroupIds": param_file["pSMSSecurityGroupId"],

diff --git a/backend/src/ml_space_lambda/training_job/lambda_functions.py b/backend/src/ml_space_lambda/training_job/lambda_functions.py
@@ -73,6 +73,7 @@ def create(event, context):
     if "AlgorithmName" in algorithm_specs and "TrainingImage" in algorithm_specs:
         del algorithm_specs["AlgorithmName"]
 
+    instance_type = resource_config["InstanceType"].removeprefix("ml.")
     training_job_definition = dict(
         TrainingJobName=training_job_name,
         HyperParameters=hyper_parameters,
@@ -87,9 +88,7 @@ def create(event, context):
             "InstanceType": resource_config["InstanceType"],
             "InstanceCount": int(resource_config["InstanceCount"]),
             "VolumeSizeInGB": int(resource_config["VolumeSizeInGB"]),
-            "VolumeKmsKeyId": (
-                "" if resource_config["InstanceType"] in kms_unsupported_instances() else param_file["pSMSKMSKeyId"]
-            ),
+            "VolumeKmsKeyId": ("" if instance_type in kms_unsupported_instances() else param_file["pSMSKMSKeyId"]),
         },
         VpcConfig={
             "SecurityGroupIds": param_file["pSMSSecurityGroupId"],