update tests

Author: bencrabtree

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

@@ -69,7 +69,7 @@ def __init__(
         region: str,
         sync_request: HubSyncRequest,
         label: Optional[str] = None,
-        thread_num: Optional[int] = 0
+        thread_num: Optional[int] = 0,
     ):
         """Multi-part S3:Copy Handler initializer.
 

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -21,7 +21,6 @@
 from botocore import exceptions
 from botocore.client import BaseClient
 from packaging.version import Version
-import tqdm
 
 from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.curated_hub.accessors import file_generator
@@ -315,9 +314,9 @@ def sync(self, model_list: List[Dict[str, str]]):
             max_workers=self._default_thread_pool_size,
             thread_name_prefix="import-models-to-curated-hub",
         ) as deploy_executor:
-                for thread_num, model in enumerate(models_to_sync):
-                    task = deploy_executor.submit(self._sync_public_model_to_hub, model, thread_num)
-                    tasks.append(task)
+            for thread_num, model in enumerate(models_to_sync):
+                task = deploy_executor.submit(self._sync_public_model_to_hub, model, thread_num)
+                tasks.append(task)
 
         # Handle failed imports
         results = futures.wait(tasks)
@@ -365,7 +364,12 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
         ).create()
 
         if len(sync_request.files) > 0:
-            MultiPartCopyHandler(thread_num=thread_num, sync_request=sync_request, region=self.region, label=dest_location.key).execute()
+            MultiPartCopyHandler(
+                thread_num=thread_num,
+                sync_request=sync_request,
+                region=self.region,
+                label=dest_location.key,
+            ).execute()
         else:
             JUMPSTART_LOGGER.warning("Nothing to copy for %s v%s", model.model_id, model.version)
 

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

@@ -163,8 +163,8 @@ def test_sync_kicks_off_parallel_syncs(
 
     mock_sync_public_models.assert_has_calls(
         [
-            mock.call(JumpStartModelInfo("mock-model-one-huggingface", "*")),
-            mock.call(JumpStartModelInfo("mock-model-two-pytorch", "1.0.2")),
+            mock.call(JumpStartModelInfo("mock-model-one-huggingface", "*"), 0),
+            mock.call(JumpStartModelInfo("mock-model-two-pytorch", "1.0.2"), 1),
         ]
     )
 
@@ -206,7 +206,7 @@ def test_sync_filters_models_that_exist_in_hub(
     hub.sync([model_one, model_two])
 
     mock_sync_public_models.assert_called_once_with(
-        JumpStartModelInfo("mock-model-one-huggingface", "*")
+        JumpStartModelInfo("mock-model-one-huggingface", "*"), 0
     )
 
 
@@ -252,8 +252,8 @@ def test_sync_updates_old_models_in_hub(
 
     mock_sync_public_models.assert_has_calls(
         [
-            mock.call(JumpStartModelInfo("mock-model-one-huggingface", "*")),
-            mock.call(JumpStartModelInfo("mock-model-two-pytorch", "1.0.2")),
+            mock.call(JumpStartModelInfo("mock-model-one-huggingface", "*"), 0),
+            mock.call(JumpStartModelInfo("mock-model-two-pytorch", "1.0.2"), 1),
         ]
     )
 
@@ -299,7 +299,7 @@ def test_sync_passes_newer_hub_models(
     hub.sync([model_one, model_two])
 
     mock_sync_public_models.assert_called_once_with(
-        JumpStartModelInfo("mock-model-one-huggingface", "*")
+        JumpStartModelInfo("mock-model-one-huggingface", "*"), 0
     )