fix: set use_spot_instances and max_wait as init params from job description (#1872)

Author: metrizable

src/sagemaker/estimator.py

@@ -871,6 +871,12 @@ class constructor
                     init_params["model_uri"] = channel["DataSource"]["S3DataSource"]["S3Uri"]
                     break
 
+        if job_details.get("EnableManagedSpotTraining", False):
+            init_params["use_spot_instances"] = True
+            max_wait = job_details.get("StoppingCondition", {}).get("MaxWaitTimeInSeconds")
+            if max_wait:
+                init_params["max_wait"] = max_wait
+
         return init_params
 
     def transformer(

tests/unit/test_estimator.py

@@ -2125,7 +2125,6 @@ def test_generic_deploy_accelerator_type(sagemaker_session):
     e.deploy(INSTANCE_COUNT, INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE)
 
     args = e.sagemaker_session.endpoint_from_production_variants.call_args[1]
-    print(args)
     assert args["name"].startswith(IMAGE_URI)
     assert args["production_variants"][0]["AcceleratorType"] == ACCELERATOR_TYPE
     assert args["production_variants"][0]["InitialInstanceCount"] == INSTANCE_COUNT
@@ -2182,7 +2181,6 @@ def test_local_mode(session_class, local_session_class):
     session_class.return_value = session
 
     e = Estimator(IMAGE_URI, ROLE, INSTANCE_COUNT, "local")
-    print(e.sagemaker_session.local_mode)
     assert e.sagemaker_session.local_mode is True
 
     e2 = Estimator(IMAGE_URI, ROLE, INSTANCE_COUNT, "local_gpu")
@@ -2248,6 +2246,25 @@ def test_prepare_init_params_from_job_description_with_algorithm_training_job():
     )
 
 
+def test_prepare_init_params_from_job_description_with_spot_training():
+    job_description = RETURNED_JOB_DESCRIPTION.copy()
+    job_description["EnableManagedSpotTraining"] = True
+    job_description["StoppingCondition"] = {
+        "MaxRuntimeInSeconds": 86400,
+        "MaxWaitTimeInSeconds": 87000,
+    }
+
+    init_params = EstimatorBase._prepare_init_params_from_job_description(
+        job_details=job_description
+    )
+
+    assert init_params["role"] == "arn:aws:iam::366:role/SageMakerRole"
+    assert init_params["instance_count"] == 1
+    assert init_params["use_spot_instances"]
+    assert init_params["max_run"] == 86400
+    assert init_params["max_wait"] == 87000
+
+
 def test_prepare_init_params_from_job_description_with_invalid_training_job():
 
     invalid_job_description = RETURNED_JOB_DESCRIPTION.copy()