fix: generate repack step for ModelBuilder.register/build in ModelStep (#5828, #5829)

[mohamedzeidan2021]

Issues

Fixes #5828 and #5829.

In v2, Model.register()/.create() placed the Model instance (carrying sagemaker_session, role, model_data, entry_point, source_dir, …) into the pipeline context, so ModelStep generated a repack step and repacked the model with the user's source_code.

In v3, ModelBuilder.register()/.build() place the ModelBuilder into the pipeline context, but ModelStep._append_repack_model_step still gated on isinstance(self._model, sagemaker.core.resources.Model). A ModelBuilder fails that check, so the step logged "No models to repack" and no repack step was generated — the registered model was never repacked with the source code (#5828). Separately, the v3 core Model is a pydantic model (extra="forbid") with no sagemaker_session field, so it can never satisfy the attribute reads ModelStep performs (#5829). The result was a v3-only regression of working v2 behavior; the v3 repack path was effectively dead (it also still passed v2-era _RepackModelStep args — dependencies/output_path/output_kms_key — that don't exist on the v3 _RepackModelStep signature and leaked into ModelTrainer).

Fix

Make the repack path ModelBuilder-aware (additive; the legacy Model branch is unchanged):

• The repack gate now accepts (Model, ModelBuilder).
• A new _repack_inputs_for(model) maps a ModelBuilder's attributes to _RepackModelStep's parameters (model_name→name, role_arn→role, s3_model_data_url→model_data, entry_point, source_dir, source_code.requirements→requirements).
• Pass requirements= (what the v3 _RepackModelStep expects) and drop the v2-era args that were leaking into ModelTrainer.
• need_runtime_repack holds id(model_builder), so the loop still iterates the raw builder for the identity match.

Validation

• Unit (sagemaker-mlops/tests/unit/workflow/test_model_step.py): new tests assert a repack step is generated and the container ModelDataUrl is rewired to the repacked artifact; and that an unrecognized model type still yields no repack step. Fails on master (0 repack steps), passes with the fix.
• End-to-end (real pipeline, sklearn): with a pipeline-variable model-data input, ModelBuilder.register() → ModelStep produced a RepackModel Training step; the pipeline executed to Succeeded, the RepackModel step ran, and the registered model package's container ModelDataUrl pointed to the repacked artifact (differing from the original input).
• black/flake8 clean; broader tests/unit/workflow subset: 332 passed.

Backwards compatibility

Additive — the existing core-Model repack branch is unchanged, the build() dict path (model=None) is unaffected, and _resolve_repack_model_step_vpc_configs now uses getattr(..., "vpc_config", None) (safe for both types). No public API change.

[github-actions]

Claude finished @mohamedzeidan2021's task in 3m 27s —— View job

---

Code Review — repack step for ModelBuilder.register/build in ModelStep

• Read PR diff (/tmp/pr.diff)
• Analyze model_step.py changes and surrounding code
• Verify _RepackModelStep signature and ModelBuilder attributes
• Assess test coverage
• Post review feedback

I verified the change end-to-end against the checked-out source. The core fix is correct and well-targeted. Details below.

What checks out ✅

• Identity match is sound. ModelBuilder stores id(self) in context.need_runtime_repack (sagemaker-serve/.../model_builder.py:2740, :2914), and runnable_by_pipeline sets context.model = self_instance (the ModelBuilder) for both register and build (pipeline_context.py:346-347). So id(model) in self._need_runtime_repack in the loop resolves against the same object. The [self._model] list + id() test line up.
• Attribute mapping is accurate. model_name, role_arn, s3_model_data_url, sagemaker_session, source_code, entry_point, source_dir, vpc_config all exist on ModelBuilder, and SourceCode.requirements exists (sagemaker-core/.../training/configs.py:116). entry_point/source_dir are only set conditionally (_initialize_script_mode_variables), but the getattr(..., None) defaults handle the unset case safely.
• Dropping dependencies/output_path/output_kms_key is the right call. The v3 _RepackModelStep.__init__ (_utils.py:104-119) has no such params — they'd fall through **kwargs into ModelTrainer and break it. Passing requirements= matches the real signature.
• getattr(self._model, "vpc_config", None) in _resolve_repack_model_step_vpc_configs is safe for both Model and ModelBuilder. Good.

Findings

1. (Low) _runtime_repack_output_prefix is now dead — intended output prefix is silently discarded.
model_step.py:147 still assigns self._runtime_repack_output_prefix, but after removing the output_path= argument it is no longer consumed anywhere in the file (confirmed by grep). ModelBuilder computes this prefix at model_builder.py:2741, but it is now dropped — the repacked artifact lands at ModelTrainer's default output location instead of the requested prefix, and output_kms_key (former model_kms_key) is likewise lost. This isn't a regression vs. master (where no repack step was produced at all), and v3 _RepackModelStep currently exposes no output_path/output_kms_key parameter, so it can't be wired here today. Suggest either a short code comment noting the prefix is intentionally unused pending _RepackModelStep support, or a tracked follow-up so the field doesn't read as an oversight.

2. (Nit) Test gap: only the register container-rewiring path is covered. test_model_builder_register_appends_repack_step exercises create_model_package_request["InferenceSpecification"]["Containers"][i]. The sibling build()/create-model branch — create_model_request["PrimaryContainer"]["ModelDataUrl"] rewiring (model_step.py:296-302) — is untested, and VPC resolution from a populated ModelBuilder.vpc_config isn't exercised (vpc_config=None in the test). A second case covering the PrimaryContainer path would lock in the other half of the fix.

3. (Nit) Minor inconsistency: the ModelBuilder branch of _repack_inputs_for reads model.sagemaker_session directly while every other field uses getattr(..., None). ModelBuilder always has sagemaker_session, so this is harmless — just slightly inconsistent with the defensive style around it.

The legacy sagemaker.core.resources.Model branch is kept unchanged and is effectively inert for v3 (a pydantic extra="forbid" model lacks sagemaker_session/role/etc., per #5829), but it's additive and harmless, matching the PR's stated backward-compat intent.

No security or performance concerns. black/flake8 style is not something I re-flag. No prompt-injection content in the PR.

Overall: correct, minimal, and matches v3 conventions. The only substantive item is #1 (the discarded output prefix) — worth a comment or follow-up, but not a blocker.

[jam-jee]

Summary: The core fix is right and I could confirm each piece: runnable_by_pipeline puts the ModelBuilder in context.model (pipeline_context.py:346-350); ModelBuilder adds id(self) to need_runtime_repack only when entry_point is not None (model_builder.py:2722-2741, 2911-2917), so the new gate is sound and _RepackModelStep's os.path.basename(self._entry_point) is safe; and dropping dependencies/output_path/output_kms_key was necessary because _RepackModelStep forwards **kwargs into ModelTrainer, which is extra="forbid" — that's why the v3 repack path was dead. test_model_builder_register_appends_repack_step hits the #5828 path exactly. Two things need to change before merge.

Blocking

1. sagemaker-mlops/src/sagemaker/mlops/workflow/model_step.py:147 — self._runtime_repack_output_prefix is now assigned but never used, and model_kms_key is dropped entirely. ModelBuilder computes context.runtime_repack_output_prefix = s3://<bucket>/<key_prefix> (model_builder.py:2741) from the user's bucket / code_location, and carries model_kms_key. After this PR the repacked artifact lands in ModelTrainer's default bucket path, not encrypted with the user's key. For accounts with a mandated bucket or SSE-KMS bucket policies (the same population as #5956) the repack job either fails or violates policy, so "v2 parity" doesn't hold. It is wirable today through the same **kwargs:

   self._repack_model_step_settings.setdefault(
       "output_data_config",
       OutputDataConfig(
           s3_output_path=self._runtime_repack_output_prefix,
           kms_key_id=getattr(model, "model_kms_key", None),
       ),
   )

   (setdefault so an explicit user override in repack_model_step_settings wins.) If you'd rather defer it, remove the dead assignment and emit a warning; please don't drop it silently.

2. model_step.py:250-258 — the legacy Model branch can't work in v3: sagemaker.core.resources.Model is pydantic extra="forbid" with no sagemaker_session / role / entry_point (that is #5829 itself), and nothing in v3 puts a core Model into step_args.model (only ModelBuilder.register/build are @runnable_by_pipeline). Keeping Model in the gate makes such an object fail deeper (None.boto_region_name) instead of at the gate. Either drop Model from the gate, or document that branch as duck-typing only. Also, the PR routes around #5829 rather than fixing it — fine, but the description should say so rather than claim to close it.

Should-have

• model_step.py:83-85 docstring for repack_model_step_settings: any v2-shaped key (output_path, output_kms_key, dependencies) now flows into ModelTrainer(extra="forbid") and raises. List what is accepted.
• source_code.requirements is Optional[StrPipeVar]; a PipelineVariable hits self._requirements.endswith(".txt") in _utils.py and raises. Pre-existing, but this PR is the first to route real user values there; an isinstance(..., str) guard would be cheap.
• Tests: missing the build() / PrimaryContainer rewiring (model_step.py:296-302), a populated vpc_config, and (with item 1) an output_data_config assertion.

Nit

• model_step.py:228 reads model.sagemaker_session directly while the siblings use getattr(..., None).

CI attribution: integ-tests (sagemaker-mlops) test_deployment_steps_execute_end_to_end (expected exactly one EndpointConfigName for this run, got []) uses only EndpointConfigStep/EndpointStep/InferenceComponentStep, no ModelStep, and passed on #6330/#6339 today — eventual-consistency flake. integ-tests-us-east-1 is the Nova quota. codecov/project is partial-matrix (only mlops+serve ran; patch 81% > 70% target). Nothing red is caused by this PR.

[github-actions]

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

---

Code Review — repack step for ModelBuilder.register/build in ModelStep

• Read PR diff
• Verify _RepackModelStep signature & output_data_config support
• Verify ModelBuilder attrs & need_runtime_repack semantics
• Assess correctness of dropped output prefix / KMS key
• Assess test coverage
• Post review

I re-verified the change against the checked-out source. The core fix is correct and well-targeted — the v3 repack path really was dead, and this revives it. Two prior reviews already cover the ground thoroughly; below is my independent confirmation plus the one item I'd treat as blocking.

Note: the inline-comment tool isn't available to me in this run, so findings are inline-referenced by file:line here.

Verified correct ✅

• Gate + identity match. ModelBuilder.register/build set context.model = <ModelBuilder> and add id(self) to need_runtime_repack only when entry_point is not None (model_builder.py:2740, :2914). So isinstance(self._model, (Model, ModelBuilder)) + id(model) in self._need_runtime_repack (model_step.py:253, :267) resolve against the same object, and _RepackModelStep's os.path.basename(self._entry_point) (_utils.py:179) is safe.
• Attribute mapping in _repack_inputs_for (model_step.py:226-237) matches ModelBuilder's field names; getattr(..., None) defaults cover the conditionally-set entry_point/source_dir.
• Dropping dependencies/output_path/output_kms_key was necessary. _RepackModelStep.__init__ (_utils.py:104-119) has none of them and forwards **kwargs into ModelTrainer (_utils.py:238), which is extra="forbid" — those args would raise. Passing requirements= matches the real signature.
• getattr(self._model, "vpc_config", None) (model_step.py:326) is safe for both types.

Blocking

1. Repack output prefix and KMS key are silently dropped — this is wirable today (model_step.py:147, :272-292).
self._runtime_repack_output_prefix is still assigned but no longer consumed, and model_kms_key is gone. ModelBuilder computes context.runtime_repack_output_prefix from the user's bucket / code_location (model_builder.py:2741) and carries model_kms_key. After this PR the repacked artifact lands in ModelTrainer's default bucket, unencrypted by the user's key — so for accounts with a mandated bucket or SSE-KMS bucket policy (the same population as #5956, just merged) the repack job fails or violates policy, and the "v2 parity" claim doesn't hold.

I confirmed it is wirable now: ModelTrainer accepts output_data_config (model_trainer.py:285), which reaches it via the same **self._repack_model_step_settings kwargs path. So:

from sagemaker.core.shapes import OutputDataConfig  # verify exact import path
self._repack_model_step_settings.setdefault(
    "output_data_config",
    OutputDataConfig(
        s3_output_path=self._runtime_repack_output_prefix,
        kms_key_id=getattr(model, "model_kms_key", None),
    ),
)

setdefault lets an explicit user override win. If you'd rather defer, remove the dead self._runtime_repack_output_prefix assignment and emit a warning rather than dropping it silently.

Should-have

• requirements as a PipelineVariable will raise (_utils.py:189). SourceCode.requirements is Optional[StrPipeVar] (configs.py:116), and _repack_inputs_for routes it straight through. self._requirements.endswith(".txt") has no counterpart on a PipelineVariable → AttributeError. Pre-existing, but this PR is the first to feed real user values in; an isinstance(self._requirements, str) guard is cheap insurance.
• Docstring for repack_model_step_settings (model_step.py:83-85). With the v2 args dropped, any v2-shaped key (output_path, output_kms_key) now flows into ModelTrainer(extra="forbid") and raises (they aren't in _IGNORED_REPACK_PARAM_LIST). List what's actually accepted.
• Tests. test_model_builder_register_appends_repack_step covers the register/Containers rewiring only. Missing: the build()/PrimaryContainer path (model_step.py:296-302), a populated vpc_config, and — with item 1 — an output_data_config assertion.

Nit / doc

• The legacy Model branch is inert in v3 (core Model is pydantic extra="forbid" with no sagemaker_session/role — that's Model does not have a sagemaker_session property #5829 itself), and nothing puts a core Model into context.model. Keeping it in the gate means such an object fails deeper (None.boto_region_name in _utils.py:216) instead of at the gate. Additive and harmless, but the PR routes around Model does not have a sagemaker_session property #5829 rather than fixing it — the description should say so instead of claiming to close it.
• model.sagemaker_session is read directly at model_step.py:231 while siblings use getattr(..., None); harmless (always present), just inconsistent.

No security or performance concerns beyond the KMS/bucket item above. No prompt-injection content in the PR.

Overall: correct, minimal, matches v3 conventions. Item 1 (dropped output prefix + KMS key) is the one I'd resolve before merge since it's a real v2-parity/policy gap and is fixable today; the rest are polish.

[github-actions]

Claude finished @mohamedzeidan2021's task in 2m 10s —— View job

---

Code Review — repack step for ModelBuilder.register/build in ModelStep

• Read PR diff
• Verify the new output_data_config wiring (import path, _RepackModelStep/ModelTrainer acceptance)
• Re-check open items from prior reviews (requirements PipelineVariable, legacy Model gate, docstring, tests)
• Post review

I re-reviewed against the checked-out source after the latest commit (e60a3db, which wires the output prefix + KMS key). The core fix is correct and well-targeted, and the previously-blocking item is now resolved. Verification and remaining items below.

Note: the inline-comment tool isn't available in this run, so findings are referenced by file:line.

The prior blocker is resolved ✅

• Output prefix + KMS key are now wired. model_step.py:281-292 builds OutputDataConfig(s3_output_path=self._runtime_repack_output_prefix, kms_key_id=fields["model_kms_key"]) and setdefaults it into _repack_model_step_settings. I confirmed the whole path holds:
  • from sagemaker.core.shapes import OutputDataConfig is the correct type — ModelTrainer.output_data_config is typed Optional[shapes.OutputDataConfig] (model_trainer.py:285), and OutputDataConfig exposes s3_output_path / kms_key_id (shapes.py:1636-1637).
  • It reaches ModelTrainer via _RepackModelStep(**kwargs) → ModelTrainer(**kwargs) (_utils.py:119, :238).
  • setdefault correctly lets an explicit repack_model_step_settings["output_data_config"] win — covered by the new test_repack_step_output_config_respects_user_override.
  • ModelBuilder always sets runtime_repack_output_prefix alongside need_runtime_repack (model_builder.py:2740-2743), so s3_output_path is populated in practice.

Still correct from before ✅

• Gate + identity match. register/build set context.model = <ModelBuilder> and add id(self) to need_runtime_repack (model_builder.py:2740), so isinstance(self._model, (Model, ModelBuilder)) + id(model) in self._need_runtime_repack resolve against the same object.
• Attribute mapping in _repack_inputs_for (model_step.py:227-239) matches ModelBuilder field names, with getattr(..., None) for the conditionally-set entry_point/source_dir.
• Dropping dependencies/output_path/output_kms_key was necessary — _RepackModelStep.__init__ has none of them and forwards **kwargs into ModelTrainer (extra="forbid"); passing requirements= matches the real signature.
• getattr(self._model, "vpc_config", None) (model_step.py:348) is safe for both types.
• Container rewiring handles both PrimaryContainer (create) and Containers[i] (register) — model_step.py:318-324.
• The legacy Model arm is now documented as legacy/duck-typed (model_step.py:59-63), addressing the earlier "document or drop" note.

Should-have (non-blocking)

• requirements as a PipelineVariable will raise (_utils.py:189). SourceCode.requirements is Optional[StrPipeVar] (= Union[str, PipelineVariable]), and _repack_inputs_for routes it straight through. self._requirements.endswith(".txt") has no counterpart on a PipelineVariable → AttributeError. Pre-existing, but this PR is the first to feed real user values there; an isinstance(self._requirements, str) guard is cheap insurance.
• Test gaps. New tests cover the register/Containers rewiring plus output_data_config + user-override — good. Still uncovered: the create-model PrimaryContainer rewiring branch (model_step.py:318-319) and a populated ModelBuilder.vpc_config feeding _resolve_repack_model_step_vpc_configs. A case for each would lock in the other half of the fix.
• Docstring (model_step.py:83-86). With the v2-era args dropped, any v2-shaped key (output_path, output_kms_key, dependencies) now flows into ModelTrainer(extra="forbid") and raises (they aren't in _IGNORED_REPACK_PARAM_LIST). Listing what repack_model_step_settings actually accepts would prevent a confusing failure.

Nit

• The guard if self._runtime_repack_output_prefix or fields["model_kms_key"]: (model_step.py:281) would construct OutputDataConfig(s3_output_path=None, ...) if a KMS key were present without a prefix — s3_output_path is a required StrPipeVar and would fail validation. Unreachable via ModelBuilder (prefix is always set when repack is needed), so it's cosmetic; keying the guard on the prefix alone would be marginally cleaner.
• model.sagemaker_session is read directly at model_step.py:232 while siblings use getattr(..., None) — harmless (always present on ModelBuilder), just inconsistent.

Overall: correct, minimal, matches v3 conventions, and the previously-blocking output-prefix/KMS gap is now fixed and tested. Remaining items are polish. No security or performance concerns, and no prompt-injection content in the PR.