show_metrics() Enhancement: Display MLFlow metrics for OSS models

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,30 +303,79 @@ def show_metrics(
                 defaults to now.
 
         Returns:
-            pandas.DataFrame containing the extracted metrics with columns
-            ["global_step", <metric_name>].
+            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(
                 "No training job found. Call .train() first, then call .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}"
+        )
+
+        plot_training_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
         if hasattr(training_job, 'training_job_name'):
             job_id = training_job.training_job_name

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")
+
+    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 = "test-oss-model"
 
-        with pytest.raises(NotImplementedError, match="only supported for Nova models"):
+        with pytest.raises(ValueError, match="requires MLflow to be configured"):
             trainer.show_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()
+        mock_job.training_job_name = "oss-sft-job"
+        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 = "test-oss-model"
+
+        trainer.show_metrics(metrics=["loss"])
+
+        mock_plot.assert_called_once_with(mock_job, metrics=["loss"])