update linter

Author: bencrabtree

src/sagemaker/jumpstart/cache.py

@@ -339,17 +339,23 @@ def _retrieval_function(
                 formatted_content=model_specs
             )
         if data_type == HubDataType.MODEL:
-            hub_name, hub_region, model_id, model_version = utils.extract_info_from_hub_content_arn(id_info)
-            hub = CuratedHub(hub_name=hub_name, hub_region=hub_region)
-            hub_content = hub.describe_model(model_id=model_id, model_version=model_version)
-            utils.emit_logs_based_on_model_specs(hub_content.content_document, self.get_region(), self._s3_client)
+            hub_name, hub_region, model_name, model_version = utils.extract_info_from_hub_content_arn(
+                id_info
+            )
+            hub = CuratedHub(hub_name=hub_name, region=hub_region)
+            hub_content = hub.describe_model(model_name=model_name, model_version=model_version)
+            utils.emit_logs_based_on_model_specs(
+                hub_content.content_document,
+                self.get_region(),
+                self._s3_client
+            )
             model_specs = JumpStartModelSpecs(hub_content.content_document, is_hub_content=True)
             return JumpStartCachedContentValue(
                 formatted_content=model_specs
             )
         if data_type == HubDataType.HUB:
             hub_name, hub_region, _, _ = utils.extract_info_from_hub_content_arn(id_info)
-            hub = CuratedHub(hub_name=hub_name, hub_region=hub_region)
+            hub = CuratedHub(hub_name=hub_name, region=hub_region)
             hub_info = hub.describe()
             return JumpStartCachedContentValue(formatted_content=hub_info)
         raise ValueError(
@@ -467,7 +473,9 @@ def get_hub_model(self, hub_model_arn: str) -> JumpStartModelSpecs:
             hub_model_arn (str): Arn for the Hub model to get specs for
         """
 
-        details, _ = self._content_cache.get(JumpStartCachedContentKey(HubDataType.MODEL, hub_model_arn))
+        details, _ = self._content_cache.get(
+            JumpStartCachedContentKey(HubDataType.MODEL, hub_model_arn)
+        )
         return details.formatted_content
 
     def get_hub(self, hub_arn: str) -> Dict[str, Any]:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -20,7 +20,7 @@
 class CuratedHub:
     """Class for creating and managing a curated JumpStart hub"""
 
-    def __init__(self, hub_name: str, region: str, session: Optional[Session]):
+    def __init__(self, hub_name: str, region: str, session: Optional[Session] = None):
         self.hub_name = hub_name
         self.region = region
         self.session = session

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -860,16 +860,12 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
 def test_jumpstart_cache_get_hub_model():
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
 
-    model_arn = (
-        "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/huggingface-mock-model-123/1.2.3"
-    )
+    model_arn = "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/huggingface-mock-model-123/1.2.3"
     assert get_spec_from_base_spec(
         model_id="huggingface-mock-model-123", version="1.2.3"
     ) == cache.get_hub_model(hub_model_arn=model_arn)
 
-    model_arn = (
-        "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/pytorch-mock-model-123/*"
-    )
-    assert get_spec_from_base_spec(model_id="pytorch-mock-model-123", version="*") == cache.get_hub_model(
-        hub_model_arn=model_arn
-    )
+    model_arn = "arn:aws:sagemaker:us-west-2:000000000000:hub-content/HubName/Model/pytorch-mock-model-123/*"
+    assert get_spec_from_base_spec(
+        model_id="pytorch-mock-model-123", version="*"
+    ) == cache.get_hub_model(hub_model_arn=model_arn)