From bc92c8f658eea0a9db5c06dc4a0b921d3c08dc03 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Fri, 23 Oct 2026 16:40:26 +0000 Subject: [PATCH] feat: expose the created job as latest_training_job and return it from ModelTrainer.train() `ModelTrainer.train()` returned `None` and kept the created job on the private `_latest_training_job`. A caller that needed the job name or ARN after `train(wait=False)` had to read the private attribute or page `ListTrainingJobs` by name prefix. The SDK's own docs, seven example notebooks, and `ModelBuilder` read the private attribute too, and `DPOTrainer` stored its job on a different attribute than the other serverless trainers, so `show_metrics()` and `stream_logs()` could not find a DPO job. Add a public `latest_training_job` property with a typed setter on `ModelTrainer` and `BaseTrainer`, return the created `TrainingJob` from `ModelTrainer.train()` for a SageMaker training job (`None` for a dry run, a pipeline capture, and local mode), store the DPO job on the shared attribute, and point the docs and error messages at the public name. --- X-AI-Prompt: How do I get the training job that ModelTrainer.train() created without reading a private attribute? X-AI-Tool: claude-code --- .../model_customization.rst | 3 +- docs/training/index.rst | 2 +- .../src/sagemaker/serve/model_builder.py | 2 +- .../src/sagemaker/train/base_trainer.py | 49 +++- .../src/sagemaker/train/dpo_trainer.py | 2 +- .../src/sagemaker/train/model_trainer.py | 43 +++- .../unit/train/test_latest_training_job.py | 215 ++++++++++++++++++ 7 files changed, 300 insertions(+), 16 deletions(-) create mode 100644 sagemaker-train/tests/unit/train/test_latest_training_job.py diff --git a/docs/model_customization/model_customization.rst b/docs/model_customization/model_customization.rst index 6e0095184d..5d36c2a6eb 100644 --- a/docs/model_customization/model_customization.rst +++ b/docs/model_customization/model_customization.rst @@ -154,6 +154,7 @@ 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( @@ -161,7 +162,7 @@ Key Features 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()** diff --git a/docs/training/index.rst b/docs/training/index.rst index e2eabfae54..f2e29da422 100644 --- a/docs/training/index.rst +++ b/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 diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 556525a45e..8587ff7c6f 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/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='') " + "trainer.latest_training_job = TrainingJob.get(training_job_name='') " "to attach a previously completed job." ) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index a60f5ba52b..00ea3339ad 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/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() 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 = '' or " - "trainer._latest_job = ''." + "to view training metrics. If training has already completed, attach " + "the job with trainer.latest_training_job = TrainingJob.get('') " + "or trainer._latest_job = ''." ) # 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 = '' or " + "If training has already completed, attach the job with " + "trainer.latest_training_job = TrainingJob.get('') or " "trainer._latest_job = ''." ) # For AgentRFTJob (MTRL), delegate to its own stream_logs which diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index b0c0cf6871..d5c64d52c9 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/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 diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 21ecee8314..23a3475b94 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/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() 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. diff --git a/sagemaker-train/tests/unit/train/test_latest_training_job.py b/sagemaker-train/tests/unit/train/test_latest_training_job.py new file mode 100644 index 0000000000..19c9089c01 --- /dev/null +++ b/sagemaker-train/tests/unit/train/test_latest_training_job.py @@ -0,0 +1,215 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""``latest_training_job`` is the public handle for the job a trainer created. + +Before this property the only handle was the private ``_latest_training_job``, so callers +that needed the job name or ARN after ``train(wait=False)`` had to read the private +attribute or page ``ListTrainingJobs`` by name prefix. ``ModelTrainer.train()`` also +returned ``None``, so the job could not be chained from the call that created it. +""" + +from __future__ import absolute_import + +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from sagemaker.core.helper.session_helper import Session +from sagemaker.core.resources import TrainingJob +from sagemaker.core.shapes import OutputDataConfig +from sagemaker.core.training.configs import Compute +from sagemaker.train.base_trainer import BaseTrainer +from sagemaker.train.dpo_trainer import DPOTrainer +from sagemaker.train.model_trainer import ModelTrainer +from sagemaker.train.sft_trainer import SFTTrainer +from sagemaker.train.common import TrainingType + +ROLE = "arn:aws:iam::123456789012:role/SageMakerRole" +IMAGE = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-algo:latest" +BUCKET = "my-bucket" + + +@pytest.fixture(scope="module", autouse=True) +def modules_session(): + with ( + patch("sagemaker.train.Session", spec=Session) as session_mock, + patch("sagemaker.train.defaults.resolve_and_validate_role", return_value=ROLE), + ): + session_instance = session_mock.return_value + session_instance.default_bucket.return_value = BUCKET + session_instance.get_caller_identity_arn.return_value = ROLE + session_instance.default_bucket_prefix = None + session_instance.boto_session = MagicMock(spec="boto3.session.Session") + session_instance.boto_region_name = "us-west-2" + yield session_instance + + +@pytest.fixture +def model_trainer(modules_session): + return ModelTrainer( + training_image=IMAGE, + role=ROLE, + sagemaker_session=modules_session, + compute=Compute(instance_type="ml.m5.xlarge", instance_count=1), + output_data_config=OutputDataConfig(s3_output_path=f"s3://{BUCKET}/output"), + ) + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_train_returns_the_created_job(mock_training_job, model_trainer): + created = TrainingJob(training_job_name="created-job") + mock_training_job.create.return_value = created + + returned = model_trainer.train(wait=False) + + assert returned is created + assert model_trainer.latest_training_job is created + assert model_trainer.latest_training_job.training_job_name == "created-job" + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_train_with_wait_returns_the_job_after_waiting( + mock_training_job, model_trainer +): + created = Mock(spec=TrainingJob) + mock_training_job.create.return_value = created + + returned = model_trainer.train(wait=True, logs=False) + + created.wait.assert_called_once_with(logs=False) + assert returned is created + + +def test_model_trainer_latest_training_job_is_none_before_train(model_trainer): + assert model_trainer.latest_training_job is None + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_dry_run_returns_none_and_creates_nothing(mock_training_job, model_trainer): + assert model_trainer.train(wait=False, dry_run=True) is None + mock_training_job.create.assert_not_called() + assert model_trainer.latest_training_job is None + + +def test_model_trainer_latest_training_job_setter_attaches_an_existing_job(model_trainer): + existing = TrainingJob(training_job_name="existing-job") + + model_trainer.latest_training_job = existing + + assert model_trainer.latest_training_job is existing + assert model_trainer.output_data.endswith("/existing-job/output/output.tar.gz") + + +def test_model_trainer_latest_training_job_setter_rejects_a_job_name(model_trainer): + with pytest.raises(TypeError, match="TrainingJob.get"): + model_trainer.latest_training_job = "existing-job" + + +class _StubTrainer(BaseTrainer): + """Minimal concrete trainer: BaseTrainer is abstract on ``train``.""" + + def train(self, input_data_config=None, wait=True, logs=True, wait_timeout=None, dry_run=False): + return None + + +@patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") +def test_base_trainer_latest_training_job_is_none_before_train(mock_get_session): + mock_get_session.return_value = Mock(sagemaker_config={}) + + trainer = _StubTrainer() + + assert trainer.latest_training_job is None + + +@patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") +def test_base_trainer_latest_training_job_setter_round_trips(mock_get_session): + mock_get_session.return_value = Mock(sagemaker_config={}) + trainer = _StubTrainer() + existing = TrainingJob(training_job_name="existing-job") + + trainer.latest_training_job = existing + + assert trainer.latest_training_job is existing + with pytest.raises(TypeError, match="TrainingJob.get"): + trainer.latest_training_job = "existing-job" + + +@patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") +def test_base_trainer_show_metrics_reads_the_public_handle(mock_get_session): + mock_get_session.return_value = Mock(sagemaker_config={}) + trainer = _StubTrainer() + + with pytest.raises(ValueError, match="trainer.latest_training_job = TrainingJob.get"): + trainer.show_metrics() + with pytest.raises(ValueError, match="trainer.latest_training_job = TrainingJob.get"): + trainer.stream_logs() + + +def _serverless_trainer_patches(module): + return [ + patch("sagemaker.train.common_utils.finetune_utils._get_beta_session"), + patch(f"{module}.validate_data_path_exists"), + patch(f"{module}._resolve_model_and_name", return_value=("test-model", "test-model")), + patch(f"{module}.TrainDefaults.get_role", return_value=ROLE), + patch( + f"{module}.TrainDefaults.get_sagemaker_session", + return_value=Mock(sagemaker_config={}), + ), + patch(f"{module}._get_unique_name", return_value="test-job-name"), + patch(f"{module}._validate_and_resolve_model_package_group", return_value="test-group"), + patch(f"{module}._create_input_data_config", return_value=[Mock()]), + patch(f"{module}._convert_input_data_to_channels", return_value=[Mock()]), + patch(f"{module}._create_output_config", return_value=Mock()), + patch(f"{module}._create_serverless_config", return_value=Mock()), + patch(f"{module}._create_mlflow_config", return_value=Mock()), + patch(f"{module}._create_model_package_config", return_value=Mock()), + ] + + +@pytest.mark.parametrize( + "trainer_cls, module", + [ + (SFTTrainer, "sagemaker.train.sft_trainer"), + (DPOTrainer, "sagemaker.train.dpo_trainer"), + ], +) +def test_serverless_trainers_expose_the_created_job(trainer_cls, module): + options = Mock() + options.to_dict.return_value = {"learning_rate": "0.001"} + created = Mock(spec=TrainingJob) + created.training_job_name = "test-job-name" + patches = _serverless_trainer_patches(module) + [ + patch( + f"{module}._get_fine_tuning_options_and_model_arn", return_value=(options, "arn", False) + ), + patch("sagemaker.core.resources.TrainingJob.create", return_value=created), + ] + for active in patches: + active.start() + try: + trainer = trainer_cls( + model="test-model", + training_type=TrainingType.LORA, + model_package_group="test-group", + training_dataset="s3://bucket/train", + ) + assert trainer.latest_training_job is None + + returned = trainer.train(wait=False) + + assert returned is created + assert trainer.latest_training_job is created + assert trainer.latest_training_job.training_job_name == "test-job-name" + finally: + for active in patches: + active.stop()