diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index b75c5a3669..3aac3f476e 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -1158,6 +1158,7 @@ def _yaml_safe_default(value): instance_count=compute.instance_count, volume_size_in_gb=compute.volume_size_in_gb, keep_alive_period_in_seconds=compute.keep_alive_period_in_seconds, + training_plan_arn=compute.training_plan_arn, ) # Build input data config (datasets resolved earlier for recipe injection) diff --git a/sagemaker-train/tests/unit/train/test_base_trainer_compute.py b/sagemaker-train/tests/unit/train/test_base_trainer_compute.py index af42e22802..c9bb4cf8aa 100644 --- a/sagemaker-train/tests/unit/train/test_base_trainer_compute.py +++ b/sagemaker-train/tests/unit/train/test_base_trainer_compute.py @@ -115,6 +115,7 @@ def test_compute_attributes_forwarded(self): instance_count=4, volume_size_in_gb=300, keep_alive_period_in_seconds=1200, + training_plan_arn=None, ) trainer.training_dataset = "s3://my-bucket/data/train/" @@ -126,6 +127,38 @@ def test_compute_attributes_forwarded(self): assert forwarded.volume_size_in_gb == 300 assert forwarded.keep_alive_period_in_seconds == 1200 + def test_training_plan_arn_forwarded(self): + trainer = _ConcreteTrainer() + trainer.compute = MagicMock( + instance_type="ml.p5.48xlarge", + instance_count=2, + volume_size_in_gb=500, + keep_alive_period_in_seconds=0, + training_plan_arn="arn:aws:sagemaker:us-west-2:123456789012:training-plan/my-plan", + ) + trainer.training_dataset = "s3://my-bucket/data/train/" + + kwargs = self._run(trainer) + + forwarded = kwargs["compute"] + assert forwarded.training_plan_arn == "arn:aws:sagemaker:us-west-2:123456789012:training-plan/my-plan" + + def test_training_plan_arn_none_when_not_set(self): + trainer = _ConcreteTrainer() + trainer.compute = MagicMock( + instance_type="ml.p4d.24xlarge", + instance_count=1, + volume_size_in_gb=30, + keep_alive_period_in_seconds=0, + training_plan_arn=None, + ) + trainer.training_dataset = "s3://my-bucket/data/train/" + + kwargs = self._run(trainer) + + forwarded = kwargs["compute"] + assert forwarded.training_plan_arn is None + def _make_hyperpod_trainer(cluster_name="my-cluster", node_count=2): """Build an SFTTrainer instance for _train_hyperpod without heavy __init__."""