linter

Author: bencrabtree

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -89,6 +89,7 @@ def _retrieve_training_artifact_key(model_specs: JumpStartModelSpecs, instance_t
 def _retrieve_model_uri(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str] = None,
     model_scope: Optional[str] = None,
     instance_type: Optional[str] = None,
     region: Optional[str] = None,
@@ -105,6 +106,8 @@ def _retrieve_model_uri(
             the model artifact S3 URI.
         model_version (str): Version of the JumpStart model for which to retrieve the model
             artifact S3 URI.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         model_scope (str): The model type, i.e. what it is used for.
             Valid values: "training" and "inference".
         instance_type (str): The ML compute instance type for the specified scope. (Default: None).
@@ -135,6 +138,7 @@ def _retrieve_model_uri(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=model_scope,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/cache.py

@@ -343,7 +343,9 @@ def _retrieval_function(
                 id_info
             )
             hub = CuratedHub(hub_name=info.hub_name, region=info.region)
-            hub_content = hub.describe_model(model_name=info.hub_content_name, model_version=info.hub_content_version)
+            hub_content = hub.describe_model(
+                model_name=info.hub_content_name, model_version=info.hub_content_version
+            )
             utils.emit_logs_based_on_model_specs(
                 hub_content.content_document,
                 self.get_region(),
@@ -467,10 +469,10 @@ def get_specs(self, model_id: str, semantic_version_str: str) -> JumpStartModelS
                 )
             )
         return specs.formatted_content
-    
+
     def get_hub_model(self, hub_model_arn: str) -> JumpStartModelSpecs:
         """Return JumpStart-compatible specs for a given Hub model
-        
+
         Args:
             hub_model_arn (str): Arn for the Hub model to get specs for
         """
@@ -479,14 +481,14 @@ def get_hub_model(self, hub_model_arn: str) -> JumpStartModelSpecs:
             JumpStartCachedContentKey(HubDataType.MODEL, hub_model_arn)
         )
         return details.formatted_content
-    
+
     def get_hub(self, hub_arn: str) -> Dict[str, Any]:
         """Return descriptive info for a given Hub
-        
+
         Args:
             hub_arn (str): Arn for the Hub to get info for
         """
-        
+
         details, _ = self._content_cache.get(JumpStartCachedContentKey(HubDataType.HUB, hub_arn))
         return details.formatted_content
 

src/sagemaker/jumpstart/curated_hub/constants.py

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