feat: jumpstart training metrics

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/integ/sagemaker/jumpstart/constants.py

@@ -42,6 +42,7 @@ def _to_s3_path(filename: str, s3_prefix: Optional[str]) -> str:
 
 TRAINING_DATASET_MODEL_DICT = {
     ("huggingface-spc-bert-base-cased", "1.0.0"): ("training-datasets/QNLI-tiny/"),
+    ("huggingface-spc-bert-base-cased", "1.2.3"): ("training-datasets/QNLI-tiny/"),
 }
 
 

tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py

@@ -13,7 +13,7 @@
 from __future__ import absolute_import
 import os
 
-from sagemaker import hyperparameters, image_uris, model_uris, script_uris
+from sagemaker import hyperparameters, metric_definitions, image_uris, model_uris, script_uris
 from sagemaker.estimator import Estimator
 from sagemaker.jumpstart.constants import (
     INFERENCE_ENTRY_POINT_SCRIPT_NAME,
@@ -35,7 +35,7 @@
 
 def test_jumpstart_transfer_learning_estimator_class(setup):
 
-    model_id, model_version = "huggingface-spc-bert-base-cased", "1.0.0"
+    model_id, model_version = "huggingface-spc-bert-base-cased", "1.2.3"
     training_instance_type = "ml.p3.2xlarge"
     inference_instance_type = "ml.p2.xlarge"
     instance_count = 1
@@ -66,6 +66,11 @@ def test_jumpstart_transfer_learning_estimator_class(setup):
 
     default_hyperparameters["epochs"] = "1"
 
+    default_metric_definitions = metric_definitions.retrieve_default(
+        model_id=model_id,
+        model_version=model_version,
+    )
+
     estimator = Estimator(
         image_uri=image_uri,
         source_dir=script_uri,
@@ -78,6 +83,7 @@ def test_jumpstart_transfer_learning_estimator_class(setup):
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         instance_count=instance_count,
         instance_type=training_instance_type,
+        metric_definitions=default_metric_definitions,
     )
 
     estimator.fit(

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/jumpstart/test_default.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,
+        )