From b0a9e928782e3971ee23310c976d9a9cca8a8e21 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Mon, 12 Oct 2026 11:58:12 -0500 Subject: [PATCH] fix(train): store the DPO training job as _latest_training_job DPOTrainer.train() saved the job as self.latest_training_job, while the other trainers (SFT, RLAIF, RLVR, BaseTrainer) use _latest_training_job and that is what the readers check: BaseTrainer.stream_logs()/show_metrics() and ModelBuilder (hasattr(model, '_latest_training_job')). After DPOTrainer.train(wait=False), stream_logs() and show_metrics() raised 'No training job found', and ModelBuilder(model=dpo_trainer) did not recognize the trained model. Set _latest_training_job, keeping the public attribute as an alias. --- sagemaker-train/src/sagemaker/train/dpo_trainer.py | 3 +++ sagemaker-train/tests/unit/train/test_dpo_trainer.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index b0c0cf6871..8978fe03c4 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/sagemaker-train/src/sagemaker/train/dpo_trainer.py @@ -433,5 +433,8 @@ def train( except TimeoutExceededError as e: logger.error("Error: %s", e) + self._latest_training_job = training_job + # Kept for backwards compatibility; the other trainers and the code that reads the + # job (stream_logs, show_metrics, ModelBuilder) use _latest_training_job. self.latest_training_job = training_job return training_job diff --git a/sagemaker-train/tests/unit/train/test_dpo_trainer.py b/sagemaker-train/tests/unit/train/test_dpo_trainer.py index 8dddd315b2..2754ab1687 100644 --- a/sagemaker-train/tests/unit/train/test_dpo_trainer.py +++ b/sagemaker-train/tests/unit/train/test_dpo_trainer.py @@ -529,6 +529,8 @@ def test_train_passes_wait_timeout( trainer.train(wait=True, wait_timeout=600) mock_wait.assert_called_once_with(mock_training_job, timeout=600, poll=5) + # BaseTrainer.stream_logs/show_metrics and ModelBuilder read _latest_training_job. + assert trainer._latest_training_job is mock_training_job @patch("sagemaker.train.common_utils.trainer_wait.wait") @patch("sagemaker.train.dpo_trainer._resolve_model_and_name")