feat: Add support for MetricDefinitions in ModelTrainer

Author: Chad Chiang

src/sagemaker/modules/configs.py

@@ -42,6 +42,7 @@
     RemoteDebugConfig,
     SessionChainingConfig,
     InstanceGroup,
+    MetricDefinition,
 )
 
 from sagemaker.modules.utils import convert_unassigned_to_none
@@ -68,6 +69,7 @@
     "Compute",
     "Networking",
     "InputData",
+    "MetricDefinition",
 ]
 
 

src/sagemaker/modules/train/model_trainer.py

@@ -66,6 +66,7 @@
     RemoteDebugConfig,
     SessionChainingConfig,
     InputData,
+    MetricDefinition,
 )
 
 from sagemaker.modules.local_core.local_container import _LocalContainer
@@ -1290,3 +1291,27 @@ def with_checkpoint_config(
         """
         self.checkpoint_config = checkpoint_config or configs.CheckpointConfig()
         return self
+    
+    def with_metric_definitions(
+        self, metric_definitions: List[MetricDefinition]
+    ) -> "ModelTrainer":  # noqa: D412
+        """Set the metric definitions for the training job.
+        Example:
+        .. code:: python
+            from sagemaker.modules.train import ModelTrainer
+            from sagemaker.modules.configs import MetricDefinition
+            metric_definitions = [
+                MetricDefinition(
+                    name="loss",
+                    regex="Loss: (.*?)",
+                )
+            ]
+            model_trainer = ModelTrainer(
+                ...
+            ).with_metric_definitions(metric_definitions)
+        Args:
+            metric_definitions (List[MetricDefinition]):
+                The metric definitions for the training job.
+        """
+        self._metric_definitions = metric_definitions
+        return self

tests/unit/sagemaker/modules/train/test_model_trainer.py

@@ -64,6 +64,7 @@
     FileSystemDataSource,
     Channel,
     DataSource,
+    MetricDefinition,
 )
 from sagemaker.modules.distributed import Torchrun, SMP, MPI
 from sagemaker.modules.train.sm_recipes.utils import _load_recipes_cfg
@@ -705,6 +706,32 @@ def test_remote_debug_config(mock_training_job, modules_session):
         )
 
 
+@patch("sagemaker.modules.train.model_trainer.TrainingJob")
+def test_metric_definitions(mock_training_job, modules_session):
+    image_uri = DEFAULT_IMAGE
+    role = DEFAULT_ROLE
+    metric_definitions = [
+        MetricDefinition(
+            name="loss",
+            regex="Loss: (.*?);",
+        )
+    ]
+
+    model_trainer = ModelTrainer(
+        training_image=image_uri, sagemaker_session=modules_session, role=role
+    ).with_metric_definitions(metric_definitions)
+
+    with patch("sagemaker.modules.train.model_trainer.Session.upload_data") as mock_upload_data:
+        mock_upload_data.return_value = "s3://dummy-bucket/dummy-prefix"
+        model_trainer.train()
+
+        mock_training_job.create.assert_called_once()
+        assert (
+            mock_training_job.create.call_args.kwargs["algorithm_specification"].metric_definitions
+            == metric_definitions
+        )
+
+
 @patch("sagemaker.modules.train.model_trainer._get_unique_name")
 @patch("sagemaker.modules.train.model_trainer.TrainingJob")
 def test_model_trainer_full_init(mock_training_job, mock_unique_name, modules_session):