fix: formatter

Author: chrstfu

src/sagemaker/jumpstart/curated_hub/accessors/file_generator.py

@@ -21,7 +21,9 @@
     HubContentDependencyType,
     S3ObjectLocation,
 )
-from sagemaker.jumpstart.curated_hub.accessors.public_model_data import PublicModelDataAccessor
+from sagemaker.jumpstart.curated_hub.accessors.public_model_data import (
+    PublicModelDataAccessor,
+)
 from sagemaker.jumpstart.types import JumpStartModelSpecs
 
 
@@ -64,7 +66,9 @@ def generate_file_infos_from_model_specs(
     )
     files = []
     for dependency in HubContentDependencyType:
-        location: S3ObjectLocation = public_model_data_accessor.get_s3_reference(dependency)
+        location: S3ObjectLocation = public_model_data_accessor.get_s3_reference(
+            dependency
+        )
         location_type = "prefix" if location.key.endswith("/") else "object"
 
         if location_type == "prefix":
@@ -89,5 +93,7 @@ def generate_file_infos_from_model_specs(
             response = s3_client.head_object(**parameters)
             size = response.get("ContentLength")
             last_updated = response.get("LastModified")
-            files.append(FileInfo(location.bucket, location.key, size, last_updated, dependency))
+            files.append(
+                FileInfo(location.bucket, location.key, size, last_updated, dependency)
+            )
     return files

src/sagemaker/jumpstart/curated_hub/accessors/public_model_data.py

@@ -73,7 +73,9 @@ def training_script_s3_reference(self):
     @property
     def default_training_dataset_s3_reference(self):
         """Retrieves s3 reference for s3 directory containing model training datasets"""
-        return S3ObjectLocation(self._get_bucket_name(), self.__get_training_dataset_prefix())
+        return S3ObjectLocation(
+            self._get_bucket_name(), self.__get_training_dataset_prefix()
+        )
 
     @property
     def demo_notebook_s3_reference(self):

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,7 +13,7 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 from concurrent import futures
-from functools import lru_cache 
+from functools import lru_cache
 from datetime import datetime
 import json
 import traceback
@@ -37,7 +37,10 @@
 from sagemaker.jumpstart.curated_hub.sync.request import HubSyncRequestFactory
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.session import Session
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION, JUMPSTART_LOGGER
+from sagemaker.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    JUMPSTART_LOGGER,
+)
 from sagemaker.jumpstart.types import (
     DescribeHubResponse,
     DescribeHubContentsResponse,
@@ -61,6 +64,7 @@
 )
 from sagemaker.utils import TagsDict
 
+
 class CuratedHub:
     """Class for creating and managing a curated JumpStart hub"""
 
@@ -98,7 +102,9 @@ def _fetch_hub_bucket_name(self) -> str:
             if hub_output_location:
                 location = create_s3_object_reference_from_uri(hub_output_location)
                 return location.bucket
-            default_bucket_name = generate_default_hub_bucket_name(self._sagemaker_session)
+            default_bucket_name = generate_default_hub_bucket_name(
+                self._sagemaker_session
+            )
             JUMPSTART_LOGGER.warning(
                 "There is not a Hub bucket associated with %s. Using %s",
                 self.hub_name,
@@ -118,7 +124,9 @@ def _generate_hub_storage_location(self, bucket_name: Optional[str] = None) -> N
         """Generates an ``S3ObjectLocation`` given a Hub name."""
         hub_bucket_name = bucket_name or self._fetch_hub_bucket_name()
         curr_timestamp = datetime.now().timestamp()
-        return S3ObjectLocation(bucket=hub_bucket_name, key=f"{self.hub_name}-{curr_timestamp}")
+        return S3ObjectLocation(
+            bucket=hub_bucket_name, key=f"{self.hub_name}-{curr_timestamp}"
+        )
 
     def create(
         self,
@@ -151,29 +159,30 @@ def describe(self) -> DescribeHubResponse:
 
         return hub_description
 
-    
     def list_models(self, clear_cache: bool = True, **kwargs) -> List[Dict[str, Any]]:
         """Lists the models in this Curated Hub.
 
-        This function caches the models in local memory 
+        This function caches the models in local memory
 
         **kwargs: Passed to invocation of ``Session:list_hub_contents``.
         """
         if clear_cache:
             self._list_hubs_cache = None
         if self._list_hubs_cache is None:
-          hub_content_summaries = self._sagemaker_session.list_hub_contents(
-            hub_name=self.hub_name, hub_content_type=HubContentType.MODEL, **kwargs
-          )
-          self._list_hubs_cache = hub_content_summaries
+            hub_content_summaries = self._sagemaker_session.list_hub_contents(
+                hub_name=self.hub_name, hub_content_type=HubContentType.MODEL, **kwargs
+            )
+            self._list_hubs_cache = hub_content_summaries
         return self._list_hubs_cache
 
     def describe_model(
         self, model_name: str, model_version: str = "*"
     ) -> DescribeHubContentsResponse:
         """Returns descriptive information about the Hub Model"""
 
-        hub_content_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
+        hub_content_description: Dict[
+            str, Any
+        ] = self._sagemaker_session.describe_hub_content(
             hub_name=self.hub_name,
             hub_content_name=model_name,
             hub_content_version=model_version,
@@ -219,10 +228,16 @@ def _populate_latest_model_version(self, model: Dict[str, str]) -> Dict[str, str
 
     def _get_jumpstart_models_in_hub(self) -> List[HubContentSummary]:
         hub_models = summary_list_from_list_api_response(self.list_models())
-        return [model for model in hub_models if get_jumpstart_model_and_version(model) is not None]
+        return [
+            model
+            for model in hub_models
+            if get_jumpstart_model_and_version(model) is not None
+        ]
 
     def _determine_models_to_sync(
-        self, model_list: List[JumpStartModelInfo], models_in_hub: Dict[str, HubContentSummary]
+        self,
+        model_list: List[JumpStartModelInfo],
+        models_in_hub: Dict[str, HubContentSummary],
     ) -> List[JumpStartModelInfo]:
         """Determines which models from `sync` params to sync into the CuratedHub.
 
@@ -286,14 +301,22 @@ def sync(self, model_list: List[Dict[str, str]]):
                     model["model_id"],
                     model["version"],
                 )
-            model_version_list.append(JumpStartModelInfo(model["model_id"], model["version"]))
+            model_version_list.append(
+                JumpStartModelInfo(model["model_id"], model["version"])
+            )
 
         js_models_in_hub = self._get_jumpstart_models_in_hub()
-        mapped_models_in_hub = {model.hub_content_name: model for model in js_models_in_hub}
+        mapped_models_in_hub = {
+            model.hub_content_name: model for model in js_models_in_hub
+        }
 
-        models_to_sync = self._determine_models_to_sync(model_version_list, mapped_models_in_hub)
+        models_to_sync = self._determine_models_to_sync(
+            model_version_list, mapped_models_in_hub
+        )
         JUMPSTART_LOGGER.warning(
-            "Syncing the following models into Hub %s: %s", self.hub_name, models_to_sync
+            "Syncing the following models into Hub %s: %s",
+            self.hub_name,
+            models_to_sync,
         )
 
         # Delete old models?
@@ -305,7 +328,9 @@ def sync(self, model_list: List[Dict[str, str]]):
             thread_name_prefix="import-models-to-curated-hub",
         ) as import_executor:
             for thread_num, model in enumerate(models_to_sync):
-                task = import_executor.submit(self._sync_public_model_to_hub, model, thread_num)
+                task = import_executor.submit(
+                    self._sync_public_model_to_hub, model, thread_num
+                )
                 tasks.append(task)
 
         # Handle failed imports
@@ -318,7 +343,9 @@ def sync(self, model_list: List[Dict[str, str]]):
                     {
                         "Exception": exception,
                         "Traceback": "".join(
-                            traceback.TracebackException.from_exception(exception).format()
+                            traceback.TracebackException.from_exception(
+                                exception
+                            ).format()
                         ),
                     }
                 )
@@ -362,7 +389,9 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
                 label=dest_location.key,
             ).execute()
         else:
-            JUMPSTART_LOGGER.warning("Nothing to copy for %s v%s", model.model_id, model.version)
+            JUMPSTART_LOGGER.warning(
+                "Nothing to copy for %s v%s", model.model_id, model.version
+            )
 
         # TODO: Tag model if specs say it is deprecated or training/inference
         # vulnerable. Update tag of HubContent ARN without version.
@@ -403,44 +432,47 @@ def _fetch_studio_specs(self, model_specs: JumpStartModelSpecs) -> Dict[str, Any
         )
         return json.loads(response["Body"].read().decode("utf-8"))
 
-
     def scan_and_tag_models(self, model_list: List[Dict[str, str]] = None) -> None:
         """Scans the Hub for JumpStart models and tags the HubContent.
-        
+
         If the scan detects a model is deprecated or vulnerable, it will tag the HubContent.
         The tags that will be added are based off the specifications in the JumpStart public hub:
         1. "deprecated_versions" -> If the public hub model is deprecated
         2. "inference_vulnerable_versions" -> If the public hub model has inference vulnerabilities
         3. "training_vulnerable_versions" -> If the public hub model has training vulnerabilities
 
-        The tag value will be a list of versions in the Curated Hub that fall under those keys. 
+        The tag value will be a list of versions in the Curated Hub that fall under those keys.
         For example, if model_a version_a is deprecated and inference is vulnerable, the
-        HubContent for `model_a` will have tags [{"deprecated_versions": [version_a]}, 
+        HubContent for `model_a` will have tags [{"deprecated_versions": [version_a]},
         {"inference_vulnerable_versions": [version_a]}]
 
-        If models are passed in, 
+        If models are passed in,
         """
-        JUMPSTART_LOGGER.info(
-            "Tagging models in hub: %s", self.hub_name
-        )
+        JUMPSTART_LOGGER.info("Tagging models in hub: %s", self.hub_name)
         if self._is_invalid_model_list_input(model_list):
             raise ValueError(
                 "Model list should be a list of objects with values 'model_id',",
                 "and optional 'version'.",
             )
-        
+
         models_to_scan = model_list if model_list else self.list_models()
-        js_models_in_hub = [model for model in models_to_scan if get_jumpstart_model_and_version(model) is not None]
+        js_models_in_hub = [
+            model
+            for model in models_to_scan
+            if get_jumpstart_model_and_version(model) is not None
+        ]
         for model in js_models_in_hub:
-            tags_to_add: List[TagsDict] = find_unsupported_flags_for_hub_content_versions(
+            tags_to_add: List[
+                TagsDict
+            ] = find_unsupported_flags_for_hub_content_versions(
                 hub_name=self.hub_name,
                 hub_content_name=model.hub_content_name,
                 region=self.region,
-                session=self._sagemaker_session
+                session=self._sagemaker_session,
             )
             tag_hub_content(
                 hub_content_arn=model.hub_content_arn,
                 tags=tags_to_add,
-                session=self._sagemaker_session
+                session=self._sagemaker_session,
             )
-        JUMPSTART_LOGGER.info("Tagging complete!")
+        JUMPSTART_LOGGER.info("Tagging complete!")

src/sagemaker/jumpstart/curated_hub/sync/request.py

@@ -31,7 +31,9 @@ class HubSyncRequest:
     destination: S3ObjectLocation
 
     def __init__(
-        self, files_to_copy: Generator[FileInfo, FileInfo, FileInfo], destination: S3ObjectLocation
+        self,
+        files_to_copy: Generator[FileInfo, FileInfo, FileInfo],
+        destination: S3ObjectLocation,
     ):
         """Contains information required to sync data into a Hub.
 
@@ -70,7 +72,9 @@ def __init__(
         # Need the file lists to be sorted for comparisons below
         self.src_files: List[FileInfo] = sorted(src_files, key=lambda x: x.location.key)
         formatted_dest_files = [self._format_dest_file(file) for file in dest_files]
-        self.dest_files: List[FileInfo] = sorted(formatted_dest_files, key=lambda x: x.location.key)
+        self.dest_files: List[FileInfo] = sorted(
+            formatted_dest_files, key=lambda x: x.location.key
+        )
 
     def _format_dest_file(self, file: FileInfo) -> FileInfo:
         """Strips HubContent data prefix from dest file name"""
@@ -143,6 +147,8 @@ def _is_same_file_name(self, src_filename: str, dest_filename: str) -> bool:
         """Determines if two files have the same file name."""
         return src_filename == dest_filename
 
-    def _is_alphabetically_earlier_file_name(self, src_filename: str, dest_filename: str) -> bool:
+    def _is_alphabetically_earlier_file_name(
+        self, src_filename: str, dest_filename: str
+    ) -> bool:
         """Determines if one filename is alphabetically earlier than another."""
         return src_filename < dest_filename

src/sagemaker/jumpstart/curated_hub/types.py

@@ -17,27 +17,36 @@
 from dataclasses import dataclass
 from datetime import datetime
 
-from sagemaker.jumpstart.types import JumpStartDataHolderType, JumpStartModelSpecs, HubContentType
+from sagemaker.jumpstart.types import (
+    JumpStartDataHolderType,
+    JumpStartModelSpecs,
+    HubContentType,
+)
+
 
 class CuratedHubUnsupportedFlag(str, Enum):
     """Enum class for Curated Hub tag names."""
+
     DEPRECATED_VERSIONS = "deprecated_versions"
     TRAINING_VULNERABLE_VERSIONS = "training_vulnerable_versions"
     INFERENCE_VULNERABLE_VERSIONS = "inference_vulnerable_versions"
 
+
 @dataclass
 class HubContentSummary:
-  """Dataclass to store HubContentSummary from List APIs."""
-  hub_content_arn: str
-  hub_content_name: str
-  hub_content_version: str
-  hub_content_type: HubContentType
-  document_schema_version: str
-  hub_content_status: str
-  creation_time: str
-  hub_content_display_name: str = None
-  hub_content_description: str = None
-  hub_content_search_keywords: List[str] = None
+    """Dataclass to store HubContentSummary from List APIs."""
+
+    hub_content_arn: str
+    hub_content_name: str
+    hub_content_version: str
+    hub_content_type: HubContentType
+    document_schema_version: str
+    hub_content_status: str
+    creation_time: str
+    hub_content_display_name: str = None
+    hub_content_description: str = None
+    hub_content_search_keywords: List[str] = None
+
 
 @dataclass
 class S3ObjectLocation:
@@ -57,6 +66,7 @@ def get_uri(self) -> str:
         """Returns the s3 URI"""
         return f"s3://{self.bucket}/{self.key}"
 
+
 @dataclass
 class JumpStartModelInfo:
     """Helper class for storing JumpStart model info."""