From 3eed72b25e05eb302bc880432a5fd2456877a40c Mon Sep 17 00:00:00 2001
From: jzhaoqwa <jzhaoqwa@amazon.com>
Date: Tue, 14 Jul 2026 13:39:25 -0700
Subject: [PATCH 1/5] 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 <jzhaoqwa@amazon.com>
Date: Fri, 17 Jul 2026 15:11:32 -0700
Subject: [PATCH 2/5] 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 <jzhaoqwa@amazon.com>
Date: Fri, 17 Jul 2026 15:47:32 -0700
Subject: [PATCH 3/5] 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}")

From e28bfc401afbf2812385a6a4870ad79299056bc5 Mon Sep 17 00:00:00 2001
From: Mohamed Zeidan <zeidmo@amazon.com>
Date: Thu, 23 Jul 2026 13:09:54 -0700
Subject: [PATCH 4/5] fix: make dependencies optional in FrameworkProcessor
 code-packaging helpers

_package_code and _pack_and_upload_code declared dependencies as a
required positional arg, but the unit tests (and any direct callers)
invoke them without it, causing:

    TypeError: _package_code() missing 1 required positional argument: 'dependencies'

Move dependencies to the end of both signatures with a None default so
it is optional, matching the public run() signature. Also apply black
formatting and remove unused imports/vars in test_processing.py.
---
 .../src/sagemaker/core/processing.py          | 33 ++++++++++++-------
 sagemaker-core/tests/unit/test_processing.py  | 30 ++++++++---------
 2 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py
index 3f48383a47..026f4e84d4 100644
--- a/sagemaker-core/src/sagemaker/core/processing.py
+++ b/sagemaker-core/src/sagemaker/core/processing.py
@@ -489,7 +489,10 @@ def _normalize_outputs(self, outputs=None):
                 # If the output's s3_uri is not an s3_uri, create one.
                 parse_result = urlparse(output.s3_output.s3_uri)
                 if parse_result.scheme != "s3":
-                    if getattr(self.sagemaker_session, "local_mode", False) and parse_result.scheme == "file":
+                    if (
+                        getattr(self.sagemaker_session, "local_mode", False)
+                        and parse_result.scheme == "file"
+                    ):
                         normalized_outputs.append(output)
                         continue
                     if _pipeline_config:
@@ -1156,10 +1159,10 @@ def _package_code(
         self,
         entry_point,
         source_dir,
-        dependencies,
         requirements,
         job_name,
         kms_key,
+        dependencies=None,
     ):
         """Package and upload code to S3.
 
@@ -1212,7 +1215,7 @@ def _package_code(
                     tar.add(item_path, arcname=item)
 
                 # Add dependency directories at the root of the archive
-                for dep_path in (dependencies or []):
+                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):
@@ -1252,7 +1255,7 @@ def run(
         job_name: Optional[str] = None,
         experiment_config: Optional[Dict[str, str]] = None,
         kms_key: Optional[str] = None,
-        entry_point: Optional[str] = None
+        entry_point: Optional[str] = None,
     ):
         """Runs a processing job.
 
@@ -1296,12 +1299,12 @@ def run(
         s3_runproc_sh, inputs, job_name = self._pack_and_upload_code(
             code,
             source_dir,
-            dependencies,
             requirements,
             job_name,
             inputs,
             kms_key,
-            entry_point
+            entry_point,
+            dependencies=dependencies,
         )
 
         # Submit a processing job.
@@ -1321,12 +1324,12 @@ def _pack_and_upload_code(
         self,
         code,
         source_dir,
-        dependencies,
         requirements,
         job_name,
         inputs,
         kms_key=None,
-        entry_point=None
+        entry_point=None,
+        dependencies=None,
     ):
         """Pack local code bundle and upload to Amazon S3."""
         if code.startswith("s3://"):
@@ -1409,7 +1412,9 @@ def _set_entrypoint(self, command, user_script_name):
         )
         self.entrypoint = self.framework_entrypoint_command + [user_script_location]
 
-    def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri, entry_point=None, source_dir=None):
+    def _create_and_upload_runproc(
+        self, user_script, kms_key, entrypoint_s3_uri, entry_point=None, source_dir=None
+    ):
         """Create runproc shell script and upload to S3 bucket."""
         from sagemaker.core.workflow.utilities import _pipeline_config, hash_object
 
@@ -1439,7 +1444,9 @@ def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri, en
 
         return s3_runproc_sh
 
-    def _generate_framework_script(self, user_script: str, entry_point: str = None, source_dir: str = None) -> str:
+    def _generate_framework_script(
+        self, user_script: str, entry_point: str = None, source_dir: str = None
+    ) -> str:
         """Generate the framework entrypoint file (as text) for a processing job."""
         if entry_point:
             return self._generate_custom_framework_script(user_script, entry_point, source_dir)
@@ -1548,11 +1555,13 @@ def _generate_custom_framework_script(
             entry_point_content = f.read()
 
         # Generate the script with embedded entry_point content
-        return dedent("""\
+        return dedent(
+            """\
             {entry_point_content}
 
             {entry_point_command} {entry_point} "$@"
-            """).format(
+            """
+        ).format(
             entry_point_content=entry_point_content,
             entry_point_command=" ".join(self.command),
             entry_point=user_script,
diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py
index dd619808ed..5ebfdd23fd 100644
--- a/sagemaker-core/tests/unit/test_processing.py
+++ b/sagemaker-core/tests/unit/test_processing.py
@@ -14,7 +14,7 @@
 import pytest
 import os
 import tempfile
-from unittest.mock import Mock, patch, MagicMock, mock_open
+from unittest.mock import Mock, patch
 from sagemaker.core.processing import (
     Processor,
     ScriptProcessor,
@@ -48,9 +48,6 @@ def mock_session():
 
 class TestProcessorNormalizeArgs:
     def test_normalize_args_with_pipeline_variable_code(self, mock_session):
-        from sagemaker.core.workflow.pipeline_context import PipelineSession
-        from sagemaker.core.workflow import is_pipeline_variable
-
         processor = Processor(
             role="arn:aws:iam::123456789012:role/SageMakerRole",
             image_uri="test-image:latest",
@@ -238,8 +235,6 @@ def test_normalize_outputs_invalid_type(self, mock_session):
             processor._normalize_outputs(["invalid"])
 
 
-
-
 class TestBugConditionFileUriReplacedInLocalMode:
     """Bug condition exploration test: file:// URIs should be preserved in local mode.
 
@@ -444,7 +439,9 @@ def test_pipeline_variable_uri_skips_normalization(self, session_local_mode_fals
             [{"key": "value"}],
         ],
     )
-    def test_non_processing_output_raises_type_error(self, invalid_output, session_local_mode_false):
+    def test_non_processing_output_raises_type_error(
+        self, invalid_output, session_local_mode_false
+    ):
         """Non-ProcessingOutput objects must raise TypeError.
 
         **Validates: Requirements 3.4**
@@ -882,9 +879,7 @@ def test_generate_framework_script_with_custom_entry_point(self, mock_session):
             entry_point_path = f.name
 
         try:
-            script = processor._generate_framework_script(
-                "train.py", entry_point=entry_point_path
-            )
+            script = processor._generate_framework_script("train.py", entry_point=entry_point_path)
             assert custom_script_content in script
             assert "python3 train.py" in script
             assert "tar -xzf sourcedir.tar.gz" not in script
@@ -1042,7 +1037,7 @@ def test_start_new_submit_success(self, mock_session):
             },
         ):
             with patch("sagemaker.core.processing.serialize", return_value={}):
-                with patch("sagemaker.core.processing.ProcessingJob") as mock_job:
+                with patch("sagemaker.core.processing.ProcessingJob"):
                     with patch(
                         "sagemaker.core.utils.code_injection.codec.transform",
                         return_value={"processing_job_name": "test-job"},
@@ -1220,7 +1215,6 @@ def test_package_code_source_dir_not_exists(self, mock_session):
                 kms_key=None,
             )
 
-
     def test_package_code_with_code_location(self, mock_session):
         processor = FrameworkProcessor(
             role="arn:aws:iam::123456789012:role/SageMakerRole",
@@ -1477,8 +1471,6 @@ def test_generate_job_name_with_invalid_chars(self, mock_session):
 
 class TestProcessorWithPipelineVariable:
     def test_get_process_args_with_pipeline_variable_role(self, mock_session):
-        from sagemaker.core.workflow import is_pipeline_variable
-
         role_var = Mock()
 
         with patch("sagemaker.core.processing.is_pipeline_variable", return_value=True):
@@ -1524,11 +1516,17 @@ def test_start_new_removes_tags_from_processing_job(self, mock_session):
                 "tags": [{"Key": "Key", "Value": "Value"}],
             },
         ):
-            with patch("sagemaker.core.processing.serialize", return_value={"tags": [{"Key": "Key", "Value": "Value"}]}):
+            with patch(
+                "sagemaker.core.processing.serialize",
+                return_value={"tags": [{"Key": "Key", "Value": "Value"}]},
+            ):
                 with patch("sagemaker.core.processing.ProcessingJob") as mock_job_class:
                     with patch(
                         "sagemaker.core.utils.code_injection.codec.transform",
-                        return_value={"processing_job_name": "test-job", "tags": [{"Key": "Key", "Value": "Value"}]},
+                        return_value={
+                            "processing_job_name": "test-job",
+                            "tags": [{"Key": "Key", "Value": "Value"}],
+                        },
                     ):
                         processor._start_new([], [], None)
                         # Verify ProcessingJob was called without tags

From 969325eacd3b4ca423e3366549b23b71be3e3a74 Mon Sep 17 00:00:00 2001
From: Mohamed Zeidan <zeidmo@amazon.com>
Date: Thu, 23 Jul 2026 14:16:26 -0700
Subject: [PATCH 5/5] fix: mount install_requirements.py for S3 source_dir
 processing jobs

When source_dir is an S3 URI, /opt/ml/processing/input/code maps to the
user's (possibly read-only) S3 location, which does not contain the
managed install_requirements.py helper. The helper was uploaded to a
separate managed prefix that was never mounted into the container, yet
the generated runproc script hardcoded

    python3 /opt/ml/processing/input/code/install_requirements.py

so any S3 source_dir bundle containing a requirements.txt failed at
runtime with a missing-file error.

Mount the managed helper via a dedicated 'aux' ProcessingInput at
/opt/ml/processing/input/aux (never writing into the user's bucket) and
parameterize the generated scripts to reference install_requirements.py
from the correct directory: /input/code for a local source_dir
(unchanged) and /input/aux for an S3 source_dir.

Update test_pack_and_upload_code_with_s3_source_dir_creates_code_input to
assert the aux mount is created (it previously locked in the broken
layout) and add regression tests verifying the generated scripts point at
the mounted helper path.

Also use a context manager for the helper file read.
---
 .../src/sagemaker/core/processing.py          | 91 +++++++++++++++----
 sagemaker-core/tests/unit/test_processing.py  | 49 +++++++++-
 2 files changed, 120 insertions(+), 20 deletions(-)

diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py
index 026f4e84d4..34385619c8 100644
--- a/sagemaker-core/src/sagemaker/core/processing.py
+++ b/sagemaker-core/src/sagemaker/core/processing.py
@@ -1028,6 +1028,14 @@ class FrameworkProcessor(ScriptProcessor):
 
     framework_entrypoint_command = ["/bin/bash"]
 
+    # Container path where the source bundle is extracted. For a local ``source_dir`` the
+    # managed code prefix (which also holds ``install_requirements.py``) is mounted here.
+    _SOURCE_CODE_CONTAINER_DIR = "/opt/ml/processing/input/code"
+    # Container path where helper scripts are mounted when ``source_dir`` is an S3 URI. In
+    # that case ``/input/code`` maps to the user's (possibly read-only) S3 location, so the
+    # managed ``install_requirements.py`` must be mounted from a separate input instead.
+    _AUX_CONTAINER_DIR = "/opt/ml/processing/input/aux"
+
     def __init__(
         self,
         image_uri: Union[str, PipelineVariable],
@@ -1350,23 +1358,42 @@ def _pack_and_upload_code(
 
         inputs = self._patch_inputs_with_payload(inputs, s3_payload)
 
-        # 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.
+        # Determine where to upload helper scripts and where install_requirements.py will be
+        # available inside the container. When source_dir is S3, s3_payload points to the
+        # user's existing (possibly read-only) location, so helpers go to a separate managed
+        # prefix which is mounted at a dedicated aux path rather than /input/code.
         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")
+            install_requirements_dir = self._AUX_CONTAINER_DIR
+            # Mount the managed helper so install_requirements.py is present in the container.
+            inputs = (inputs or []) + [
+                ProcessingInput(
+                    input_name="aux",
+                    s3_input=ProcessingS3Input(
+                        s3_uri=install_req_s3_uri,
+                        local_path=self._AUX_CONTAINER_DIR,
+                        s3_data_type="S3Prefix",
+                        s3_input_mode="File",
+                    ),
+                )
+            ]
         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")
+            # For a local source_dir the managed code prefix (with install_requirements.py) is
+            # mounted at /input/code alongside the extracted bundle.
+            install_requirements_dir = self._SOURCE_CODE_CONTAINER_DIR
 
         # Upload install_requirements helper
         import sagemaker.core.utils.install_requirements as _ir_mod
 
         evaluated_kms_key = kms_key if kms_key else self.output_kms_key
+        with open(_ir_mod.__file__, "r") as _ir_file:
+            _ir_body = _ir_file.read()
         s3.S3Uploader.upload_string_as_file_body(
-            body=open(_ir_mod.__file__, "r").read(),
+            body=_ir_body,
             desired_s3_uri=install_req_s3_uri,
             kms_key=evaluated_kms_key,
             sagemaker_session=self.sagemaker_session,
@@ -1374,7 +1401,12 @@ def _pack_and_upload_code(
 
         script = os.path.basename(code)
         s3_runproc_sh = self._create_and_upload_runproc(
-            script, evaluated_kms_key, entrypoint_s3_uri, entry_point, source_dir
+            script,
+            evaluated_kms_key,
+            entrypoint_s3_uri,
+            entry_point,
+            source_dir,
+            install_requirements_dir,
         )
 
         return s3_runproc_sh, inputs, job_name
@@ -1413,13 +1445,21 @@ def _set_entrypoint(self, command, user_script_name):
         self.entrypoint = self.framework_entrypoint_command + [user_script_location]
 
     def _create_and_upload_runproc(
-        self, user_script, kms_key, entrypoint_s3_uri, entry_point=None, source_dir=None
+        self,
+        user_script,
+        kms_key,
+        entrypoint_s3_uri,
+        entry_point=None,
+        source_dir=None,
+        install_requirements_dir=None,
     ):
         """Create runproc shell script and upload to S3 bucket."""
         from sagemaker.core.workflow.utilities import _pipeline_config, hash_object
 
         if _pipeline_config and _pipeline_config.pipeline_name:
-            runproc_file_str = self._generate_framework_script(user_script, entry_point, source_dir)
+            runproc_file_str = self._generate_framework_script(
+                user_script, entry_point, source_dir, install_requirements_dir
+            )
             runproc_file_hash = hash_object(runproc_file_str)
             s3_uri = s3.s3_path_join(
                 self._s3_code_prefix(),
@@ -1436,7 +1476,9 @@ def _create_and_upload_runproc(
             )
         else:
             s3_runproc_sh = s3.S3Uploader.upload_string_as_file_body(
-                self._generate_framework_script(user_script, entry_point, source_dir),
+                self._generate_framework_script(
+                    user_script, entry_point, source_dir, install_requirements_dir
+                ),
                 desired_s3_uri=entrypoint_s3_uri,
                 kms_key=kms_key,
                 sagemaker_session=self.sagemaker_session,
@@ -1445,25 +1487,33 @@ def _create_and_upload_runproc(
         return s3_runproc_sh
 
     def _generate_framework_script(
-        self, user_script: str, entry_point: str = None, source_dir: str = None
+        self,
+        user_script: str,
+        entry_point: str = None,
+        source_dir: str = None,
+        install_requirements_dir: str = None,
     ) -> str:
         """Generate the framework entrypoint file (as text) for a processing job."""
         if entry_point:
-            return self._generate_custom_framework_script(user_script, entry_point, source_dir)
+            return self._generate_custom_framework_script(
+                user_script, entry_point, source_dir, install_requirements_dir
+            )
+
+        install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_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/
-            
+
             # Debug: List files before extraction
             echo "Files in /opt/ml/processing/input/code/ before extraction:"
             ls -la
-            
+
             # Extract source code
             if [ -f sourcedir.tar.gz ]; then
                 tar -xzf sourcedir.tar.gz
@@ -1478,18 +1528,23 @@ def _generate_framework_script(
                 # Some py3 containers has typing, which may breaks pip install
                 pip uninstall --yes typing
 
-                python3 /opt/ml/processing/input/code/install_requirements.py requirements.txt
+                python3 {install_requirements_dir}/install_requirements.py requirements.txt
             fi
 
             {entry_point_command} {entry_point} "$@"
         """
         ).format(
+            install_requirements_dir=install_requirements_dir,
             entry_point_command=" ".join(self.command),
             entry_point=user_script,
         )
 
     def _generate_custom_framework_script(
-        self, user_script: str, entry_point: str, source_dir: str = None
+        self,
+        user_script: str,
+        entry_point: str,
+        source_dir: str = None,
+        install_requirements_dir: str = None,
     ) -> str:
         """
         Generate a custom framework script with a user-provided entrypoint embedded.
@@ -1502,6 +1557,8 @@ def _generate_custom_framework_script(
             entry_point (str): Path to the custom entrypoint script file
             source_dir (str): Path to the source directory. If provided and entry_point
                 is relative, it will be combined with source_dir.
+            install_requirements_dir (str): Container directory that holds
+                ``install_requirements.py`` (default: the extracted source code dir).
 
         Returns:
             str: The generated script content
@@ -1510,6 +1567,7 @@ def _generate_custom_framework_script(
         # Instead, generate a script that executes the entry_point from the extracted
         # source bundle on the container.
         if self._is_s3_uri(source_dir):
+            install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_DIR
             return dedent(
                 """\
                 #!/bin/bash
@@ -1529,7 +1587,7 @@ def _generate_custom_framework_script(
 
                 if [[ -f 'requirements.txt' ]]; then
                     pip uninstall --yes typing
-                    python3 /opt/ml/processing/input/code/install_requirements.py requirements.txt
+                    python3 {install_requirements_dir}/install_requirements.py requirements.txt
                 fi
 
                 # Execute custom entrypoint
@@ -1539,6 +1597,7 @@ def _generate_custom_framework_script(
                 {entry_point_command} {user_script} "$@"
             """
             ).format(
+                install_requirements_dir=install_requirements_dir,
                 entry_point=entry_point,
                 entry_point_command=" ".join(self.command),
                 user_script=user_script,
diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py
index 5ebfdd23fd..ca811424d1 100644
--- a/sagemaker-core/tests/unit/test_processing.py
+++ b/sagemaker-core/tests/unit/test_processing.py
@@ -1910,12 +1910,20 @@ def test_pack_and_upload_code_with_s3_source_dir_creates_code_input(self, mock_s
                 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"
+            # Should have a 'code' input pointing to the S3 source dir plus an 'aux' input
+            # that mounts the managed install_requirements.py helper into the container.
+            assert len(result_inputs) == 2
+            inputs_by_name = {i.input_name: i for i in result_inputs}
+
+            code_input = inputs_by_name["code"]
             assert "s3://my-bucket/code/" in code_input.s3_input.s3_uri
 
+            aux_input = inputs_by_name["aux"]
+            assert aux_input.s3_input.local_path == "/opt/ml/processing/input/aux"
+            assert aux_input.s3_input.s3_uri.endswith("install_requirements.py")
+            # The helper must NOT come from the user's (possibly read-only) S3 location.
+            assert not aux_input.s3_input.s3_uri.startswith("s3://my-bucket/code/")
+
     # --- _generate_custom_framework_script with S3 source_dir ---
 
     def test_generate_custom_framework_script_with_s3_source_dir(self, mock_session):
@@ -1933,6 +1941,39 @@ def test_generate_custom_framework_script_with_s3_source_dir(self, mock_session)
         assert "./setup.sh" in script
         assert "python train.py" in script
 
+    def test_s3_source_dir_script_references_mounted_install_requirements(self, mock_session):
+        """The install_requirements.py path in the generated script must match the aux mount.
+
+        For an S3 source_dir, /opt/ml/processing/input/code maps to the user's location and
+        does NOT contain install_requirements.py; the helper is mounted at the aux path. The
+        generated script must invoke it from there, not from the code dir.
+        """
+        processor = self._make_processor(mock_session)
+        aux_dir = FrameworkProcessor._AUX_CONTAINER_DIR
+
+        # Both the plain framework script and the custom-entrypoint script must point at
+        # the aux mount when told the helper lives there.
+        plain = processor._generate_framework_script(
+            user_script="train.py",
+            source_dir="s3://my-bucket/code/sourcedir.tar.gz",
+            install_requirements_dir=aux_dir,
+        )
+        custom = processor._generate_custom_framework_script(
+            user_script="train.py",
+            entry_point="setup.sh",
+            source_dir="s3://my-bucket/code/sourcedir.tar.gz",
+            install_requirements_dir=aux_dir,
+        )
+        for script in (plain, custom):
+            assert f"{aux_dir}/install_requirements.py" in script
+            assert "/opt/ml/processing/input/code/install_requirements.py" not in script
+
+    def test_local_source_dir_script_references_code_dir_install_requirements(self, mock_session):
+        """For a local source_dir the helper is co-located under /input/code (unchanged)."""
+        processor = self._make_processor(mock_session)
+        script = processor._generate_framework_script(user_script="train.py")
+        assert "/opt/ml/processing/input/code/install_requirements.py" in script
+
     def test_generate_custom_framework_script_with_local_source_dir(self, mock_session):
         processor = self._make_processor(mock_session)