ModelBuilder silently drops `source_code` (no repack) for `image_uri` / `ModelTrainer` builds

[Alex-Mesmer]

PySDK Version

• PySDK V2 (2.x)
• PySDK V3 (3.x)

Describe the bug

ModelBuilder accepts a source_code argument, documented as "Source code configuration for
custom inference code." When building a model from an existing script-mode container
(image_uri) together with a trained model artifact (a model data S3 URI), I expect
ModelBuilder.build() to repack the provided source_code into the model artifact — i.e. bake
code/ into model.tar.gz — exactly the way the classic
sagemaker.model.Model(model_data=…, entry_point=…, source_dir=…) did, and the way pipelines did
via _RepackModelStep.

A self-contained artifact is required for register(): a registered model package is persisted and
served on its own, so it cannot depend on an external SAGEMAKER_SUBMIT_DIRECTORY pointing at a
transient S3 source tarball.

Instead, for any image_uri-based build with no model/inference_spec, ModelBuilder
classifies the build as passthrough and silently discards source_code:
_build_for_passthrough() resets self.source_dir = None and self.entry_point = None, so
is_repack() returns False and no code is ever packaged. The resulting artifact contains the model
data but none of the inference code, and register() produces a model package that fails to
serve (no model_fn/input_fn/predict_fn/output_fn).

This is existing behavior in v2 of this SDK so this is a regression.

To reproduce

Case 1 — image_uri + model URI + source_code (silent drop):

from sagemaker.serve import ModelBuilder
from sagemaker.core.training.configs import SourceCode

mb = ModelBuilder(
 image_uri="<script-mode framework or custom (BYO) container>",
 model_path="s3://my-bucket/.../model.tar.gz", # trained artifact — the "model URI"
 source_code=SourceCode( # custom inference code
 source_dir="./code",
 entry_script="inference.py", # defines model_fn/input_fn/predict_fn/output_fn
 ),
 role_arn=role,
 sagemaker_session=session,
)
model = mb.build()

# BUG: the model.tar.gz backing `model` does NOT contain code/inference.py.
# `source_code` was dropped and no repack occurred. Registering `model` yields a
# package with no serving code.

Case 2 — ModelTrainer (the natural train→serve flow) + source_code:

mb = ModelBuilder(
 model=trainer, # ModelTrainer carrying the trained artifact
 source_code=SourceCode(source_dir="./code", entry_script="inference.py"),
 image_uri="<script-mode container>",
 role_arn=role,
 sagemaker_session=session,
)
mb.build()
# raises: ValueError("InferenceSpec is required when using ModelTrainer, ...")
# => source_code cannot be combined with a ModelTrainer at all.

Expected behavior

When a model artifact (a model URI - via model_path, a ModelTrainer, or a TrainingJob) is
supplied alongside source_code, ModelBuilder.build() should repack the source code into the
model artifact and produce a self-contained model.tar.gz (code under code/), mirroring the
classic Model + _RepackModelStep behavior. source_code must not be silently ignored.

Screenshots or logs
If applicable, add screenshots or logs to help explain your problem.

System information
A description of your system. Please provide:

• SageMaker Python SDK version: sagemaker==3.15.1, sagemaker-core==2.16.0
• Framework name (eg. PyTorch) or algorithm (eg. KMeans):
• Framework version: any script-mode framework or custom (BYO) container — e.g. the managed SKLearn serving image (first-party) or an in-house image (non-first-party). Reproduces for both.
• Python version: 3.12
• CPU or GPU:
• Custom Docker image (Y/N):

Additional context

Root cause / code references

All references are to sagemaker/serve/model_builder.py in sagemaker==3.15.1.

1. source_code is honored initially. _initialize_script_mode_variables() maps it onto the
   script-mode attributes:

   • L1362-1376: self.entry_point = self.source_code.entry_script;
     self.source_dir = self.source_code.source_dir.

2. _build_validations() forces passthrough for image-only builds (both first- and
   non-first-party images):

   • L1564-1572: image_uri + is_1p_image_uri(...) + no model + no inference_spec
     → self._passthrough = True.
   • L1574-1582: image_uri + not is_1p_image_uri(...) + no model + no inference_spec
     → self._passthrough = True.

3. _build_for_passthrough() then discards the source code:

   • L1603-1605:

     # Ensure no script-mode artifacts are injected for passthrough
     self.source_dir = None
     self.entry_point = None

4. is_repack() therefore returns False, so _upload_code(..., repack=True) →
   repack_model(...) never runs:

   • L1924-1925: if self.source_dir is None or self.entry_point is None: return False.

The repack code path exists (_upload_code(..., repack=True) at L1932, calling repack_model(...)),
but it is unreachable for these inputs.

ModelTrainer-specific manifestation

• model=<ModelTrainer> requires inference_spec — _build_validations() raises
  "InferenceSpec is required when using ModelTrainer, ..." at L1527-1536. So source_code
  cannot be combined with a ModelTrainer.
• Even when inference_spec is supplied with a ModelTrainer, repack is explicitly disabled —
  is_repack() short-circuits at L1927-1928:

  if isinstance(self.model, ModelTrainer) and self.inference_spec:
   return False

  That routes serving through the multi-model-server / serve.pkl (cloudpickle) path, which does
  not repack a multi-file source_dir into the artifact.

Net: there is no combination of ModelBuilder inputs that produces
"model data URI + inference source code → repacked, self-contained model.tar.gz" for a script-mode
framework or custom image. This blocks migrating classic Model / PipelineModel registration
(which relied on the repack) to ModelBuilder.

Suggested fix

For image_uri-based builds, honor source_code instead of nulling it in
_build_for_passthrough(): when both a model artifact and source_code are present, take the
repack path (_upload_code(..., repack=True)) so the code is baked into the artifact. Equivalently,
allow source_code + a model artifact (script-mode) with a ModelTrainer without requiring an
InferenceSpec, so the classic train → repack → register flow remains expressible in v3.

[d-vesely]

@lucasjia-aws the RepackModelStep is still regressed in comparison to v2, where a register step would generate a repack step, if custom inference code was provided. See related issue #5828 .

[Alex-Mesmer]

@lucasjia-aws the RepackModelStep is still regressed in comparison to v2, where a register step would generate a repack step, if custom inference code was provided. See related issue #5828 .

The team said that the fix for this has been released in v3.18 of the sdk. I haven't had time to validate it yet since I have moved to other work in the mean time but that is a good place to start.