test: unskip sm-core integ tests (aws#5856)

* test: unskip sm-core integ tests

* fix: resolve intelligent defaults bug in ImageRetriever by replacing ineffective locals() assignment with args dict

* fix: use PascalCase for config attribute lookup to match schema keys

* test: update expected sagemaker-distribution image version from 3.0.0 to 3.2.0

* fix: ignore unknown fields in DefaultPayloadsModel to handle hub schema additions

* fix: port kms_utils from SageMakerHulkPythonSDK and enable test_advanced_job_setting

- Add tests/integ/kms_utils.py with get_or_create_kms_key utility ported
  from SageMakerHulkPythonSDK, adapted for sagemaker-core import paths
- Uncomment kms_utils import in test_decorator.py to fix NameError in
  s3_kms_key fixture

* fix: unskip test_decorator_with_spark_job with conditional skipif and fix test bugs

- Add skipif for Python versions without Spark image (only 3.9/3.12 supported)
- Fix Properties set literal to dict: {"spark.app.name": "remote-spark-test"}
- Fix spark.conf.get() to use string key instead of attribute reference

* Revert "fix: ignore unknown fields in DefaultPayloadsModel to handle hub schema additions"

* update kms helper name

Author: lucasjia-aws

sagemaker-core/src/sagemaker/core/image_retriever/image_retriever.py

@@ -25,9 +25,15 @@
     config_for_framework,
 )
 from sagemaker.core.workflow.utilities import override_pipeline_parameter_var
-from sagemaker.core.config.config_schema import IMAGE_RETRIEVER, MODULES, SAGEMAKER, _simple_path
+from sagemaker.core.config.config_schema import IMAGE_RETRIEVER, MODULES, PYTHON_SDK, SAGEMAKER, _simple_path
 from sagemaker.core.config.config_manager import SageMakerConfig
 
+
+def _to_pascal_case(name):
+    """Convert snake_case to PascalCase."""
+    camel = to_camel_case(name)
+    return camel[0].upper() + camel[1:] if camel else camel
+
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
 PYTORCH_FRAMEWORK = "pytorch"
@@ -114,11 +120,25 @@ def retrieve_hugging_face_uri(
             if name in CONFIGURABLE_ATTRIBUTES and not val:
                 default_value = ImageRetriever._config.resolve_value_from_config(
                     config_path=_simple_path(
-                        SAGEMAKER, MODULES, IMAGE_RETRIEVER, to_camel_case(name)
+                        SAGEMAKER, PYTHON_SDK, MODULES, IMAGE_RETRIEVER, _to_pascal_case(name)
                     )
                 )
                 if default_value is not None:
-                    locals()[name] = default_value
+                    args[name] = default_value
+
+        # Apply resolved defaults back to local variables
+        version = args.get("version", version)
+        py_version = args.get("py_version", py_version)
+        instance_type = args.get("instance_type", instance_type)
+        accelerator_type = args.get("accelerator_type", accelerator_type)
+        image_scope = args.get("image_scope", image_scope)
+        container_version = args.get("container_version", container_version)
+        distributed = args.get("distributed", distributed)
+        base_framework_version = args.get("base_framework_version", base_framework_version)
+        training_compiler_config = args.get("training_compiler_config", training_compiler_config)
+        sdk_version = args.get("sdk_version", sdk_version)
+        inference_tool = args.get("inference_tool", inference_tool)
+        serverless_inference_config = args.get("serverless_inference_config", serverless_inference_config)
 
         if training_compiler_config:
             final_image_scope = image_scope
@@ -503,11 +523,28 @@ def retrieve(
             if name in CONFIGURABLE_ATTRIBUTES and not val:
                 default_value = ImageRetriever._config.resolve_value_from_config(
                     config_path=_simple_path(
-                        SAGEMAKER, MODULES, IMAGE_RETRIEVER, to_camel_case(name)
+                        SAGEMAKER, PYTHON_SDK, MODULES, IMAGE_RETRIEVER, _to_pascal_case(name)
                     )
                 )
                 if default_value is not None:
-                    locals()[name] = default_value
+                    args[name] = default_value
+
+        # Apply resolved defaults back to local variables
+        version = args.get("version", version)
+        py_version = args.get("py_version", py_version)
+        instance_type = args.get("instance_type", instance_type)
+        accelerator_type = args.get("accelerator_type", accelerator_type)
+        image_scope = args.get("image_scope", image_scope)
+        container_version = args.get("container_version", container_version)
+        distributed = args.get("distributed", distributed)
+        smp = args.get("smp", smp)
+        base_framework_version = args.get("base_framework_version", base_framework_version)
+        training_compiler_config = args.get("training_compiler_config", training_compiler_config)
+        model_id = args.get("model_id", model_id)
+        model_version = args.get("model_version", model_version)
+        sdk_version = args.get("sdk_version", sdk_version)
+        inference_tool = args.get("inference_tool", inference_tool)
+        serverless_inference_config = args.get("serverless_inference_config", serverless_inference_config)
 
         for name, val in args.items():
             if is_pipeline_variable(val):

sagemaker-core/tests/integ/image_retriever/test_image_retriever.py

@@ -94,7 +94,7 @@ def test_retrieve_base_python_image_uri():
     assert image_uri == "236514542706.dkr.ecr.us-west-2.amazonaws.com/sagemaker-base-python-310:1.0"
 
 
-@pytest.mark.skip(reason="Test is failing due to locals()[name] = default_value in Image Retriever")
+# @pytest.mark.skip(reason="Test is failing due to locals()[name] = default_value in Image Retriever")
 @patch.object(SageMakerConfig, "resolve_value_from_config")
 def test_retrieve_image_uri_intelligent_default(mock_load_config):
     def custom_return(config_path=None, **kwargs):
@@ -116,5 +116,5 @@ def custom_return(config_path=None, **kwargs):
     )
     assert (
         image_uri
-        == "053634841547.dkr.ecr.us-west-1.amazonaws.com/sagemaker-distribution-prod:3.0.0-gpu"
+        == "053634841547.dkr.ecr.us-west-1.amazonaws.com/sagemaker-distribution-prod:3.2.0-gpu"
     )

sagemaker-core/tests/integ/integ_test_kms_helpers.py

@@ -0,0 +1,138 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""KMS test helpers for integration tests.
+
+Ported from SageMakerHulkPythonSDK/tests/integ/kms_utils.py.
+
+NOTE: KMS keys created by these helpers use a fixed alias and are intentionally
+reused across test runs rather than deleted after each run. This is because KMS
+keys have a mandatory 7-day minimum deletion window (schedule_key_deletion), so
+per-run create/delete is not practical. The persistent shared key approach avoids
+accumulating orphaned keys and unnecessary costs.
+"""
+from __future__ import absolute_import
+
+import json
+
+from sagemaker.core.common_utils import aws_partition, sts_regional_endpoint
+
+PRINCIPAL_TEMPLATE = (
+    '["{account_id}", "{role_arn}", '
+    '"arn:{partition}:iam::{account_id}:role/{sagemaker_role}"] '
+)
+
+KEY_ALIAS = "SageMakerTestKMSKey"
+POLICY_NAME = "default"
+KEY_POLICY = """
+{{
+  "Version": "2012-10-17",
+  "Id": "{id}",
+  "Statement": [
+    {{
+      "Sid": "Enable IAM User Permissions",
+      "Effect": "Allow",
+      "Principal": {{
+        "AWS": {principal}
+      }},
+      "Action": "kms:*",
+      "Resource": "*"
+    }}
+  ]
+}}
+"""
+
+
+def _get_kms_key_arn(kms_client, alias):
+    try:
+        response = kms_client.describe_key(KeyId="alias/" + alias)
+        return response["KeyMetadata"]["Arn"]
+    except kms_client.exceptions.NotFoundException:
+        return None
+
+
+def _get_kms_key_id(kms_client, alias):
+    try:
+        response = kms_client.describe_key(KeyId="alias/" + alias)
+        return response["KeyMetadata"]["KeyId"]
+    except kms_client.exceptions.NotFoundException:
+        return None
+
+
+def _create_kms_key(
+    kms_client, account_id, region, role_arn=None, sagemaker_role="SageMakerRole", alias=KEY_ALIAS
+):
+    if role_arn:
+        principal = PRINCIPAL_TEMPLATE.format(
+            partition=aws_partition(region),
+            account_id=account_id,
+            role_arn=role_arn,
+            sagemaker_role=sagemaker_role,
+        )
+    else:
+        principal = '"{account_id}"'.format(account_id=account_id)
+
+    response = kms_client.create_key(
+        Policy=KEY_POLICY.format(
+            id=POLICY_NAME, principal=principal, sagemaker_role=sagemaker_role
+        ),
+        Description="KMS key for SageMaker Python SDK integ tests",
+    )
+    key_arn = response["KeyMetadata"]["Arn"]
+
+    if alias:
+        kms_client.create_alias(AliasName="alias/" + alias, TargetKeyId=key_arn)
+    return key_arn
+
+
+def _add_role_to_policy(
+    kms_client, account_id, role_arn, region, alias=KEY_ALIAS, sagemaker_role="SageMakerRole"
+):
+    key_id = _get_kms_key_id(kms_client, alias)
+    policy = kms_client.get_key_policy(KeyId=key_id, PolicyName=POLICY_NAME)
+    policy = json.loads(policy["Policy"])
+    principal = policy["Statement"][0]["Principal"]["AWS"]
+
+    if role_arn not in principal or sagemaker_role not in principal:
+        principal = PRINCIPAL_TEMPLATE.format(
+            partition=aws_partition(region),
+            account_id=account_id,
+            role_arn=role_arn,
+            sagemaker_role=sagemaker_role,
+        )
+
+        kms_client.put_key_policy(
+            KeyId=key_id,
+            PolicyName=POLICY_NAME,
+            Policy=KEY_POLICY.format(id=POLICY_NAME, principal=principal),
+        )
+
+
+def get_or_create_kms_key(
+    sagemaker_session, role_arn=None, alias=KEY_ALIAS, sagemaker_role="SageMakerRole"
+):
+    kms_client = sagemaker_session.boto_session.client("kms")
+    kms_key_arn = _get_kms_key_arn(kms_client, alias)
+
+    region = sagemaker_session.boto_region_name
+    sts_client = sagemaker_session.boto_session.client(
+        "sts", region_name=region, endpoint_url=sts_regional_endpoint(region)
+    )
+    account_id = sts_client.get_caller_identity()["Account"]
+
+    if kms_key_arn is None:
+        return _create_kms_key(kms_client, account_id, region, role_arn, sagemaker_role, alias)
+
+    if role_arn:
+        _add_role_to_policy(kms_client, account_id, role_arn, region, alias, sagemaker_role)
+
+    return kms_key_arn

sagemaker-core/tests/integ/remote_function/test_decorator.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+import sys
 import time
 from typing import Union
 
@@ -40,7 +41,7 @@
 from sagemaker.core.common_utils import unique_name_from_base
 from tests.integ.s3_utils import assert_s3_files_exist
 
-# from tests.integ.kms_utils import get_or_create_kms_key  # TODO: provide KMS utils
+from tests.integ.integ_test_kms_helpers import get_or_create_kms_key
 import os
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "data")
@@ -122,7 +123,7 @@ def divide(x, y):
 
 
 # TODO: add VPC settings, update SageMakerRole with KMS permissions
-@pytest.mark.skip
+# @pytest.mark.skip
 def test_advanced_job_setting(
     sagemaker_session, dummy_container_without_error, cpu_instance_type, s3_kms_key
 ):
@@ -573,7 +574,10 @@ def my_func():
     assert client_error_message in str(error)
 
 
-@pytest.mark.skip
+@pytest.mark.skipif(
+    sys.version_info[:2] not in [(3, 9), (3, 12)],
+    reason="SageMaker Spark image only available for Python 3.9 and 3.12",
+)
 def test_decorator_with_spark_job(sagemaker_session, cpu_instance_type):
     @remote(
         role=ROLE,
@@ -584,7 +588,7 @@ def test_decorator_with_spark_job(sagemaker_session, cpu_instance_type):
             configuration=[
                 {
                     "Classification": "spark-defaults",
-                    "Properties": {"spark.app.name", "remote-spark-test"},
+                    "Properties": {"spark.app.name": "remote-spark-test"},
                 }
             ]
         ),
@@ -594,12 +598,12 @@ def test_spark_transform():
 
         spark = SparkSession.builder.getOrCreate()
 
-        assert spark.conf.get(spark.app.name) == "remote-spark-test"
+        assert spark.conf.get("spark.app.name") == "remote-spark-test"
 
     test_spark_transform()
 
 
-@pytest.mark.skip
+# @pytest.mark.skip
 def test_decorator_auto_capture(sagemaker_session, auto_capture_test_container):
     """
     This test runs a docker container. The Container invocation will execute a python script