fix: rename signing key

Author: Namrata Madan

requirements/extras/test_requirements.txt

@@ -13,4 +13,5 @@ graphene
 typing_extensions>=4.9.0
 tensorflow>=2.16.2,<=2.19.0
 build
-docker>=5.0.2,<8.0.0
+docker>=5.0.2,<8.0
+filelock>=3.0.0

sagemaker-core/pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "cloudpickle>=2.0.0",
     "paramiko>=2.11.0",
     "tblib>=1.7.0",
+    "cryptography>=46.0.0",
 ]
 requires-python = ">=3.9"
 classifiers = [

sagemaker-core/src/sagemaker/core/remote_function/core/stored_function.py

@@ -55,7 +55,7 @@ def __init__(
         self,
         sagemaker_session: Session,
         s3_base_uri: str,
-        hmac_key: str,
+        signing_key: str,
         s3_kms_key: str = None,
         context: Context = Context(),
     ):
@@ -66,13 +66,13 @@ def __init__(
                 AWS service calls are delegated to.
             s3_base_uri: the base uri to which serialized artifacts will be uploaded.
             s3_kms_key: KMS key used to encrypt artifacts uploaded to S3.
-            hmac_key: Key used to encrypt serialized and deserialized function and arguments.
+            signing_key: Key used to encrypt serialized and deserialized function and arguments.
             context: Build or run context of a pipeline step.
         """
         self.sagemaker_session = sagemaker_session
         self.s3_base_uri = s3_base_uri
         self.s3_kms_key = s3_kms_key
-        self.hmac_key = hmac_key
+        self.signing_key = signing_key
         self.context = context
 
         # For pipeline steps, function code is at: base/step_name/build_timestamp/
@@ -114,7 +114,7 @@ def save(self, func, *args, **kwargs):
             sagemaker_session=self.sagemaker_session,
             s3_uri=s3_path_join(self.func_upload_path, FUNCTION_FOLDER),
             s3_kms_key=self.s3_kms_key,
-            private_key=self.hmac_key,
+            private_key=self.signing_key,
         )
 
         logger.info(
@@ -126,7 +126,7 @@ def save(self, func, *args, **kwargs):
             obj=(args, kwargs),
             sagemaker_session=self.sagemaker_session,
             s3_uri=s3_path_join(self.func_upload_path, ARGUMENTS_FOLDER),
-            signing_key=self.hmac_key,
+            signing_key=self.signing_key,
             s3_kms_key=self.s3_kms_key,
         )
 
@@ -144,7 +144,7 @@ def save_pipeline_step_function(self, serialized_data):
         )
         serialization._upload_payload_and_metadata_to_s3_signed(
             bytes_to_upload=serialized_data.func,
-            private_key=self.hmac_key,
+            private_key=self.signing_key,
             s3_uri=s3_path_join(self.func_upload_path, FUNCTION_FOLDER),
             sagemaker_session=self.sagemaker_session,
             s3_kms_key=self.s3_kms_key,
@@ -156,7 +156,7 @@ def save_pipeline_step_function(self, serialized_data):
         )
         serialization._upload_payload_and_metadata_to_s3_signed(
             bytes_to_upload=serialized_data.args,
-            private_key=self.hmac_key,
+            private_key=self.signing_key,
             s3_uri=s3_path_join(self.func_upload_path, ARGUMENTS_FOLDER),
             sagemaker_session=self.sagemaker_session,
             s3_kms_key=self.s3_kms_key,
@@ -172,7 +172,7 @@ def load_and_invoke(self) -> Any:
         func = serialization.deserialize_func_from_s3(
             sagemaker_session=self.sagemaker_session,
             s3_uri=s3_path_join(self.func_upload_path, FUNCTION_FOLDER),
-            public_key_pem=self.hmac_key,
+            public_key_pem=self.signing_key,
         )
 
         logger.info(
@@ -182,7 +182,7 @@ def load_and_invoke(self) -> Any:
         args, kwargs = serialization.deserialize_obj_from_s3(
             sagemaker_session=self.sagemaker_session,
             s3_uri=s3_path_join(self.func_upload_path, ARGUMENTS_FOLDER),
-            verification_key=self.hmac_key,
+            verification_key=self.signing_key,
         )
 
         logger.info("Resolving pipeline variables")

sagemaker-core/src/sagemaker/core/remote_function/invoke_function.py

@@ -98,7 +98,7 @@ def _load_pipeline_context(args) -> Context:
 
 
 def _execute_remote_function(
-    sagemaker_session, s3_base_uri, s3_kms_key, run_in_context, hmac_key, context
+    sagemaker_session, s3_base_uri, s3_kms_key, run_in_context, signing_key, context
 ):
     """Execute stored remote function"""
     from sagemaker.core.remote_function.core.stored_function import StoredFunction
@@ -107,7 +107,7 @@ def _execute_remote_function(
         sagemaker_session=sagemaker_session,
         s3_base_uri=s3_base_uri,
         s3_kms_key=s3_kms_key,
-        hmac_key=hmac_key,
+        signing_key=signing_key,
         context=context,
     )
 
@@ -138,15 +138,15 @@ def main(sys_args=None):
         run_in_context = args.run_in_context
         pipeline_context = _load_pipeline_context(args)
 
-        hmac_key = os.getenv("REMOTE_FUNCTION_SECRET_KEY")
+        signing_key = os.getenv("REMOTE_FUNCTION_SECRET_KEY")
 
         sagemaker_session = _get_sagemaker_session(region)
         _execute_remote_function(
             sagemaker_session=sagemaker_session,
             s3_base_uri=s3_base_uri,
             s3_kms_key=s3_kms_key,
             run_in_context=run_in_context,
-            hmac_key=hmac_key,
+            signing_key=signing_key,
             context=pipeline_context,
         )
 

sagemaker-core/src/sagemaker/core/remote_function/job.py

@@ -842,19 +842,19 @@ def _get_default_spark_image(session):
 class _Job:
     """Helper class that interacts with the SageMaker training service."""
 
-    def __init__(self, job_name: str, s3_uri: str, sagemaker_session: Session, hmac_key: str):
+    def __init__(self, job_name: str, s3_uri: str, sagemaker_session: Session, verification_key: str):
         """Initialize a _Job object.
 
         Args:
             job_name (str): The training job name.
             s3_uri (str): The training job output S3 uri.
             sagemaker_session (Session): SageMaker boto session.
-            hmac_key (str): Remote function secret key.
+            verification_key (str): Remote function secret key.
         """
         self.job_name = job_name
         self.s3_uri = s3_uri
         self.sagemaker_session = sagemaker_session
-        self.hmac_key = hmac_key
+        self.verification_key = verification_key
         self._last_describe_response = None
 
     @staticmethod
@@ -870,9 +870,9 @@ def from_describe_response(describe_training_job_response, sagemaker_session):
         """
         job_name = describe_training_job_response["TrainingJobName"]
         s3_uri = describe_training_job_response["OutputDataConfig"]["S3OutputPath"]
-        hmac_key = describe_training_job_response["Environment"]["REMOTE_FUNCTION_SECRET_KEY"]
+        verification_key = describe_training_job_response["Environment"]["REMOTE_FUNCTION_SECRET_KEY"]
 
-        job = _Job(job_name, s3_uri, sagemaker_session, hmac_key)
+        job = _Job(job_name, s3_uri, sagemaker_session, verification_key)
         job._last_describe_response = describe_training_job_response
         return job
 
@@ -965,15 +965,15 @@ def compile(
             stored_function = StoredFunction(
                 sagemaker_session=job_settings.sagemaker_session,
                 s3_base_uri=s3_base_uri,
-                hmac_key=private_key,
+                signing_key=private_key,
                 s3_kms_key=job_settings.s3_kms_key,
             )
             stored_function.save(func, *func_args, **func_kwargs)
         else:
             stored_function = StoredFunction(
                 sagemaker_session=job_settings.sagemaker_session,
                 s3_base_uri=s3_base_uri,
-                hmac_key=private_key,
+                signing_key=private_key,
                 s3_kms_key=job_settings.s3_kms_key,
                 context=Context(
                     step_name=step_compilation_context.step_name,

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

@@ -21,6 +21,7 @@
 from contextlib import contextmanager
 
 import docker
+import filelock
 import pytest
 from docker.errors import BuildError
 
@@ -149,45 +150,52 @@ def gpu_instance_type():
 
 
 @pytest.fixture(scope="session")
-def dummy_container_without_error(sagemaker_session, compatible_python_version):
-    ecr_uri = _build_container(sagemaker_session, compatible_python_version, DOCKERFILE_TEMPLATE)
-    return ecr_uri
+def dummy_container_without_error(sagemaker_session, compatible_python_version, sagemaker_sdk_tar_path, tmp_path_factory):
+    return _build_container_once(
+        "dummy_container_without_error", sagemaker_session, compatible_python_version,
+        DOCKERFILE_TEMPLATE, sagemaker_sdk_tar_path, tmp_path_factory,
+    )
 
 
 @pytest.fixture(scope="session")
-def dummy_container_with_user_and_workdir(sagemaker_session, compatible_python_version):
-    ecr_uri = _build_container(
-        sagemaker_session,
-        compatible_python_version,
-        DOCKERFILE_TEMPLATE_WITH_USER_AND_WORKDIR,
+def dummy_container_with_user_and_workdir(sagemaker_session, compatible_python_version, sagemaker_sdk_tar_path, tmp_path_factory):
+    return _build_container_once(
+        "dummy_container_with_user_and_workdir", sagemaker_session, compatible_python_version,
+        DOCKERFILE_TEMPLATE_WITH_USER_AND_WORKDIR, sagemaker_sdk_tar_path, tmp_path_factory,
     )
-    return ecr_uri
 
 
 @pytest.fixture(scope="session")
-def dummy_container_incompatible_python_runtime(sagemaker_session, incompatible_python_version):
-    ecr_uri = _build_container(sagemaker_session, incompatible_python_version, DOCKERFILE_TEMPLATE)
-    return ecr_uri
+def dummy_container_incompatible_python_runtime(sagemaker_session, incompatible_python_version, sagemaker_sdk_tar_path, tmp_path_factory):
+    return _build_container_once(
+        "dummy_container_incompatible_python_runtime", sagemaker_session, incompatible_python_version,
+        DOCKERFILE_TEMPLATE, sagemaker_sdk_tar_path, tmp_path_factory,
+    )
 
 
 @pytest.fixture(scope="session")
-def dummy_container_with_conda(sagemaker_session, compatible_python_version):
-    ecr_uri = _build_container(
-        sagemaker_session, compatible_python_version, DOCKERFILE_TEMPLATE_WITH_CONDA
+def dummy_container_with_conda(sagemaker_session, compatible_python_version, sagemaker_sdk_tar_path, tmp_path_factory):
+    return _build_container_once(
+        "dummy_container_with_conda", sagemaker_session, compatible_python_version,
+        DOCKERFILE_TEMPLATE_WITH_CONDA, sagemaker_sdk_tar_path, tmp_path_factory,
     )
-    return ecr_uri
 
 
 @pytest.fixture(scope="session")
-def auto_capture_test_container(sagemaker_session):
-    ecr_uri = _build_auto_capture_client_container("3.10", AUTO_CAPTURE_CLIENT_DOCKER_TEMPLATE)
-    return ecr_uri
+def auto_capture_test_container(sagemaker_session, sagemaker_sdk_tar_path, tmp_path_factory):
+    return _build_container_once(
+        "auto_capture_test_container", sagemaker_session, "3.10",
+        AUTO_CAPTURE_CLIENT_DOCKER_TEMPLATE, sagemaker_sdk_tar_path, tmp_path_factory,
+        is_auto_capture=True,
+    )
 
 
 @pytest.fixture(scope="session")
-def spark_test_container(sagemaker_session):
-    ecr_uri = _build_container(sagemaker_session, "3.9", DOCKERFILE_TEMPLATE)
-    return ecr_uri
+def spark_test_container(sagemaker_session, sagemaker_sdk_tar_path, tmp_path_factory):
+    return _build_container_once(
+        "spark_test_container", sagemaker_session, "3.9",
+        DOCKERFILE_TEMPLATE, sagemaker_sdk_tar_path, tmp_path_factory,
+    )
 
 
 @pytest.fixture(scope="session")
@@ -208,6 +216,27 @@ def conda_env_yml():
         os.remove(conda_yml_file_name)
 
 
+@pytest.fixture(scope="session")
+def sagemaker_sdk_tar_path(tmp_path_factory):
+    """Build the sagemaker-core sdist once and share it across all xdist workers.
+
+    Uses a file lock so only one worker runs the build; others wait and reuse
+    the already-built tar.gz from the shared temp directory.
+    """
+    # tmp_path_factory.getbasetemp().parent is shared across all xdist workers
+    root_tmp = tmp_path_factory.getbasetemp().parent
+    tar_dir = root_tmp / "sagemaker_sdk_tar"
+    tar_dir.mkdir(exist_ok=True)
+    lock_file = root_tmp / "sagemaker_sdk_tar.lock"
+
+    with filelock.FileLock(str(lock_file)):
+        existing = list(tar_dir.glob("*.tar.gz"))
+        if not existing:
+            _generate_sagemaker_sdk_tar(str(tar_dir))
+            existing = list(tar_dir.glob("*.tar.gz"))
+    return str(existing[0])
+
+
 def _tmpdir():
     """Create a temporary directory context manager."""
     import tempfile
@@ -222,7 +251,33 @@ def _tmpdir():
 _tmpdir = contextmanager(_tmpdir)
 
 
-def _build_container(sagemaker_session, py_version, docker_template):
+def _build_container_once(
+    fixture_name, sagemaker_session, py_version, docker_template, sdk_tar_path,
+    tmp_path_factory, is_auto_capture=False,
+):
+    """Build and push a container image exactly once across all xdist workers.
+
+    Uses a file lock keyed by fixture_name so parallel workers wait for the
+    first worker to finish, then reuse the ECR URI written to a shared file.
+    """
+    root_tmp = tmp_path_factory.getbasetemp().parent
+    uri_file = root_tmp / f"{fixture_name}.ecr_uri"
+    lock_file = root_tmp / f"{fixture_name}.lock"
+
+    with filelock.FileLock(str(lock_file)):
+        if uri_file.exists():
+            return uri_file.read_text().strip()
+        if is_auto_capture:
+            ecr_uri = _build_auto_capture_client_container(
+                py_version, docker_template, sdk_tar_path
+            )
+        else:
+            ecr_uri = _build_container(sagemaker_session, py_version, docker_template, sdk_tar_path)
+        uri_file.write_text(ecr_uri)
+    return ecr_uri
+
+
+def _build_container(sagemaker_session, py_version, docker_template, sdk_tar_path):
     """Build a dummy test container locally and push to ECR."""
     region = sagemaker_session.boto_region_name
     image_tag = f"{py_version.replace('.', '-')}-{sagemaker_timestamp()}"
@@ -231,7 +286,8 @@ def _build_container(sagemaker_session, py_version, docker_template):
 
     with _tmpdir() as tmpdir:
         print("building docker image locally in ", tmpdir)
-        source_archive = _generate_sagemaker_sdk_tar(tmpdir)
+        source_archive = os.path.basename(sdk_tar_path)
+        shutil.copy2(sdk_tar_path, os.path.join(tmpdir, source_archive))
         with open(os.path.join(tmpdir, "Dockerfile"), "w") as file:
             content = docker_template.format(py_version=py_version, source_archive=source_archive)
             print(f"Dockerfile contents: \n{content}\n")
@@ -267,10 +323,11 @@ def _build_container(sagemaker_session, py_version, docker_template):
     return ecr_image
 
 
-def _build_auto_capture_client_container(py_version, docker_template):
+def _build_auto_capture_client_container(py_version, docker_template, sdk_tar_path):
     """Build a test docker container for auto_capture tests."""
     with _tmpdir() as tmpdir:
-        source_archive = _generate_sdk_tar_with_public_version(tmpdir)
+        source_archive = os.path.basename(sdk_tar_path)
+        shutil.copy2(sdk_tar_path, os.path.join(tmpdir, source_archive))
         _move_auto_capture_test_file(tmpdir)
         with open(os.path.join(tmpdir, "Dockerfile"), "w") as file:
             content = docker_template.format(py_version=py_version, source_archive=source_archive)
@@ -304,7 +361,11 @@ def _ecr_image_uri(account, region, image_name, tag):
 def _generate_sagemaker_sdk_tar(destination_folder):
     """Run build to generate the SDK tar file."""
     command = f"python -m build --sdist -o {destination_folder}"
-    result = subprocess.run(command, shell=True, check=True, capture_output=True)
+    try:
+        subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error when building sagemaker-core sdist: {e.stderr}")
+        raise
     destination_folder_contents = os.listdir(destination_folder)
     source_archive = [f for f in destination_folder_contents if f.endswith("tar.gz")][0]
     return source_archive

sagemaker-core/tests/unit/remote_function/test_job.py

@@ -146,7 +146,7 @@ def test_init(self, mock_session):
         job = _Job("test-job", "s3://bucket/output", mock_session, "test-key")
         assert job.job_name == "test-job"
         assert job.s3_uri == "s3://bucket/output"
-        assert job.hmac_key == "test-key"
+        assert job.verification_key == "test-key"
 
     def test_from_describe_response(self, mock_session):
         """Test creating _Job from describe response."""
@@ -158,7 +158,7 @@ def test_from_describe_response(self, mock_session):
         job = _Job.from_describe_response(response, mock_session)
         assert job.job_name == "test-job"
         assert job.s3_uri == "s3://bucket/output"
-        assert job.hmac_key == "test-key"
+        assert job.verification_key == "test-key"
 
     def test_describe_returns_cached_response(self, mock_session):
         """Test that describe returns cached response for completed jobs."""

sagemaker-core/tests/unit/remote_function/test_job_comprehensive.py

@@ -149,8 +149,7 @@ def test_from_describe_response(self, mock_session):
         job = _Job.from_describe_response(response, mock_session)
         assert job.job_name == "test-job"
         assert job.s3_uri == "s3://bucket/output"
-        assert job.hmac_key == "test-key"
-        assert job._last_describe_response == response
+        assert job.verification_key == "test-key"
 
     def test_describe_cached_completed(self, mock_session):
         """Test lines 865-871: describe with cached completed job."""