aws / aws/sagemaker-python-sdk
ModelBuilder silently drops `source_code` (no repack) for `image_uri` / `ModelTrainer` builds
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 1.3k
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 35
Description
**PySDK Version**
- [ ] PySDK V2 (2.x)
- [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):**
```python
from sagemaker.serve import ModelBuilder
from sagemaker.core.training.configs import SourceCode
mb = ModelBuilder(
image_uri="",
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`:**
```python
mb = ModelBuilder(
model=trainer, # ModelTrainer carrying the trained artifact
source_code=SourceCode(source_dir="./code", entry_script="inference.py"),
image_uri="",
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`:
```python
# 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=` **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`:
```python
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.
Contributor guide
Research direction
Start in sagemaker/serve/model_builder.py by tracing _build_validations(), _build_for_passthrough(), is_repack(), and _upload_code(..., repack=True); reproduce the image_uri and ModelTrainer cases described in the issue. Done means source_code is not discarded, model artifacts are repacked with code/ and inference.py, and the resulting artifact supports registration and serving.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100