Fix source_dir in FrameworkProcessor

[zhaoqizqwang]

fix(processing): Support S3 URIs for FrameworkProcessor source_dir

Restore v2 behavior where source_dir accepts S3 URIs pointing to
tar.gz archives or S3 prefixes. The v3 rewrite only handled local
directories, causing ValueError when customers passed S3 paths.

Add _is_s3_uri, _resolve_s3_source_dir, and _resolve_helper_scripts_prefix
helpers to cleanly route S3 vs local source_dir through _package_code and
_pack_and_upload_code. Update _generate_custom_framework_script to handle
the case where entry_point cannot be read locally.

[zhaoqizqwang]

Tested with real use case

def _upload_source_dir_to_s3(session, local_source_dir, s3_prefix):
    """Package local source_dir as sourcedir.tar.gz and upload to S3.

    This simulates what a customer would have pre-staged in S3.
    """
    bucket = session.default_bucket()
    s3_key = f"{s3_prefix}/sourcedir.tar.gz"

    with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
        with tarfile.open(tmp.name, "w:gz") as tar:
            for item in os.listdir(local_source_dir):
                item_path = os.path.join(local_source_dir, item)
                tar.add(item_path, arcname=item)

        s3_client = session.boto_session.client("s3")
        s3_client.upload_file(tmp.name, bucket, s3_key)
        os.unlink(tmp.name)

    return f"s3://{bucket}/{s3_key}"


def test_framework_processor_with_s3_source_dir(sagemaker_session, role):
    """FrameworkProcessor.run() with S3 source_dir should complete successfully.

    This test:
    1. Packages a local source dir as sourcedir.tar.gz
    2. Uploads it to S3 (simulating a pre-staged archive)
    3. Runs FrameworkProcessor with the S3 URI as source_dir
    4. Verifies the processing job completes successfully
    """
    region = sagemaker_session.boto_region_name
    bucket = sagemaker_session.default_bucket()
    s3_prefix = "integ-test-s3-source-dir"
    processing_job_name = "s3-srcdir-{}".format(strftime("%d-%H-%M-%S", gmtime()))
    output_destination = f"s3://{bucket}/{s3_prefix}/output"

    # Path to our test source code
    local_source_dir = os.path.join(
        os.path.dirname(__file__), "code", "s3_source_dir_processing"
    )

    try:
        # Upload source_dir to S3 as a tar.gz archive
        s3_source_dir_uri = _upload_source_dir_to_s3(
            sagemaker_session, local_source_dir, f"{s3_prefix}/code"
        )

        # Use a Python image (no framework-specific dependencies needed)
        image_uri = get_training_image_uri(
            region=region,
            framework="pytorch",
            framework_version="1.13",
            py_version="py39",
            instance_type="ml.m5.xlarge",
        )

        processor = FrameworkProcessor(
            image_uri=image_uri,
            role=role,
            instance_type="ml.m5.xlarge",
            instance_count=1,
            sagemaker_session=sagemaker_session,
        )

        # This is the critical call — source_dir is an S3 URI.
        # Before the fix, this would raise:
        #   ValueError: source_dir does not exist: /...mangled-path.../s3:/bucket/...
        processor.run(
            code="process.py",
            source_dir=s3_source_dir_uri,
            job_name=processing_job_name,
            outputs=[
                ProcessingOutput(
                    output_name="output",
                    s3_output=ProcessingS3Output(
                        s3_uri=output_destination,
                        local_path="/opt/ml/processing/output",
                        s3_upload_mode="EndOfJob",
                    ),
                ),
            ],
            wait=False,
        )

        # Wait for job to complete (10 min timeout)
        job = processor.latest_job
        timeout = 600
        start_time = time.time()

        while time.time() - start_time < timeout:
            job.refresh()
            status = job.processing_job_status

            if status == "Completed":
                break
            elif status in ["Failed", "Stopped"]:
                pytest.fail(
                    f"Processing job {status}. "
                    f"Job name: {processing_job_name}"
                )
            time.sleep(30)
        else:
            pytest.fail(f"Processing job timed out after {timeout} seconds")

        assert status == "Completed"

    finally:
        # Cleanup S3 resources
        s3 = boto3.resource("s3")
        bucket_obj = s3.Bucket(bucket)
        bucket_obj.objects.filter(Prefix=f"{s3_prefix}/").delete()

[github-actions]

Claude finished @mohamedzeidan2021's task in 5m 1s —— View job

---

PR Review: Fix source_dir in FrameworkProcessor

• Read PR diff and changed files
• Trace _package_code / _pack_and_upload_code / script-generation flow
• Check backward compatibility and public interface changes
• Assess test coverage
• Post review

I reviewed the S3 source_dir support restoration. The overall approach is sound and the refactor into _is_s3_uri / _resolve_s3_source_dir / _resolve_helper_scripts_prefix reads cleanly. I found one correctness bug and a couple of smaller items.

---

🔴 High — install_requirements.py is not mounted in the container for S3 source_dir

sagemaker-core/src/sagemaker/core/processing.py:1356-1362

When source_dir is an S3 URI, helper scripts are deliberately uploaded to a separate managed prefix (_resolve_helper_scripts_prefix, i.e. <code_prefix>/<job_name>/source/):

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")

But the code input mounted at /opt/ml/processing/input/code/ points at the user's S3 location (_patch_inputs_with_payload uses s3_payload, which for S3 is the user's sourcedir.tar.gz dir — processing.py:1391). runproc.sh is fine because it's passed as code=s3_runproc_sh and gets mounted separately as the entrypoint input. However, the generated S3 script references:

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

install_requirements.py was uploaded to the managed helper prefix, which is not mounted at /opt/ml/processing/input/code/. So any S3 source_dir whose bundle contains a requirements.txt will fail at runtime with a missing-file error inside the container. (The requirements.txt itself comes from the extracted tar, so the if [[ -f 'requirements.txt' ]] branch does fire.)

The test test_pack_and_upload_code_with_s3_source_dir actually asserts the helpers do not live under the user prefix, so it locks in this broken layout rather than catching it. Consider uploading install_requirements.py into the user's S3 code dir (co-located so it mounts alongside sourcedir.tar.gz), or adding the helper prefix as an additional mounted input, and add a test that exercises the requirements path end-to-end for S3 source_dir.

Fix this →

---

🟡 Medium — modules/configs.py docstring change may overstate behavior

sagemaker-core/src/sagemaker/core/modules/configs.py:93-97

This PR only changes processing code, but it edits the SourceCode.source_dir docstring (used by the training/ModelTrainer path) to add "When an S3 URI is provided, the source code is used directly from S3 without local packaging." There's no corresponding code change in this PR to the training/SourceCode consumers, so this claim isn't demonstrably backed here. If training already supported S3 (it's phrased differently in training/configs.py:98-100), the "without local packaging" detail may not be accurate for that path. I'd either drop this file from the PR or confirm the training path behaves as documented.

---

🟢 Minor / nits

• Unclosed file handles (processing.py:1233, 1369): open(tmp.name, "rb").read() and open(_ir_mod.__file__, "r").read() leak the handle until GC. This matches the pre-existing style, so not blocking — worth a with if you touch these lines.
• _resolve_s3_source_dir extension check: .endswith(".tar.gz") won't match .tgz archives. Probably fine given the documented contract ("must point to a tar.gz file named sourcedir.tar.gz"), just noting the mismatch between the strict prefix-append behavior and the broader "S3 prefix" wording in the docstring.

---

✅ What looks good

• _is_s3_uri correctly handles None/empty/uppercase; good unit coverage.
• New dependencies param is additive and keyword-only at call sites — no break to existing callers. The explicit ValueError when combining dependencies with an S3 source_dir is a nice guardrail.
• Helper-script routing keeps user S3 buckets read-only (doesn't write into the customer's prefix), which is the right instinct — just needs the mount fix above to actually work with requirements.

Overall: solid restoration of v2 behavior; the S3 + requirements mounting bug is the one thing I'd want fixed before merge.

[github-actions]

Claude finished @mohamedzeidan2021's task in 8m 49s —— View job

---

PR Review: Fix source_dir in FrameworkProcessor

• Read PR diff and changed files
• Trace _package_code / _pack_and_upload_code / script-generation flow
• Verify the previously-flagged install_requirements.py mount bug
• Check backward compatibility and public interface changes
• Assess test coverage

Note: the inline-comment tool isn't available to me in this run, so findings are consolidated below with file:line references.

I re-reviewed the S3 source_dir restoration. The overall approach is clean and the previously-reported high-severity bug is now fixed. One backward-compatibility concern and a few minor items remain.

---

✅ Previously-flagged install_requirements.py bug is now resolved

The earlier review flagged that, for an S3 source_dir, install_requirements.py was uploaded to a managed helper prefix but the generated script referenced it under /opt/ml/processing/input/code/, where it wasn't mounted. The current diff addresses this:

• _pack_and_upload_code (processing.py:1365-1387) now mounts the helper as a dedicated aux input at _AUX_CONTAINER_DIR (/opt/ml/processing/input/aux) and threads install_requirements_dir through to script generation.
• Both _generate_framework_script and _generate_custom_framework_script now interpolate {install_requirements_dir}/install_requirements.py (processing.py:1533, 1596), so the path matches the mount.
• test_s3_source_dir_script_references_mounted_install_requirements locks this in.

The S3Prefix mount of the exact object key .../install_requirements.py won't accidentally pull in the sibling runproc.sh (different key string), and runproc.sh is mounted separately as the entrypoint — so the layout is correct.

---

🟡 Medium — new dependencies param shifts positional argument order in run()

processing.py:1252-1257

def run(self, code, source_dir=None, dependencies=None, requirements=None, inputs=None, ...)

dependencies is inserted between source_dir and requirements. This restores the v2 parameter ordering, but relative to the current v3 run() (where requirements was the 3rd positional param) it is a breaking change for positional callers: existing code like processor.run("proc.py", "src/", "requirements.txt") will now bind "requirements.txt" to dependencies instead of requirements, silently changing behavior.

Since the stated goal is v2 parity this may be intentional, but it's worth either (a) making the new params keyword-only with a *, marker to prevent positional misbinding, or (b) appending dependencies after requirements, and calling out the ordering change in the PR description/changelog. Also consider a unit test asserting the ValueError guardrail fires when dependencies is combined with an S3 source_dir (processing.py:1201-1205) — that path is currently untested.

---

🟢 Minor / nits

• modules/configs.py:93-97 — This PR only changes processing, but edits the training-path SourceCode.source_dir docstring to add "used directly from S3 without local packaging." training/configs.py:98-100 already documented S3 support with different wording; the "without local packaging" claim isn't backed by any code change here. Consider dropping this hunk or confirming the training path actually behaves that way.
• _resolve_s3_source_dir (processing.py:1148-1159) — .endswith(".tar.gz") won't match .tgz. Fine given the documented sourcedir.tar.gz contract, just noting the mismatch with the looser "S3 prefix" wording in the _package_code docstring.
• Duplicated bash — the S3 custom-entrypoint branch (processing.py:1590-1607) duplicates the extract + requirements-install block from _generate_framework_script. Not blocking; a shared helper would reduce drift.
• Custom-entrypoint semantics differ for S3 — for a local source_dir the entry_point content is embedded and replaces the extraction logic; for S3 the generated script extracts first, then runs ./{entry_point}, then python {user_script}. Reasonable given the file can't be read locally, but the behavioral difference is worth a comment so future readers don't assume parity.

---

What looks good

• _is_s3_uri correctly handles None/empty/uppercase, with solid unit coverage.
• Keeping the user's S3 bucket read-only and routing helpers to a separate managed prefix is the right instinct — and the aux mount now makes it actually work with requirements.
• Good end-to-end and unit coverage for both the tar.gz and prefix forms; the integ test exercises cross-file imports from the S3 bundle.

Overall: the correctness bug from the prior round is fixed. My only real ask is deciding intentionally on the run() positional-arg ordering (keyword-only vs. append) so v3 positional callers don't break silently.