fix: workflow passing spot training param to training job (#1599)

Co-authored-by: Chuyang Deng <[email protected]>
Co-authored-by: Lauren Yu <[email protected]>

Author: 3 people

src/sagemaker/workflow/airflow.py

@@ -181,6 +181,9 @@ def training_base_config(estimator, inputs=None, job_name=None, mini_batch_size=
     if job_config["vpc_config"] is not None:
         train_config["VpcConfig"] = job_config["vpc_config"]
 
+    if estimator.train_use_spot_instances:
+        train_config["EnableManagedSpotTraining"] = True
+
     if estimator.hyperparameters() is not None:
         hyperparameters = {str(k): str(v) for (k, v) in estimator.hyperparameters().items()}
 

tests/unit/test_airflow.py

@@ -105,6 +105,7 @@ def test_byo_training_config_all_args(sagemaker_session):
         model_uri="{{ model_uri }}",
         model_channel_name="{{ model_chanel }}",
         sagemaker_session=sagemaker_session,
+        train_use_spot_instances=True,
     )
 
     byo.set_hyperparameters(epochs=32, feature_dim=1024, mini_batch_size=256)
@@ -155,6 +156,7 @@ def test_byo_training_config_all_args(sagemaker_session):
             "Subnets": ["{{ subnet }}"],
             "SecurityGroupIds": ["{{ security_group_ids }}"],
         },
+        "EnableManagedSpotTraining": True,
         "HyperParameters": {"epochs": "32", "feature_dim": "1024", "mini_batch_size": "256"},
         "Tags": [{"{{ key }}": "{{ value }}"}],
     }