fix: more linting :/

chrstfu · chrstfu · commit d21fd41d16a0 · 2024-03-18T22:28:43.000Z
diff --git a/src/sagemaker/jumpstart/artifacts/environment_variables.py b/src/sagemaker/jumpstart/artifacts/environment_variables.py
@@ -133,7 +133,6 @@ def _retrieve_default_environment_variables(
             )
 
             if gated_model_env_var is None and model_specs.is_gated_model():
-
                 possible_env_vars: Set[str] = {
                     retrieve_gated_env_var_for_instance_type(instance_type)
                     for instance_type in model_specs.supported_training_instance_types
diff --git a/src/sagemaker/jumpstart/artifacts/model_packages.py b/src/sagemaker/jumpstart/artifacts/model_packages.py
@@ -84,7 +84,6 @@ def _retrieve_model_package_arn(
     )
 
     if scope == JumpStartScriptScope.INFERENCE:
-
         instance_specific_arn: Optional[str] = (
             model_specs.hosting_instance_type_variants.get_model_package_arn(
                 region=region, instance_type=instance_type
@@ -155,7 +154,6 @@ def _retrieve_model_package_model_artifact_s3_uri(
     """
 
     if scope == JumpStartScriptScope.TRAINING:
-
         if region is None:
             region = JUMPSTART_DEFAULT_REGION_NAME
 
diff --git a/src/sagemaker/jumpstart/artifacts/model_uris.py b/src/sagemaker/jumpstart/artifacts/model_uris.py
@@ -149,7 +149,6 @@ def _retrieve_model_uri(
     model_artifact_key: str
 
     if model_scope == JumpStartScriptScope.INFERENCE:
-
         is_prepacked = not model_specs.use_inference_script_uri()
 
         model_artifact_key = (
@@ -159,7 +158,6 @@ def _retrieve_model_uri(
         )
 
     elif model_scope == JumpStartScriptScope.TRAINING:
-
         model_artifact_key = _retrieve_training_artifact_key(model_specs, instance_type)
 
     default_jumpstart_bucket: str = (
diff --git a/src/sagemaker/jumpstart/artifacts/resource_requirements.py b/src/sagemaker/jumpstart/artifacts/resource_requirements.py
@@ -134,7 +134,6 @@ def _retrieve_default_resources(
     }
 
     if is_dynamic_container_deployment_supported:
-
         all_resource_requirement_kwargs = {}
 
         for (
diff --git a/src/sagemaker/jumpstart/cache.py b/src/sagemaker/jumpstart/cache.py
@@ -34,7 +34,6 @@
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     MODEL_TYPE_TO_MANIFEST_MAP,
     MODEL_TYPE_TO_SPECS_MAP,
-    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.exceptions import (
     get_wildcard_model_version_msg,
@@ -178,9 +177,7 @@ def set_manifest_file_s3_key(
         }
         property_name = file_mapping.get(file_type)
         if not property_name:
-            raise ValueError(
-                self._file_type_error_msg(file_type, manifest_only=True)
-            )
+            raise ValueError(self._file_type_error_msg(file_type, manifest_only=True))
         if key != property_name:
             setattr(self, property_name, key)
             self.clear()
@@ -193,9 +190,7 @@ def get_manifest_file_s3_key(
             return self._manifest_file_s3_key
         if file_type == JumpStartS3FileType.PROPRIETARY_MANIFEST:
             return self._proprietary_manifest_s3_key
-        raise ValueError(
-            self._file_type_error_msg(file_type, manifest_only=True)
-        )
+        raise ValueError(self._file_type_error_msg(file_type, manifest_only=True))
 
     def set_s3_bucket_name(self, s3_bucket_name: str) -> None:
         """Set s3 bucket used for cache."""
@@ -248,7 +243,8 @@ def _model_id_retrieval_function(
         sm_version = utils.get_sagemaker_version()
         manifest = self._content_cache.get(
             JumpStartCachedContentKey(
-                MODEL_TYPE_TO_MANIFEST_MAP[model_type], self._manifest_file_s3_map[model_type])
+                MODEL_TYPE_TO_MANIFEST_MAP[model_type], self._manifest_file_s3_map[model_type]
+            )
         )[0].formatted_content
 
         versions_compatible_with_sagemaker = [
@@ -265,7 +261,8 @@ def _model_id_retrieval_function(
             return JumpStartVersionedModelId(model_id, sm_compatible_model_version)
 
         versions_incompatible_with_sagemaker = [
-            Version(header.version) for header in manifest.values()  # type: ignore
+            Version(header.version)
+            for header in manifest.values()  # type: ignore
             if header.model_id == model_id
         ]
         sm_incompatible_model_version = self._select_version(
@@ -295,9 +292,7 @@ def _model_id_retrieval_function(
             raise KeyError(error_msg)
 
         error_msg = f"Unable to find model manifest for '{model_id}' with version '{version}'. "
-        error_msg += (
-            f"Visit {MODEL_ID_LIST_WEB_URL} for updated list of models. "
-        )
+        error_msg += f"Visit {MODEL_ID_LIST_WEB_URL} for updated list of models. "
 
         other_model_id_version = None
         if model_type == JumpStartModelType.OPEN_WEIGHTS:
@@ -306,19 +301,17 @@ def _model_id_retrieval_function(
             )  # all versions here are incompatible with sagemaker
         elif model_type == JumpStartModelType.PROPRIETARY:
             all_possible_model_id_version = [
-                header.version for header in manifest.values()  # type: ignore
+                header.version
+                for header in manifest.values()  # type: ignore
                 if header.model_id == model_id
             ]
             other_model_id_version = (
-                None
-                if not all_possible_model_id_version
-                else all_possible_model_id_version[0]
+                None if not all_possible_model_id_version else all_possible_model_id_version[0]
             )
 
         if other_model_id_version is not None:
             error_msg += (
-                f"Consider using model ID '{model_id}' with version "
-                f"'{other_model_id_version}'."
+                f"Consider using model ID '{model_id}' with version " f"'{other_model_id_version}'."
             )
         else:
             possible_model_ids = [header.model_id for header in manifest.values()]  # type: ignore
@@ -360,15 +353,15 @@ def _get_json_file_and_etag_from_s3(self, key: str) -> Tuple[Union[dict, list],
 
     def _is_local_metadata_mode(self) -> bool:
         """Returns True if the cache should use local metadata mode, based off env variables."""
-        return (ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE in os.environ
-                and os.path.isdir(os.environ[ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE])
-                and ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE in os.environ
-                and os.path.isdir(os.environ[ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE]))
+        return (
+            ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE in os.environ
+            and os.path.isdir(os.environ[ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE])
+            and ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE in os.environ
+            and os.path.isdir(os.environ[ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE])
+        )
 
     def _get_json_file(
-        self,
-        key: str,
-        filetype: JumpStartS3FileType
+        self, key: str, filetype: JumpStartS3FileType
     ) -> Tuple[Union[dict, list], Optional[str]]:
         """Returns json file either from s3 or local file system.
 
@@ -392,21 +385,19 @@ def _get_json_md5_hash(self, key: str):
         return self._s3_client.head_object(Bucket=self.s3_bucket_name, Key=key)["ETag"]
 
     def _get_json_file_from_local_override(
-        self,
-        key: str,
-        filetype: JumpStartS3FileType
+        self, key: str, filetype: JumpStartS3FileType
     ) -> Union[dict, list]:
         """Reads json file from local filesystem and returns data."""
         if filetype == JumpStartS3FileType.OPEN_WEIGHT_MANIFEST:
-            metadata_local_root = (
-                os.environ[ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE]
-            )
+            metadata_local_root = os.environ[
+                ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE
+            ]
         elif filetype == JumpStartS3FileType.OPEN_WEIGHT_SPECS:
             metadata_local_root = os.environ[ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE]
         else:
             raise ValueError(f"Unsupported file type for local override: {filetype}")
         file_path = os.path.join(metadata_local_root, key)
-        with open(file_path, 'r') as f:
+        with open(file_path, "r") as f:
             data = json.load(f)
         return data
 
@@ -443,17 +434,15 @@ def _retrieval_function(
                 formatted_content=utils.get_formatted_manifest(formatted_body),
                 md5_hash=etag,
             )
-        
+
         if data_type in {
             JumpStartS3FileType.OPEN_WEIGHT_SPECS,
             JumpStartS3FileType.PROPRIETARY_SPECS,
         }:
             formatted_body, _ = self._get_json_file(id_info, data_type)
             model_specs = JumpStartModelSpecs(formatted_body)
             utils.emit_logs_based_on_model_specs(model_specs, self.get_region(), self._s3_client)
-            return JumpStartCachedContentValue(
-                formatted_content=model_specs
-            )
+            return JumpStartCachedContentValue(formatted_content=model_specs)
 
         if data_type == HubContentType.MODEL:
             hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
@@ -463,21 +452,15 @@ def _retrieval_function(
                 hub_name=hub_name,
                 hub_content_name=model_name,
                 hub_content_version=model_version,
-                hub_content_type=data_type
+                hub_content_type=data_type,
             )
 
             model_specs = JumpStartModelSpecs(
                 DescribeHubContentsResponse(hub_model_description), is_hub_content=True
             )
 
-            utils.emit_logs_based_on_model_specs(
-                model_specs,
-                self.get_region(),
-                self._s3_client
-            )
-            return JumpStartCachedContentValue(
-                formatted_content=model_specs
-            )
+            utils.emit_logs_based_on_model_specs(model_specs, self.get_region(), self._s3_client)
+            return JumpStartCachedContentValue(formatted_content=model_specs)
 
         if data_type == HubType.HUB:
             hub_name, _, _, _ = hub_utils.get_info_from_hub_resource_arn(id_info)
@@ -487,9 +470,7 @@ def _retrieval_function(
                 formatted_content=DescribeHubResponse(hub_description)
             )
 
-        raise ValueError(
-            self._file_type_error_msg(data_type)
-        )
+        raise ValueError(self._file_type_error_msg(data_type))
 
     def get_manifest(
         self,
@@ -498,7 +479,8 @@ def get_manifest(
         """Return entire JumpStart models manifest."""
         manifest_dict = self._content_cache.get(
             JumpStartCachedContentKey(
-                MODEL_TYPE_TO_MANIFEST_MAP[model_type], self._manifest_file_s3_map[model_type])
+                MODEL_TYPE_TO_MANIFEST_MAP[model_type], self._manifest_file_s3_map[model_type]
+            )
         )[0].formatted_content
         manifest = list(manifest_dict.values())  # type: ignore
         return manifest
@@ -555,16 +537,14 @@ def _select_version(
         except InvalidSpecifier:
             raise KeyError(f"Bad semantic version: {version_str}")
         available_versions_filtered = list(spec.filter(available_versions))
-        return (
-            str(max(available_versions_filtered)) if available_versions_filtered != [] else None
-        )
+        return str(max(available_versions_filtered)) if available_versions_filtered != [] else None
 
     def _get_header_impl(
         self,
         model_id: str,
         semantic_version_str: str,
         attempt: int = 0,
-        model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS
+        model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
     ) -> JumpStartModelHeader:
         """Lower-level function to return header.
 
@@ -587,7 +567,8 @@ def _get_header_impl(
 
         manifest = self._content_cache.get(
             JumpStartCachedContentKey(
-                MODEL_TYPE_TO_MANIFEST_MAP[model_type], self._manifest_file_s3_map[model_type])
+                MODEL_TYPE_TO_MANIFEST_MAP[model_type], self._manifest_file_s3_map[model_type]
+            )
         )[0].formatted_content
 
         try:
@@ -603,7 +584,7 @@ def get_specs(
         self,
         model_id: str,
         version_str: str,
-        model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS
+        model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
     ) -> JumpStartModelSpecs:
         """Return specs for a given JumpStart model ID and semantic version.
 
@@ -616,16 +597,12 @@ def get_specs(
         header = self.get_header(model_id, version_str, model_type)
         spec_key = header.spec_key
         specs, cache_hit = self._content_cache.get(
-            JumpStartCachedContentKey(
-                MODEL_TYPE_TO_SPECS_MAP[model_type], spec_key
-            )
+            JumpStartCachedContentKey(MODEL_TYPE_TO_SPECS_MAP[model_type], spec_key)
         )
 
         if not cache_hit and "*" in version_str:
             JUMPSTART_LOGGER.warning(
-                get_wildcard_model_version_msg(
-                    header.model_id, version_str, header.version
-                )
+                get_wildcard_model_version_msg(header.model_id, version_str, header.version)
             )
         return specs.formatted_content
 
diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py
@@ -734,7 +734,6 @@ def attach(
         """
 
         if model_id is None:
-
             model_id, model_version = get_model_id_version_from_training_job(
                 training_job_name=training_job_name, sagemaker_session=sagemaker_session
             )
diff --git a/src/sagemaker/jumpstart/factory/model.py b/src/sagemaker/jumpstart/factory/model.py
@@ -330,7 +330,6 @@ def _add_entry_point_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartMod
         tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
         sagemaker_session=kwargs.sagemaker_session,
     ):
-
         entry_point = entry_point or INFERENCE_ENTRY_POINT_SCRIPT_NAME
 
     kwargs.entry_point = entry_point
diff --git a/src/sagemaker/jumpstart/notebook_utils.py b/src/sagemaker/jumpstart/notebook_utils.py
@@ -379,7 +379,6 @@ def _generate_jumpstart_model_versions(  # pylint: disable=redefined-builtin
     is_model_type_filter = SpecialSupportedFilterKeys.MODEL_TYPE in all_keys
 
     def evaluate_model(model_manifest: JumpStartModelHeader) -> Optional[Tuple[str, str]]:
-
         copied_filter = copy.deepcopy(filter)
 
         manifest_specs_cached_values: Dict[str, Union[bool, int, float, str, dict, list]] = {}
diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py
@@ -743,7 +743,6 @@ def resolve_estimator_sagemaker_config_field(
     # JumpStart Estimators have certain default field values. We want
     # sagemaker config values to take priority over the model-specific defaults.
     if field_name == "enable_network_isolation":
-
         resolved_val = resolve_value_from_config(
             direct_input=None,
             config_path=TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
@@ -754,7 +753,6 @@ def resolve_estimator_sagemaker_config_field(
         return resolved_val if resolved_val is not None else field_val
 
     if field_name == "encrypt_inter_container_traffic":
-
         resolved_val = resolve_value_from_config(
             direct_input=None,
             config_path=TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
@@ -868,12 +866,13 @@ def generate_studio_spec_file_prefix(model_id: str, model_version: str) -> str:
     """Returns the Studio Spec file prefix given a model ID and version."""
     return f"studio_models/{model_id}/studio_specs_v{model_version}.json"
 
+
 def extract_info_from_hub_content_arn(
     arn: str,
 ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
     """Extracts hub_name, content_name, and content_version from a HubContentArn"""
 
-    match = re.match(constants.HUB_MODEL_ARN_REGEX, arn)
+    match = re.match(constants.HUB_CONTENT_ARN_REGEX, arn)
     if match:
         hub_name = match.group(4)
         hub_region = match.group(2)

Original file line number	Diff line number	Diff line change
`@@ -133,7 +133,6 @@ def _retrieve_default_environment_variables(`
`133`	`133`	`)`
`134`	`134`
`135`	`135`	`if gated_model_env_var is None and model_specs.is_gated_model():`
`136`		`-`
`137`	`136`	`possible_env_vars: Set[str] = {`
`138`	`137`	`retrieve_gated_env_var_for_instance_type(instance_type)`
`139`	`138`	`for instance_type in model_specs.supported_training_instance_types`
Original file line number	Diff line number	Diff line change
`@@ -134,7 +134,6 @@ def _retrieve_default_resources(`
`134`	`134`	`}`
`135`	`135`
`136`	`136`	`if is_dynamic_container_deployment_supported:`
`137`		`-`
`138`	`137`	`all_resource_requirement_kwargs = {}`
`139`	`138`
`140`	`139`	`for (`
Original file line number	Diff line number	Diff line change
`@@ -734,7 +734,6 @@ def attach(`
`734`	`734`	`"""`
`735`	`735`
`736`	`736`	`if model_id is None:`
`737`		`-`
`738`	`737`	`model_id, model_version = get_model_id_version_from_training_job(`
`739`	`738`	`training_job_name=training_job_name, sagemaker_session=sagemaker_session`
`740`	`739`	`)`