rebase with master

Author: bencrabtree

src/sagemaker/jumpstart/cache.py

@@ -40,7 +40,6 @@
     get_wildcard_model_version_msg,
     get_wildcard_proprietary_model_version_msg,
 )
-from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
     JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
     JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
@@ -480,6 +479,7 @@ def _retrieval_function(
             return JumpStartCachedContentValue(
                 formatted_content=model_specs
             )
+
         if data_type == HubType.HUB:
             hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
             response: Dict[str, Any] = self._sagemaker_session.describe_hub(hub_name=hub_name)

src/sagemaker/jumpstart/types.py

@@ -15,7 +15,6 @@
 from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, List, Optional, Set, Union
-from sagemaker.session import Session
 from sagemaker.utils import get_instance_type_family, format_tags, Tags
 from sagemaker.enums import EndpointType
 from sagemaker.model_metrics import ModelMetrics

src/sagemaker/remote_function/runtime_environment/runtime_environment_manager.py

@@ -24,8 +24,6 @@
 import dataclasses
 import json
 
-import sagemaker
-
 
 class _UTCFormatter(logging.Formatter):
     """Class that overrides the default local time provider in log formatter."""

src/sagemaker/utils.py

@@ -22,7 +22,6 @@
 import random
 import re
 import shutil
-import sys
 import tarfile
 import tempfile
 import time

tests/unit/sagemaker/jumpstart/test_accessors.py

@@ -137,6 +137,7 @@ def test_jumpstart_proprietary_models_cache_get(mock_cache):
         > 0
     )
 
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._cache")
 def test_jumpstart_models_cache_get_model_specs(mock_cache):
     mock_cache.get_specs = Mock()

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -28,7 +28,6 @@
     JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY,
     JumpStartModelsCache,
 )
-from sagemaker.session_settings import SessionSettings
 from sagemaker.jumpstart.constants import (
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,

tests/unit/sagemaker/jumpstart/utils.py

@@ -253,6 +253,13 @@ def patched_retrieval_function(
                 model_type=JumpStartModelType.PROPRIETARY,
             )
         )
+
+    if datatype == HubContentType.MODEL:
+        _, _, _, model_name, model_version = id_info.split("/")
+        return JumpStartCachedContentValue(
+            formatted_content=get_spec_from_base_spec(model_id=model_name, version=model_version)
+        )
+
     # TODO: Implement
     if datatype == HubType.HUB:
         return None