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
+cryptography>=46.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/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."""

sagemaker-core/tox.ini

@@ -6,7 +6,6 @@
 [tox]
 isolated_build = true
 envlist = black-format,flake8,pylint,docstyle,sphinx,doc8,twine,py39,py310,py311,py312
-
 skip_missing_interpreters = False
 
 [flake8]

sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py

@@ -1157,7 +1157,7 @@ def get_function_step_result(
         return deserialize_obj_from_s3(
             sagemaker_session=sagemaker_session,
             s3_uri=s3_uri,
-            hmac_key=describe_training_job_response["Environment"]["REMOTE_FUNCTION_SECRET_KEY"],
+            verification_key=describe_training_job_response["Environment"]["REMOTE_FUNCTION_SECRET_KEY"],
         )
 
     raise RemoteFunctionError(_ERROR_MSG_OF_STEP_INCOMPLETE)