feature: adding resourcekey and tags for api in config for intelligent defaults

Adding ResourceKey to the Config schema
for the createEdgePackaging Job API

Adding a new entry into the Config Schema
for CreateEndpoint API with the addition of tags attribute

Author: akrishna1995

src/sagemaker/config/__init__.py

@@ -48,6 +48,7 @@
     COMPILATION_JOB_VPC_CONFIG_PATH,
     COMPILATION_JOB,
     EDGE_PACKAGING_ROLE_ARN_PATH,
+    EDGE_PACKAGING_RESOURCE_KEY_PATH,
     EDGE_PACKAGING_OUTPUT_CONFIG_PATH,
     EDGE_PACKAGING_JOB,
     TRANSFORM_JOB,
@@ -69,10 +70,13 @@
     MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH,
     ENDPOINT_CONFIG_PRODUCTION_VARIANTS_PATH,
     KMS_KEY_ID,
+    RESOURCE_KEY,
     ENDPOINT_CONFIG_KMS_KEY_ID_PATH,
     ENDPOINT_CONFIG,
     ENDPOINT_CONFIG_DATA_CAPTURE_PATH,
     ENDPOINT_CONFIG_ASYNC_INFERENCE_PATH,
+    ENDPOINT,
+    ENDPOINT_TAGS_PATH,
     SAGEMAKER,
     FEATURE_GROUP,
     TAGS,

src/sagemaker/config/config_schema.py

@@ -18,6 +18,7 @@
 ENABLE_NETWORK_ISOLATION = "EnableNetworkIsolation"
 VOLUME_KMS_KEY_ID = "VolumeKmsKeyId"
 KMS_KEY_ID = "KmsKeyId"
+RESOURCE_KEY = "ResourceKey"
 ROLE_ARN = "RoleArn"
 TAGS = "Tags"
 KEY = "Key"
@@ -78,6 +79,7 @@
 MODEL = "Model"
 MONITORING_SCHEDULE = "MonitoringSchedule"
 ENDPOINT_CONFIG = "EndpointConfig"
+ENDPOINT = "Endpoint"
 AUTO_ML_JOB = "AutoMLJob"
 COMPILATION_JOB = "CompilationJob"
 CUSTOM_PARAMETERS = "CustomParameters"
@@ -131,6 +133,7 @@ def _simple_path(*args: str):
 )
 EDGE_PACKAGING_OUTPUT_CONFIG_PATH = _simple_path(SAGEMAKER, EDGE_PACKAGING_JOB, OUTPUT_CONFIG)
 EDGE_PACKAGING_ROLE_ARN_PATH = _simple_path(SAGEMAKER, EDGE_PACKAGING_JOB, ROLE_ARN)
+EDGE_PACKAGING_RESOURCE_KEY_PATH = _simple_path(SAGEMAKER, EDGE_PACKAGING_JOB, RESOURCE_KEY)
 ENDPOINT_CONFIG_DATA_CAPTURE_KMS_KEY_ID_PATH = _simple_path(
     SAGEMAKER, ENDPOINT_CONFIG, DATA_CAPTURE_CONFIG, KMS_KEY_ID
 )
@@ -145,6 +148,7 @@ def _simple_path(*args: str):
     SAGEMAKER, ENDPOINT_CONFIG, ASYNC_INFERENCE_CONFIG, OUTPUT_CONFIG, KMS_KEY_ID
 )
 ENDPOINT_CONFIG_KMS_KEY_ID_PATH = _simple_path(SAGEMAKER, ENDPOINT_CONFIG, KMS_KEY_ID)
+ENDPOINT_TAGS_PATH = _simple_path(SAGEMAKER, ENDPOINT, TAGS)
 FEATURE_GROUP_ONLINE_STORE_CONFIG_PATH = _simple_path(SAGEMAKER, FEATURE_GROUP, ONLINE_STORE_CONFIG)
 FEATURE_GROUP_OFFLINE_STORE_CONFIG_PATH = _simple_path(
     SAGEMAKER, FEATURE_GROUP, OFFLINE_STORE_CONFIG
@@ -746,6 +750,15 @@ def _simple_path(*args: str):
                         TAGS: {"$ref": "#/definitions/tags"},
                     },
                 },
+                # Endpoint
+                # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateEndpoint.html
+                ENDPOINT: {
+                    TYPE: OBJECT,
+                    ADDITIONAL_PROPERTIES: False,
+                    PROPERTIES: {
+                        TAGS: {"$ref": "#/definitions/tags"}
+                    }
+                },
                 # Endpoint Config
                 # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateEndpointConfig.html
                 # Note: there is a separate API for creating Endpoints.
@@ -992,6 +1005,7 @@ def _simple_path(*args: str):
                             ADDITIONAL_PROPERTIES: False,
                             PROPERTIES: {KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"}},
                         },
+                        RESOURCE_KEY: {"$ref": "#/definitions/kmsKeyId"},
                         ROLE_ARN: {"$ref": "#/definitions/roleArn"},
                         TAGS: {"$ref": "#/definitions/tags"},
                     },

src/sagemaker/model.py

@@ -34,6 +34,7 @@
     EDGE_PACKAGING_KMS_KEY_ID_PATH,
     EDGE_PACKAGING_ROLE_ARN_PATH,
     MODEL_CONTAINERS_PATH,
+    EDGE_PACKAGING_RESOURCE_KEY_PATH,
     MODEL_VPC_CONFIG_PATH,
     MODEL_ENABLE_NETWORK_ISOLATION_PATH,
     MODEL_EXECUTION_ROLE_ARN_PATH,
@@ -959,6 +960,7 @@ def package_for_edge(
         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)
         if role is not None:
             role = self.sagemaker_session.expand_role(role)
         config = self._edge_packaging_job_config(

src/sagemaker/session.py

@@ -70,6 +70,7 @@
     COMPILATION_JOB,
     EDGE_PACKAGING_ROLE_ARN_PATH,
     EDGE_PACKAGING_OUTPUT_CONFIG_PATH,
+    EDGE_PACKAGING_RESOURCE_KEY_PATH,
     EDGE_PACKAGING_JOB,
     TRANSFORM_JOB,
     TRANSFORM_JOB_ENVIRONMENT_PATH,
@@ -95,6 +96,7 @@
     ENDPOINT_CONFIG,
     ENDPOINT_CONFIG_DATA_CAPTURE_PATH,
     ENDPOINT_CONFIG_ASYNC_INFERENCE_PATH,
+    ENDPOINT_TAGS_PATH,
     SAGEMAKER,
     FEATURE_GROUP,
     TAGS,
@@ -2503,6 +2505,7 @@ def package_model_for_edge(
             "EdgePackagingJobName": job_name,
             "CompilationJobName": compilation_job_name,
         }
+        resource_key = resolve_value_from_config(resource_key, EDGE_PACKAGING_RESOURCE_KEY_PATH, sagemaker_session=self)
         tags = _append_project_tags(tags)
         tags = self._append_sagemaker_config_tags(
             tags, "{}.{}.{}".format(SAGEMAKER, EDGE_PACKAGING_JOB, TAGS)
@@ -4018,6 +4021,8 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
             config_name (str): Name of the Amazon SageMaker endpoint configuration to deploy.
             wait (bool): Whether to wait for the endpoint deployment to complete before returning
                 (default: True).
+            tags (list[dict[str, str]]): A list of key-value pairs for tagging the endpoint
+                (default: None).
 
         Returns:
             str: Name of the Amazon SageMaker ``Endpoint`` created.
@@ -4026,6 +4031,9 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
 
         tags = tags or []
         tags = _append_project_tags(tags)
+        tags = self._append_sagemaker_config_tags(
+            tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_TAGS_PATH, TAGS)
+        )
 
         self.sagemaker_client.create_endpoint(
             EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags
@@ -4481,6 +4489,7 @@ def endpoint_from_model_data(
         model_vpc_config=None,
         accelerator_type=None,
         data_capture_config=None,
+        tags=None,
     ):
         """Create and deploy to an ``Endpoint`` using existing model data stored in S3.
 
@@ -4512,14 +4521,23 @@ def endpoint_from_model_data(
             data_capture_config (sagemaker.model_monitor.DataCaptureConfig): Specifies
                 configuration related to Endpoint data capture for use with
                 Amazon SageMaker Model Monitoring. Default: None.
+            tags (list[dict[str, str]]): A list of key-value pairs for tagging the endpoint
+                (default: None).
 
         Returns:
             str: Name of the ``Endpoint`` that is created.
         """
         model_environment_vars = model_environment_vars or {}
         name = name or name_from_image(image_uri)
         model_vpc_config = vpc_utils.sanitize(model_vpc_config)
-
+        endpoint_config_tags = _append_project_tags(tags)
+        endpoint_tags = _append_project_tags(tags)
+        endpoint_tags = self._append_sagemaker_config_tags(
+            endpoint_tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_TAGS_PATH, TAGS)
+        )
+        endpoint_config_tags = self._append_sagemaker_config_tags(
+            endpoint_config_tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_CONFIG, TAGS)
+        )
         primary_container = container_def(
             image_uri=image_uri,
             model_data_url=model_s3_location,
@@ -4542,12 +4560,13 @@ def endpoint_from_model_data(
                 instance_type=instance_type,
                 accelerator_type=accelerator_type,
                 data_capture_config_dict=data_capture_config_dict,
+                tags=endpoint_config_tags
             )
         )
 
         # to make change backwards compatible
         response = _create_resource(
-            lambda: self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+            lambda: self.create_endpoint(endpoint_name=name, config_name=name, Tags=endpoint_tags, wait=wait)
         )
         if not response:
             raise ValueError(
@@ -4612,12 +4631,18 @@ def endpoint_from_production_variants(
             if supports_kms
             else kms_key
         )
-        tags = _append_project_tags(tags)
-        tags = self._append_sagemaker_config_tags(
-            tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_CONFIG, TAGS)
+
+        endpoint_config_tags = _append_project_tags(tags)
+        endpoint_tags = _append_project_tags(tags)
+
+        endpoint_tags = self._append_sagemaker_config_tags(
+            endpoint_tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_TAGS_PATH, TAGS)
         )
-        if tags:
-            config_options["Tags"] = tags
+        endpoint_config_tags = self._append_sagemaker_config_tags(
+            endpoint_config_tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_CONFIG, TAGS)
+        )
+        if endpoint_config_tags:
+            config_options["Tags"] = endpoint_config_tags
         if kms_key:
             config_options["KmsKeyId"] = kms_key
         if data_capture_config_dict is not None:
@@ -4638,7 +4663,7 @@ def endpoint_from_production_variants(
         LOGGER.info("Creating endpoint-config with name %s", name)
         self.sagemaker_client.create_endpoint_config(**config_options)
 
-        return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
+        return self.create_endpoint(endpoint_name=name, config_name=name, tags=endpoint_tags, wait=wait)
 
     def expand_role(self, role):
         """Expand an IAM role name into an ARN.

tests/data/config/config.yaml

@@ -53,6 +53,10 @@ SageMaker:
             Subnets:
               - 'subnet-1234'
         RoleArn: 'arn:aws:iam::555555555555:role/IMRole'
+  Endpoint:
+    Tags:
+      - Key: "tag1"
+        Value: "tagValue1"
   EndpointConfig:
     AsyncInferenceConfig:
       OutputConfig:
@@ -171,4 +175,28 @@ SageMaker:
   EdgePackagingJob:
     OutputConfig:
       KmsKeyId: 'kmskeyid1'
-    RoleArn: 'arn:aws:iam::555555555555:role/IMRole'
+    RoleArn: 'arn:aws:iam::555555555555:role/IMRole'
+    ResourceKey: 'kmskeyid1'
+  PythonSDK:
+    Modules:
+      RemoteFunction:
+        Dependencies: "./requirements.txt"
+        EnvironmentVariables:
+          "var1": "value1"
+          "var2": "value2"
+        ImageUri: "123456789012.dkr.ecr.us-west-2.amazonaws.com/myimage:latest"
+        IncludeLocalWorkDir: true
+        InstanceType: "ml.m5.xlarge"
+        JobCondaEnvironment: "some_conda_env"
+        RoleArn: "arn:aws:iam::555555555555:role/IMRole"
+        S3KmsKeyId: "kmskeyid1"
+        S3RootUri: "s3://my-bucket/key"
+        Tags:
+          - Key: "tag1"
+            Value: "tagValue1"
+        VolumeKmsKeyId: "kmskeyid2"
+        VpcConfig:
+          SecurityGroupIds:
+            - 'sg123'
+          Subnets:
+            - 'subnet-1234'

tests/unit/__init__.py

@@ -39,6 +39,8 @@
     COMPILATION_JOB,
     OUTPUT_CONFIG,
     EDGE_PACKAGING_JOB,
+    RESOURCE_KEY,
+    ENDPOINT,
     ENDPOINT_CONFIG,
     DATA_CAPTURE_CONFIG,
     PRODUCTION_VARIANTS,
@@ -145,6 +147,7 @@
             OUTPUT_CONFIG: {
                 KMS_KEY_ID: "configKmsKeyId",
             },
+            RESOURCE_KEY: "kmskeyid1",
             ROLE_ARN: "arn:aws:iam::111111111111:role/ConfigRole",
             TAGS: [{KEY: "some-tag", VALUE: "value-for-tag"}],
         },
@@ -173,6 +176,15 @@
     },
 }
 
+SAGEMAKER_CONFIG_ENDPOINT = {
+    SCHEMA_VERSION: "1.0",
+    SAGEMAKER: {
+        ENDPOINT: {
+            TAGS: [{KEY: "some-tag", VALUE: "value-for-tag"}],
+        }
+    }
+}
+
 SAGEMAKER_CONFIG_AUTO_ML = {
     SCHEMA_VERSION: "1.0",
     SAGEMAKER: {