aws / aws/sagemaker-python-sdk

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

Open
#6,105 2 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.