intelligent defaults - custom parameters and small fixes (aws#845)

* fix: sagemaker-config - S3 session, tuning tags, config schema test side-effects

* feature: sagemaker-config - support for custom parameters in config schema

---------

Co-authored-by: Ruban Hussain <[email protected]>

Author: rubanh

src/sagemaker/config/config_schema.py

@@ -70,9 +70,11 @@
 ENDPOINT_CONFIG = "EndpointConfig"
 AUTO_ML = "AutoML"
 COMPILATION_JOB = "CompilationJob"
+CUSTOM_PARAMETERS = "CustomParameters"
 PIPELINE = "Pipeline"
 TRANSFORM_JOB = "TransformJob"
 PROPERTIES = "properties"
+PATTERN_PROPERTIES = "patternProperties"
 TYPE = "type"
 OBJECT = "object"
 ADDITIONAL_PROPERTIES = "additionalProperties"
@@ -104,7 +106,6 @@ def _simple_path(*args: str):
     SAGEMAKER, TRAINING_JOB, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
 )
 
-
 SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {
     "$schema": "https://json-schema.org/draft/2020-12/schema",
     TYPE: OBJECT,
@@ -231,6 +232,13 @@ def _simple_path(*args: str):
             "enum": ["1.0"],
             "description": "The schema version of the document.",
         },
+        CUSTOM_PARAMETERS: {
+            TYPE: OBJECT,
+            ADDITIONAL_PROPERTIES: False,
+            PATTERN_PROPERTIES: {
+                "^[\w\s\d_.:/=+\-@]+$": {TYPE: "string"},
+            }
+        },
         SAGEMAKER: {
             TYPE: OBJECT,
             ADDITIONAL_PROPERTIES: False,

src/sagemaker/session.py

@@ -41,7 +41,6 @@
     ENDPOINT_CONFIG,
     FEATURE_GROUP,
     KEY,
-    HYPER_PARAMETER_TUNING_JOB,
     SAGEMAKER,
     MODEL,
     MONITORING_SCHEDULE,
@@ -219,7 +218,7 @@ def _initialize(
         if sagemaker_config:
             self.sagemaker_config = sagemaker_config
         else:
-            self.sagemaker_config = SageMakerConfig(s3_resource=self.boto_session)
+            self.sagemaker_config = SageMakerConfig(s3_resource=self.boto_session.resource("s3"))
 
     @property
     def boto_region_name(self):
@@ -2642,9 +2641,6 @@ def tune(  # noqa: C901
             tune_request["WarmStartConfig"] = warm_start_config
 
         tags = _append_project_tags(tags)
-        tags = self._append_sagemaker_config_tags(
-            tags, "{}.{}.{}".format(SAGEMAKER, HYPER_PARAMETER_TUNING_JOB, TAGS)
-        )
         if tags is not None:
             tune_request["Tags"] = tags
 
@@ -2747,9 +2743,6 @@ def _get_tuning_request(
             tune_request["WarmStartConfig"] = warm_start_config
 
         tags = _append_project_tags(tags)
-        tags = self._append_sagemaker_config_tags(
-            tags, "{}.{}.{}".format(SAGEMAKER, HYPER_PARAMETER_TUNING_JOB, TAGS)
-        )
         if tags is not None:
             tune_request["Tags"] = tags
 

tests/unit/sagemaker/config/conftest.py

@@ -17,22 +17,22 @@
 from mock import MagicMock
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def base_config_with_schema():
     return {"SchemaVersion": "1.0"}
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_vpc_config():
     return {"SecurityGroupIds": ["sg123"], "Subnets": ["subnet-1234"]}
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_iam_role_arn():
     return "arn:aws:iam::366666666666:role/IMRole"
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_feature_group_config(valid_iam_role_arn):
     s3_storage_config = {"KmsKeyId": "somekmskeyid"}
     security_storage_config = {"KmsKeyId": "someotherkmskeyid"}
@@ -45,15 +45,15 @@ def valid_feature_group_config(valid_iam_role_arn):
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_edge_packaging_config(valid_iam_role_arn):
     return {
         "OutputConfig": {"KmsKeyId": "somekeyid"},
         "RoleArn": valid_iam_role_arn,
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_model_config(valid_iam_role_arn, valid_vpc_config):
     return {
         "EnableNetworkIsolation": True,
@@ -62,7 +62,7 @@ def valid_model_config(valid_iam_role_arn, valid_vpc_config):
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_model_package_config(valid_iam_role_arn):
     transform_job_definition = {
         "TransformOutput": {"KmsKeyId": "somerandomkmskeyid"},
@@ -75,7 +75,7 @@ def valid_model_package_config(valid_iam_role_arn):
     return {"ValidationSpecification": validation_specification}
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_processing_job_config(valid_iam_role_arn, valid_vpc_config):
     network_config = {"EnableNetworkIsolation": True, "VpcConfig": valid_vpc_config}
     dataset_definition = {
@@ -94,7 +94,7 @@ def valid_processing_job_config(valid_iam_role_arn, valid_vpc_config):
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_training_job_config(valid_iam_role_arn, valid_vpc_config):
     return {
         "EnableNetworkIsolation": True,
@@ -105,12 +105,12 @@ def valid_training_job_config(valid_iam_role_arn, valid_vpc_config):
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_pipeline_config(valid_iam_role_arn):
     return {"RoleArn": valid_iam_role_arn}
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_compilation_job_config(valid_iam_role_arn, valid_vpc_config):
     return {
         "OutputConfig": {"KmsKeyId": "somekmskey"},
@@ -119,7 +119,7 @@ def valid_compilation_job_config(valid_iam_role_arn, valid_vpc_config):
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_transform_job_config():
     return {
         "DataCaptureConfig": {"KmsKeyId": "somekmskey"},
@@ -128,7 +128,7 @@ def valid_transform_job_config():
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_automl_config(valid_iam_role_arn, valid_vpc_config):
     return {
         "AutoMLJobConfig": {
@@ -139,7 +139,7 @@ def valid_automl_config(valid_iam_role_arn, valid_vpc_config):
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_endpointconfig_config():
     return {
         "AsyncInferenceConfig": {"OutputConfig": {"KmsKeyId": "somekmskey"}},
@@ -150,7 +150,7 @@ def valid_endpointconfig_config():
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_monitoring_schedule_config(valid_iam_role_arn, valid_vpc_config):
     network_config = {"EnableNetworkIsolation": True, "VpcConfig": valid_vpc_config}
     return {
@@ -165,7 +165,7 @@ def valid_monitoring_schedule_config(valid_iam_role_arn, valid_vpc_config):
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def valid_config_with_all_the_scopes(
     valid_feature_group_config,
     valid_monitoring_schedule_config,
@@ -196,11 +196,11 @@ def valid_config_with_all_the_scopes(
     }
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def s3_resource_mock():
     return MagicMock(name="s3")
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture()
 def get_data_dir():
     return os.path.join(os.path.dirname(__file__), "..", "..", "..", "data", "config")

tests/unit/sagemaker/config/test_config_schema.py

@@ -38,7 +38,7 @@ def test_invalid_schema_version():
 
 
 def test_valid_config_with_all_the_features(
-    base_config_with_schema, valid_config_with_all_the_scopes
+        base_config_with_schema, valid_config_with_all_the_scopes
 ):
     _validate_config(base_config_with_schema, valid_config_with_all_the_scopes)
 
@@ -137,3 +137,31 @@ def test_invalid_feature_group_schema(base_config_with_schema):
     config["SageMaker"] = {"FeatureGroup": feature_group_config}
     with pytest.raises(exceptions.ValidationError):
         validate(config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
+
+
+def test_valid_custom_parameters_schema(base_config_with_schema):
+    config = base_config_with_schema
+    config["CustomParameters"] = {
+        "custom_key": "custom_value",
+        "CustomKey": "CustomValue",
+        "custom key": "custom value",
+        "custom-key": "custom-value",
+        "custom0123 key0123": "custom0123 value0123",
+    }
+    validate(config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
+
+
+def test_invalid_custom_parameters_schema(base_config_with_schema):
+    config = base_config_with_schema
+
+    config["CustomParameters"] = {"^&": "custom_value"}
+    with pytest.raises(exceptions.ValidationError):
+        validate(config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
+
+    config["CustomParameters"] = {"custom_key": 476}
+    with pytest.raises(exceptions.ValidationError):
+        validate(config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)
+
+    config["CustomParameters"] = {"custom_key": {"custom_key": "custom_value"}}
+    with pytest.raises(exceptions.ValidationError):
+        validate(config, SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA)