feat: expose the created job as latest_training_job and return it from ModelTrainer.train()

docs/model_customization/model_customization.rst

@@ -154,14 +154,15 @@ Key Features
 
     # Re-attach (HyperPod — needs cluster name for log group resolution)
     from sagemaker.train import SFTTrainer
+    from sagemaker.core.resources import TrainingJob
     from sagemaker.core.training.configs import HyperPodCompute
 
     trainer = SFTTrainer(
         model="nova-textgeneration-micro",
         training_dataset="s3://unused",
         compute=HyperPodCompute(cluster_name="my-cluster", instance_type="ml.p5.48xlarge")
     )
-    trainer._latest_training_job = "my-hp-job"
+    trainer.latest_training_job = TrainingJob.get("my-hp-job")
     df = trainer.show_metrics()
 
 **Monitoring: stream_logs()**

docs/training/index.rst

@@ -434,7 +434,7 @@ Provide an ordered list of candidate instance types and the platform launches th
 
 .. code-block:: python
 
-   training_job = model_trainer._latest_training_job
+   training_job = model_trainer.latest_training_job
    training_job.refresh()
 
    resource_config = training_job.resource_config

sagemaker-serve/src/sagemaker/serve/model_builder.py

@@ -3429,7 +3429,7 @@ def _build_single_modelbuilder(
                 raise ValueError(
                     "The trainer passed to ModelBuilder does not have a completed training job. "
                     "Either call trainer.train() first, or manually set "
-                    "trainer._latest_training_job = TrainingJob.get(training_job_name='<job-name>') "
+                    "trainer.latest_training_job = TrainingJob.get(training_job_name='<job-name>') "
                     "to attach a previously completed job."
                 )
 

sagemaker-train/src/sagemaker/train/base_trainer.py

@@ -120,7 +120,35 @@ class BaseTrainer(ABC):
     input_data_config: Optional[List[Union[Channel, InputData]]] = None
     environment: Optional[Dict[str, str]] = None
     training_image: Optional[str] = None
-    latest_training_job: Optional[TrainingJob] = None
+
+    @property
+    def latest_training_job(self) -> Optional[TrainingJob]:
+        """The ``TrainingJob`` created by the most recent ``train()`` call.
+
+        ``None`` until ``train()`` has submitted a job (a ``dry_run`` or a
+        ``PipelineSession`` capture submits nothing). The resource is bound to the
+        trainer's session and region, so ``refresh()``, ``wait()``, ``stop()`` and the
+        ``training_job_name`` / ``training_job_arn`` attributes work without a second
+        lookup. ``show_metrics()`` and ``stream_logs()`` read this same job.
+
+        Assign a ``TrainingJob`` to re-attach a trainer to a job it did not create::
+
+            trainer.latest_training_job = TrainingJob.get("my-job")
+
+        Returns:
+            Optional[TrainingJob]: The latest created training job, or ``None``.
+        """
+        return getattr(self, "_latest_training_job", None)
+
+    @latest_training_job.setter
+    def latest_training_job(self, training_job: Optional[TrainingJob]) -> None:
+        if training_job is not None and not isinstance(training_job, TrainingJob):
+            raise TypeError(
+                "latest_training_job must be a sagemaker.core.resources.TrainingJob or None, "
+                f"got {type(training_job).__name__}. Use TrainingJob.get(<name>) to attach "
+                "an existing job."
+            )
+        self._latest_training_job = training_job
 
     @classmethod
     @_telemetry_emitter(
@@ -186,6 +214,7 @@ def __init__(
         self.disable_output_compression = disable_output_compression
         self.notification_rule_arn = None
         self.source_code = None
+        self._latest_training_job: Optional[TrainingJob] = None
 
         # Set up notifications if configured
         if notifications:
@@ -518,17 +547,17 @@ def show_metrics(
             ValueError: If no training job has been run yet, no logs/metrics
                 are found, or MLflow is not configured for OSS models.
         """
-        # Resolve the job reference. Prefer _latest_training_job (CreateTrainingJob),
+        # Resolve the job reference. Prefer latest_training_job (CreateTrainingJob),
         # fall back to _latest_job (generic CreateJob API used by MTRL).
-        resolved_job = getattr(self, "_latest_training_job", None)
+        resolved_job = self.latest_training_job
         if resolved_job is None:
             latest_job = getattr(self, "_latest_job", None)
             if latest_job is None:
                 raise ValueError(
                     "No training job found. Call .train() first, then call .show_metrics() "
-                    "to view training metrics. If training has already completed, set the "
-                    "job name directly via trainer._latest_training_job = '<job-name>' or "
-                    "trainer._latest_job = '<job-name>'."
+                    "to view training metrics. If training has already completed, attach "
+                    "the job with trainer.latest_training_job = TrainingJob.get('<job-name>') "
+                    "or trainer._latest_job = '<job-name>'."
                 )
             # For AgentRFTJob (MTRL), delegate to its own get_training_metrics()
             # which fetches metrics from MLflow via the Job API (not TrainingJob).
@@ -787,17 +816,17 @@ def stream_logs(
         Raises:
             ValueError: If no training job has been run yet.
         """
-        # Resolve the job reference. Prefer _latest_training_job (CreateTrainingJob),
+        # Resolve the job reference. Prefer latest_training_job (CreateTrainingJob),
         # fall back to _latest_job (generic CreateJob API used by MTRL).
-        resolved_job = getattr(self, "_latest_training_job", None)
+        resolved_job = self.latest_training_job
         if resolved_job is None:
             latest_job = getattr(self, "_latest_job", None)
             if latest_job is None:
                 raise ValueError(
                     "No training job found. Call .train(wait=False) first, "
                     "then call .stream_logs() to stream logs in real-time. "
-                    "If training has already completed, set the job name directly via "
-                    "trainer._latest_training_job = '<job-name>' or "
+                    "If training has already completed, attach the job with "
+                    "trainer.latest_training_job = TrainingJob.get('<job-name>') or "
                     "trainer._latest_job = '<job-name>'."
                 )
             # For AgentRFTJob (MTRL), delegate to its own stream_logs which

sagemaker-train/src/sagemaker/train/dpo_trainer.py

@@ -433,5 +433,5 @@ def train(
             except TimeoutExceededError as e:
                 logger.error("Error: %s", e)
 
-        self.latest_training_job = training_job
+        self._latest_training_job = training_job
         return training_job

sagemaker-train/src/sagemaker/train/model_trainer.py

@@ -352,6 +352,35 @@ class ModelTrainer(BaseModel):
 
     config_mgr: SageMakerConfig = SageMakerConfig()
 
+    @property
+    def latest_training_job(self) -> Optional[resources.TrainingJob]:
+        """The ``TrainingJob`` created by the most recent ``train()`` call.
+
+        ``None`` until ``train()`` has submitted a job (a ``dry_run`` or a
+        ``PipelineSession`` capture submits nothing). The resource is bound to the
+        session and region the job was created with, so ``refresh()``, ``wait()``,
+        ``stop()`` and the ``training_job_name`` / ``training_job_arn`` attributes
+        work without a second lookup.
+
+        Assign a ``TrainingJob`` to re-attach a trainer to a job it did not create::
+
+            trainer.latest_training_job = TrainingJob.get("my-job")
+
+        Returns:
+            Optional[TrainingJob]: The latest created training job, or ``None``.
+        """
+        return self._latest_training_job
+
+    @latest_training_job.setter
+    def latest_training_job(self, training_job: Optional[resources.TrainingJob]) -> None:
+        if training_job is not None and not isinstance(training_job, resources.TrainingJob):
+            raise TypeError(
+                "latest_training_job must be a sagemaker.core.resources.TrainingJob or None, "
+                f"got {type(training_job).__name__}. Use TrainingJob.get(<name>) to attach "
+                "an existing job."
+            )
+        self._latest_training_job = training_job
+
     @property
     def output_data(self) -> str:
         """The S3 URI of the training job's ``output.tar.gz`` archive.
@@ -915,7 +944,7 @@ def train(
         wait: Optional[bool] = True,
         logs: Optional[bool] = True,
         dry_run: bool = False,
-    ):
+    ) -> Optional[resources.TrainingJob]:
         """Train a model using AWS SageMaker.
 
         Args:
@@ -934,6 +963,13 @@ def train(
                 If True, runs all validation (configuration, input resolution, hyperparameters)
                 without submitting a job. Returns None on success, raises on validation failure.
                 Defaults to False.
+
+        Returns:
+            Optional[TrainingJob]: The created ``sagemaker.core.resources.TrainingJob`` for a
+                SageMaker training job, bound to the trainer's session and region. ``None`` for
+                a ``dry_run``, a ``PipelineSession`` capture (the step runs the job later), and
+                local-mode training. The same object is available afterwards as
+                ``latest_training_job``.
         """
         training_request = self._create_training_job_args(input_data_config=input_data_config)
 
@@ -980,7 +1016,7 @@ def train(
                 training_job.wait(logs=logs)
             if logs and not wait:
                 logger.warning(
-                    "Not displaing the training container logs as 'wait' is set to False."
+                    "Not displaying the training container logs as 'wait' is set to False."
                 )
 
         else:
@@ -1007,6 +1043,9 @@ def train(
             local_container.train(wait)
         if self._temp_code_dir is not None:
             self._temp_code_dir.cleanup()
+        if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB:
+            return self._latest_training_job
+        return None
 
     def _resolve_staging_bucket(self) -> tuple[str, str]:
         """Resolve the S3 bucket and key prefix for staging training artifacts.