fix: SDK defaults formatting changes

Making changes wrt formatting to
fix tox issues

Author: Ashwin Krishna

src/sagemaker/config/config_schema.py

@@ -171,14 +171,20 @@ def _simple_path(*args: str):
 )
 AUTO_ML_JOB_CONFIG_PATH = _simple_path(SAGEMAKER, AUTO_ML_JOB, AUTO_ML_JOB_CONFIG)
 MONITORING_JOB_DEFINITION_PREFIX = _simple_path(
-    SAGEMAKER, MONITORING_SCHEDULE, MONITORING_SCHEDULE_CONFIG, MONITORING_JOB_DEFINITION
+    SAGEMAKER,
+    MONITORING_SCHEDULE,
+    MONITORING_SCHEDULE_CONFIG,
+    MONITORING_JOB_DEFINITION,
 )
 MONITORING_JOB_ENVIRONMENT_PATH = _simple_path(MONITORING_JOB_DEFINITION_PREFIX, ENVIRONMENT)
 MONITORING_JOB_OUTPUT_KMS_KEY_ID_PATH = _simple_path(
     MONITORING_JOB_DEFINITION_PREFIX, MONITORING_OUTPUT_CONFIG, KMS_KEY_ID
 )
 MONITORING_JOB_VOLUME_KMS_KEY_ID_PATH = _simple_path(
-    MONITORING_JOB_DEFINITION_PREFIX, MONITORING_RESOURCES, CLUSTER_CONFIG, VOLUME_KMS_KEY_ID
+    MONITORING_JOB_DEFINITION_PREFIX,
+    MONITORING_RESOURCES,
+    CLUSTER_CONFIG,
+    VOLUME_KMS_KEY_ID,
 )
 MONITORING_JOB_NETWORK_CONFIG_PATH = _simple_path(MONITORING_JOB_DEFINITION_PREFIX, NETWORK_CONFIG)
 MONITORING_JOB_ENABLE_NETWORK_ISOLATION_PATH = _simple_path(
@@ -292,7 +298,11 @@ def _simple_path(*args: str):
     SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, VPC_CONFIG, SECURITY_GROUP_IDS
 )
 REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = _simple_path(
-    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
+    SAGEMAKER,
+    PYTHON_SDK,
+    MODULES,
+    REMOTE_FUNCTION,
+    ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
 )
 MONITORING_SCHEDULE_INTER_CONTAINER_ENCRYPTION_PATH = _simple_path(
     SAGEMAKER,
@@ -472,7 +482,11 @@ def _simple_path(*args: str):
             "maxProperties": 48,
         },
         # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_S3DataSource.html#sagemaker-Type-S3DataSource-S3Uri
-        "s3Uri": {TYPE: "string", "pattern": "^(https|s3)://([^/]+)/?(.*)$", "maxLength": 1024},
+        "s3Uri": {
+            TYPE: "string",
+            "pattern": "^(https|s3)://([^/]+)/?(.*)$",
+            "maxLength": 1024,
+        },
         # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_AlgorithmSpecification.html#sagemaker-Type-AlgorithmSpecification-ContainerEntrypoint
         "preExecutionCommand": {TYPE: "string", "pattern": r".*"},
         # Regex based on https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_PipelineDefinitionS3Location.html
@@ -755,9 +769,7 @@ def _simple_path(*args: str):
                 ENDPOINT: {
                     TYPE: OBJECT,
                     ADDITIONAL_PROPERTIES: False,
-                    PROPERTIES: {
-                        TAGS: {"$ref": "#/definitions/tags"}
-                    }
+                    PROPERTIES: {TAGS: {"$ref": "#/definitions/tags"}},
                 },
                 # Endpoint Config
                 # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateEndpointConfig.html

src/sagemaker/model.py

@@ -50,7 +50,10 @@
 from sagemaker.predictor import PredictorBase
 from sagemaker.serverless import ServerlessInferenceConfig
 from sagemaker.transformer import Transformer
-from sagemaker.jumpstart.utils import add_jumpstart_tags, get_jumpstart_base_name_if_jumpstart_model
+from sagemaker.jumpstart.utils import (
+    add_jumpstart_tags,
+    get_jumpstart_base_name_if_jumpstart_model,
+)
 from sagemaker.utils import (
     unique_name_from_base,
     update_container_with_inference_params,
@@ -63,15 +66,25 @@
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline, PipelineSession
-from sagemaker.inference_recommender.inference_recommender_mixin import InferenceRecommenderMixin
+from sagemaker.inference_recommender.inference_recommender_mixin import (
+    InferenceRecommenderMixin,
+)
 
 LOGGER = logging.getLogger("sagemaker")
 
 NEO_ALLOWED_FRAMEWORKS = set(
     ["mxnet", "tensorflow", "keras", "pytorch", "onnx", "xgboost", "tflite"]
 )
 
-NEO_IOC_TARGET_DEVICES = ["ml_c4", "ml_c5", "ml_m4", "ml_m5", "ml_p2", "ml_p3", "ml_g4dn"]
+NEO_IOC_TARGET_DEVICES = [
+    "ml_c4",
+    "ml_c5",
+    "ml_m4",
+    "ml_m5",
+    "ml_p2",
+    "ml_p3",
+    "ml_g4dn",
+]
 
 NEO_MULTIVERSION_UNSUPPORTED = [
     "imx8mplus",
@@ -300,7 +313,9 @@ def __init__(
         self._base_name = None
         self.sagemaker_session = sagemaker_session
         self.role = resolve_value_from_config(
-            role, MODEL_EXECUTION_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session
+            role,
+            MODEL_EXECUTION_ROLE_ARN_PATH,
+            sagemaker_session=self.sagemaker_session,
         )
         self.vpc_config = resolve_value_from_config(
             vpc_config, MODEL_VPC_CONFIG_PATH, sagemaker_session=self.sagemaker_session
@@ -585,7 +600,9 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
         local_code = utils.get_config_value("local.local_code", self.sagemaker_session.config)
 
         bucket, key_prefix = s3.determine_bucket_and_prefix(
-            bucket=self.bucket, key_prefix=key_prefix, sagemaker_session=self.sagemaker_session
+            bucket=self.bucket,
+            key_prefix=key_prefix,
+            sagemaker_session=self.sagemaker_session,
         )
 
         if (self.sagemaker_session.local_mode and local_code) or self.entry_point is None:
@@ -633,7 +650,8 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
             else:
                 repacked_model_data = "s3://" + "/".join([bucket, key_prefix, "model.tar.gz"])
                 self.uploaded_code = fw_utils.UploadedCode(
-                    s3_prefix=repacked_model_data, script_name=os.path.basename(self.entry_point)
+                    s3_prefix=repacked_model_data,
+                    script_name=os.path.basename(self.entry_point),
                 )
 
             LOGGER.info(
@@ -693,7 +711,11 @@ def enable_network_isolation(self):
         return False if not self._enable_network_isolation else self._enable_network_isolation
 
     def _create_sagemaker_model(
-        self, instance_type=None, accelerator_type=None, tags=None, serverless_inference_config=None
+        self,
+        instance_type=None,
+        accelerator_type=None,
+        tags=None,
+        serverless_inference_config=None,
     ):
         """Create a SageMaker Model Entity
 
@@ -734,10 +756,14 @@ def _create_sagemaker_model(
         self._init_sagemaker_session_if_does_not_exist(instance_type)
         # Depending on the instance type, a local session (or) a session is initialized.
         self.role = resolve_value_from_config(
-            self.role, MODEL_EXECUTION_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session
+            self.role,
+            MODEL_EXECUTION_ROLE_ARN_PATH,
+            sagemaker_session=self.sagemaker_session,
         )
         self.vpc_config = resolve_value_from_config(
-            self.vpc_config, MODEL_VPC_CONFIG_PATH, sagemaker_session=self.sagemaker_session
+            self.vpc_config,
+            MODEL_VPC_CONFIG_PATH,
+            sagemaker_session=self.sagemaker_session,
         )
         self._enable_network_isolation = resolve_value_from_config(
             self._enable_network_isolation,
@@ -955,12 +981,16 @@ def package_for_edge(
             job_name = f"packaging{self._compilation_job_name[11:]}"
         self._init_sagemaker_session_if_does_not_exist(None)
         s3_kms_key = resolve_value_from_config(
-            s3_kms_key, EDGE_PACKAGING_KMS_KEY_ID_PATH, sagemaker_session=self.sagemaker_session
+            s3_kms_key,
+            EDGE_PACKAGING_KMS_KEY_ID_PATH,
+            sagemaker_session=self.sagemaker_session,
         )
         role = resolve_value_from_config(
             role, EDGE_PACKAGING_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session
         )
-        resource_key = resolve_value_from_config(resource_key, EDGE_PACKAGING_RESOURCE_KEY_PATH, sagemaker_session=self)
+        resource_key = resolve_value_from_config(
+            resource_key, EDGE_PACKAGING_RESOURCE_KEY_PATH, sagemaker_session=self
+        )
         if role is not None:
             role = self.sagemaker_session.expand_role(role)
         config = self._edge_packaging_job_config(
@@ -1066,7 +1096,9 @@ def compile(
 
         self._init_sagemaker_session_if_does_not_exist(target_instance_family)
         role = resolve_value_from_config(
-            role, COMPILATION_JOB_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session
+            role,
+            COMPILATION_JOB_ROLE_ARN_PATH,
+            sagemaker_session=self.sagemaker_session,
         )
         if not role:
             # Originally IAM role was a required parameter.
@@ -1231,10 +1263,14 @@ def deploy(
         self._init_sagemaker_session_if_does_not_exist(instance_type)
         # Depending on the instance type, a local session (or) a session is initialized.
         self.role = resolve_value_from_config(
-            self.role, MODEL_EXECUTION_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session
+            self.role,
+            MODEL_EXECUTION_ROLE_ARN_PATH,
+            sagemaker_session=self.sagemaker_session,
         )
         self.vpc_config = resolve_value_from_config(
-            self.vpc_config, MODEL_VPC_CONFIG_PATH, sagemaker_session=self.sagemaker_session
+            self.vpc_config,
+            MODEL_VPC_CONFIG_PATH,
+            sagemaker_session=self.sagemaker_session,
         )
         self._enable_network_isolation = resolve_value_from_config(
             self._enable_network_isolation,
@@ -1243,7 +1279,9 @@ def deploy(
         )
 
         tags = add_jumpstart_tags(
-            tags=tags, inference_model_uri=self.model_data, inference_script_uri=self.source_dir
+            tags=tags,
+            inference_model_uri=self.model_data,
+            inference_script_uri=self.source_dir,
         )
 
         if self.role is None:
@@ -1291,7 +1329,9 @@ def deploy(
         compiled_model_suffix = None if is_serverless else "-".join(instance_type.split(".")[:-1])
         if self._is_compiled_model and not is_serverless:
             self._ensure_base_name_if_needed(
-                image_uri=self.image_uri, script_uri=self.source_dir, model_uri=self.model_data
+                image_uri=self.image_uri,
+                script_uri=self.source_dir,
+                model_uri=self.model_data,
             )
             if self._base_name is not None:
                 self._base_name = "-".join((self._base_name, compiled_model_suffix))
@@ -1668,7 +1708,12 @@ class ModelPackage(Model):
     """A SageMaker ``Model`` that can be deployed to an ``Endpoint``."""
 
     def __init__(
-        self, role=None, model_data=None, algorithm_arn=None, model_package_arn=None, **kwargs
+        self,
+        role=None,
+        model_data=None,
+        algorithm_arn=None,
+        model_package_arn=None,
+        **kwargs,
     ):
         """Initialize a SageMaker ModelPackage.