feature: support checkpoint to be passed from estimator (#2849)

Co-authored-by: marckarp <[email protected]>

Author: marckarp

src/sagemaker/workflow/airflow.py

@@ -195,6 +195,11 @@ def training_base_config(estimator, inputs=None, job_name=None, mini_batch_size=
     if s3_operations:
         train_config["S3Operations"] = s3_operations
 
+    if (estimator.checkpoint_local_path is not None) & (estimator.checkpoint_s3_uri is not None):
+        train_config["CheckpointConfig"] = {
+            "LocalPath": estimator.checkpoint_local_path,
+            "S3Uri": estimator.checkpoint_s3_uri,
+        }
     return train_config
 
 

tests/unit/test_airflow.py

@@ -262,7 +262,7 @@ def test_framework_training_config_all_args(retrieve_image_uri, sagemaker_sessio
         py_version="py3",
         framework_version="1.15.2",
         role="{{ role }}",
-        instance_count="{{ instance_count }}",
+        instance_count=1,
         instance_type="ml.c4.2xlarge",
         volume_size="{{ volume_size }}",
         volume_kms_key="{{ volume_kms_key }}",
@@ -276,6 +276,8 @@ def test_framework_training_config_all_args(retrieve_image_uri, sagemaker_sessio
         security_group_ids=["{{ security_group_ids }}"],
         metric_definitions=[{"Name": "{{ name }}", "Regex": "{{ regex }}"}],
         sagemaker_session=sagemaker_session,
+        checkpoint_local_path="{{ checkpoint_local_path }}",
+        checkpoint_s3_uri="{{ checkpoint_s3_uri }}",
     )
 
     data = "{{ training_data }}"
@@ -294,7 +296,7 @@ def test_framework_training_config_all_args(retrieve_image_uri, sagemaker_sessio
         "TrainingJobName": "{{ base_job_name }}-%s" % TIME_STAMP,
         "StoppingCondition": {"MaxRuntimeInSeconds": "{{ max_run }}"},
         "ResourceConfig": {
-            "InstanceCount": "{{ instance_count }}",
+            "InstanceCount": 1,
             "InstanceType": "ml.c4.2xlarge",
             "VolumeSizeInGB": "{{ volume_size }}",
             "VolumeKmsKeyId": "{{ volume_kms_key }}",
@@ -338,6 +340,10 @@ def test_framework_training_config_all_args(retrieve_image_uri, sagemaker_sessio
                 }
             ]
         },
+        "CheckpointConfig": {
+            "LocalPath": "{{ checkpoint_local_path }}",
+            "S3Uri": "{{ checkpoint_s3_uri }}",
+        },
     }
     assert config == expected_config