fix: improve notebook utils logging and add model-specific info messages

Author: evakravi

src/sagemaker/jumpstart/constants.py

@@ -13,6 +13,7 @@
 """This module stores constants related to SageMaker JumpStart."""
 from __future__ import absolute_import
 import logging
+import os
 from typing import Dict, Set, Type
 import boto3
 from sagemaker.base_deserializers import BaseDeserializer, JSONDeserializer
@@ -33,6 +34,8 @@
 from sagemaker.session import Session
 
 
+ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING = "DISABLE_JUMPSTART_LOGGING"
+
 JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = set(
     [
         JumpStartLaunchedRegionInfo(
@@ -209,6 +212,19 @@
 
 JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart")
 
+# disable logging if env var is set
+JUMPSTART_LOGGER.addHandler(
+    type(
+        "",
+        (logging.StreamHandler,),
+        {
+            "emit": lambda self, *args, **kwargs: logging.StreamHandler.emit(self, *args, **kwargs)
+            if not os.environ.get(ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING)
+            else None
+        },
+    )()
+)
+
 try:
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session(
         boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)

src/sagemaker/jumpstart/filters.py

@@ -45,7 +45,6 @@ class SpecialSupportedFilterKeys(str, Enum):
 
     TASK = "task"
     FRAMEWORK = "framework"
-    SUPPORTED_MODEL = "supported_model"
 
 
 FILTER_OPERATOR_STRING_MAPPINGS = {
@@ -74,7 +73,6 @@ class SpecialSupportedFilterKeys(str, Enum):
     [
         SpecialSupportedFilterKeys.TASK,
         SpecialSupportedFilterKeys.FRAMEWORK,
-        SpecialSupportedFilterKeys.SUPPORTED_MODEL,
     ]
 )
 

src/sagemaker/jumpstart/notebook_utils.py

@@ -15,10 +15,14 @@
 import copy
 
 from functools import cmp_to_key
+import os
 from typing import Any, Generator, List, Optional, Tuple, Union, Set, Dict
 from packaging.version import Version
 from sagemaker.jumpstart import accessors
-from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
+from sagemaker.jumpstart.constants import (
+    ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING,
+    JUMPSTART_DEFAULT_REGION_NAME,
+)
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.filters import (
     SPECIAL_SUPPORTED_FILTER_KEYS,
@@ -281,126 +285,160 @@ def _generate_jumpstart_model_versions(  # pylint: disable=redefined-builtin
             results. (Default: False).
     """
 
-    if isinstance(filter, str):
-        filter = Identity(filter)
+    class _ModelSearchContext:
+        """Context manager for conducting model searches."""
 
-    models_manifest_list = accessors.JumpStartModelsAccessor._get_manifest(region=region)
-    manifest_keys = set(models_manifest_list[0].__slots__)
+        def __init__(self):
+            """Initialize context manager."""
 
-    all_keys: Set[str] = set()
+            self.old_disable_js_logging_env_var_value = os.environ.get(
+                ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING
+            )
 
-    model_filters: Set[ModelFilter] = set()
+        def __enter__(self, *args, **kwargs):
+            """Enter context.
 
-    for operator in _model_filter_in_operator_generator(filter):
-        model_filter = operator.unresolved_value
-        key = model_filter.key
-        all_keys.add(key)
-        model_filters.add(model_filter)
+            JumpStart logs get disabled to avoid excessive logging.
+            """
 
-    for key in all_keys:
-        if "." in key:
-            raise NotImplementedError(f"No support for multiple level metadata indexing ('{key}').")
+            os.environ[ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING] = "true"
 
-    metadata_filter_keys = all_keys - SPECIAL_SUPPORTED_FILTER_KEYS
+        def __exit__(self, *args, **kwargs):
+            """Exit context.
 
-    required_manifest_keys = manifest_keys.intersection(metadata_filter_keys)
-    possible_spec_keys = metadata_filter_keys - manifest_keys
+            Restore JumpStart logging settings, and reset cache so
+            new logs would appear for models previously searched.
+            """
 
-    unrecognized_keys: Set[str] = set()
+            if self.old_disable_js_logging_env_var_value:
+                os.environ[
+                    ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING
+                ] = self.old_disable_js_logging_env_var_value
+            else:
+                os.environ.pop(ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING, None)
+            accessors.JumpStartModelsAccessor.reset_cache()
 
-    is_task_filter = SpecialSupportedFilterKeys.TASK in all_keys
-    is_framework_filter = SpecialSupportedFilterKeys.FRAMEWORK in all_keys
-    is_supported_model_filter = SpecialSupportedFilterKeys.SUPPORTED_MODEL in all_keys
+    with _ModelSearchContext():
 
-    for model_manifest in models_manifest_list:
+        if isinstance(filter, str):
+            filter = Identity(filter)
 
-        copied_filter = copy.deepcopy(filter)
+        models_manifest_list = accessors.JumpStartModelsAccessor._get_manifest(region=region)
+        manifest_keys = set(models_manifest_list[0].__slots__)
 
-        manifest_specs_cached_values: Dict[str, Union[bool, int, float, str, dict, list]] = {}
+        all_keys: Set[str] = set()
 
-        model_filters_to_resolved_values: Dict[ModelFilter, BooleanValues] = {}
+        model_filters: Set[ModelFilter] = set()
 
-        for val in required_manifest_keys:
-            manifest_specs_cached_values[val] = getattr(model_manifest, val)
+        for operator in _model_filter_in_operator_generator(filter):
+            model_filter = operator.unresolved_value
+            key = model_filter.key
+            all_keys.add(key)
+            model_filters.add(model_filter)
 
-        if is_task_filter:
-            manifest_specs_cached_values[
-                SpecialSupportedFilterKeys.TASK
-            ] = extract_framework_task_model(model_manifest.model_id)[1]
+        for key in all_keys:
+            if "." in key:
+                raise NotImplementedError(
+                    f"No support for multiple level metadata indexing ('{key}')."
+                )
 
-        if is_framework_filter:
-            manifest_specs_cached_values[
-                SpecialSupportedFilterKeys.FRAMEWORK
-            ] = extract_framework_task_model(model_manifest.model_id)[0]
+        metadata_filter_keys = all_keys - SPECIAL_SUPPORTED_FILTER_KEYS
 
-        if is_supported_model_filter:
-            manifest_specs_cached_values[SpecialSupportedFilterKeys.SUPPORTED_MODEL] = Version(
-                model_manifest.min_version
-            ) <= Version(get_sagemaker_version())
+        required_manifest_keys = manifest_keys.intersection(metadata_filter_keys)
+        possible_spec_keys = metadata_filter_keys - manifest_keys
 
-        _populate_model_filters_to_resolved_values(
-            manifest_specs_cached_values,
-            model_filters_to_resolved_values,
-            model_filters,
-        )
+        unrecognized_keys: Set[str] = set()
 
-        _put_resolved_booleans_into_filter(copied_filter, model_filters_to_resolved_values)
+        is_task_filter = SpecialSupportedFilterKeys.TASK in all_keys
+        is_framework_filter = SpecialSupportedFilterKeys.FRAMEWORK in all_keys
 
-        copied_filter.eval()
+        for model_manifest in models_manifest_list:
 
-        if copied_filter.resolved_value in [BooleanValues.TRUE, BooleanValues.FALSE]:
-            if copied_filter.resolved_value == BooleanValues.TRUE:
-                yield (model_manifest.model_id, model_manifest.version)
-            continue
+            copied_filter = copy.deepcopy(filter)
 
-        if copied_filter.resolved_value == BooleanValues.UNEVALUATED:
-            raise RuntimeError(
-                "Filter expression in unevaluated state after using values from model manifest. "
-                "Model ID and version that is failing: "
-                f"{(model_manifest.model_id, model_manifest.version)}."
+            manifest_specs_cached_values: Dict[str, Union[bool, int, float, str, dict, list]] = {}
+
+            model_filters_to_resolved_values: Dict[ModelFilter, BooleanValues] = {}
+
+            for val in required_manifest_keys:
+                manifest_specs_cached_values[val] = getattr(model_manifest, val)
+
+            if is_task_filter:
+                manifest_specs_cached_values[
+                    SpecialSupportedFilterKeys.TASK
+                ] = extract_framework_task_model(model_manifest.model_id)[1]
+
+            if is_framework_filter:
+                manifest_specs_cached_values[
+                    SpecialSupportedFilterKeys.FRAMEWORK
+                ] = extract_framework_task_model(model_manifest.model_id)[0]
+
+            if Version(model_manifest.min_version) > Version(get_sagemaker_version()):
+                continue
+
+            _populate_model_filters_to_resolved_values(
+                manifest_specs_cached_values,
+                model_filters_to_resolved_values,
+                model_filters,
             )
-        copied_filter_2 = copy.deepcopy(filter)
 
-        model_specs = accessors.JumpStartModelsAccessor.get_model_specs(
-            region=region,
-            model_id=model_manifest.model_id,
-            version=model_manifest.version,
-        )
+            _put_resolved_booleans_into_filter(copied_filter, model_filters_to_resolved_values)
 
-        model_specs_keys = set(model_specs.__slots__)
+            copied_filter.eval()
 
-        unrecognized_keys -= model_specs_keys
-        unrecognized_keys_for_single_spec = possible_spec_keys - model_specs_keys
-        unrecognized_keys.update(unrecognized_keys_for_single_spec)
+            if copied_filter.resolved_value in [BooleanValues.TRUE, BooleanValues.FALSE]:
+                if copied_filter.resolved_value == BooleanValues.TRUE:
+                    yield (model_manifest.model_id, model_manifest.version)
+                continue
 
-        for val in possible_spec_keys:
-            if hasattr(model_specs, val):
-                manifest_specs_cached_values[val] = getattr(model_specs, val)
+            if copied_filter.resolved_value == BooleanValues.UNEVALUATED:
+                raise RuntimeError(
+                    "Filter expression in unevaluated state after using "
+                    "values from model manifest. Model ID and version that "
+                    f"is failing: {(model_manifest.model_id, model_manifest.version)}."
+                )
+            copied_filter_2 = copy.deepcopy(filter)
 
-        _populate_model_filters_to_resolved_values(
-            manifest_specs_cached_values,
-            model_filters_to_resolved_values,
-            model_filters,
-        )
-        _put_resolved_booleans_into_filter(copied_filter_2, model_filters_to_resolved_values)
+            model_specs = accessors.JumpStartModelsAccessor.get_model_specs(
+                region=region,
+                model_id=model_manifest.model_id,
+                version=model_manifest.version,
+            )
 
-        copied_filter_2.eval()
+            model_specs_keys = set(model_specs.__slots__)
 
-        if copied_filter_2.resolved_value != BooleanValues.UNEVALUATED:
-            if copied_filter_2.resolved_value == BooleanValues.TRUE or (
-                BooleanValues.UNKNOWN and list_incomplete_models
-            ):
-                yield (model_manifest.model_id, model_manifest.version)
-            continue
+            unrecognized_keys -= model_specs_keys
+            unrecognized_keys_for_single_spec = possible_spec_keys - model_specs_keys
+            unrecognized_keys.update(unrecognized_keys_for_single_spec)
 
-        raise RuntimeError(
-            "Filter expression in unevaluated state after using values from model specs. "
-            "Model ID and version that is failing: "
-            f"{(model_manifest.model_id, model_manifest.version)}."
-        )
+            for val in possible_spec_keys:
+                if hasattr(model_specs, val):
+                    manifest_specs_cached_values[val] = getattr(model_specs, val)
+
+            _populate_model_filters_to_resolved_values(
+                manifest_specs_cached_values,
+                model_filters_to_resolved_values,
+                model_filters,
+            )
+            _put_resolved_booleans_into_filter(copied_filter_2, model_filters_to_resolved_values)
+
+            copied_filter_2.eval()
+
+            if copied_filter_2.resolved_value != BooleanValues.UNEVALUATED:
+                if copied_filter_2.resolved_value == BooleanValues.TRUE or (
+                    BooleanValues.UNKNOWN and list_incomplete_models
+                ):
+                    yield (model_manifest.model_id, model_manifest.version)
+                continue
+
+            raise RuntimeError(
+                "Filter expression in unevaluated state after using values from model specs. "
+                "Model ID and version that is failing: "
+                f"{(model_manifest.model_id, model_manifest.version)}."
+            )
 
-    if len(unrecognized_keys) > 0:
-        raise RuntimeError(f"Unrecognized keys: {str(unrecognized_keys)}")
+        if len(unrecognized_keys) > 0:
+            raise RuntimeError(f"Unrecognized keys: {str(unrecognized_keys)}")
 
 
 def get_model_url(

src/sagemaker/jumpstart/types.py

@@ -730,6 +730,7 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "training_dependencies",
         "training_vulnerabilities",
         "deprecated",
+        "info_message",
         "deprecated_message",
         "deprecate_warn_message",
         "default_inference_instance_type",
@@ -801,6 +802,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.deprecated: bool = bool(json_obj["deprecated"])
         self.deprecated_message: Optional[str] = json_obj.get("deprecated_message")
         self.deprecate_warn_message: Optional[str] = json_obj.get("deprecate_warn_message")
+        self.info_message: Optional[str] = json_obj.get("info_message")
         self.default_inference_instance_type: Optional[str] = json_obj.get(
             "default_inference_instance_type"
         )

src/sagemaker/jumpstart/utils.py

@@ -511,6 +511,9 @@ def emit_logs_based_on_model_specs(
     if model_specs.deprecate_warn_message:
         constants.JUMPSTART_LOGGER.warning(model_specs.deprecate_warn_message)
 
+    if model_specs.info_message:
+        constants.JUMPSTART_LOGGER.info(model_specs.info_message)
+
     if model_specs.inference_vulnerable or model_specs.training_vulnerable:
         constants.JUMPSTART_LOGGER.warning(
             "Using vulnerable JumpStart model '%s' and version '%s'.",

tests/unit/sagemaker/jumpstart/constants.py

@@ -6221,6 +6221,7 @@
         "ml.c5.2xlarge",
     ],
     "hosting_use_script_uri": True,
+    "info_message": None,
     "metrics": [{"Regex": "val_accuracy: ([0-9\\.]+)", "Name": "pytorch-ic:val-accuracy"}],
     "model_kwargs": {"some-model-kwarg-key": "some-model-kwarg-value"},
     "deploy_kwargs": {"some-model-deploy-kwarg-key": "some-model-deploy-kwarg-value"},

tests/unit/sagemaker/jumpstart/test_notebook_utils.py

@@ -383,34 +383,27 @@ def test_list_jumpstart_models_region(
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
     @patch("sagemaker.jumpstart.notebook_utils.get_sagemaker_version")
-    def test_list_jumpstart_models_unsupported_models(
+    @patch("sagemaker.jumpstart.notebook_utils.accessors.JumpStartModelsAccessor.reset_cache")
+    @patch.dict("os.environ", {})
+    @patch("logging.StreamHandler.emit")
+    @patch("sagemaker.jumpstart.constants.JUMPSTART_LOGGER.propagate", False)
+    def test_list_jumpstart_models_disables_logging_resets_cache(
         self,
+        patched_emit: Mock,
+        patched_reset_cache: Mock,
         patched_get_sagemaker_version: Mock,
         patched_get_model_specs: Mock,
         patched_get_manifest: Mock,
     ):
         patched_get_model_specs.side_effect = get_prototype_model_spec
         patched_get_manifest.side_effect = get_prototype_manifest
 
-        patched_get_sagemaker_version.return_value = "0.0.0"
+        patched_get_sagemaker_version.return_value = "3.0.0"
 
-        assert [] == list_jumpstart_models("supported_model == True")
-        patched_get_model_specs.assert_not_called()
-        assert [] == list_jumpstart_models(
-            And("supported_model == True", "training_supported in [False, True]")
-        )
-        patched_get_model_specs.assert_not_called()
-
-        assert [] != list_jumpstart_models("supported_model == False")
-
-        patched_get_sagemaker_version.return_value = "999999.0.0"
-
-        assert [] != list_jumpstart_models("supported_model == True")
-
-        patched_get_model_specs.reset_mock()
+        list_jumpstart_models("deprecate_warn_message is blah")
 
-        assert [] != list_jumpstart_models("training_supported in [False, True]")
-        patched_get_model_specs.assert_called()
+        patched_emit.assert_not_called()
+        patched_reset_cache.assert_called_once()
 
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")