From 469698a186baf44b8e9b0d89394d35c7143fe840 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 28 Sep 2026 14:24:59 -0700 Subject: [PATCH 1/2] feat: include training_plan_arn in remote function classes (#5157) --- .../sagemaker/core/remote_function/client.py | 12 +++ .../src/sagemaker/core/remote_function/job.py | 8 ++ .../tests/unit/remote_function/test_client.py | 25 ++++- .../tests/unit/remote_function/test_job.py | 94 +++++++++++++++++++ 4 files changed, 138 insertions(+), 1 deletion(-) diff --git a/sagemaker-core/src/sagemaker/core/remote_function/client.py b/sagemaker-core/src/sagemaker/core/remote_function/client.py index a45d644288..bc8a15c899 100644 --- a/sagemaker-core/src/sagemaker/core/remote_function/client.py +++ b/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() diff --git a/sagemaker-core/src/sagemaker/core/remote_function/job.py b/sagemaker-core/src/sagemaker/core/remote_function/job.py index efdcf7db4f..eef2370a30 100644 --- a/sagemaker-core/src/sagemaker/core/remote_function/job.py +++ b/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 diff --git a/sagemaker-core/tests/unit/remote_function/test_client.py b/sagemaker-core/tests/unit/remote_function/test_client.py index 167233cab9..2d9abd2e82 100644 --- a/sagemaker-core/tests/unit/remote_function/test_client.py +++ b/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""" diff --git a/sagemaker-core/tests/unit/remote_function/test_job.py b/sagemaker-core/tests/unit/remote_function/test_job.py index 39e2d56d28..7145f01abf 100644 --- a/sagemaker-core/tests/unit/remote_function/test_job.py +++ b/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.""" From 6b6c47fe5ce3ec7c21293c916f86c3002dd031e2 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Tue, 29 Sep 2026 13:26:06 -0700 Subject: [PATCH 2/2] test: add training_plan_arn to remote-decorator field consistency guard (#5157) --- .../feature_store/feature_processor/test_feature_scheduler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/test_feature_scheduler.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/test_feature_scheduler.py index 2cc113282a..ee45785c3b 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/test_feature_scheduler.py +++ b/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(