From 3eed72b25e05eb302bc880432a5fd2456877a40c Mon Sep 17 00:00:00 2001 From: jzhaoqwa Date: Tue, 14 Jul 2026 13:39:25 -0700 Subject: [PATCH 1/3] fix(iam): scope repo-level ECR actions to prevent false deny in preflight validation Split ecr_policy statements in training, serving, and hyperpod role types so that only ecr:GetAuthorizationToken (account-level) remains under Resource: "*". The repository-level actions (BatchGetImage, GetDownloadUrlForLayer, BatchCheckLayerAvailability) are now scoped to arn:aws:ecr:*:*:repository/*, which excludes them from _get_smoke_test_actions. This prevents SimulatePrincipalPolicy from returning implicitDeny for roles that correctly scope ECR permissions to specific repo ARNs (least privilege), fixing the regression that blocked deploys/training/pipelines for those customers. --- .../src/sagemaker/core/helper/iam_policies.py | 37 +++-- .../unit/helper/test_iam_role_resolver.py | 136 ++++++++++++++++++ 2 files changed, 164 insertions(+), 9 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/helper/iam_policies.py b/sagemaker-core/src/sagemaker/core/helper/iam_policies.py index f47b1dda83..f583b3b586 100644 --- a/sagemaker-core/src/sagemaker/core/helper/iam_policies.py +++ b/sagemaker-core/src/sagemaker/core/helper/iam_policies.py @@ -49,15 +49,26 @@ "Version": "2012-10-17", "Statement": [ { + # GetAuthorizationToken is an account-level action that + # does not target a specific repository — Resource: "*" is + # correct and safe to simulate without ResourceArns. + "Effect": "Allow", + "Action": ["ecr:GetAuthorizationToken"], + "Resource": "*", + }, + { + # These are repository-level actions. Scoping to + # repository/* prevents false implicitDeny results from + # SimulatePrincipalPolicy when the customer's real policy + # correctly scopes them to specific repo ARNs. "Effect": "Allow", "Action": [ "ecr:GetDownloadUrlForLayer", "ecr:BatchGetImage", "ecr:BatchCheckLayerAvailability", - "ecr:GetAuthorizationToken", ], - "Resource": "*", - } + "Resource": "arn:aws:ecr:*:*:repository/*", + }, ], }, "cloudwatch_logs_policy": { @@ -341,16 +352,20 @@ "ecr_policy": { "Version": "2012-10-17", "Statement": [ + { + "Effect": "Allow", + "Action": ["ecr:GetAuthorizationToken"], + "Resource": "*", + }, { "Effect": "Allow", "Action": [ "ecr:GetDownloadUrlForLayer", "ecr:BatchGetImage", "ecr:BatchCheckLayerAvailability", - "ecr:GetAuthorizationToken", ], - "Resource": "*", - } + "Resource": "arn:aws:ecr:*:*:repository/*", + }, ], }, "cloudwatch_logs_policy": { @@ -515,16 +530,20 @@ "ecr_policy": { "Version": "2012-10-17", "Statement": [ + { + "Effect": "Allow", + "Action": ["ecr:GetAuthorizationToken"], + "Resource": "*", + }, { "Effect": "Allow", "Action": [ "ecr:GetDownloadUrlForLayer", "ecr:BatchGetImage", "ecr:BatchCheckLayerAvailability", - "ecr:GetAuthorizationToken", ], - "Resource": "*", - } + "Resource": "arn:aws:ecr:*:*:repository/*", + }, ], }, "cloudwatch_logs_policy": { diff --git a/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py b/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py index 2ebcc35c31..48fcdf3be0 100644 --- a/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py +++ b/sagemaker-core/tests/unit/helper/test_iam_role_resolver.py @@ -14,6 +14,7 @@ HYPERPOD_CLI_CONNECT_ACTIONS, _load_policy_config, _get_required_actions, + _get_smoke_test_actions, _replace_placeholders, _get_boto_session, _simulate_denied_actions, @@ -262,6 +263,10 @@ def paginate(**kwargs): assert not any(a.startswith("sagemaker-mlflow:") for a in simulated) assert "sagemaker:DescribeHubContent" not in simulated assert "s3:GetObject" not in simulated # S3 is scoped to S3_PLACEHOLDER + # Repository-level ECR actions are scoped and excluded from the gate. + assert "ecr:BatchGetImage" not in simulated + assert "ecr:GetDownloadUrlForLayer" not in simulated + assert "ecr:BatchCheckLayerAvailability" not in simulated # ...but *-resource smoke-test actions are included. assert "cloudwatch:PutMetricData" in simulated assert "ecr:GetAuthorizationToken" in simulated @@ -634,6 +639,137 @@ def test_all_role_types_have_source_account_placeholder(self): ), f"{role_type} trust policy missing aws:SourceAccount placeholder" +class TestEcrPolicyScopedCorrectly: + """Regression tests for ECR policy scoping (V2287033493). + + Repository-level ECR actions (BatchGetImage, GetDownloadUrlForLayer, + BatchCheckLayerAvailability) must be scoped to repository ARNs, NOT + Resource: "*". Only GetAuthorizationToken is account-level and belongs + under "*". When all four are under "*", _get_smoke_test_actions includes + the repo-level ones and SimulatePrincipalPolicy (called without ResourceArns) + returns implicitDeny for roles with least-privilege ECR policies — a false + positive that blocks deploys/training/pipelines. + """ + + ECR_REPO_ACTIONS = { + "ecr:GetDownloadUrlForLayer", + "ecr:BatchGetImage", + "ecr:BatchCheckLayerAvailability", + } + + @pytest.mark.parametrize("role_type", ["training", "serving", "hyperpod"]) + def test_ecr_repo_actions_not_in_smoke_test(self, role_type): + """Repository-level ECR actions must NOT appear in the smoke test set.""" + smoke_actions = set(_get_smoke_test_actions(role_type)) + overlap = smoke_actions & self.ECR_REPO_ACTIONS + assert not overlap, ( + f"role_type={role_type}: repo-level ECR actions {overlap} should not be " + f"in smoke test (would cause false implicitDeny)" + ) + + @pytest.mark.parametrize("role_type", ["training", "serving", "hyperpod"]) + def test_ecr_get_authorization_token_in_smoke_test(self, role_type): + """GetAuthorizationToken is account-level and SHOULD be in the smoke test.""" + smoke_actions = set(_get_smoke_test_actions(role_type)) + assert "ecr:GetAuthorizationToken" in smoke_actions + + @pytest.mark.parametrize("role_type", ["training", "serving", "hyperpod"]) + def test_ecr_repo_actions_scoped_to_repository_arn(self, role_type): + """Repository-level ECR actions must have a repository/* resource scope.""" + config = _load_policy_config() + ecr_stmts = config[role_type]["policies"]["ecr_policy"]["Statement"] + # Find the statement containing BatchGetImage + repo_stmt = None + for stmt in ecr_stmts: + actions = stmt.get("Action", []) + if isinstance(actions, str): + actions = [actions] + if "ecr:BatchGetImage" in actions: + repo_stmt = stmt + break + assert repo_stmt is not None, "No statement with ecr:BatchGetImage found" + resource = repo_stmt["Resource"] + assert resource == "arn:aws:ecr:*:*:repository/*", ( + f"Expected repository/* scope, got: {resource}" + ) + + @pytest.mark.parametrize("role_type", ["training", "serving", "hyperpod"]) + def test_ecr_get_authorization_token_resource_is_wildcard(self, role_type): + """GetAuthorizationToken must remain under Resource: '*'.""" + config = _load_policy_config() + ecr_stmts = config[role_type]["policies"]["ecr_policy"]["Statement"] + auth_stmt = None + for stmt in ecr_stmts: + actions = stmt.get("Action", []) + if isinstance(actions, str): + actions = [actions] + if "ecr:GetAuthorizationToken" in actions: + auth_stmt = stmt + break + assert auth_stmt is not None, "No statement with ecr:GetAuthorizationToken found" + assert auth_stmt["Resource"] == "*" + + @pytest.mark.parametrize("role_type", ["training", "serving", "hyperpod"]) + def test_all_ecr_actions_still_in_required_actions(self, role_type): + """All four ECR actions must remain in the full required actions list.""" + all_actions = set(_get_required_actions(role_type)) + expected = self.ECR_REPO_ACTIONS | {"ecr:GetAuthorizationToken"} + assert expected.issubset(all_actions), ( + f"Missing ECR actions from required set: {expected - all_actions}" + ) + + def test_least_privilege_ecr_role_passes_validation(self): + """A role with ECR permissions scoped to specific repos must not be blocked. + + This is the actual customer scenario from V2287033493: the customer's role + grants BatchGetImage/GetDownloadUrlForLayer/BatchCheckLayerAvailability on + specific repository ARNs, not '*'. The smoke test should only simulate + GetAuthorizationToken (which the customer grants on '*'), so the validation + should pass. + """ + mock_session, mock_iam, _ = _make_session( + "arn:aws:sts::123456789012:assumed-role/MyTrainingRole/sess" + ) + mock_iam.get_role.return_value = { + "Role": { + "Arn": "arn:aws:iam::123456789012:role/MyTrainingRole", + "AssumeRolePolicyDocument": _trusted_doc(), + } + } + + # Simulate: all smoke-test actions are allowed (customer has them on *) + captured = {} + + def paginate(**kwargs): + captured["actions"] = kwargs.get("ActionNames", []) + return [ + { + "EvaluationResults": [ + {"EvalActionName": a, "EvalDecision": "allowed"} + for a in kwargs["ActionNames"] + ] + } + ] + + paginator = MagicMock() + paginator.paginate.side_effect = paginate + mock_iam.get_paginator.return_value = paginator + + # Should succeed without raising + result = resolve_and_validate_role( + provided_role=None, role_type="training", sagemaker_session=mock_session + ) + assert result == "arn:aws:iam::123456789012:role/MyTrainingRole" + + # Verify repo-level ECR actions were NOT simulated + simulated = set(captured["actions"]) + assert "ecr:BatchGetImage" not in simulated + assert "ecr:GetDownloadUrlForLayer" not in simulated + assert "ecr:BatchCheckLayerAvailability" not in simulated + # But GetAuthorizationToken WAS simulated + assert "ecr:GetAuthorizationToken" in simulated + + class TestReplacePlaceholders: """Tests for the pure-data _replace_placeholders() helper.""" From e3a75e6213727cd9af5138cc335a428540dd2d02 Mon Sep 17 00:00:00 2001 From: jzhaoqwa Date: Fri, 17 Jul 2026 15:11:32 -0700 Subject: [PATCH 2/3] fix: support s3 uri in FrameworkProcessor --- .../src/sagemaker/core/modules/configs.py | 5 +- .../src/sagemaker/core/processing.py | 91 ++++++- sagemaker-core/tests/unit/test_processing.py | 233 ++++++++++++++++++ 3 files changed, 320 insertions(+), 9 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/modules/configs.py b/sagemaker-core/src/sagemaker/core/modules/configs.py index 09e2671894..a23c5e14c4 100644 --- a/sagemaker-core/src/sagemaker/core/modules/configs.py +++ b/sagemaker-core/src/sagemaker/core/modules/configs.py @@ -91,7 +91,10 @@ class SourceCode(BaseConfig): Parameters: source_dir (Optional[str]): - The local directory containing the source code to be used in the training job container. + The local directory, S3 URI, or path to a tar.gz file stored locally or in S3 + that contains the source code to be used in the training job container. + When an S3 URI is provided, the source code is used directly from S3 + without local packaging. requirements (Optional[str]): The path within ``source_dir`` to a ``requirements.txt`` file. If specified, the listed requirements will be installed in the training job container. diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 7ceeb9c19f..f751e7081f 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -1129,6 +1129,29 @@ def _s3_code_prefix(self): self.sagemaker_session.default_bucket_prefix or "", ) + @staticmethod + def _is_s3_uri(path: Optional[str]) -> bool: + """Check whether the given path is an S3 URI.""" + return bool(path) and path.lower().startswith("s3://") + + def _resolve_s3_source_dir(self, source_dir: str) -> str: + """Resolve an S3 source_dir to a sourcedir.tar.gz URI. + + If the URI already points to a .tar.gz file, return it unchanged. + Otherwise treat it as an S3 prefix and append sourcedir.tar.gz. + """ + if source_dir.lower().endswith(".tar.gz"): + return source_dir + return source_dir.rstrip("/") + "/sourcedir.tar.gz" + + def _resolve_helper_scripts_prefix(self, job_name: str) -> str: + """Return an S3 prefix for uploading helper scripts (runproc.sh, install_requirements.py).""" + return s3.s3_path_join( + self._s3_code_prefix(), + job_name, + "source", + ) + def _package_code( self, entry_point, @@ -1137,10 +1160,20 @@ def _package_code( job_name, kms_key, ): - """Package and upload code to S3.""" + """Package and upload code to S3. + + If source_dir is an S3 URI, it is used directly as the code payload + (no local packaging or upload is performed). The S3 URI should point to + either a tar.gz archive or an S3 prefix containing the source code. + """ import tarfile import tempfile + # S3 source_dir: use it directly without local packaging. + # This restores v2 behavior where S3 paths with tar.gz archives were supported. + if self._is_s3_uri(source_dir): + return self._resolve_s3_source_dir(source_dir) + # If source_dir is not provided, use the directory containing entry_point if source_dir is None: if os.path.isabs(entry_point): @@ -1158,12 +1191,10 @@ def _package_code( # Create tar.gz with source_dir contents with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp: with tarfile.open(tmp.name, "w:gz") as tar: - # Add all files from source_dir to the root of the tar for item in os.listdir(source_dir): item_path = os.path.join(source_dir, item) tar.add(item_path, arcname=item) - # Upload to S3 s3_uri = s3.s3_path_join( self._s3_code_prefix(), job_name, @@ -1171,7 +1202,6 @@ def _package_code( "sourcedir.tar.gz", ) - # Upload the tar file directly to S3 s3.S3Uploader.upload_string_as_file_body( body=open(tmp.name, "rb").read(), desired_s3_uri=s3_uri, @@ -1283,12 +1313,20 @@ def _pack_and_upload_code( inputs = self._patch_inputs_with_payload(inputs, s3_payload) - entrypoint_s3_uri = s3_payload.replace("sourcedir.tar.gz", "runproc.sh") + # Determine where to upload helper scripts. + # When source_dir is S3, s3_payload points to the user's existing location, + # so helpers go to a separate managed prefix. + if self._is_s3_uri(source_dir): + helper_prefix = self._resolve_helper_scripts_prefix(job_name) + entrypoint_s3_uri = s3.s3_path_join(helper_prefix, "runproc.sh") + install_req_s3_uri = s3.s3_path_join(helper_prefix, "install_requirements.py") + else: + entrypoint_s3_uri = s3_payload.replace("sourcedir.tar.gz", "runproc.sh") + install_req_s3_uri = s3_payload.replace("sourcedir.tar.gz", "install_requirements.py") - # Upload the CodeArtifact-aware install_requirements script alongside the source code + # Upload install_requirements helper import sagemaker.core.utils.install_requirements as _ir_mod - install_req_s3_uri = s3_payload.replace("sourcedir.tar.gz", "install_requirements.py") evaluated_kms_key = kms_key if kms_key else self.output_kms_key s3.S3Uploader.upload_string_as_file_body( body=open(_ir_mod.__file__, "r").read(), @@ -1298,7 +1336,6 @@ def _pack_and_upload_code( ) script = os.path.basename(code) - evaluated_kms_key = kms_key if kms_key else self.output_kms_key s3_runproc_sh = self._create_and_upload_runproc( script, evaluated_kms_key, entrypoint_s3_uri, entry_point, source_dir ) @@ -1428,6 +1465,44 @@ def _generate_custom_framework_script( Returns: str: The generated script content """ + # When source_dir is an S3 URI, we cannot read the entry_point file locally. + # Instead, generate a script that executes the entry_point from the extracted + # source bundle on the container. + if self._is_s3_uri(source_dir): + return dedent( + """\ + #!/bin/bash + + # Exit on any error. SageMaker uses error code to mark failed job. + set -e + + cd /opt/ml/processing/input/code/ + + # Extract source code + if [ -f sourcedir.tar.gz ]; then + tar -xzf sourcedir.tar.gz + else + echo "ERROR: sourcedir.tar.gz not found!" + exit 1 + fi + + if [[ -f 'requirements.txt' ]]; then + pip uninstall --yes typing + python3 /opt/ml/processing/input/code/install_requirements.py requirements.txt + fi + + # Execute custom entrypoint + chmod +x {entry_point} + ./{entry_point} + + {entry_point_command} {user_script} "$@" + """ + ).format( + entry_point=entry_point, + entry_point_command=" ".join(self.command), + user_script=user_script, + ) + # Resolve the full path to the entry_point file if source_dir and not os.path.isabs(entry_point): full_entry_point_path = os.path.join(source_dir, entry_point) diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index f5349faea0..dd619808ed 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -1763,3 +1763,236 @@ def test_latest_job_updated_after_run(self, mock_session): assert processor.latest_job == mock_job2 assert len(processor.jobs) == 2 + + +class TestFrameworkProcessorS3SourceDir: + """Tests for FrameworkProcessor S3 source_dir support (v2 parity regression fix).""" + + def _make_processor(self, mock_session, **kwargs): + return 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, + **kwargs, + ) + + # --- _is_s3_uri helper --- + + def test_is_s3_uri_with_s3_path(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("s3://bucket/path") is True + + def test_is_s3_uri_with_uppercase(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("S3://bucket/path") is True + + def test_is_s3_uri_with_local_path(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("/local/path") is False + + def test_is_s3_uri_with_none(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri(None) is False + + def test_is_s3_uri_with_empty_string(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("") is False + + # --- _resolve_s3_source_dir helper --- + + def test_resolve_s3_source_dir_with_tar_gz(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._resolve_s3_source_dir("s3://bucket/code/sourcedir.tar.gz") + assert result == "s3://bucket/code/sourcedir.tar.gz" + + def test_resolve_s3_source_dir_with_prefix(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._resolve_s3_source_dir("s3://bucket/code/") + assert result == "s3://bucket/code/sourcedir.tar.gz" + + def test_resolve_s3_source_dir_with_prefix_no_trailing_slash(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._resolve_s3_source_dir("s3://bucket/code") + assert result == "s3://bucket/code/sourcedir.tar.gz" + + # --- _package_code with S3 source_dir --- + + def test_package_code_with_s3_tar_gz(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + kms_key=None, + ) + assert result == "s3://my-bucket/code/sourcedir.tar.gz" + + def test_package_code_with_s3_prefix(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code/", + requirements=None, + job_name="test-job", + kms_key=None, + ) + assert result == "s3://my-bucket/code/sourcedir.tar.gz" + + def test_package_code_with_s3_prefix_no_trailing_slash(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code", + requirements=None, + job_name="test-job", + kms_key=None, + ) + assert result == "s3://my-bucket/code/sourcedir.tar.gz" + + def test_package_code_with_s3_does_not_upload(self, mock_session): + """S3 source_dir should not trigger any S3 upload.""" + processor = self._make_processor(mock_session) + with patch("sagemaker.core.s3.S3Uploader.upload_string_as_file_body") as mock_upload: + processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + kms_key=None, + ) + mock_upload.assert_not_called() + + # --- _pack_and_upload_code with S3 source_dir --- + + def test_pack_and_upload_code_with_s3_source_dir(self, mock_session): + processor = self._make_processor(mock_session) + + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ) as mock_upload: + result_uri, result_inputs, result_job_name = processor._pack_and_upload_code( + code="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + inputs=None, + kms_key=None, + ) + + # Should have uploaded install_requirements.py and runproc.sh + assert mock_upload.call_count == 2 + upload_uris = [ + call.kwargs.get("desired_s3_uri") or call.args[1] + for call in mock_upload.call_args_list + ] + assert any("install_requirements.py" in uri for uri in upload_uris) + assert any("runproc.sh" in uri for uri in upload_uris) + + # Helpers should go to the managed prefix, not the user's S3 location + for uri in upload_uris: + assert not uri.startswith("s3://my-bucket/code/") + + def test_pack_and_upload_code_with_s3_source_dir_creates_code_input(self, mock_session): + processor = self._make_processor(mock_session) + + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ): + _, result_inputs, _ = processor._pack_and_upload_code( + code="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + inputs=None, + kms_key=None, + ) + + # Should have a 'code' input pointing to the S3 source dir + assert len(result_inputs) == 1 + code_input = result_inputs[0] + assert code_input.input_name == "code" + assert "s3://my-bucket/code/" in code_input.s3_input.s3_uri + + # --- _generate_custom_framework_script with S3 source_dir --- + + def test_generate_custom_framework_script_with_s3_source_dir(self, mock_session): + processor = self._make_processor(mock_session) + script = processor._generate_custom_framework_script( + user_script="train.py", + entry_point="setup.sh", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + ) + + # Should not try to read local files — generates a container script instead + assert "#!/bin/bash" in script + assert "tar -xzf sourcedir.tar.gz" in script + assert "chmod +x setup.sh" in script + assert "./setup.sh" in script + assert "python train.py" in script + + def test_generate_custom_framework_script_with_local_source_dir(self, mock_session): + processor = self._make_processor(mock_session) + + with tempfile.TemporaryDirectory() as tmpdir: + entry_point_path = os.path.join(tmpdir, "setup.sh") + with open(entry_point_path, "w") as f: + f.write("#!/bin/bash\necho setup") + + script = processor._generate_custom_framework_script( + user_script="train.py", + entry_point="setup.sh", + source_dir=tmpdir, + ) + + # Should have read and embedded the local entry_point file content + assert "echo setup" in script + + # --- _generate_framework_script (no custom entry_point) --- + + def test_generate_framework_script_contains_extraction(self, mock_session): + processor = self._make_processor(mock_session) + script = processor._generate_framework_script(user_script="train.py") + + assert "tar -xzf sourcedir.tar.gz" in script + assert "python train.py" in script + + # --- Full run() integration with S3 source_dir --- + + def test_run_with_s3_source_dir(self, mock_session): + """End-to-end: run() with S3 source_dir should not raise ValueError.""" + processor = self._make_processor(mock_session) + mock_job = Mock() + + with patch.object(processor, "_start_new", return_value=mock_job): + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ): + processor.run( + code="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + wait=False, + ) + assert processor.latest_job == mock_job + + def test_run_with_s3_source_dir_prefix(self, mock_session): + """run() with S3 prefix (no .tar.gz) should also work.""" + processor = self._make_processor(mock_session) + mock_job = Mock() + + with patch.object(processor, "_start_new", return_value=mock_job): + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ): + processor.run( + code="train.py", + source_dir="s3://my-bucket/code/", + wait=False, + ) + assert processor.latest_job == mock_job From f83dea447fe2cd01d6ccda678938e9bc33cabd4f Mon Sep 17 00:00:00 2001 From: jzhaoqwa Date: Fri, 17 Jul 2026 15:47:32 -0700 Subject: [PATCH 3/3] add local dependencies to source_dir --- .../src/sagemaker/core/processing.py | 38 ++++++++++++++++++- .../code/s3_source_dir_processing/helpers.py | 6 +++ .../code/s3_source_dir_processing/process.py | 27 +++++++++++++ 3 files changed, 69 insertions(+), 2 deletions(-) create mode 100644 sagemaker-mlops/tests/integ/code/s3_source_dir_processing/helpers.py create mode 100644 sagemaker-mlops/tests/integ/code/s3_source_dir_processing/process.py diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index f751e7081f..3f48383a47 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -1156,6 +1156,7 @@ def _package_code( self, entry_point, source_dir, + dependencies, requirements, job_name, kms_key, @@ -1165,6 +1166,15 @@ def _package_code( If source_dir is an S3 URI, it is used directly as the code payload (no local packaging or upload is performed). The S3 URI should point to either a tar.gz archive or an S3 prefix containing the source code. + + Args: + entry_point (str): Path to the entry point script. + source_dir (str): Local directory, S3 URI, or None. + dependencies (list[str]): Additional local directories to include + in the tar.gz bundle (default: None). Not supported with S3 source_dir. + requirements (str): Path to requirements.txt relative to source_dir. + job_name (str): Processing job name (used in S3 key). + kms_key (str): KMS key for S3 upload encryption. """ import tarfile import tempfile @@ -1172,6 +1182,11 @@ def _package_code( # S3 source_dir: use it directly without local packaging. # This restores v2 behavior where S3 paths with tar.gz archives were supported. if self._is_s3_uri(source_dir): + if dependencies: + raise ValueError( + "dependencies is not supported when source_dir is an S3 URI. " + "Bundle dependencies into the S3 tar.gz archive instead." + ) return self._resolve_s3_source_dir(source_dir) # If source_dir is not provided, use the directory containing entry_point @@ -1188,13 +1203,22 @@ def _package_code( if not os.path.exists(source_dir): raise ValueError(f"source_dir does not exist: {source_dir}") - # Create tar.gz with source_dir contents + # Create tar.gz with source_dir contents + dependencies with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp: with tarfile.open(tmp.name, "w:gz") as tar: + # Add all files from source_dir for item in os.listdir(source_dir): item_path = os.path.join(source_dir, item) tar.add(item_path, arcname=item) + # Add dependency directories at the root of the archive + for dep_path in (dependencies or []): + if not os.path.isabs(dep_path): + dep_path = os.path.abspath(dep_path) + if not os.path.exists(dep_path): + 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, @@ -1218,6 +1242,7 @@ def run( self, code: str, source_dir: Optional[str] = None, + dependencies: Optional[List[str]] = None, requirements: Optional[str] = None, inputs: Optional[List[ProcessingInput]] = None, outputs: Optional[List["ProcessingOutput"]] = None, @@ -1236,7 +1261,13 @@ def run( framework script to run. source_dir (str): Path (absolute, relative or an S3 URI) to a directory with any other processing source code dependencies aside from the entry - point file (default: None). + point file (default: None). If ``source_dir`` is an S3 URI, it must + point to a tar.gz file named ``sourcedir.tar.gz``. + dependencies (list[str]): A list of paths to directories (absolute or + relative) with any additional libraries that will be exported to the + container (default: None). The library folders are copied into the + same tar.gz bundle as source_dir. Not supported when source_dir is + an S3 URI. requirements (str): Path to a requirements.txt file relative to source_dir (default: None). inputs (list[:class:`~sagemaker.processing.ProcessingInput`]): Input files for @@ -1265,6 +1296,7 @@ def run( s3_runproc_sh, inputs, job_name = self._pack_and_upload_code( code, source_dir, + dependencies, requirements, job_name, inputs, @@ -1289,6 +1321,7 @@ def _pack_and_upload_code( self, code, source_dir, + dependencies, requirements, job_name, inputs, @@ -1306,6 +1339,7 @@ def _pack_and_upload_code( s3_payload = self._package_code( entry_point=code, source_dir=source_dir, + dependencies=dependencies, requirements=requirements, job_name=job_name, kms_key=kms_key, diff --git a/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/helpers.py b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/helpers.py new file mode 100644 index 0000000000..46711e1d95 --- /dev/null +++ b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/helpers.py @@ -0,0 +1,6 @@ +"""Helper module to verify cross-file imports work from S3 source_dir.""" + + +def get_greeting(name: str) -> str: + """Return a greeting string.""" + return f"Hello from S3 source_dir, {name}!" diff --git a/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/process.py b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/process.py new file mode 100644 index 0000000000..a215c079de --- /dev/null +++ b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/process.py @@ -0,0 +1,27 @@ +"""Simple processing script for S3 source_dir integ test. + +This script validates that: +1. It can be executed from an S3-based source_dir +2. It can import from a sibling module in the same source bundle +""" +import os +import json + +from helpers import get_greeting + + +if __name__ == "__main__": + output_dir = "/opt/ml/processing/output" + os.makedirs(output_dir, exist_ok=True) + + result = { + "status": "success", + "greeting": get_greeting("integration-test"), + "source_dir_type": "s3", + } + + output_path = os.path.join(output_dir, "result.json") + with open(output_path, "w") as f: + json.dump(result, f) + + print(f"Processing complete. Output written to {output_path}")