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

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

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

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

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

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

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

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