Add support for MetricDefinitions in ModelTrainer

For commit: 0215512

Author: zhaoqizqwang

sagemaker-core/src/sagemaker/core/modules/configs.py

@@ -45,6 +45,7 @@
     InstanceGroup,
     TensorBoardOutputConfig,
     CheckpointConfig,
+    MetricDefinition,
 )
 
 
@@ -70,6 +71,7 @@
     "Compute",
     "Networking",
     "InputData",
+    "MetricDefinition",
 ]
 
 from sagemaker.core.modules.utils import convert_unassigned_to_none

sagemaker-train/src/sagemaker/train/model_trainer.py

@@ -64,6 +64,7 @@
     RemoteDebugConfig,
     SessionChainingConfig,
     InputData,
+    MetricDefinition,
 )
 
 from sagemaker.train.distributed import Torchrun, DistributedConfig
@@ -244,6 +245,7 @@ class ModelTrainer(BaseModel):
     _infra_check_config: Optional[InfraCheckConfig] = PrivateAttr(default=None)
     _session_chaining_config: Optional[SessionChainingConfig] = PrivateAttr(default=None)
     _remote_debug_config: Optional[RemoteDebugConfig] = PrivateAttr(default=None)
+    _metric_definitions: Optional[List[MetricDefinition]] = PrivateAttr(default=None)
 
     # Private Attributes for Recipes
     _temp_recipe_train_dir: Optional[TemporaryDirectory] = PrivateAttr(default=None)
@@ -654,6 +656,7 @@ def train(
             training_image_config=self.training_image_config,
             container_entrypoint=container_entrypoint,
             container_arguments=container_arguments,
+            metric_definitions=self._metric_definitions,
         )
 
         resource_config = self.compute._to_resource_config()
@@ -1496,3 +1499,29 @@ 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

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

@@ -64,6 +64,7 @@
     FileSystemDataSource,
     Channel,
     DataSource,
+    MetricDefinition,
 )
 from sagemaker.train.distributed import Torchrun, SMP, MPI
 from sagemaker.train.sm_recipes.utils import _load_recipes_cfg
@@ -830,6 +831,7 @@ def mock_upload_data(path, bucket, key_prefix):
             training_input_mode=training_input_mode,
             training_image=training_image,
             algorithm_name=None,
+            metric_definitions=None,
             container_entrypoint=DEFAULT_ENTRYPOINT,
             container_arguments=DEFAULT_ARGUMENTS,
             training_image_config=training_image_config,
@@ -1321,3 +1323,27 @@ def test_input_merge(mock_training_job, modules_session):
             input_mode="File",
         ),
     ]
+
+@patch("sagemaker.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.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
+        )