run linter

Author: Eli Davidson

src/sagemaker/jumpstart/utils.py

@@ -1631,7 +1631,9 @@ def get_draft_model_content_bucket(provider: Dict, region: str) -> str:
         return get_jumpstart_content_bucket(region=region)
     return neo_bucket
 
+
 def get_latest_version(versions: List[str]) -> Optional[str]:
-    try: return None if not versions else max(versions, key=Version)
+    try:
+        return None if not versions else max(versions, key=Version)
     except InvalidVersion as e:
         return max(versions)

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -1124,9 +1124,10 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
         ]
     )
 
+
 @patch.object(JumpStartModelsCache, "_retrieval_function")
 def test_jumpstart_cache_handles_versioning_correctly_for_open_source_weights(
-   retrieval_function: Mock
+    retrieval_function: Mock
 ):
     sm_version = Version(utils.get_sagemaker_version())
     new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
@@ -1150,7 +1151,7 @@ def test_jumpstart_cache_handles_versioning_correctly_for_open_source_weights(
             "spec_key": "spec_key"
         }
     )
-    
+
     manifest_dict = {}
     for header in manifest:
         header_obj = JumpStartModelHeader(header)
@@ -1163,15 +1164,16 @@ def test_jumpstart_cache_handles_versioning_correctly_for_open_source_weights(
     key = JumpStartVersionedModelId("test-model", "*")
 
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
-    result = cache._get_open_weight_manifest_key_from_model_id( key = key, value = None )
+    result = cache._get_open_weight_manifest_key_from_model_id(key=key, value=None)
 
     assert_key = JumpStartVersionedModelId("test-model", "2.16.0")
 
     assert result == assert_key
 
+
 @patch.object(JumpStartModelsCache, "_retrieval_function")
 def test_jumpstart_cache_handles_versioning_correctly_for_proprietary_weights(
-   retrieval_function: Mock
+    retrieval_function: Mock
 ):
     sm_version = Version(utils.get_sagemaker_version())
     new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
@@ -1195,7 +1197,7 @@ def test_jumpstart_cache_handles_versioning_correctly_for_proprietary_weights(
             "spec_key": "spec_key"
         }
     )
-    
+
     manifest_dict = {}
     for header in manifest:
         header_obj = JumpStartModelHeader(header)
@@ -1208,7 +1210,7 @@ def test_jumpstart_cache_handles_versioning_correctly_for_proprietary_weights(
     key = JumpStartVersionedModelId("test-model", "*")
 
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
-    result = cache._get_proprietary_manifest_key_from_model_id( key = key, value = None )
+    result = cache._get_proprietary_manifest_key_from_model_id(key=key, value=None)
 
     assert_key = JumpStartVersionedModelId("test-model", "2.16.0")
 
@@ -1217,7 +1219,7 @@ def test_jumpstart_cache_handles_versioning_correctly_for_proprietary_weights(
 
 @patch.object(JumpStartModelsCache, "_retrieval_function")
 def test_jumpstart_cache_handles_versioning_correctly_non_sem_ver(
-   retrieval_function: Mock
+    retrieval_function: Mock
 ):
     sm_version = Version(utils.get_sagemaker_version())
     new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
@@ -1232,7 +1234,7 @@ def test_jumpstart_cache_handles_versioning_correctly_non_sem_ver(
         }
         for version in versions
     ]
-    
+
     manifest_dict = {}
     for header in manifest:
         header_obj = JumpStartModelHeader(header)
@@ -1245,8 +1247,8 @@ def test_jumpstart_cache_handles_versioning_correctly_non_sem_ver(
     key = JumpStartVersionedModelId("test-model", "*")
 
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
-    result = cache._get_open_weight_manifest_key_from_model_id( key = key, value = None )
+    result = cache._get_open_weight_manifest_key_from_model_id(key=key, value=None)
 
     assert_key = JumpStartVersionedModelId("test-model", "abc")
 
-    assert result == assert_key
+    assert result == assert_key

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -2143,18 +2143,23 @@ def test_add_instance_rate_stats_to_benchmark_metrics_client_ex(
 def test_has_instance_rate_stat(stats, expected):
     assert utils.has_instance_rate_stat(stats) is expected
 
+
 def test_get_latest_version():
     assert utils.get_latest_version(["2.9.1", "2.16.0", "1.0.0"]) == "2.16.0"
 
+
 def test_get_latest_version_empty_list_is_none():
-    assert utils.get_latest_version([]) == None
+    assert utils.get_latest_version([]) is None
+
 
 def test_get_latest_version_none_is_none():
-    assert utils.get_latest_version(None) == None
+    assert utils.get_latest_version(None) is None
+
 
 def test_get_latest_version_with_invalid_sem_ver():
     assert utils.get_latest_version(["2.9.1", "2.16.0", "1.0.0", "abc"]) == "abc"
 
+
 @pytest.mark.parametrize(
     "data, expected",
     [(None, []), ([], []), (get_base_deployment_configs_metadata(), get_base_deployment_configs())],