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")