feat: include training_plan_arn in remote function classes (#5157)

sagemaker-core/src/sagemaker/core/remote_function/client.py

@@ -99,6 +99,7 @@ def remote(
     use_torchrun: bool = False,
     use_mpirun: bool = False,
     nproc_per_node: Optional[int] = None,
+    training_plan_arn: str = None,
 ):
     """Decorator for running the annotated function as a SageMaker training job.
 
@@ -234,6 +235,10 @@ def remote(
           latency than over using SageMaker managed warm pools alone by caching the package source
           downloaded in the previous runs.
 
+        training_plan_arn (str): The Amazon Resource Name (ARN) of the training plan to use for
+          this job. Training plans provide reserved capacity for training jobs. Defaults to
+          ``None``.
+
         max_retry_attempts (int): The max number of times the job is retried on
           ``InternalServerFailure`` Error from SageMaker service. Defaults to 1.
 
@@ -340,6 +345,7 @@ def _remote(func):
             use_torchrun=use_torchrun,
             use_mpirun=use_mpirun,
             nproc_per_node=nproc_per_node,
+            training_plan_arn=training_plan_arn,
         )
 
         @functools.wraps(func)
@@ -558,6 +564,7 @@ def __init__(
         use_torchrun: bool = False,
         use_mpirun: bool = False,
         nproc_per_node: Optional[int] = None,
+        training_plan_arn: str = None,
     ):
         """Constructor for RemoteExecutor
 
@@ -690,6 +697,10 @@ def __init__(
               up latency than over using SageMaker managed warm pools alone by caching the package
               source downloaded in the previous runs.
 
+            training_plan_arn (str): The Amazon Resource Name (ARN) of the training plan to use
+              for this job. Training plans provide reserved capacity for training jobs. Defaults
+              to ``None``.
+
             max_parallel_jobs (int): Maximum number of jobs that run in parallel. Defaults to 1.
 
             max_retry_attempts (int): The max number of times the job is retried on
@@ -812,6 +823,7 @@ def __init__(
             use_torchrun=use_torchrun,
             use_mpirun=use_mpirun,
             nproc_per_node=nproc_per_node,
+            training_plan_arn=training_plan_arn,
         )
 
         self._state_condition = threading.Condition()

sagemaker-core/src/sagemaker/core/remote_function/job.py

@@ -419,6 +419,7 @@ def __init__(
         use_torchrun: bool = False,
         use_mpirun: bool = False,
         nproc_per_node: Optional[int] = None,
+        training_plan_arn: str = None,
     ):
         """Initialize a _JobSettings instance which configures the remote job.
 
@@ -544,6 +545,10 @@ def __init__(
               latency than over using SageMaker managed warm pools alone by caching the package
               source downloaded in the previous runs.
 
+            training_plan_arn (str): The Amazon Resource Name (ARN) of the training plan to use
+              for this job. Training plans provide reserved capacity for training jobs. Defaults
+              to ``None``.
+
             max_retry_attempts (int, PipelineVariable): The max number of times the job is retried
               on ``InternalServerFailure`` Error from SageMaker service. Defaults to 1.
 
@@ -727,6 +732,7 @@ def __init__(
 
         self.instance_count = instance_count
         self.volume_size = volume_size
+        self.training_plan_arn = training_plan_arn
         self.max_runtime_in_seconds = max_runtime_in_seconds
         self.max_retry_attempts = max_retry_attempts
         self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
@@ -1158,6 +1164,8 @@ def compile(
             resource_config["VolumeKmsKeyId"] = job_settings.volume_kms_key
         if job_settings.keep_alive_period_in_seconds is not None:
             resource_config["KeepAlivePeriodInSeconds"] = job_settings.keep_alive_period_in_seconds
+        if job_settings.training_plan_arn is not None:
+            resource_config["TrainingPlanArn"] = job_settings.training_plan_arn
 
         request_dict["ResourceConfig"] = resource_config
 

sagemaker-core/tests/unit/remote_function/test_client.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 
 import pytest
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 from collections import deque
 
 from sagemaker.core.remote_function.client import (
+    remote,
     RemoteExecutor,
     _submit_worker,
     _polling_worker,
@@ -26,6 +27,28 @@
     _FINISHED,
 )
 
+TRAINING_PLAN_ARN = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/test-plan"
+
+
+class TestTrainingPlanArnForwarding:
+    """Ensure training_plan_arn is forwarded into _JobSettings by both entry points."""
+
+    @patch("sagemaker.core.remote_function.client._JobSettings")
+    def test_remote_decorator_forwards_training_plan_arn(self, mock_job_settings):
+        @remote(instance_type="ml.m5.xlarge", training_plan_arn=TRAINING_PLAN_ARN)
+        def my_func():
+            pass
+
+        _, kwargs = mock_job_settings.call_args
+        assert kwargs["training_plan_arn"] == TRAINING_PLAN_ARN
+
+    @patch("sagemaker.core.remote_function.client._JobSettings")
+    def test_remote_executor_forwards_training_plan_arn(self, mock_job_settings):
+        RemoteExecutor(instance_type="ml.m5.xlarge", training_plan_arn=TRAINING_PLAN_ARN)
+
+        _, kwargs = mock_job_settings.call_args
+        assert kwargs["training_plan_arn"] == TRAINING_PLAN_ARN
+
 
 class TestConstants:
     """Test module constants"""

sagemaker-core/tests/unit/remote_function/test_job.py

@@ -111,6 +111,28 @@ def test_init_without_instance_type_raises_error(self, mock_session):
         with pytest.raises(ValueError, match="instance_type is a required parameter"):
             _JobSettings(sagemaker_session=mock_session, image_uri="test-image")
 
+    def test_init_stores_training_plan_arn(self, mock_session):
+        """training_plan_arn is stored on the settings when provided."""
+        training_plan_arn = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/test-plan"
+        settings = _JobSettings(
+            sagemaker_session=mock_session,
+            image_uri="test-image",
+            instance_type="ml.m5.xlarge",
+            role="arn:aws:iam::123456789012:role/test",
+            training_plan_arn=training_plan_arn,
+        )
+        assert settings.training_plan_arn == training_plan_arn
+
+    def test_init_training_plan_arn_defaults_to_none(self, mock_session):
+        """training_plan_arn defaults to None when not provided."""
+        settings = _JobSettings(
+            sagemaker_session=mock_session,
+            image_uri="test-image",
+            instance_type="ml.m5.xlarge",
+            role="arn:aws:iam::123456789012:role/test",
+        )
+        assert settings.training_plan_arn is None
+
     @patch.dict(os.environ, {"SAGEMAKER_INTERNAL_IMAGE_URI": "custom-image"})
     def test_get_default_image_from_env(self, mock_session):
         """Test getting default image from environment variable."""
@@ -859,6 +881,7 @@ def test_compile_basic(self, mock_input_config, mock_stored_func, mock_session):
         job_settings.job_conda_env = None
         job_settings.spark_config = None
         job_settings.dependencies = None
+        job_settings.training_plan_arn = None
 
         def test_func():
             pass
@@ -868,6 +891,77 @@ def test_func():
         assert result["TrainingJobName"] == "test-job"
         assert result["RoleArn"] == "arn:aws:iam::123456789012:role/test"
 
+    @staticmethod
+    def _base_job_settings(mock_session):
+        """Build a mock ``_JobSettings`` with the minimum attrs ``compile`` reads."""
+        job_settings = Mock()
+        job_settings.max_runtime_in_seconds = 3600
+        job_settings.max_wait_time_in_seconds = None
+        job_settings.max_retry_attempts = 1
+        job_settings.role = "arn:aws:iam::123456789012:role/test"
+        job_settings.tags = None
+        job_settings.s3_kms_key = None
+        job_settings.disable_output_compression = False
+        job_settings.volume_size = 30
+        job_settings.instance_count = 1
+        job_settings.instance_type = "ml.m5.xlarge"
+        job_settings.volume_kms_key = None
+        job_settings.keep_alive_period_in_seconds = None
+        job_settings.enable_network_isolation = False
+        job_settings.encrypt_inter_container_traffic = False
+        job_settings.vpc_config = None
+        job_settings.use_spot_instances = False
+        job_settings.environment_variables = {}
+        job_settings.image_uri = "test-image"
+        job_settings.sagemaker_session = mock_session
+        job_settings.use_torchrun = False
+        job_settings.use_mpirun = False
+        job_settings.nproc_per_node = None
+        job_settings.job_conda_env = None
+        job_settings.spark_config = None
+        job_settings.dependencies = None
+        job_settings.training_plan_arn = None
+        return job_settings
+
+    @patch("sagemaker.core.remote_function.job.StoredFunction")
+    @patch("sagemaker.core.remote_function.job._generate_input_data_config")
+    def test_compile_with_training_plan_arn(
+        self, mock_input_config, mock_stored_func, mock_session
+    ):
+        """TrainingPlanArn is added to ResourceConfig when training_plan_arn is set."""
+        mock_input_config.return_value = []
+        mock_stored_func.return_value.save = Mock()
+
+        training_plan_arn = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/test-plan"
+        job_settings = self._base_job_settings(mock_session)
+        job_settings.training_plan_arn = training_plan_arn
+
+        def test_func():
+            pass
+
+        result = _Job.compile(job_settings, "test-job", "s3://bucket", test_func, (), {})
+
+        assert result["ResourceConfig"]["TrainingPlanArn"] == training_plan_arn
+
+    @patch("sagemaker.core.remote_function.job.StoredFunction")
+    @patch("sagemaker.core.remote_function.job._generate_input_data_config")
+    def test_compile_without_training_plan_arn(
+        self, mock_input_config, mock_stored_func, mock_session
+    ):
+        """TrainingPlanArn is absent from ResourceConfig when training_plan_arn is None."""
+        mock_input_config.return_value = []
+        mock_stored_func.return_value.save = Mock()
+
+        job_settings = self._base_job_settings(mock_session)
+        job_settings.training_plan_arn = None
+
+        def test_func():
+            pass
+
+        result = _Job.compile(job_settings, "test-job", "s3://bucket", test_func, (), {})
+
+        assert "TrainingPlanArn" not in result["ResourceConfig"]
+
 
 class TestJobStart:
     """Test _Job.start method."""

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/test_feature_scheduler.py

@@ -974,6 +974,7 @@ def test_remote_decorator_fields_consistency(get_execution_role, session):
         "use_torchrun",
         "use_mpirun",
         "nproc_per_node",
+        "training_plan_arn",
     }
 
     job_settings = _JobSettings(