fix: pass training_plan_arn through serverful compute reconstruction (#6258)

The serverful path in _train_serverful_smtj() reconstructs TrainingJobCompute
with only 4 fields, silently dropping training_plan_arn. This causes customers
using Flexible Training Plans with high-level trainers (SFTTrainer, DPOTrainer,
etc.) to have their plan ARN ignored.

Add training_plan_arn to the reconstruction and update tests.

Author: mrpic-amzn

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)

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__."""