From 8d3bf7c1ea228ef43cf2e779923167706282d2ee Mon Sep 17 00:00:00 2001 From: Ealynn Hsu Date: Fri, 10 Jul 2026 15:40:44 +0000 Subject: [PATCH 1/3] show_metrics() Enhancement: Display MLFlow metrics for OSS models --- .../src/sagemaker/train/base_trainer.py | 82 ++++++++++++++++--- .../common_utils/test_cloudwatch_metrics.py | 30 +++++-- 2 files changed, 94 insertions(+), 18 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index ab353c1612..93d6a1b23f 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -32,6 +32,7 @@ _validate_hyperparameter_values, _get_smhp_replicas_enum, ) +from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics from sagemaker.train.common_utils.mlflow_config_utils import resolve_mlflow_tracking_fields from sagemaker.train.common_utils.validator import validate_hyperpod_compute from sagemaker.train.common_utils.cloudwatch_metrics import fetch_and_plot_metrics, _get_smhp_log_group @@ -283,7 +284,11 @@ def show_metrics( start_time: Optional[Any] = None, end_time: Optional[Any] = None, ) -> Any: - """Plot training metrics extracted from CloudWatch logs using matplotlib. + """Plot training metrics from CloudWatch logs (Nova) or MLflow (OSS). + + For Nova models, parses CloudWatch logs for training_loss, lr, and reward_score. + For non-Nova (OSS) models, pulls metrics from MLflow (requires mlflow_resource_arn + to be configured on the trainer or auto-resolved). Args: metrics: Optional list of metric names to plot. If None, plots all @@ -298,22 +303,14 @@ def show_metrics( defaults to now. Returns: - pandas.DataFrame containing the extracted metrics with columns - ["global_step", ]. + pandas.DataFrame containing the extracted metrics. Raises: NotImplementedError: If the training technique does not support metric extraction (e.g., DPO). - ValueError: If no training job has been run yet, or no logs/metrics - are found. + ValueError: If no training job has been run yet, no logs/metrics + are found, or MLflow is not configured for OSS models. """ - # Gate to Nova models only - model_name = getattr(self, '_model_name', None) - if model_name and not _is_nova_model(model_name): - raise NotImplementedError( - "show_metrics() is currently only supported for Nova models. " - ) - # Validate that we have a training job to get metrics from if not hasattr(self, '_latest_training_job') or self._latest_training_job is None: raise ValueError( @@ -321,8 +318,67 @@ def show_metrics( "to view training metrics." ) - # Resolve job ID + # Route based on model type + model_name = getattr(self, '_model_name', None) + is_nova = _is_nova_model(model_name) if model_name else False + + if is_nova: + return self._show_metrics_cloudwatch(metrics, starting_step, ending_step, start_time, end_time) + else: + return self._show_metrics_mlflow(metrics, starting_step, ending_step) + + def _show_metrics_mlflow( + self, + metrics: Optional[List[str]] = None, + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, + ) -> None: + """Pull and plot training metrics from MLflow for non-Nova models.""" + training_job = self._latest_training_job + + # Resolve the TrainingJob object if it's a string + if isinstance(training_job, str): + logger.info(f"Resolving training job: {training_job}") + training_job = TrainingJob.get(training_job_name=training_job) + + # Validate MLflow is configured + mlflow_config = getattr(training_job, 'mlflow_config', None) + if not mlflow_config or not getattr(mlflow_config, 'mlflow_resource_arn', None): + raise ValueError( + "show_metrics() for non-Nova models requires MLflow to be configured. " + "Either pass mlflow_resource_arn when creating the trainer, or ensure " + "your account has an MLflow app set up." + ) + + mlflow_details = getattr(training_job, 'mlflow_details', None) + if not mlflow_details or not getattr(mlflow_details, 'mlflow_run_id', None): + raise ValueError( + "No MLflow run ID found on the training job. " + "MLflow metrics are only available after the job completes. " + "If the job is still running, wait for it to finish and try again. " + f"MLflow app ARN: {mlflow_config.mlflow_resource_arn}" + ) + + logger.info( + f"Fetching metrics from MLflow app: {mlflow_config.mlflow_resource_arn}, " + f"run: {mlflow_details.mlflow_run_id}" + ) + + from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics as _plot_mlflow_metrics + _plot_mlflow_metrics(training_job, metrics=metrics) + + def _show_metrics_cloudwatch( + self, + metrics: Optional[List[str]] = None, + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, + start_time: Optional[Any] = None, + end_time: Optional[Any] = None, + ) -> Any: + """Parse and plot training metrics from CloudWatch logs (Nova models).""" training_job = self._latest_training_job + + # Resolve job ID if hasattr(training_job, 'training_job_name'): job_id = training_job.training_job_name elif isinstance(training_job, str): diff --git a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py index aa880cfc85..cf135fef68 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py @@ -272,10 +272,30 @@ def test_smtj_stops_on_completed(self, mock_handler_cls, mock_get_job): mock_get_job.assert_called() - def test_show_metrics_rejects_oss_models(self): - """show_metrics() raises NotImplementedError for non-Nova models.""" - trainer = self._make_trainer(latest_job="some-job") - trainer._model_name = "test-oss-model" - with pytest.raises(NotImplementedError, match="only supported for Nova models"): + def test_show_metrics_oss_without_mlflow_raises(self): + """show_metrics() raises ValueError for non-Nova models without MLflow configured.""" + trainer = self._make_trainer(latest_job=MagicMock( + training_job_name="some-job", + mlflow_config=None, + mlflow_details=None, + )) + trainer._model_name = "meta-textgeneration-llama-3-2-1b-instruct" + + with pytest.raises(ValueError, match="requires MLflow to be configured"): trainer.show_metrics() + + @patch("sagemaker.train.common_utils.metrics_visualizer.plot_training_metrics") + def test_show_metrics_oss_with_mlflow_delegates(self, mock_plot): + """show_metrics() for OSS models with MLflow configured calls plot_training_metrics.""" + mock_job = MagicMock() + mock_job.training_job_name = "oss-sft-job" + mock_job.mlflow_config.mlflow_resource_arn = "arn:aws:sagemaker:us-east-1:123:mlflow-app/app-123" + mock_job.mlflow_details.mlflow_run_id = "run-abc123" + + trainer = self._make_trainer(latest_job=mock_job) + trainer._model_name = "meta-textgeneration-llama-3-2-1b-instruct" + + trainer.show_metrics(metrics=["loss"]) + + mock_plot.assert_called_once_with(mock_job, metrics=["loss"]) From e0380ad3213af4568202c43c265924c086acdce9 Mon Sep 17 00:00:00 2001 From: Ealynn Hsu Date: Fri, 10 Jul 2026 16:00:52 +0000 Subject: [PATCH 2/3] Update unit tests --- .../unit/train/common_utils/test_cloudwatch_metrics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py index cf135fef68..cd02d56596 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py @@ -280,7 +280,7 @@ def test_show_metrics_oss_without_mlflow_raises(self): mlflow_config=None, mlflow_details=None, )) - trainer._model_name = "meta-textgeneration-llama-3-2-1b-instruct" + trainer._model_name = "test-oss-model" with pytest.raises(ValueError, match="requires MLflow to be configured"): trainer.show_metrics() @@ -290,11 +290,11 @@ def test_show_metrics_oss_with_mlflow_delegates(self, mock_plot): """show_metrics() for OSS models with MLflow configured calls plot_training_metrics.""" mock_job = MagicMock() mock_job.training_job_name = "oss-sft-job" - mock_job.mlflow_config.mlflow_resource_arn = "arn:aws:sagemaker:us-east-1:123:mlflow-app/app-123" + mock_job.mlflow_config.mlflow_resource_arn = "arn:aws:sagemaker:us-east-1:012345678910:mlflow-app/app-123" mock_job.mlflow_details.mlflow_run_id = "run-abc123" trainer = self._make_trainer(latest_job=mock_job) - trainer._model_name = "meta-textgeneration-llama-3-2-1b-instruct" + trainer._model_name = "test-oss-model" trainer.show_metrics(metrics=["loss"]) From 7c46a4f542bc14b390f9219902136ce50b0717dc Mon Sep 17 00:00:00 2001 From: Ealynn Hsu Date: Fri, 10 Jul 2026 17:38:28 +0000 Subject: [PATCH 3/3] Address code comments --- sagemaker-train/src/sagemaker/train/base_trainer.py | 6 ++---- .../unit/train/common_utils/test_cloudwatch_metrics.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index 93d6a1b23f..ff3bdcd751 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -364,8 +364,7 @@ def _show_metrics_mlflow( f"run: {mlflow_details.mlflow_run_id}" ) - from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics as _plot_mlflow_metrics - _plot_mlflow_metrics(training_job, metrics=metrics) + plot_training_metrics(training_job, metrics=metrics) def _show_metrics_cloudwatch( self, @@ -376,9 +375,8 @@ def _show_metrics_cloudwatch( end_time: Optional[Any] = None, ) -> Any: """Parse and plot training metrics from CloudWatch logs (Nova models).""" + training_job = self._latest_training_job - - # Resolve job ID if hasattr(training_job, 'training_job_name'): job_id = training_job.training_job_name elif isinstance(training_job, str): diff --git a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py index cd02d56596..a67fd31c1d 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py @@ -285,7 +285,7 @@ def test_show_metrics_oss_without_mlflow_raises(self): with pytest.raises(ValueError, match="requires MLflow to be configured"): trainer.show_metrics() - @patch("sagemaker.train.common_utils.metrics_visualizer.plot_training_metrics") + @patch("sagemaker.train.base_trainer.plot_training_metrics") def test_show_metrics_oss_with_mlflow_delegates(self, mock_plot): """show_metrics() for OSS models with MLflow configured calls plot_training_metrics.""" mock_job = MagicMock()