feat: jumpstart training metrics

Author: evakravi

src/sagemaker/jumpstart/artifacts.py

@@ -12,8 +12,9 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart ECR and S3 URIs."""
 from __future__ import absolute_import
+from copy import deepcopy
 import os
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 from sagemaker import image_uris
 from sagemaker.jumpstart.constants import (
     ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE,
@@ -363,3 +364,32 @@ def _retrieve_default_environment_variables(
     for environment_variable in model_specs.inference_environment_variables:
         default_environment_variables[environment_variable.name] = str(environment_variable.default)
     return default_environment_variables
+
+
+def _retrieve_default_training_metric_definitions(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+) -> Optional[List[Dict[str, str]]]:
+    """Retrieves the default training metric definitions for the model.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the default training metric definitions.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            default training metric definitions.
+        region (Optional[str]): Region for which to retrieve default training metric
+            definitions.
+
+    Returns:
+        list: the default training metric definitions to use for the model or None.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
+        region=region, model_id=model_id, version=model_version
+    )
+
+    return deepcopy(model_specs.metrics) if model_specs.metrics else None

src/sagemaker/jumpstart/types.py

@@ -292,6 +292,7 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "training_dependencies",
         "training_vulnerabilities",
         "deprecated",
+        "metrics",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -328,6 +329,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.training_dependencies: List[str] = json_obj["training_dependencies"]
         self.training_vulnerabilities: List[str] = json_obj["training_vulnerabilities"]
         self.deprecated: bool = bool(json_obj["deprecated"])
+        self.metrics: Optional[List[Dict[str, str]]] = json_obj.get("metrics", None)
 
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(

src/sagemaker/metric_definitions.py

@@ -0,0 +1,52 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Accessors to retrieve metric definition for training jobs."""
+
+from __future__ import absolute_import
+
+import logging
+from typing import Dict, Optional, List
+
+from sagemaker.jumpstart import utils as jumpstart_utils
+from sagemaker.jumpstart import artifacts
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve_default(
+    region: Optional[str] = None,
+    model_id: Optional[str] = None,
+    model_version: Optional[str] = None,
+) -> Optional[List[Dict[str, str]]]:
+    """Retrieves the default training metric definitions for the model matching the given arguments.
+
+    Args:
+        region (str): The AWS Region for which to retrieve the default default training metric
+            definitions. Defaults to ``None``.
+        model_id (str): The model ID of the model for which to
+            retrieve the default training metric definitions. (Default: None).
+        model_version (str): The version of the model for which to retrieve the
+            default training metric definitions. (Default: None).
+    Returns:
+        list: The default metric definitions to use for the model or None.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
+        raise ValueError(
+            "Must specify `model_id` and `model_version` when retrieving default training "
+            "metric definitions."
+        )
+
+    return artifacts._retrieve_default_training_metric_definitions(model_id, model_version, region)

tests/unit/sagemaker/jumpstart/constants.py

@@ -1183,6 +1183,7 @@
     "training_dependencies": [],
     "training_vulnerabilities": [],
     "deprecated": False,
+    "metrics": [{"Regex": "val_accuracy: ([0-9\\.]+)", "Name": "pytorch-ic:val-accuracy"}],
 }
 
 BASE_HEADER = {

tests/unit/sagemaker/metric_definitions/__init__.py

@@ -0,0 +1,80 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+
+from mock.mock import patch
+import pytest
+
+from sagemaker import metric_definitions
+
+from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_default_metric_definitions(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_spec_from_base_spec
+
+    model_id = "pytorch-ic-mobilenet-v2"
+    region = "us-west-2"
+
+    definitions = metric_definitions.retrieve_default(
+        region=region,
+        model_id=model_id,
+        model_version="*",
+    )
+    assert definitions == [
+        {"Regex": "val_accuracy: ([0-9\\.]+)", "Name": "pytorch-ic:val-accuracy"}
+    ]
+
+    patched_get_model_specs.assert_called_once_with(region=region, model_id=model_id, version="*")
+
+    patched_get_model_specs.reset_mock()
+
+    definitions = metric_definitions.retrieve_default(
+        region=region,
+        model_id=model_id,
+        model_version="1.*",
+    )
+    assert definitions == [
+        {"Regex": "val_accuracy: ([0-9\\.]+)", "Name": "pytorch-ic:val-accuracy"}
+    ]
+
+    patched_get_model_specs.assert_called_once_with(region=region, model_id=model_id, version="1.*")
+
+    patched_get_model_specs.reset_mock()
+
+    with pytest.raises(KeyError):
+        metric_definitions.retrieve_default(
+            region=region,
+            model_id="blah",
+            model_version="*",
+        )
+
+    with pytest.raises(ValueError):
+        metric_definitions.retrieve_default(
+            region="mars-south-1",
+            model_id=model_id,
+            model_version="*",
+        )
+
+    with pytest.raises(ValueError):
+        metric_definitions.retrieve_default(
+            model_version="*",
+        )
+
+    with pytest.raises(ValueError):
+        metric_definitions.retrieve_default(
+            model_id=model_id,
+        )