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,
@@ -453,7 +452,9 @@ def _retrieval_function(
             formatted_body, _ = self._get_json_file(id_info, data_type)
             model_specs = JumpStartModelSpecs(formatted_body)
             utils.emit_logs_based_on_model_specs(model_specs, self.get_region(), self._s3_client)
-            return JumpStartCachedContentValue(formatted_content=model_specs)
+            return JumpStartCachedContentValue(
+                formatted_content=model_specs
+            )
 
         if data_type == HubContentType.MODEL:
             hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
@@ -478,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/instance_types/jumpstart/test_instance_types.py

@@ -206,6 +206,7 @@ def test_jumpstart_instance_types_from_hub(patched_get_model_specs):
         model_id=model_id,
         version=model_version,
         s3_client=mock_client,
+        model_type=JumpStartModelType.OPEN_WEIGHTS,
     )
 
     patched_get_model_specs.reset_mock()
@@ -227,6 +228,7 @@ def test_jumpstart_instance_types_from_hub(patched_get_model_specs):
         model_id=model_id,
         version=model_version,
         s3_client=mock_client,
+        model_type=JumpStartModelType.OPEN_WEIGHTS,
     )
 
     patched_get_model_specs.reset_mock()
@@ -253,6 +255,7 @@ def test_jumpstart_instance_types_from_hub(patched_get_model_specs):
         model_id=model_id,
         version=model_version,
         s3_client=mock_client,
+        model_type=JumpStartModelType.OPEN_WEIGHTS,
     )
 
     patched_get_model_specs.reset_mock()
@@ -282,6 +285,7 @@ def test_jumpstart_instance_types_from_hub(patched_get_model_specs):
         model_id=model_id,
         version=model_version,
         s3_client=mock_client,
+        model_type=JumpStartModelType.OPEN_WEIGHTS,
     )
 
     patched_get_model_specs.reset_mock()

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -292,7 +292,7 @@ def test_prepacked(
     @mock.patch("sagemaker.jumpstart.factory.model._retrieve_model_deploy_kwargs")
     @mock.patch("sagemaker.jumpstart.factory.estimator._retrieve_estimator_fit_kwargs")
     @mock.patch("sagemaker.jumpstart.curated_hub.utils.construct_hub_arn_from_name")
-    @mock.patch("sagemaker.jumpstart.estimator.is_valid_model_id")
+    @mock.patch("sagemaker.jumpstart.estimator.validate_model_id_and_get_type")
     @mock.patch("sagemaker.jumpstart.factory.model.Session")
     @mock.patch("sagemaker.jumpstart.factory.estimator.Session")
     @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
@@ -309,7 +309,7 @@ def test_hub_model(
         mock_get_model_specs: mock.Mock,
         mock_session_estimator: mock.Mock,
         mock_session_model: mock.Mock,
-        mock_is_valid_model_id: mock.Mock,
+        mock_validate_model_id_and_get_type: mock.Mock,
         mock_construct_hub_arn_from_name: mock.Mock,
         mock_retrieve_estimator_fit_kwargs: mock.Mock,
         mock_retrieve_model_deploy_kwargs: mock.Mock,
@@ -321,7 +321,7 @@ def test_hub_model(
         mock_get_caller_identity.return_value = "123456789123"
         mock_estimator_deploy.return_value = default_predictor
 
-        mock_is_valid_model_id.return_value = True
+        mock_validate_model_id_and_get_type.return_value = JumpStartModelType.OPEN_WEIGHTS
 
         model_id, _ = "pytorch-hub-model-1", "*"
         hub_arn = "arn:aws:sagemaker:us-west-2:123456789123:hub/my-mock-hub"

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()
@@ -147,7 +148,9 @@ def test_jumpstart_models_cache_get_model_specs(mock_cache):
     accessors.JumpStartModelsAccessor.get_model_specs(
         region=region, model_id=model_id, version=version
     )
-    mock_cache.get_specs.assert_called_once_with(model_id=model_id, semantic_version_str=version)
+    mock_cache.get_specs.assert_called_once_with(
+        model_id=model_id, version_str=version, model_type=JumpStartModelType.OPEN_WEIGHTS
+    )
     mock_cache.get_hub_model.assert_not_called()
 
     accessors.JumpStartModelsAccessor.get_model_specs(

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,
@@ -559,8 +558,8 @@ def test_jumpstart_proprietary_cache_accepts_input_parameters():
     )
     assert cache.get_region() == region
     assert cache.get_bucket() == bucket
-    assert cache._s3_cache._max_cache_items == max_s3_cache_items
-    assert cache._s3_cache._expiration_horizon == s3_cache_expiration_horizon
+    assert cache._content_cache._max_cache_items == max_s3_cache_items
+    assert cache._content_cache._expiration_horizon == s3_cache_expiration_horizon
     assert (
         cache._proprietary_model_id_manifest_key_cache._max_cache_items
         == max_semantic_version_cache_items

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