fix(serve): propagate additional_model_data_sources for JumpStart models - #6151
Open
e-davidson wants to merge 8 commits into
Open
fix(serve): propagate additional_model_data_sources for JumpStart models#6151e-davidson wants to merge 8 commits into
e-davidson wants to merge 8 commits into
Conversation
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.
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 01:46 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 01:46 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 01:46 — with
GitHub Actions
Error
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.
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 13:20 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 13:20 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 13:20 — with
GitHub Actions
Error
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.
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 14:29 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 14:29 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 14:29 — with
GitHub Actions
Error
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.
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 14:59 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 14:59 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 14:59 — with
GitHub Actions
Error
… 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.
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 15:32 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 15:32 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 15:32 — with
GitHub Actions
Error
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.
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 18:04 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 18:04 — with
GitHub Actions
Error
e-davidson
had a problem deploying
to
manual-approval
August 4, 2026 18:04 — with
GitHub Actions
Error
…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.
e-davidson
requested a deployment
to
manual-approval
August 4, 2026 19:38 — with
GitHub Actions
Waiting
e-davidson
requested a deployment
to
manual-approval
August 4, 2026 19:38 — with
GitHub Actions
Waiting
e-davidson
requested a deployment
to
manual-approval
August 4, 2026 19:38 — with
GitHub Actions
Waiting
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Issue #, if available: N/A
Description of changes:
ModelBuilder._build_for_jumpstart()copiedimage_uri,env,enable_network_isolation,model_reference_arn, andmodel_datafrom theresolved JumpStart init_kwargs, but never propagated
additional_model_data_sources. Any JumpStart model whose spec declares anextra 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_sourcesfrom init_kwargs ontothe builder so
_prepare_container_def_baseattaches it to the CreateModelrequest. 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 twoadjustments, mirroring exactly how
container_deftreats the primary model'sModelDataSource:HostingEulaKeyis always removed. It is not part of the CreateModel APIshape, so request validation would reject it client-side.
accept_eulais set (TrueorFalse), it is folded into eachsource's
S3DataSource.ModelAccessConfig = {"AcceptEula": <value>}, thesame rule
container_defapplies to the primary model. Whenaccept_eulais 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
ModelAccessConfigwithAcceptEulatrue. This is already how a gatedprimary model behaves in this SDK today, and additional sources now behave
identically: one
accept_eulaflag, uniformly applied, server-enforced.Testing done:
Seven end-to-end tests in
TestJumpStartAdditionalModelDataSourcesUserFlow(
sagemaker-serve/tests/unit/servers/test_model_builder_servers.py). Eachtest drives the documented single-call customer API:
and asserts on the CreateModel request itself: the
container_defscapturedfrom
sagemaker_session.create_model. This is what boto receives, notbuilder 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-v2v3.0.25: public model, no additional data sourcesopenai-reasoning-gpt-oss-20bv3.38.0: public model whose default(
lmi-optimized) config carries an ungated EAGLE speculative-decodingsource. This is the model that surfaced this bug
meta-textgeneration-llama-3-1-70bv2.16.2: gated model whoselmi-optimizedconfig carries a gateddraft_modelsource(
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,
HostingEulaKeyabsence, EULA folding). The fixture stays the single sourceof 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)JumpStartModeltest suite patched: the spec-fetch boundary(
JumpStartModelsAccessor.get_model_specsroutes by model id into thecaptured spec files (the v2
PROTOTYPICAL_MODEL_SPECS_DICTpattern) and_get_manifestserves captured headers) and the AWS boundary (the mocksession whose
create_modelcall is the assertion target, IAM rolevalidation, artifact staging, and the post-create describe). Everything in
between runs for real: JumpStart model-id detection,
JumpStartModelSpecsparsing, inference-config resolution, the
get_init_kwargsfactory (image URIresolution, content-bucket injection, PascalCase shaping), the propagation and
EULA fold in
_build_for_jumpstart, and_prepare_container_def_baseassembling the CreateModel request.
Case inventory (assertions on the CreateModel request):
AdditionalModelDataSourceskey inthe request at all
(
s3://jumpstart-cache-prod-us-west-2/lmi-eagle-heads/...) reachesCreateModel with no
ModelAccessConfig. This is the original bug scenarioaccept_eula=Trueapplies uniformly:ModelAccessConfig={"AcceptEula": True}is folded onto the primaryModelDataSourceand the additionalsource alike, matching the primary-model rule
ModelAccessConfig={"AcceptEula": True}on both the gated source (reals3://jumpstart-private-cache-prod-us-west-2/...URI, noHostingEulaKey)and the primary
ModelDataSourceaccept_eula: the request goes to the wirewith no
ModelAccessConfigon the source; the service rejects it with itsEULA validation error, exactly as it does for a gated primary model
accept_eula=False: the requestcarries
ModelAccessConfig={"AcceptEula": False}and the service rejects itlmi) config sends noAdditionalModelDataSources. Real config resolution gates which sourcesapply
Full serve test file: 63 passed (7 end-to-end tests + 7 pre-existing
TestBuildForJumpStarttests + the rest of the suite, unaffected). Thesagemaker-core session_helper suite covering
container_defalso 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.