Skip to content

fix(serve): propagate additional_model_data_sources for JumpStart models - #6151

Open
e-davidson wants to merge 8 commits into
aws:masterfrom
e-davidson:fix/propagate-additional-model-data-sources
Open

fix(serve): propagate additional_model_data_sources for JumpStart models#6151
e-davidson wants to merge 8 commits into
aws:masterfrom
e-davidson:fix/propagate-additional-model-data-sources

Conversation

@e-davidson

@e-davidson e-davidson commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Issue #, if available: N/A

Description of changes:

ModelBuilder._build_for_jumpstart() copied image_uri, env,
enable_network_isolation, model_reference_arn, and model_data from the
resolved JumpStart init_kwargs, but never propagated
additional_model_data_sources. Any JumpStart model whose spec declares an
extra S3 data channel (e.g. an EAGLE speculative decoding draft model, a LoRA
adapter, or a tokenizer override) therefore had that channel silently dropped
from the CreateModel call, and the container failed at runtime trying to read
artifacts from a path that was never mounted.

This change propagates additional_model_data_sources from init_kwargs onto
the builder so _prepare_container_def_base attaches it to the CreateModel
request. The change is confined to
sagemaker-serve/src/sagemaker/serve/model_builder_servers.py.

EULA handling for additional data sources

The JumpStart factory returns these sources already API-shaped, plus one
internal metadata key: HostingEulaKey. The propagation applies two
adjustments, mirroring exactly how container_def treats the primary model's
ModelDataSource:

  • HostingEulaKey is always removed. It is not part of the CreateModel API
    shape, so request validation would reject it client-side.
  • When accept_eula is set (True or False), it is folded into each
    source's S3DataSource.ModelAccessConfig = {"AcceptEula": <value>}, the
    same rule container_def applies to the primary model. When accept_eula
    is unset, sources pass through without a ModelAccessConfig.

There is intentionally no client-side gatedness inference and no client-side
raise. Gated-source enforcement is the service's responsibility: the SageMaker
control plane determines gatedness from the artifact's bucket and rejects
CreateModel with an EULA validation error when a gated source lacks
ModelAccessConfig with AcceptEula true. This is already how a gated
primary model behaves in this SDK today, and additional sources now behave
identically: one accept_eula flag, uniformly applied, server-enforced.

Testing done:

Seven end-to-end tests in TestJumpStartAdditionalModelDataSourcesUserFlow
(sagemaker-serve/tests/unit/servers/test_model_builder_servers.py). Each
test drives the documented single-call customer API:

builder = ModelBuilder.from_jumpstart_config(
    jumpstart_config=JumpStartConfig(
        model_id="meta-textgeneration-llama-3-1-70b",
        accept_eula=True,
        inference_config_name="lmi-optimized",
    ),
    compute=Compute(instance_type="ml.p4d.24xlarge"),
    role_arn=..., sagemaker_session=...,
)
builder.build()

and asserts on the CreateModel request itself: the container_defs captured
from sagemaker_session.create_model. This is what boto receives, not
builder internals.

The spec fixtures under tests/unit/servers/data/jumpstart_specs/ are real,
unmodified specs captured from the production JumpStart content bucket
(jumpstart-cache-prod-us-west-2):

  • pytorch-ic-mobilenet-v2 v3.0.25: public model, no additional data sources
  • openai-reasoning-gpt-oss-20b v3.38.0: public model whose default
    (lmi-optimized) config carries an ungated EAGLE speculative-decoding
    source. This is the model that surfaced this bug
  • meta-textgeneration-llama-3-1-70b v2.16.2: gated model whose
    lmi-optimized config carries a gated draft_model source
    (hosting_eula_key)

Expected values are derived from the fixture files, not hand-copied: a test
helper reads the raw source from the captured spec (channel name, compression,
data type, key prefix) and applies the expected transformation declaratively
(snake_case to PascalCase key mapping, content-bucket resolution,
HostingEulaKey absence, EULA folding). The fixture stays the single source
of truth for data; only the behavior under test is spelled out in the test,
which avoids both duplication and tautological assertions.

Patching is centralized once in setUp, on the same seams the legacy (v2)
JumpStartModel test suite patched: the spec-fetch boundary
(JumpStartModelsAccessor.get_model_specs routes by model id into the
captured spec files (the v2 PROTOTYPICAL_MODEL_SPECS_DICT pattern) and
_get_manifest serves captured headers) and the AWS boundary (the mock
session whose create_model call is the assertion target, IAM role
validation, artifact staging, and the post-create describe). Everything in
between runs for real: JumpStart model-id detection, JumpStartModelSpecs
parsing, inference-config resolution, the get_init_kwargs factory (image URI
resolution, content-bucket injection, PascalCase shaping), the propagation and
EULA fold in _build_for_jumpstart, and _prepare_container_def_base
assembling the CreateModel request.

Case inventory (assertions on the CreateModel request):

  • model without additional sources: no AdditionalModelDataSources key in
    the request at all
  • ungated additional source, no EULA ever set: the real EAGLE source
    (s3://jumpstart-cache-prod-us-west-2/lmi-eagle-heads/...) reaches
    CreateModel with no ModelAccessConfig. This is the original bug scenario
  • accept_eula=True applies uniformly: ModelAccessConfig={"AcceptEula": True} is folded onto the primary ModelDataSource and the additional
    source alike, matching the primary-model rule
  • gated additional source with accepted EULA:
    ModelAccessConfig={"AcceptEula": True} on both the gated source (real
    s3://jumpstart-private-cache-prod-us-west-2/... URI, no HostingEulaKey)
    and the primary ModelDataSource
  • gated additional source without accept_eula: the request goes to the wire
    with no ModelAccessConfig on the source; the service rejects it with its
    EULA validation error, exactly as it does for a gated primary model
  • gated additional source with explicit accept_eula=False: the request
    carries ModelAccessConfig={"AcceptEula": False} and the service rejects it
  • unselected config: the gated model on its default (lmi) config sends no
    AdditionalModelDataSources. Real config resolution gates which sources
    apply

Full serve test file: 63 passed (7 end-to-end tests + 7 pre-existing
TestBuildForJumpStart tests + the rest of the suite, unaffected). The
sagemaker-core session_helper suite covering container_def also passes:
103 tests.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

ModelBuilder._build_for_jumpstart() copied image_uri, env,
enable_network_isolation, model_reference_arn, and model_data from the
resolved JumpStart init_kwargs, but never propagated
additional_model_data_sources. Any JumpStart model whose spec declares
an extra S3 data channel (e.g. an EAGLE speculative decoding draft
model, a LoRA adapter, or a tokenizer override) therefore had that
channel silently dropped from the CreateModel call, and the container
failed at runtime trying to read artifacts from a path that was never
mounted.

Propagate additional_model_data_sources from init_kwargs onto the
builder so _prepare_container_def_base attaches it to the CreateModel
request. Each source is filtered to only the keys the CreateModel API
accepts (ChannelName, S3DataSource); JumpStart specs may include extra
metadata keys such as HostingEulaKey that the API rejects.

Adds unit tests for both the propagation+filtering path and the no-op
case when the spec declares no additional data sources.
A source carrying a truthy HostingEulaKey is gated and requires license
acceptance, exactly like the base model. container_def folds accept_eula
into the base model's ModelAccessConfig but passes additional sources
through untouched, so the acceptance is translated onto each gated
source here: fold ModelAccessConfig={AcceptEula: True} into its
S3DataSource when accept_eula is set, raise the same error the base
model raises when it is not. Ungated sources just drop the
JumpStart-internal HostingEulaKey the CreateModel API rejects.

Test coverage mirrors the v2 helper's case inventory: multiple public
sources, ungated ignores acceptance, mixed public+gated, gated accept,
gated missing/false acceptance raises, plus a raw-spec-JSON end-to-end
test driving the real JumpStartModelSpecs parsing and factory shaping.
Replace the mock-harness additional_model_data_sources tests with
user-flow tests that exercise the documented user contract: construct a
real ModelBuilder, set accept_eula, and call build(). Only the spec
fetch boundary (JumpStartModelsAccessor.get_model_specs / _get_manifest)
and the AWS boundary (_prepare_for_mode / _create_model, mock session)
are mocked, mirroring how master-v2 tested JumpStartModel.

Raw spec JSON fixtures live in data/jumpstart_specs.json as a base spec
plus overlay variants (public model without/with public sources, gated
sources, mixed), reusable across tests. Everything between the two
boundaries runs for real: JumpStart detection, JumpStartModelSpecs
parsing, the get_init_kwargs factory shaping, and the EULA translation
in _build_for_jumpstart.
Replace the synthetic base-spec-plus-variants fixture with real,
unmodified specs captured from the production JumpStart content bucket,
mirroring master-v2's PROTOTYPICAL_MODEL_SPECS_DICT pattern: the spec
accessor routes by model id into one captured spec file per model, and
each test calls the SDK exactly like a customer does.

Fixture models (tests/unit/servers/data/jumpstart_specs/):
- pytorch-ic-mobilenet-v2 v3.0.25: public, no additional data sources
- openai-reasoning-gpt-oss-20b v3.38.0: ungated EAGLE speculative
  decoding source in its default (lmi-optimized) config -- the model
  that surfaced the propagation bug
- meta-textgeneration-llama-3-1-70b v2.16.2: gated draft_model source
  (hosting_eula_key) in its lmi-optimized config, selected in tests via
  the public set_deployment_config API

All patching is centralized in setUp (spec fetch, manifest, and the
AWS boundary), so test bodies contain only public API calls:
ModelBuilder(model=<real id>, instance_type=...), accept_eula, and
build(). Real inference-config resolution now runs in every test.
… API

Restructure the additional-model-data-source tests on three axes:

1. Test names describe the case under test (gated source with accepted
   EULA sends ModelAccessConfig, unselected config sources do not leak,
   etc.), not the fixture model.

2. Each test drives the documented single-call customer API:
   ModelBuilder.from_jumpstart_config(JumpStartConfig(model_id,
   accept_eula, inference_config_name), compute=Compute(...)) followed
   by build(). No attribute pokes and no separate config-selection call.

3. Assertions target the CreateModel request itself: _create_model is
   no longer patched, so the real _create_sagemaker_model and
   _prepare_container_def_base run, and every test inspects the
   container_defs captured from sagemaker_session.create_model. This
   covers the full downstream impact: AdditionalModelDataSources
   presence/absence/shape, ModelAccessConfig on both the gated
   additional source and the primary ModelDataSource, and
   create_model never being called when the EULA gate raises.
Eli Davidson added 2 commits August 4, 2026 18:00
Replace the hand-copied expected-source class constants with a helper
that reads the captured spec fixture and applies the expected
transformation declaratively (PascalCase key mapping, content-bucket
resolution, HostingEulaKey absence, EULA folding). The fixture stays
the single source of truth for data values; only the behavior under
test remains spelled out in the test, avoiding both duplication and
tautological assertions.
Move the additional-model-data-source EULA handling out of the
JumpStart-specific block in _build_for_jumpstart and into
_apply_accept_eula_to_additional_model_data_sources next to
container_def, the single place the primary model's accept_eula
already becomes ModelAccessConfig. Both the primary model and the
additional sources are now driven by the same accept_eula flag in
adjacent code, and every container-assembly call site (including the
deployment-config path that never passes through _build_for_jumpstart)
gets the gated-source enforcement. _build_for_jumpstart is reduced to
pure propagation. No behavior change on the JumpStart path: all 63
serve tests pass unchanged, plus 103 core session_helper tests.
…pstart

Scope the EULA handling down to mirror the primary model exactly: strip
the JumpStart-internal HostingEulaKey (not part of the CreateModel API
shape) and fold accept_eula into each source's ModelAccessConfig when
set. No client-side gatedness check or raise: the control plane
determines gatedness from the artifact's bucket and rejects CreateModel
with its own EULA validation error, for the primary model and
additional sources alike.

This reverts the sagemaker-core half of the previous commit:
session_helper.py returns to upstream and the whole change lives in the
JumpStart builder that produces the spec-derived sources. Tests assert
the new contract on the CreateModel request: acceptance folded when
accept_eula is set (True or False), no ModelAccessConfig when unset.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant