From c53c1bb72dd415e7a047dd1ab2931511213ae225 Mon Sep 17 00:00:00 2001 From: hippi345 <13539685+hippi345@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:42:39 +0000 Subject: [PATCH 1/2] fix: use pipeline code_hash for FrameworkProcessor sourcedir upload FrameworkProcessor._package_code() now uses _pipeline_config.code_hash in the S3 key when available, matching the behavior of ScriptProcessor._upload_code(). This enables pipeline cache hits for FrameworkProcessor jobs with source_dir. Fixes #6493 --- .../src/sagemaker/core/processing.py | 23 +++++++++--- sagemaker-core/tests/unit/test_processing.py | 37 +++++++++++++++++++ 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 9315b746d5..bf094f494b 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -1336,12 +1336,23 @@ def _package_code( raise ValueError(f"Dependency path does not exist: {dep_path}") tar.add(dep_path, arcname=os.path.basename(dep_path)) - s3_uri = s3.s3_path_join( - self._s3_code_prefix(), - job_name, - "source", - "sourcedir.tar.gz", - ) + # Use pipeline code_hash for content-addressed caching when available + from sagemaker.core.workflow.utilities import _pipeline_config + + if _pipeline_config and _pipeline_config.code_hash: + s3_uri = s3.s3_path_join( + self._s3_code_prefix(), + _pipeline_config.pipeline_name, + "source", + _pipeline_config.code_hash, + ) + else: + s3_uri = s3.s3_path_join( + self._s3_code_prefix(), + job_name, + "source", + "sourcedir.tar.gz", + ) with open(tmp.name, "rb") as tar_file: tar_bytes = tar_file.read() diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index 22cdabfc7a..6dfdbc56ff 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -1411,6 +1411,43 @@ def checking_unlink(path, *args, **kwargs): observed.get("still_open") == [] ), "temp tar.gz handle was still open when os.unlink was called" + def test_package_code_with_pipeline_config(self, mock_session): + """Test that _package_code uses pipeline code_hash when available. + + When _pipeline_config.code_hash is set, _package_code should use it + in the S3 key (similar to ScriptProcessor._upload_code) for pipeline + cache hits. This test verifies the fix for issue #6493. + """ + processor = FrameworkProcessor( + role="arn:aws:iam::123456789012:role/SageMakerRole", + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + entry_point = os.path.join(tmpdir, "train.py") + with open(entry_point, "w") as f: + f.write("print('training')") + + with patch("sagemaker.core.workflow.utilities._pipeline_config") as mock_config: + mock_config.pipeline_name = "test-pipeline" + mock_config.code_hash = "abc123" + + result = processor._package_code( + entry_point=entry_point, + source_dir=tmpdir, + requirements=None, + job_name="test-job", + kms_key=None, + ) + + # Should use code_hash in the S3 path, not job_name + assert "test-pipeline" in result + assert "abc123" in result + assert "test-job" not in result + class TestFrameworkProcessorRun: def test_run_with_s3_code(self, mock_session): From 5cab44bd96b28e31a981ed58b969d94e073e15a8 Mon Sep 17 00:00:00 2001 From: hippi345 <13539685+hippi345@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:22:30 +0000 Subject: [PATCH 2/2] fix: keep sourcedir.tar.gz under pipeline code_hash prefix Preserve the sourcedir.tar.gz filename under the pipeline code-hash prefix so _pack_and_upload_code helper replace() URIs stay distinct. Add a regression test covering _pack_and_upload_code with two hashes. --- .../src/sagemaker/core/processing.py | 1 + sagemaker-core/tests/unit/test_processing.py | 69 ++++++++++++++++++- 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index bf094f494b..912c0d1075 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -1345,6 +1345,7 @@ def _package_code( _pipeline_config.pipeline_name, "source", _pipeline_config.code_hash, + "sourcedir.tar.gz", ) else: s3_uri = s3.s3_path_join( diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index 6dfdbc56ff..8ce04bd148 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -1443,9 +1443,11 @@ def test_package_code_with_pipeline_config(self, mock_session): kms_key=None, ) - # Should use code_hash in the S3 path, not job_name + # Should use code_hash in the S3 path, not job_name, and keep + # sourcedir.tar.gz so helper-script replace() URIs stay distinct. assert "test-pipeline" in result assert "abc123" in result + assert "sourcedir.tar.gz" in result assert "test-job" not in result @@ -1668,6 +1670,71 @@ def test_pack_and_upload_code_explicit_kms_key_wins(self, mock_session): upload_kms = [call.kwargs.get("kms_key") for call in mock_upload.call_args_list] assert upload_kms and all(k == explicit_kms for k in upload_kms) + def test_pack_and_upload_code_with_pipeline_config(self, mock_session): + """Pipeline code_hash prefix must keep sourcedir.tar.gz so helper uploads do not collide. + + ``_pack_and_upload_code`` derives runproc.sh and install_requirements.py URIs + with ``s3_payload.replace("sourcedir.tar.gz", ...)``. If the payload URI is + only the hash (no sourcedir.tar.gz filename), those replacements are no-ops + and the helper scripts overwrite the source archive. Two different hashes + must also produce different S3 keys (issue #6493). + """ + processor = FrameworkProcessor( + role="arn:aws:iam::123456789012:role/SageMakerRole", + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + ) + + def pack_with_hash(code_hash): + with tempfile.TemporaryDirectory() as tmpdir: + entry_point = os.path.join(tmpdir, "train.py") + with open(entry_point, "w") as f: + f.write("print('training')") + + with patch("sagemaker.core.workflow.utilities._pipeline_config") as mock_config: + mock_config.pipeline_name = "test-pipeline" + mock_config.code_hash = code_hash + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body" + ) as mock_upload: + mock_upload.side_effect = ( + lambda body, desired_s3_uri=None, **kwargs: desired_s3_uri + ) + processor._pack_and_upload_code( + code=entry_point, + source_dir=tmpdir, + requirements=None, + job_name="test-job", + inputs=None, + kms_key=None, + ) + upload_uris = [ + call.kwargs.get("desired_s3_uri") for call in mock_upload.call_args_list + ] + + sourcedir_uris = [uri for uri in upload_uris if uri and "sourcedir.tar.gz" in uri] + helper_uris = [ + uri + for uri in upload_uris + if uri and ("install_requirements.py" in uri or uri.endswith("/runproc.sh")) + ] + assert sourcedir_uris, "sourcedir.tar.gz must remain in the pipeline hash S3 key" + payload_uri = sourcedir_uris[0] + assert code_hash in payload_uri + assert "test-pipeline" in payload_uri + assert "test-job" not in payload_uri + assert helper_uris, "expected helper-script uploads" + assert payload_uri not in helper_uris + return payload_uri + + uri_a = pack_with_hash("abc123") + uri_b = pack_with_hash("def456") + assert uri_a != uri_b + assert "abc123" in uri_a + assert "def456" in uri_b + class TestProcessingInputOutputHelpers: def test_processing_input_with_app_managed(self):