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.