docstrings and linting

Author: bencrabtree

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

@@ -14,7 +14,6 @@
 from __future__ import absolute_import
 from typing import Any, Dict, List
 
-from datetime import datetime
 from botocore.client import BaseClient
 
 from sagemaker.jumpstart.curated_hub.types import (
@@ -90,7 +89,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

@@ -17,8 +17,8 @@
 from sagemaker.jumpstart.curated_hub.types import (
     HubContentDependencyType,
     S3ObjectLocation,
-    create_s3_object_reference_from_uri,
 )
+from sagemaker.jumpstart.curated_hub.utils import create_s3_object_reference_from_uri
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.types import JumpStartModelSpecs
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -46,12 +46,12 @@
 from sagemaker.jumpstart.curated_hub.utils import (
     create_hub_bucket_if_it_does_not_exist,
     generate_default_hub_bucket_name,
+    create_s3_object_reference_from_uri,
 )
 from sagemaker.jumpstart.curated_hub.types import (
     HubContentDocument_v2,
     JumpStartModelInfo,
     S3ObjectLocation,
-    create_s3_object_reference_from_uri,
 )
 
 
@@ -302,7 +302,7 @@ def sync(self, model_list: List[Dict[str, str]]):
             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["name"]: model for model in js_models_in_hub }
+        mapped_models_in_hub = {model["name"]: model for model in js_models_in_hub}
 
         models_to_sync = self._determine_models_to_sync(model_version_list, mapped_models_in_hub)
         JUMPSTART_LOGGER.warning(
@@ -316,9 +316,9 @@ def sync(self, model_list: List[Dict[str, str]]):
         with futures.ThreadPoolExecutor(
             max_workers=self._default_thread_pool_size,
             thread_name_prefix="import-models-to-curated-hub",
-        ) as deploy_executor:
+        ) as import_executor:
             for thread_num, model in enumerate(models_to_sync):
-                task = deploy_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
@@ -353,7 +353,7 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
 
         dest_location = S3ObjectLocation(
             bucket=self.hub_storage_location.bucket,
-            key=f"{self.hub_storage_location.key}/{model.model_id}/{model.version}"
+            key=f"{self.hub_storage_location.key}/curated_models/{model.model_id}/{model.version}",
         )
         src_files = file_generator.generate_file_infos_from_model_specs(
             model_specs, studio_specs, self.region, self._s3_client

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

@@ -35,16 +35,21 @@ def __init__(
     ):
         """Contains information required to sync data into a Hub.
 
-        Returns:
-            :var: files (List[FileInfo]): Files that shoudl be synced.
-            :var: destination (S3ObjectLocation): Location to which to sync the files.
+        Attrs:
+            files (List[FileInfo]): Files that should be synced.
+            destination (S3ObjectLocation): Location to which to sync the files.
         """
         self.files = list(files_to_copy)
         self.destination = destination
 
 
 class HubSyncRequestFactory:
-    """Generates a ``HubSyncRequest`` which is required to sync data into a Hub."""
+    """Generates a ``HubSyncRequest`` which is required to sync data into a Hub.
+
+    Creates a ``HubSyncRequest`` class containing:
+            :var: files (List[FileInfo]): Files that should be synced.
+            :var: destination (S3ObjectLocation): Location to which to sync the files.
+    """
 
     def __init__(
         self,
@@ -59,11 +64,6 @@ def __init__(
             src_files (List[FileInfo]): List of files to sync to destination bucket
             dest_files (List[FileInfo]): List of files already in destination bucket
             destination (S3ObjectLocation): S3 destination for copied data
-
-        Returns:
-            ``HubSyncRequest`` class containing:
-                :var: files (List[FileInfo]): Files that shoudl be synced.
-                :var: destination (S3ObjectLocation): Location to which to sync the files.
         """
         self.comparator = comparator
         self.destination = destination

src/sagemaker/jumpstart/curated_hub/types.py

@@ -18,7 +18,6 @@
 from datetime import datetime
 
 from sagemaker.jumpstart.types import JumpStartDataHolderType, JumpStartModelSpecs
-from sagemaker.s3_utils import parse_s3_url
 
 
 @dataclass
@@ -40,16 +39,6 @@ def get_uri(self) -> str:
         return f"s3://{self.bucket}/{self.key}"
 
 
-def create_s3_object_reference_from_uri(s3_uri: str) -> S3ObjectLocation:
-    """Utiity to help generate an S3 object reference"""
-    bucket, key = parse_s3_url(s3_uri)
-
-    return S3ObjectLocation(
-        bucket=bucket,
-        key=key,
-    )
-
-
 @dataclass
 class JumpStartModelInfo:
     """Helper class for storing JumpStart model info."""

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -14,6 +14,8 @@
 from __future__ import absolute_import
 import re
 from typing import Optional
+from sagemaker.jumpstart.curated_hub.types import S3ObjectLocation
+from sagemaker.s3_utils import parse_s3_url
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
 from sagemaker.jumpstart.types import (
@@ -131,6 +133,16 @@ def generate_default_hub_bucket_name(
     return f"sagemaker-hubs-{region}-{account_id}"
 
 
+def create_s3_object_reference_from_uri(s3_uri: str) -> S3ObjectLocation:
+    """Utiity to help generate an S3 object reference"""
+    bucket, key = parse_s3_url(s3_uri)
+
+    return S3ObjectLocation(
+        bucket=bucket,
+        key=key,
+    )
+
+
 def create_hub_bucket_if_it_does_not_exist(
     bucket_name: Optional[str] = None,
     sagemaker_session: Session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,

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

@@ -32,6 +32,7 @@
 
 FAKE_TIME = datetime.datetime(1997, 8, 14, 00, 00, 00)
 
+
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name="boto_session")
@@ -80,7 +81,9 @@ def test_create_with_no_bucket_name(
     hub_search_keywords,
     tags,
 ):
-    storage_location = S3ObjectLocation("sagemaker-hubs-us-east-1-123456789123",f"{hub_name}-{FAKE_TIME.timestamp()}")
+    storage_location = S3ObjectLocation(
+        "sagemaker-hubs-us-east-1-123456789123", f"{hub_name}-{FAKE_TIME.timestamp()}"
+    )
     mock_generate_hub_storage_location.return_value = storage_location
     create_hub = {"HubArn": f"arn:aws:sagemaker:us-east-1:123456789123:hub/{hub_name}"}
     sagemaker_session.create_hub = Mock(return_value=create_hub)
@@ -133,7 +136,7 @@ def test_create_with_bucket_name(
     hub_search_keywords,
     tags,
 ):
-    storage_location = S3ObjectLocation(hub_bucket_name,f"{hub_name}-{FAKE_TIME.timestamp()}")
+    storage_location = S3ObjectLocation(hub_bucket_name, f"{hub_name}-{FAKE_TIME.timestamp()}")
     mock_generate_hub_storage_location.return_value = storage_location
     create_hub = {"HubArn": f"arn:aws:sagemaker:us-east-1:123456789123:hub/{hub_name}"}
     sagemaker_session.create_hub = Mock(return_value=create_hub)