fix linter

Author: bencrabtree

src/sagemaker/jumpstart/cache.py

@@ -30,6 +30,7 @@
     MODEL_ID_LIST_WEB_URL,
 )
 from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
+from sagemaker.jumpstart.curated_hub.utils import get_info_from_hub_resource_arn
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
     JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
@@ -339,7 +340,7 @@ def _retrieval_function(
                 formatted_content=model_specs
             )
         if data_type == HubContentType.MODEL:
-            info = utils.get_info_from_hub_resource_arn(
+            info = get_info_from_hub_resource_arn(
                 id_info
             )
             hub = CuratedHub(hub_name=info.hub_name, region=info.region)
@@ -356,7 +357,7 @@ def _retrieval_function(
                 formatted_content=model_specs
             )
         if data_type == HubContentType.HUB:
-            info = utils.get_info_from_hub_resource_arn(
+            info = get_info_from_hub_resource_arn(
                 id_info
             )
             hub = CuratedHub(hub_name=info.hub_name, region=info.region)

src/sagemaker/jumpstart/curated_hub/types.py

@@ -16,6 +16,7 @@
 
 from sagemaker.jumpstart.types import JumpStartDataHolderType
 
+
 class HubArnExtractedInfo(JumpStartDataHolderType):
     """Data class for info extracted from Hub arn."""
 

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -21,6 +21,7 @@
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
 
+
 def get_info_from_hub_resource_arn(
     arn: str,
 ) -> HubArnExtractedInfo:

src/sagemaker/jumpstart/types.py

@@ -1746,4 +1746,3 @@ def __init__(
         self.nearest_model_name = nearest_model_name
         self.data_input_configuration = data_input_configuration
         self.skip_model_validation = skip_model_validation
-

src/sagemaker/jumpstart/utils.py

@@ -14,7 +14,6 @@
 from __future__ import absolute_import
 import logging
 import os
-import re
 from typing import Any, Dict, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 import boto3
@@ -42,7 +41,7 @@
 )
 from sagemaker.session import Session
 from sagemaker.config import load_sagemaker_config
-from sagemaker.utils import aws_partition, resolve_value_from_config, TagsDict
+from sagemaker.utils import resolve_value_from_config, TagsDict
 from sagemaker.workflow import is_pipeline_variable
 
 
@@ -830,4 +829,3 @@ def get_jumpstart_model_id_version_from_resource_arn(
             model_version = model_version_from_tag
 
     return model_id, model_version
-

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

@@ -17,6 +17,7 @@
 from sagemaker.jumpstart.curated_hub import utils
 from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
 
+
 def test_get_info_from_hub_resource_arn():
     model_arn = (
         "arn:aws:sagemaker:us-west-2:000000000000:hub-content/MockHub/Model/my-mock-model/1.0.2"
@@ -113,7 +114,9 @@ def test_generate_hub_arn_for_estimator_init_kwargs():
     )
 
     assert (
-        utils.generate_hub_arn_for_estimator_init_kwargs(hub_name, "us-east-1", session=mock_default_session)
+        utils.generate_hub_arn_for_estimator_init_kwargs(
+            hub_name, "us-east-1", session=mock_default_session
+        )
         == "arn:aws:sagemaker:us-east-1:123456789123:hub/my-hub-name"
     )
 
@@ -127,16 +130,24 @@ def test_generate_hub_arn_for_estimator_init_kwargs():
         == "arn:aws:sagemaker:us-east-2:000000000000:hub/my-hub-name"
     )
 
-    assert utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, session=mock_default_session) == hub_arn
-
     assert (
-        utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, "us-east-1", session=mock_default_session)
+        utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, session=mock_default_session)
         == hub_arn
     )
 
     assert (
-        utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, "us-east-1", mock_custom_session) == hub_arn
+        utils.generate_hub_arn_for_estimator_init_kwargs(
+            hub_arn, "us-east-1", session=mock_default_session
+        )
+        == hub_arn
     )
 
-    assert utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, None, mock_custom_session) == hub_arn
+    assert (
+        utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, "us-east-1", mock_custom_session)
+        == hub_arn
+    )
 
+    assert (
+        utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, None, mock_custom_session)
+        == hub_arn
+    )