linting

aws · Mar 18, 2024 · 1af132e · 1af132e
1 parent b50c557
commit 1af132e
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 23 deletions.
diff --git a/src/sagemaker/jumpstart/curated_hub/curated_hub.py b/src/sagemaker/jumpstart/curated_hub/curated_hub.py
@@ -394,6 +394,7 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
         self._sagemaker_session.import_hub_content(
             document_schema_version=HubContentDocument_v2.SCHEMA_VERSION,
             hub_content_name=model.model_id,
+            hub_content_version=model.version,
             hub_name=self.hub_name,
             hub_content_document=hub_content_document,
             hub_content_type=HubContentType.MODEL,

diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
@@ -1434,8 +1434,6 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "model_version",
         "model_type",
         "hub_arn",
-        "model_type",
-        "hub_arn",
         "region",
         "tolerate_deprecated_model",
         "tolerate_vulnerable_model",

diff --git a/tests/unit/sagemaker/jumpstart/test_accessors.py b/tests/unit/sagemaker/jumpstart/test_accessors.py
@@ -83,7 +83,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, version_str=version, model_type=JumpStartModelType.OPEN_WEIGHTS)
+    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(

diff --git a/tests/unit/sagemaker/jumpstart/utils.py b/tests/unit/sagemaker/jumpstart/utils.py
@@ -254,26 +254,6 @@ def patched_retrieval_function(
             )
         )
 
-    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
-
-    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
-
     raise ValueError(f"Bad value for datatype: {datatype}")