fix: keep source_code for a pipeline-variable model URI and stop repack mutating caller source

[mohamedzeidan2021]

Issue #, if available: Fixes #6450, Fixes #6481

Description

Two coupled defects on the ModelBuilder pipeline repack path.

#6450 — source_code silently dropped when the model URI is a PipelineVariable

ModelBuilder._build_for_passthrough() recognised a model artifact only when the value was a str starting with s3://:

if self.model_path and self.model_path.startswith("s3://"):
    model_artifact_uri = self.model_path
elif isinstance(self.s3_model_data_url, str) and self.s3_model_data_url.startswith("s3://"):
    model_artifact_uri = self.s3_model_data_url

When the artifact is produced by an upstream pipeline step the value is a PipelineVariable, which fails that isinstance check. model_artifact_uri stayed None, the repack branch was skipped, and the fall-through cleared source_dir, entry_point, dependencies and git_config. The CreateModel request carried the pipeline expression with "Environment": {} and no repack was scheduled, so a script-mode container could never find the entry script. This is the recurrence of #6105 (fixed by #6112) on the path that fix did not cover.

The SDK already handles this case correctly one level down: _upload_code(repack=True) detects a PipelineVariable artifact and defers the repack to pipeline execution by registering need_runtime_repack and setting SAGEMAKER_SUBMIT_DIRECTORY=/opt/ml/model/code. The fix routes through that existing machinery rather than adding new logic — one elif is_pipeline_variable(self.s3_model_data_url) branch, five lines. The pipeline variable object itself is left intact (not stringified, not replaced), model_path still takes precedence as it does today, and a concrete s3:// URI keeps its current eager-repack behaviour.

#6481 — _RepackModelStep mutates the caller's source directory

_prepare_for_repacking() wrote _repack_model.py and _repack_script_launcher.sh into the caller-provided source_dir, once in the constructor and again on every arguments access. For an s3:// source_dir it downloaded the archive, added the files, and _save_model()'d the result back over the original S3 object.

The helper files now go into a copy the SDK owns:

• local source_dir → shutil.copytree into a temporary directory; the caller's directory is never written to
• s3:// source_dir → the staged archive is uploaded under a new sibling key (sourcedir.tar.gz → sourcedir-sagemaker-repack.tar.gz). Keeping the caller's bucket and prefix means the staged object inherits the same access and encryption policies; the original is read but never overwritten.
• source_dir=None → unchanged; _establish_source_dir() already builds an SDK-owned directory, so the files go straight into it.

Staging is now idempotent (guarded on self._staged_source_dir), so repeated serialization neither re-stages nor re-uploads. self._source_dir keeps the caller's value, so the source_dir hyperparameter and the repack script's view of it are byte-for-byte unchanged.

Not done deliberately: the issue suggests deleting the temporary copy after serialization. There is no point at which the SDK knows serialization is finished, and removing the directory before ModelTrainer uploads it would break the repack. This matches the existing _establish_source_dir() behaviour, which also leaves its mkdtemp() in place.

Why the two are sequenced together

#6481 is latent today only because #6450 blocks that repack path for ModelBuilder. Landing the #6450 fix on its own would ship code that writes into customer-owned source directories and overwrites their original S3 objects on the newly-reachable path. They are not independent, so they are in one PR.

Relationship to #6338 — please read before merging

PR #6338 ("fix: generate repack step for ModelBuilder.register/build in ModelStep") covers the other half of #6450, and this PR deliberately does not duplicate or conflict with it. The two touch disjoint files: #6338 changes only sagemaker-mlops/.../workflow/model_step.py; this PR changes sagemaker-serve/.../serve/model_builder.py and sagemaker-mlops/.../workflow/_utils.py. #6338's diff applies cleanly on top of this branch and vice versa.

Both halves are required, and I verified this against #6338's actual diff rather than by inspection:

configuration	source_code kept	need_runtime_repack	repack step	container ModelDataUrl
master	no	set()	0	raw pipeline variable
master + #6338 only	no	set()	0	raw pipeline variable
master + this PR only	yes	{id(builder)}	0 ("No models to repack")	raw pipeline variable
master + #6338 + this PR	yes	{id(builder)}	1	Steps.<...>.ModelArtifacts.S3ModelArtifacts

So: #6338 alone does not fix #6450 — _build_for_passthrough() drops source_code before _upload_code() ever runs, so need_runtime_repack stays empty and there is nothing for ModelStep to act on. This PR alone does not fully fix #6450 either, because ModelStep._append_repack_model_step() gates on isinstance(self._model, sagemaker.core.resources.Model), which the captured ModelBuilder can never satisfy. With only this PR merged the container would get SAGEMAKER_SUBMIT_DIRECTORY=/opt/ml/model/code without a repack step producing that directory. #6450 should therefore be closed only once both are in; please merge #6338 first or together with this PR. I have left #6338 untouched rather than re-implementing its change here.

The #6481 half is fully independent of #6338 and safe to merge on its own — and #6338 makes it more urgent, since it opens the same _RepackModelStep(source_dir=<caller's SourceCode.source_dir>) path for ModelBuilder.register(). I confirmed the non-mutation fix covers that path: with #6338 applied, the caller's directory still contains only inference.py after ModelStep construction and the helper files are in the staged copy.

On the reporter's suggested patches

#6450 — adopted in spirit, not literally. The reporter correctly identified both the isinstance(..., str) check and that _upload_code() already has the deferred-repack logic, but suggested no concrete diff. I went with the minimal routing change rather than new repack code. I did not take the issue's alternative ("raise a clear error if pipeline-variable repacking is not supported") because the capability already exists and only the dispatch was missing.

#6481 — adopted the staging approach, rejected two specifics.

• Rejected "remove the temporary copy after serialization" (see above — no safe hook, and it would break the upload).
• Rejected creating the staging copy on each _prepare_for_repacking() call, which is what a literal reading implies. arguments is accessed more than once during pipeline compilation, so a fresh mkdtemp() per call would make the compiled definition non-deterministic and re-upload the S3 archive every time. Staging is guarded to run at most once.

I also checked for the shadowed-class hazard that bit #6474. While prototyping I had translated _RepackModelStep's output_path/output_kms_key into ModelTrainer's output_data_config and hit exactly that class of trap: sagemaker.core.shapes.OutputDataConfig defaults unset fields to Unassigned() while sagemaker.core.training.configs.OutputDataConfig defaults them to None, and only the latter lets TrainDefaults.get_output_data_config() fill in its CompressionType default. That code is not in this PR — #6338 owns it — so I am only noting it for #6338's author. The practical effect is nil (CompressionType is simply omitted and the service defaults it to GZIP; I validated the resulting request against the live shape, see below).

Checked every caller of the code I changed: _inject_repack_script_and_launcher, _prepare_for_repacking, _establish_source_dir and _RepackModelStep have no callers outside _utils.py, model_step.py and their unit tests. _build_for_passthrough is called only from build() and its tests. No public API change; _RepackModelStep's signature is unchanged.

Testing

All of this is offline — there are no working AWS credentials in my environment, so there is no live AWS round-trip in any of the evidence below. The repack training job and the pipeline execution were never run; CreateModel/CreateTrainingJob were validated against the service contract but never submitted.

Reproduced both defects against the installed sagemaker==3.25.0 (role resolver stubbed, PipelineSession outbound edges stubbed):

• ModelBuilder silently drops source_code when model data is a PipelineVariable #6450, s3_model_data_url=ParameterString(...): source_dir=None, entry_point=None, "Environment": {}, need_runtime_repack=set(), 0 repack steps. The concrete-s3:// control reached repack_model() and produced the correct script-mode env vars — confirming the fix is scoped to the pipeline-variable path.
• _RepackModelStep modifies the caller's source directory #6481, local source_dir: ['inference.py'] → ['_repack_model.py', '_repack_script_launcher.sh', 'inference.py'] after the constructor, and again after re-cleaning and touching arguments.
• _RepackModelStep modifies the caller's source directory #6481, s3:// source_dir: _save_model target was s3://caller-bucket/caller/prefix/sourcedir.tar.gz — the caller's original object. After the fix: .../sourcedir-sagemaker-repack.tar.gz, original untouched.

Service-contract validation. With #6338 applied alongside, I compiled a full Pipeline definition and validated both generated requests against the real botocore input shapes (botocore.validate.ParamValidator against operation_model("CreateModel").input_shape and CreateTrainingJob, with the deliberately-trimmed ModelName/TrainingJobName restored): both VALID. The repack step's training channel carries {"Get": "Parameters.ModelDataUrl"} and the model container's ModelDataUrl resolves to {"Get": "Steps.CreateModel-RepackModel-....ModelArtifacts.S3ModelArtifacts"}.

New unit tests — 6 added, in the files that own the changed code.

sagemaker-serve/tests/unit/test_model_builder_methods.py, in the existing TestBuildForPassthroughSourceCodeRepack class (the one added for #6105):

• test_repacks_and_preserves_script_mode_vars_with_pipeline_variable
• test_pipeline_variable_repack_branch_makes_is_repack_true
• test_pipeline_variable_without_source_code_still_clears_vars (guard: pure image passthrough unchanged)

sagemaker-mlops/tests/unit/workflow/test_utils.py:

• test_caller_source_dir_not_modified — constructor and two consecutive arguments accesses
• test_staging_happens_once — repeated arguments does not re-stage
• test_s3_source_dir_staged_to_new_key — upload target differs from the original URI

I also updated test_inject_repack_script_local_source_dir, which asserted the buggy behaviour (that the helper files land in the caller's directory). It now asserts they land in the staged copy and that the caller's directory is clean.

The existing concrete-s3:// tests (test_repacks_and_preserves_script_mode_vars_with_s3_model_data_url, test_bridges_model_path_to_s3_model_data_url, test_repack_branch_makes_is_repack_true) are unmodified and still pass, which is the "no change to the working path" check.

Negative check (fix reverted, new tests kept). Restored the pristine 3.25.0 files and re-ran:

sagemaker-mlops/tests/unit/workflow/test_utils.py      6 failed, 16 passed
  FAILED test_inject_repack_script_local_source_dir
  FAILED test_caller_source_dir_not_modified
  FAILED test_s3_source_dir_staged_to_new_key
  FAILED test_staging_happens_once
  (+ 2 since-removed output_data_config tests)
sagemaker-serve/tests/unit/test_model_builder_methods.py -k Passthrough   2 failed, 13 passed
  FAILED test_repacks_and_preserves_script_mode_vars_with_pipeline_variable
  FAILED test_pipeline_variable_repack_branch_makes_is_repack_true

With the fix applied, all pass. (test_pipeline_variable_without_source_code_still_clears_vars passes both ways by design — it is a no-regression guard, not a bug witness.)

Regression runs, fix applied vs. reverted.

sagemaker-mlops/tests/unit: 835 passed (baseline on master: 832 passed; +3 = the new tests). 28 collection errors in both runs, byte-identical, all ModuleNotFoundError: No module named 'pyspark' under feature_store/feature_processor.

sagemaker-serve/tests/unit: 70 failed, 2073 passed, 25 skipped with the fix; 72 failed, 2071 passed, 25 skipped reverted. Diffing the two FAILED lists, the only delta is the 2 new #6450 tests flipping fail → pass; nothing newly broken. All 70 are pre-existing and environmental: 42 ClientError (ExpiredToken) GetObject, 23 ClientError (ExpiredToken) GetCallerIdentity, 5 ModuleNotFoundError: No module named 'tensorflow'.

Lint. black --line-length 100 --check, flake8 --max-line-length 120 and pydocstyle clean on all 5 changed files. pylint --rcfile=.pylintrc: no new findings on either source file — the only diff in the message lists is a line-number shift on a pre-existing E0203 in model_builder.py. Per-file scores improved (_utils.py 7.31 → 7.79, model_builder.py 9.84 → 9.85). Package-level scores are unchanged from master in both directions (sagemaker-mlops 9.84, sagemaker-serve 9.89); both sit below the --fail-under=9.9 gate on unmodified master in my sandbox because pylint runs there without the full tox dependency set and emits no-name-in-module noise, so I could not reproduce the real CI gate locally.

Integration test. Extended sagemaker-serve/tests/integ/test_passthrough_source_code_repack_integration.py with test_build_preserves_source_code_with_pipeline_variable_model_uri, which the existing file did not cover (it only exercises a concrete S3 URI). It builds under a real PipelineSession with a ParameterString model URI and asserts source_code survives, need_runtime_repack is registered, no eager repack upload happens, and the script-mode env vars are wired. It asserts only the ModelBuilder-side contract so it is valid on this branch alone; the end-to-end repack-step assertion belongs with #6338. I could not execute it — it needs live credentials and a SageMaker execution role, and I have neither. It needs a reviewer with a test account to run it.

Merge Checklist

Put an x in the boxes that apply.

General

• I have read the CONTRIBUTING doc
• I certify that the changes I'm introducing will be backward compatible, and I have discussed concerns about this, if any, with the Python SDK team
• I used the commit message format described in CONTRIBUTING
• I have passed the region in to all S3 and STS clients that I've initialized as part of this change
• I have updated any necessary documentation, including READMEs and API docs (if appropriate)

Tests

• I have added tests that prove my fix is effective or that my feature works (if appropriate)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes
• I have checked that my tests are not configured for a specific region or account (if appropriate)
• I have used unique_name_from_base to create resource names in integ tests (if appropriate) — the new integ test uses uuid.uuid4().hex[:8], matching the existing tests in that file, and creates no SageMaker resources (build() under a PipelineSession only); happy to switch it if you prefer

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

🤖 AI code review

This PR fixes two coupled defects on the ModelBuilder pipeline repack path: (#6450) source_code was dropped when the model artifact URI is a PipelineVariable because _build_for_passthrough only recognized concrete s3:// strings, and (#6481) _RepackModelStep mutated the caller-provided source_dir (writing repack helper files locally and overwriting the original S3 object). The fix routes the PipelineVariable case through the existing _upload_code(repack=True) machinery, and stages an SDK-owned copy of the source (temp dir for local, a new sibling S3 key for s3://) with idempotent, once-only staging. The changes are correct, consistent with the existing repack script's handling of the source_dir hyperparameter (used only as a truthiness flag; actual code arrives via the SourceCode channel), and well covered by new unit and integration tests. No blocking issues found.

_{Reviewed commit aaf749e. Automated review; verify before acting.}

[aviruthen]

Looks good, let's make sure the relevant tests pass for this though