linting

Author: malavhs

src/sagemaker/jumpstart/model.py

@@ -1044,7 +1044,7 @@ def _get_deployment_configs(
                 sagemaker_session=self.sagemaker_session,
                 region=self.region,
                 model_version=self.model_version,
-                hub_arn=self.hub_arn
+                hub_arn=self.hub_arn,
             )
 
             deployment_config_metadata = DeploymentConfigMetadata(

tests/integ/sagemaker/jumpstart/conftest.py

@@ -48,11 +48,14 @@ def _setup():
     test_hub_description = "PySDK Integ Test Private Hub"
     os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID: test_suit_id})
     os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME: test_hub_name})
-    hub = Hub(hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session())
+    hub = Hub(
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
     hub.create(description=test_hub_description)
     describe_hub_response = hub.describe()
     JUMPSTART_LOGGER.info(f"Describe Hub {describe_hub_response}")
 
+
 def _teardown():
     print("Tearing down...")
 
@@ -133,38 +136,35 @@ def _teardown():
 
     # delete private hubs
     _delete_hubs(sagemaker_session)
-    
+
 
 def _delete_hubs(sagemaker_session):
-    #list Hubs created by PySDK integration tests
+    # list Hubs created by PySDK integration tests
     list_hub_response = sagemaker_session.list_hubs(name_contains=HUB_NAME_PREFIX)
 
-    for hub in list_hub_response['HubSummaries']:
-        if hub['HubName'] != SM_JUMPSTART_PUBLIC_HUB_NAME:
-            #delete all hub contents first
-            _delete_hub_contents(sagemaker_session, hub['HubName'])
+    for hub in list_hub_response["HubSummaries"]:
+        if hub["HubName"] != SM_JUMPSTART_PUBLIC_HUB_NAME:
+            # delete all hub contents first
+            _delete_hub_contents(sagemaker_session, hub["HubName"])
             JUMPSTART_LOGGER.info(f"Deleting {hub['HubName']}")
-            sagemaker_session.delete_hub(hub['HubName'])
+            sagemaker_session.delete_hub(hub["HubName"])
 
 
 def _delete_hub_contents(sagemaker_session, test_hub_name):
-    #list hub_contents for the given hub
+    # list hub_contents for the given hub
     list_hub_content_response = sagemaker_session.list_hub_contents(
-        hub_name=test_hub_name, 
-        hub_content_type=HubContentType.MODEL_REFERENCE.value
+        hub_name=test_hub_name, hub_content_type=HubContentType.MODEL_REFERENCE.value
     )
     JUMPSTART_LOGGER.info(f"Listing HubContents {list_hub_content_response}")
 
-    #delete hub_contents for the given hub
-    for models in list_hub_content_response['HubContentSummaries']:
+    # delete hub_contents for the given hub
+    for models in list_hub_content_response["HubContentSummaries"]:
         sagemaker_session.delete_hub_content_reference(
-            hub_name=test_hub_name, 
+            hub_name=test_hub_name,
             hub_content_type=HubContentType.MODEL_REFERENCE.value,
-            hub_content_name=models['HubContentName']
+            hub_content_name=models["HubContentName"],
         )
 
-    
-
 
 @pytest.fixture(scope="session", autouse=True)
 def setup(request):

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

@@ -50,14 +50,16 @@
     "catboost-regression-model",
 }
 
+
 @pytest.fixture(scope="module")
 def add_models():
     # Create Model References to test in Hub
-    hub_instance = Hub(hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session())
+    hub_instance = Hub(
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
     for model in TEST_MODEL_IDS:
-        hub_instance.create_model_reference(
-            model_arn = get_public_hub_model_arn(hub_instance, model)
-        )
+        hub_instance.create_model_reference(model_arn=get_public_hub_model_arn(hub_instance, model))
+
 
 def test_jumpstart_hub_model(setup, add_models):
 
@@ -70,14 +72,15 @@ def test_jumpstart_hub_model(setup, add_models):
         model_id=model_id,
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
-        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
     )
 
     # uses ml.m5.4xlarge instance
     model.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
     )
 
+
 def test_jumpstart_hub_gated_model(setup, add_models):
 
     model_id = "meta-textgeneration-llama-3-2-1b"
@@ -86,7 +89,7 @@ def test_jumpstart_hub_gated_model(setup, add_models):
         model_id=model_id,
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
-        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
     )
 
     # uses ml.g6.xlarge instance
@@ -104,6 +107,7 @@ def test_jumpstart_hub_gated_model(setup, add_models):
 
     assert response is not None
 
+
 def test_jumpstart_gated_model_inference_component_enabled(setup, add_models):
 
     model_id = "meta-textgeneration-llama-2-7b"
@@ -115,14 +119,14 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models):
     sagemaker_session = get_sm_session()
 
     hub_arn = generate_hub_arn_for_init_kwargs(
-                hub_name=hub_name, region=region, session=sagemaker_session
-            )
+        hub_name=hub_name, region=region, session=sagemaker_session
+    )
 
     model = JumpStartModel(
         model_id=model_id,
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=sagemaker_session,
-        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
     )
 
     # uses ml.g5.2xlarge instance
@@ -136,7 +140,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models):
         endpoint_name=model.endpoint_name,
         sagemaker_session=sagemaker_session,
         tolerate_vulnerable_model=True,
-        hub_arn=hub_arn
+        hub_arn=hub_arn,
     )
 
     payload = {
@@ -149,13 +153,13 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models):
     assert response is not None
 
     model = JumpStartModel.attach(
-        predictor.endpoint_name, 
-        sagemaker_session=sagemaker_session, 
-        hub_name=hub_name)
+        predictor.endpoint_name, sagemaker_session=sagemaker_session, hub_name=hub_name
+    )
     assert model.model_id == model_id
     assert model.endpoint_name == predictor.endpoint_name
     assert model.inference_component_name == predictor.component_name
 
+
 def test_instatiating_model(setup, add_models):
 
     model_id = "catboost-regression-model"
@@ -166,10 +170,9 @@ def test_instatiating_model(setup, add_models):
         model_id=model_id,
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
-        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
     )
 
     elapsed_time = time.perf_counter() - start_time
 
     assert elapsed_time <= MAX_INIT_TIME_SECONDS
-

tests/integ/sagemaker/jumpstart/private_hub/test_hub.py

@@ -14,27 +14,29 @@
     HUB_NAME_PREFIX,
 )
 
+
 @pytest.fixture
 def hub_instance():
-    HUB_NAME=f"{HUB_NAME_PREFIX}-{get_test_suite_id()}"
+    HUB_NAME = f"{HUB_NAME_PREFIX}-{get_test_suite_id()}"
     hub = Hub(HUB_NAME, sagemaker_session=get_sm_session())
     yield hub
 
+
 def test_private_hub(setup, hub_instance):
-    #Createhub
+    # Createhub
     create_hub_response = hub_instance.create(
-      description="This is a Test Private Hub.",
-      display_name="malavhs Test hub",
-      search_keywords=["jumpstart-sdk-integ-test"],
+        description="This is a Test Private Hub.",
+        display_name="malavhs Test hub",
+        search_keywords=["jumpstart-sdk-integ-test"],
     )
 
-    #Create Hub Verifications
+    # Create Hub Verifications
     assert create_hub_response is not None
 
-    #Describe Hub
+    # Describe Hub
     hub_description = hub_instance.describe()
     assert hub_description is not None
 
-    #Delete Hub
+    # Delete Hub
     delete_hub_response = hub_instance.delete()
-    assert delete_hub_response is not None 
+    assert delete_hub_response is not None

tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py

@@ -6,9 +6,7 @@
 from tests.integ.sagemaker.jumpstart.utils import (
     get_sm_session,
 )
-from tests.integ.sagemaker.jumpstart.utils import (
-    get_public_hub_model_arn
-)
+from tests.integ.sagemaker.jumpstart.utils import get_public_hub_model_arn
 from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME,
 )
@@ -18,21 +16,21 @@
 def test_hub_model_reference(setup):
     model_id = "meta-textgenerationneuron-llama-3-2-1b-instruct"
 
-    hub_instance = Hub(hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session())
+    hub_instance = Hub(
+        hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
+    )
 
-    #Create Model Reference
+    # Create Model Reference
     create_model_response = hub_instance.create_model_reference(
-        model_arn = get_public_hub_model_arn(hub_instance, model_id)
+        model_arn=get_public_hub_model_arn(hub_instance, model_id)
     )
-    assert create_model_response is not None  
+    assert create_model_response is not None
 
-    #Describe Model
-    describe_model_response = hub_instance.describe_model(
-        model_name = model_id
-    )
+    # Describe Model
+    describe_model_response = hub_instance.describe_model(model_name=model_id)
     assert describe_model_response is not None
     assert type(describe_model_response) == DescribeHubContentResponse
 
-    #Delete Model Reference
+    # Delete Model Reference
     delete_model_response = hub_instance.delete_model_reference(model_name=model_id)
-    assert delete_model_response is not None  
+    assert delete_model_response is not None

tests/integ/sagemaker/jumpstart/utils.py

@@ -115,16 +115,20 @@ def get_tabular_data(data_filename: str) -> Tuple[pd.DataFrame, pd.DataFrame]:
 def download_file(local_download_path, s3_bucket, s3_key, s3_client) -> None:
     s3_client.download_file(s3_bucket, s3_key, local_download_path)
 
+
 def get_public_hub_model_arn(hub: Hub, model_id: str) -> str:
     filter_value = f"model_id == {model_id}"
     response = hub.list_sagemaker_public_hub_models(filter=filter_value)
 
     models = response["hub_content_summaries"]
     while response["next_token"]:
-        response = hub.list_sagemaker_public_hub_models(filter=filter_value, next_token=response["next_token"])
+        response = hub.list_sagemaker_public_hub_models(
+            filter=filter_value, next_token=response["next_token"]
+        )
         models.extend(response["hub_content_summaries"])
 
-    return models[0]['hub_content_arn']
+    return models[0]["hub_content_arn"]
+
 
 class EndpointInvoker:
     def __init__(