fix: linter

Author: chrstfu

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

@@ -66,9 +66,7 @@ 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":
@@ -93,7 +91,5 @@ 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,9 +73,7 @@ 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,12 +13,10 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 from concurrent import futures
-from functools import lru_cache
 from datetime import datetime
 import json
 import traceback
 from typing import Optional, Dict, List, Any
-
 import boto3
 from botocore import exceptions
 from botocore.client import BaseClient
@@ -102,9 +100,7 @@ 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,
@@ -124,9 +120,7 @@ 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,
@@ -180,9 +174,7 @@ def describe_model(
     ) -> 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,
@@ -228,11 +220,7 @@ 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,
@@ -301,18 +289,12 @@ 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,
@@ -328,9 +310,7 @@ 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
@@ -343,9 +323,7 @@ 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()
                         ),
                     }
                 )
@@ -389,9 +367,7 @@ 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.
@@ -457,14 +433,10 @@ def scan_and_tag_models(self, model_list: List[Dict[str, str]] = None) -> None:
 
         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
+            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,

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

@@ -72,9 +72,7 @@ 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"""
@@ -147,8 +145,6 @@ 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/utils.py

@@ -19,8 +19,7 @@
 from sagemaker.s3_utils import parse_s3_url
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
-from typing import Optional, Dict, List, Any, Set
-from botocore.exceptions import ClientError
+from typing import Optional, Dict, List, Any
 from sagemaker.jumpstart.types import HubContentType, HubArnExtractedInfo
 from sagemaker.jumpstart.curated_hub.types import (
     CuratedHubUnsupportedFlag,
@@ -29,13 +28,10 @@
 )
 from sagemaker.jumpstart import constants
 from sagemaker.jumpstart import utils
-from sagemaker.session import Session
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.curated_hub.constants import (
     JUMPSTART_HUB_MODEL_ID_TAG_PREFIX,
     JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX,
-    TASK_TAG_PREFIX,
-    FRAMEWORK_TAG_PREFIX,
 )
 from sagemaker.utils import format_tags, TagsDict
 
@@ -95,9 +91,7 @@ def construct_hub_arn_from_name(
     return f"arn:{partition}:sagemaker:{region}:{account_id}:hub/{hub_name}"
 
 
-def construct_hub_model_arn_from_inputs(
-    hub_arn: str, model_name: str, version: str
-) -> str:
+def construct_hub_model_arn_from_inputs(hub_arn: str, model_name: str, version: str) -> str:
     """Constructs a HubContent model arn from the Hub name, model name, and model version."""
 
     info = get_info_from_hub_resource_arn(hub_arn)
@@ -127,9 +121,7 @@ def generate_hub_arn_for_init_kwargs(
         if match:
             hub_arn = hub_name
         else:
-            hub_arn = construct_hub_arn_from_name(
-                hub_name=hub_name, region=region, session=session
-            )
+            hub_arn = construct_hub_arn_from_name(hub_name=hub_name, region=region, session=session)
     return hub_arn
 
 
@@ -185,21 +177,18 @@ def create_hub_bucket_if_it_does_not_exist(
     return bucket_name
 
 
-def tag_hub_content(
-    hub_content_arn: str, tags: List[TagsDict], session: Session
-) -> None:
+def tag_hub_content(hub_content_arn: str, tags: List[TagsDict], session: Session) -> None:
     session.add_tags(ResourceArn=hub_content_arn, Tags=tags)
-    JUMPSTART_LOGGER.info(
-        f"Added tags to HubContentArn %s: %s", hub_content_arn, TagsDict
-    )
+    JUMPSTART_LOGGER.info("Added tags to HubContentArn %s: %s", hub_content_arn, TagsDict)
 
 
 def find_unsupported_flags_for_hub_content_versions(
     hub_name: str, hub_content_name: str, region: str, session: Session
 ) -> List[TagsDict]:
     """Finds the JumpStart public hub model for a HubContent and calculates relevant tags.
 
-    Since tags are the same for all versions of a HubContent, these tags will map from the key to a list of versions impacted.
+    Since tags are the same for all versions of a HubContent,
+    these tags will map from the key to a list of versions impacted.
     For example, if certain public hub model versions are deprecated,
     this utility will return a `deprecated` tag mapped to the deprecated versions for the HubContent.
     """
@@ -242,7 +231,8 @@ def find_unsupported_flags_for_model_version(
     """Finds relevant CuratedHubTags for a version of a JumpStart public hub model.
 
     For example, if the public hub model is deprecated, this utility will return a `deprecated` tag.
-    Since tags are the same for all versions of a HubContent, these tags will map from the key to a list of versions impacted.
+    Since tags are the same for all versions of a HubContent,
+    these tags will map from the key to a list of versions impacted.
     """
     flags_to_add: List[CuratedHubUnsupportedFlag] = []
     jumpstart_model_specs = utils.verify_model_region_and_return_specs(
@@ -293,14 +283,10 @@ def get_jumpstart_model_and_version(
     jumpstart_model_version = jumpstart_model_version_tag[
         len(JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX) :
     ]
-    return JumpStartModelInfo(
-        model_id=jumpstart_model_id, version=jumpstart_model_version
-    )
+    return JumpStartModelInfo(model_id=jumpstart_model_id, version=jumpstart_model_version)
 
 
-def summary_from_list_api_response(
-    hub_content_summary: Dict[str, Any]
-) -> HubContentSummary:
+def summary_from_list_api_response(hub_content_summary: Dict[str, Any]) -> HubContentSummary:
     return HubContentSummary(
         hub_content_arn=hub_content_summary.get("HubContentArn"),
         hub_content_name=hub_content_summary.get("HubContentName"),

tests/unit/sagemaker/jumpstart/curated_hub/test_curated_hub.py

@@ -21,7 +21,7 @@
 from sagemaker.jumpstart.curated_hub.types import JumpStartModelInfo, S3ObjectLocation, HubContentSummary
 from sagemaker.jumpstart.types import JumpStartModelSpecs
 from tests.unit.sagemaker.jumpstart.constants import BASE_SPEC
-from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, HubContentType
+from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
 
 
 REGION = "us-east-1"
@@ -552,4 +552,4 @@ def test_determine_models_to_sync(sagemaker_session):
     }
     # Old model_one, same model_two
     res = hub._determine_models_to_sync([model_one, model_two], js_model_map)
-    assert res == [model_one]
+    assert res == [model_one]