feature: Support SageMakerTrainingPlan for training jobs (#1544)

Co-authored-by: Tritin Truong <[email protected]>

Author: 2 people

src/sagemaker/estimator.py

@@ -185,6 +185,7 @@ def __init__(
         disable_output_compression: bool = False,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
+        training_plan: Optional[Union[str, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -554,6 +555,8 @@ def __init__(
                 Specifies whether RemoteDebug is enabled for the training job.
             enable_session_tag_chaining (bool or PipelineVariable): Optional.
                 Specifies whether SessionTagChaining is enabled for the training job.
+            training_plan (str or PipelineVariable): Optional.
+                Specifies which training plan arn to use for the training job
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -762,8 +765,7 @@ def __init__(
 
         self.tensorboard_output_config = tensorboard_output_config
 
-        self.debugger_rule_configs = None
-        self.collection_configs = None
+        self.debugger_rule_configs, self.collection_configs = None, None
 
         self.enable_sagemaker_metrics = enable_sagemaker_metrics
 
@@ -774,6 +776,7 @@ def __init__(
             sagemaker_session=self.sagemaker_session,
         )
 
+        self.profiler_rule_configs, self.profiler_rules = None, None
         self.profiler_config = profiler_config
         self.disable_profiler = resolve_value_from_config(
             direct_input=disable_profiler,
@@ -796,8 +799,6 @@ def __init__(
         ) or _instance_type_supports_profiler(self.instance_type):
             self.disable_profiler = True
 
-        self.profiler_rule_configs = None
-        self.profiler_rules = None
         self.debugger_rules = None
         self.disable_output_compression = disable_output_compression
         validate_source_code_input_against_pipeline_variables(
@@ -807,6 +808,8 @@ def __init__(
             enable_network_isolation=self._enable_network_isolation,
         )
 
+        self.training_plan = training_plan
+
         # Internal flag
         self._is_output_path_set_from_default_bucket_and_prefix = False
 
@@ -1960,6 +1963,9 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
                 "KeepAlivePeriodInSeconds"
             ]
 
+        if "TrainingPlanArn" in job_details["ResourceConfig"]:
+            init_params["training_plan"] = job_details["ResourceConfig"]["TrainingPlanArn"]
+
         has_hps = "HyperParameters" in job_details
         init_params["hyperparameters"] = job_details["HyperParameters"] if has_hps else {}
 
@@ -2840,6 +2846,7 @@ def __init__(
         enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
+        training_plan: Optional[Union[str, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -3205,6 +3212,8 @@ def __init__(
                 Specifies whether RemoteDebug is enabled for the training job
             enable_session_tag_chaining (bool or PipelineVariable): Optional.
                  Specifies whether SessionTagChaining is enabled for the training job
+            training_plan (str or PipelineVariable): Optional.
+                Specifies which training plan arn to use for the training job
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
@@ -3258,6 +3267,7 @@ def __init__(
             disable_output_compression=disable_output_compression,
             enable_remote_debug=enable_remote_debug,
             enable_session_tag_chaining=enable_session_tag_chaining,
+            training_plan=training_plan,
             **kwargs,
         )
 

src/sagemaker/job.py

@@ -83,6 +83,7 @@ def _load_config(inputs, estimator, expand_role=True, validate_uri=True):
             estimator.volume_size,
             estimator.volume_kms_key,
             estimator.keep_alive_period_in_seconds,
+            estimator.training_plan,
         )
         stop_condition = _Job._prepare_stop_condition(estimator.max_run, estimator.max_wait)
         vpc_config = estimator.get_vpc_config()
@@ -294,6 +295,7 @@ def _prepare_resource_config(
         volume_size,
         volume_kms_key,
         keep_alive_period_in_seconds,
+        training_plan,
     ):
         """Placeholder docstring"""
         resource_config = {
@@ -319,6 +321,8 @@ def _prepare_resource_config(
                 )
             resource_config["InstanceCount"] = instance_count
             resource_config["InstanceType"] = instance_type
+        if training_plan is not None:
+            resource_config["TrainingPlanArn"] = training_plan
 
         return resource_config
 

src/sagemaker/jumpstart/estimator.py

@@ -115,6 +115,7 @@ def __init__(
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         config_name: Optional[str] = None,
         enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
+        training_plan: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Initializes a ``JumpStartEstimator``.
 
@@ -511,6 +512,8 @@ def __init__(
                 Name of the training configuration to apply to the Estimator. (Default: None).
             enable_session_tag_chaining (bool or PipelineVariable): Optional.
                 Specifies whether SessionTagChaining is enabled for the training job
+            training_plan (str or PipelineVariable): Optional.
+                Specifies which training plan arn to use for the training job
 
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
@@ -599,6 +602,7 @@ def _validate_model_id_and_get_type_hook():
             enable_remote_debug=enable_remote_debug,
             config_name=config_name,
             enable_session_tag_chaining=enable_session_tag_chaining,
+            training_plan=training_plan,
         )
 
         self.hub_arn = estimator_init_kwargs.hub_arn

src/sagemaker/jumpstart/factory/estimator.py

@@ -144,6 +144,7 @@ def get_init_kwargs(
     enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
     config_name: Optional[str] = None,
     enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
+    training_plan: Optional[Union[str, PipelineVariable]] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -205,6 +206,7 @@ def get_init_kwargs(
         enable_remote_debug=enable_remote_debug,
         config_name=config_name,
         enable_session_tag_chaining=enable_session_tag_chaining,
+        training_plan=training_plan,
     )
 
     estimator_init_kwargs, orig_session = _set_temp_sagemaker_session_if_not_set(

src/sagemaker/jumpstart/types.py

@@ -2406,6 +2406,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "hub_content_type",
         "model_reference_arn",
         "specs",
+        "training_plan",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2479,6 +2480,7 @@ def __init__(
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         config_name: Optional[str] = None,
         enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
+        training_plan: Optional[Union[str, PipelineVariable]] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -2541,6 +2543,7 @@ def __init__(
         self.enable_remote_debug = enable_remote_debug
         self.config_name = config_name
         self.enable_session_tag_chaining = enable_session_tag_chaining
+        self.training_plan = training_plan
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):

src/sagemaker/session.py

@@ -2470,6 +2470,75 @@ def describe_training_job(self, job_name):
         """
         return self.sagemaker_client.describe_training_job(TrainingJobName=job_name)
 
+    def describe_training_plan(self, training_plan_name):
+        """Calls the DescribeTrainingPlan API for the given training plan and returns the response.
+
+        Args:
+            training_plan_name (str): The name of the training plan to describe.
+
+        Returns:
+            dict: A dictionary response with the training plan description.
+        """
+        return self.sagemaker_client.describe_training_plan(TrainingPlanName=training_plan_name)
+
+    def list_training_plans(
+        self,
+        filters=None,
+        requested_start_time_after=None,
+        requested_start_time_before=None,
+        start_time_after=None,
+        start_time_before=None,
+        sort_order=None,
+        sort_by=None,
+        max_results=None,
+        next_token=None,
+    ):
+        """Calls the ListrTrainingPlan API for the given filters and returns the response.
+
+        Args:
+            filters (dict): A dictionary of key-value pairs used to filter the training plans.
+                Default to None.
+            requested_start_time_after (datetime): A timestamp that filters the results
+                to only include training plans with a requested start time after this timestamp.
+            requested_start_time_before (datetime): A timestamp that filters the results
+                to only include training plans with a requested start time before this timestamp.
+            start_time_after (datetime): A timestamp that filters the results
+                to only include training plans with an actual start time after this timestamp.
+            start_time_before (datetime): A timestamp that filters the results
+                to only include training plans with an actual start time before this timestamp.
+            sort_order (str): The order that the training plans will be listed in result.
+                Default to None.
+            sort_by (str): The value that the training plans will be sorted by.
+                Default to None.
+            max_results (int): The number of candidates will be listed in results,
+                between 1 and 100. Default to None. If None, will return all the training_plans.
+            next_token (str): The pagination token. Default to None.
+
+        Returns:
+            dict: A dictionary containing the following keys:
+                - "TrainingPlanSummaries": A list of dictionaries, where each dictionary represents
+                  a training plan.
+                - "NextToken": A token to retrieve the next set of results, if there are more
+                  than the maximum number of results returned.
+        """
+        list_training_plan_args = {}
+
+        def check_object(key, value):
+            if value is not None:
+                list_training_plan_args[key] = value
+
+        check_object("Filters", filters)
+        check_object("SortBy", sort_by)
+        check_object("SortOrder", sort_order)
+        check_object("RequestedStartTimeAfter", requested_start_time_after)
+        check_object("RequestedStartTimeBefore", requested_start_time_before)
+        check_object("StartTimeAfter", start_time_after)
+        check_object("StartTimeBefore", start_time_before)
+        check_object("NextToken", next_token)
+        check_object("MaxResults", max_results)
+
+        return self.sagemaker_client.list_training_plans(**list_training_plan_args)
+
     def auto_ml(
         self,
         input_config,

tests/unit/test_estimator.py

@@ -89,6 +89,7 @@
 INSTANCE_COUNT = 1
 INSTANCE_TYPE = "c4.4xlarge"
 KEEP_ALIVE_PERIOD_IN_SECONDS = 1800
+TRAINING_PLAN = "arn:aws:sagemaker:us-west-2:336:training-plan/test_training_plan"
 ACCELERATOR_TYPE = "ml.eia.medium"
 ROLE = "DummyRole"
 IMAGE_URI = "fakeimage"
@@ -861,6 +862,23 @@ def test_framework_with_keep_alive_period(sagemaker_session):
     assert args["resource_config"]["KeepAlivePeriodInSeconds"] == KEEP_ALIVE_PERIOD_IN_SECONDS
 
 
+def test_framework_with_training_plan(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.m4.xlarge", 2),
+        ],
+        training_plan=TRAINING_PLAN,
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args["resource_config"]["TrainingPlanArn"] == TRAINING_PLAN
+
+
 def test_framework_with_both_training_repository_config(sagemaker_session):
     f = DummyFramework(
         entry_point=SCRIPT_PATH,

tests/unit/test_job.py

@@ -31,6 +31,7 @@
 INSTANCE_COUNT = 1
 INSTANCE_TYPE = "c4.4xlarge"
 KEEP_ALIVE_PERIOD = 1800
+TRAINING_PLAN = "arn:aws:sagemaker:us-west-2:336:training-plan/test_training_plan"
 INSTANCE_GROUP = InstanceGroup("group", "ml.c4.xlarge", 1)
 VOLUME_SIZE = 1
 MAX_RUNTIME = 1
@@ -633,7 +634,13 @@ def test_prepare_output_config_kms_key_none():
 
 def test_prepare_resource_config():
     resource_config = _Job._prepare_resource_config(
-        INSTANCE_COUNT, INSTANCE_TYPE, None, VOLUME_SIZE, None, None
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        None,
+        VOLUME_SIZE,
+        None,
+        None,
+        None,
     )
 
     assert resource_config == {
@@ -643,9 +650,35 @@ def test_prepare_resource_config():
     }
 
 
+def test_prepare_resource_config_with_training_plan():
+    resource_config = _Job._prepare_resource_config(
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        None,
+        VOLUME_SIZE,
+        VOLUME_KMS_KEY,
+        None,
+        TRAINING_PLAN,
+    )
+
+    assert resource_config == {
+        "InstanceCount": INSTANCE_COUNT,
+        "InstanceType": INSTANCE_TYPE,
+        "VolumeSizeInGB": VOLUME_SIZE,
+        "VolumeKmsKeyId": VOLUME_KMS_KEY,
+        "TrainingPlanArn": TRAINING_PLAN,
+    }
+
+
 def test_prepare_resource_config_with_keep_alive_period():
     resource_config = _Job._prepare_resource_config(
-        INSTANCE_COUNT, INSTANCE_TYPE, None, VOLUME_SIZE, VOLUME_KMS_KEY, KEEP_ALIVE_PERIOD
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        None,
+        VOLUME_SIZE,
+        VOLUME_KMS_KEY,
+        KEEP_ALIVE_PERIOD,
+        None,
     )
 
     assert resource_config == {
@@ -659,7 +692,13 @@ def test_prepare_resource_config_with_keep_alive_period():
 
 def test_prepare_resource_config_with_volume_kms():
     resource_config = _Job._prepare_resource_config(
-        INSTANCE_COUNT, INSTANCE_TYPE, None, VOLUME_SIZE, VOLUME_KMS_KEY, None
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        None,
+        VOLUME_SIZE,
+        VOLUME_KMS_KEY,
+        None,
+        None,
     )
 
     assert resource_config == {
@@ -678,6 +717,7 @@ def test_prepare_resource_config_with_heterogeneous_cluster():
         VOLUME_SIZE,
         None,
         None,
+        None,
     )
 
     assert resource_config == {
@@ -698,6 +738,7 @@ def test_prepare_resource_config_with_instance_groups_instance_type_instance_cou
             VOLUME_SIZE,
             None,
             None,
+            None,
         )
     assert "instance_count and instance_type cannot be set when instance_groups is set" in str(
         error
@@ -713,6 +754,7 @@ def test_prepare_resource_config_with_instance_groups_instance_type_instance_cou
             VOLUME_SIZE,
             None,
             None,
+            None,
         )
     assert "instance_count and instance_type must be set if instance_groups is not set" in str(
         error