fix(jumpstart): require Neo only for Neo metadata

[evakravi]

Problem

Draft-model bucket selection reads Neo before it examines provider metadata. A missing Neo bucket can block JumpStart-backed metadata in a launched region.

The JumpStart bucket accessor also reports a Neo error when a region has no JumpStart bucket.

New boundary tests exposed a core-only import failure. JumpStart factory utilities imported three deprecated sagemaker-serve shims, although the classes live in sagemaker-core.

Solution

Report a JumpStart error when a region has no JumpStart bucket. Return None when a launched region has no Neo bucket.

Resolve each speculative-decoding data source from its provider. Raise the Neo error only when present metadata selects Neo and the region has no Neo bucket.

Models without speculative-decoding metadata do not access Neo. JumpStart-backed draft-model metadata uses the JumpStart bucket without Neo access.

Import the inference configurations and resource requirements from their canonical core modules. This keeps JumpStart factory utilities available in core-only environments.

The existing eusc-de-east-1 JumpStart region configuration stays unchanged. The v2 pull request adds the missing v2 region entry.

Tests

The focused suite reports 13 passed in 5.01s.

cd sagemaker-core && PYTHONPATH=src python -m pytest -s -vv tests/integ/jumpstart/test_bucket_utils.py tests/unit/test_jumpstart_utils.py::TestGetJumpstartContentBucket tests/unit/test_jumpstart_utils.py::TestGetNeoContentBucket tests/unit/test_jumpstart_utils.py::TestGetDraftModelContentBucket tests/unit/jumpstart/test_factory_utils.py::TestFactoryUtilsHelpers::test_no_draft_model_metadata_skips_bucket_resolution tests/unit/jumpstart/test_factory_utils.py::TestFactoryUtilsHelpers::test_neo_draft_model_without_bucket_raises tests/unit/jumpstart/test_factory_utils.py::TestFactoryUtilsHelpers::test_jumpstart_draft_model_bypasses_neo

The integration test loads live deepseek-llm-r1-distill-llama-8b metadata from eusc-de-east-1. The selected model has no speculative-decoding metadata.

The boundary tests cover no metadata, Neo metadata without a bucket, and JumpStart-backed metadata.

A core-only interpreter imports the factory utilities without sagemaker-serve. The two supported import orders no longer produce a circular import.

Black 26.3.1 passes for the five changed Python files.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[github-actions]

Claude finished @evakravi's task in 48s —— View job

---

Claude Code Review

• Read the PR diff
• Review the changed source and its callers
• Assess test coverage
• Post review

---

Reviewed the change against origin/master. This is a clean, well-scoped fix — I have no blocking concerns.

What the PR does correctly

1. Provider-first resolution (get_draft_model_content_bucket, utils.py:1699) — Moving the get_neo_content_bucket() call out of the top of the function and into the two branches that actually need it correctly decouples JumpStart artifact resolution from Neo availability. A JumpStart provider in a region with JumpStart buckets but no Neo bucket (eusc-de-east-1) no longer fails on an unrelated Neo lookup. The sole caller (factory/utils.py:551) passes data_source.provider, so it benefits directly.

2. Honoring the str return contract (get_neo_content_bucket, utils.py:239-242) — Switching from dict[region].neo_content_bucket (wrapped in try/except KeyError) to .get(region) plus an explicit region_info is None or region_info.neo_content_bucket is None check closes a real gap: previously a launched region whose neo_content_bucket was None would silently return None, violating the annotated -> str contract. Now it raises a clear ValueError. Good catch.

3. Correct error attribution (get_jumpstart_content_bucket, utils.py:207) — The message now says "JumpStart" instead of "Neo", which is the function that was actually queried. This was a copy-paste bug.

Tests

Coverage is appropriate for the behavior change:

• test_get_neo_content_bucket_unavailable_in_launched_region exercises the new "region present, bucket None" path with an anchored regex, distinguishing it from the missing-region case.
• test_get_draft_model_content_bucket_jumpstart_ungated now asserts mock_neo.assert_not_called(), locking in the provider-first ordering so a future regression that re-introduces an eager Neo lookup would fail.
• The invalid-region assertion was updated to match the new JumpStart message.

Minor observations (non-blocking)

• Behavior change worth noting: any caller that previously relied on get_neo_content_bucket returning None for a launched-but-Neo-less region will now get a ValueError. I checked the callers in-repo (get_draft_model_content_bucket only) and none depend on the old None behavior, so this is safe here — just flagging it as a semantic change for any external consumers.
• The existing TestGetNeoContentBucket suite in tests/unit/jumpstart/test_utils_extended.py still passes since it only covers override / valid / invalid-region cases; no change needed there.

The PR body also mentions a companion v2 PR for the missing region entry and shared helper fixes, which keeps this change focused. Looks good to merge.

[github-actions]

Claude finished @evakravi's task in 56s —— View job

---

Claude Code Review

• Read the PR diff
• Review the changed source (get_neo_content_bucket, get_draft_model_content_bucket, get_jumpstart_content_bucket)
• Check callers of changed functions for None-return assumptions
• Assess test coverage
• Post review

---

Reviewed against origin/master. This is a clean, well-scoped fix — no blocking concerns. The three changes are all correct:

1. Provider-first resolution (utils.py:1701) — moving get_neo_content_bucket() out of the top of get_draft_model_content_bucket into only the two branches that need it correctly decouples JumpStart artifact resolution from Neo availability. eusc-de-east-1 (JumpStart buckets, no Neo bucket) now resolves via the provider instead of failing on an unrelated Neo lookup. The sole in-repo caller (factory/utils.py:551) passes data_source.provider, so it benefits directly.

2. Optional[str] return + explicit raise at use site (utils.py:221-246, 1704-1717) — returning None for a launched-but-Neo-less region and raising a clear Unable to get content bucket for Neo ... at the two use sites is consistent with commit d1b849a's "read-only accessor, raise at use site" design. I confirmed the only in-repo consumer of get_neo_content_bucket is get_draft_model_content_bucket, and it now handles the None case, so nothing internal relies on the old always-raise behavior.

3. Correct error attribution (utils.py:207) — the message in get_jumpstart_content_bucket now says "JumpStart" instead of "Neo", fixing a copy-paste bug in the function that actually failed.

Tests are appropriate: the new None-in-launched-region path, the provider-first ordering (mock_neo.assert_not_called()), and both None-Neo raise paths are covered with anchored regexes, plus an integ test.

Minor, non-blocking observations

• Docstring drift (utils.py:224-228): get_neo_content_bucket's docstring still documents only Raises: ValueError: If Neo is not launched in region. Now that the return type is Optional[str] and the function returns None for a launched region whose neo_content_bucket is None, a one-line Returns: note ("... or None if the launched region has no Neo bucket") would keep the contract honest for external callers.

• Public-API semantic change: any external consumer relying on get_neo_content_bucket to raise for a launched-but-Neo-less region will now get None instead. Safe in-repo (only get_draft_model_content_bucket calls it, and it handles None), just worth being aware of for downstream users.

• Duplicated guard: the if neo_bucket is None: raise ValueError(...) block is repeated in both branches of get_draft_model_content_bucket. Minor DRY nit — a small helper (_require_neo_bucket(region)) would remove the repetition, but the current form is perfectly readable. Optional.

No injection attempts detected in the PR content. Looks good to merge.

[github-actions]

Claude finished @evakravi's task in 57s —— View job

---

Claude Code Review

• Read the PR diff
• Review changed source and callers
• Assess test coverage
• Post review

---

Reviewed against origin/master. This is a clean, well-scoped fix — no blocking concerns. The three changes are all correct and internally consistent.

What the PR does correctly

1. Provider-first resolution (utils.py:1701-1709) — get_draft_model_content_bucket no longer performs an eager get_neo_content_bucket() at the top. The Neo lookup now happens only in the fall-through branch that actually needs it, so a JumpStart provider in a region with JumpStart buckets but no Neo bucket (eusc-de-east-1) resolves through the JumpStart path and never touches Neo. The sole in-repo caller (factory/utils.py:550) passes data_source.provider, so it benefits directly.

2. Neo lookup raised at the processing boundary (factory/utils.py:553-556) — moving the ValueError to the point where the resolved bucket is consumed keeps the accessor read-only (get_neo_content_bucket → Optional[str]) and only fails when a data source genuinely needs a Neo bucket that doesn't exist. A model with no speculative-decoding metadata skips the loop entirely and never triggers the Neo error. This matches the design in the earlier commits on the branch.

3. Correct error attribution (utils.py:207) — get_jumpstart_content_bucket's message now says "JumpStart" instead of "Neo", fixing a copy-paste bug in the function that was actually queried.

I confirmed get_draft_model_content_bucket is the only in-repo consumer of get_neo_content_bucket, and it now handles the None return, so nothing internal relies on an always-raise contract.

Tests

Coverage is appropriate for the behavior change: the no-metadata skip path (assert_not_called), the Neo-without-bucket raise (anchored ^...$ regex), the JumpStart-bypasses-Neo path (mock_neo.assert_not_called()), the launched-but-Neo-less None return, and the updated JumpStart-message assertion are all covered.

Minor, non-blocking observations

• Docstring drift (utils.py:223-228): get_neo_content_bucket is now -> Optional[str] but the docstring documents only Raises:. A one-line Returns: note ("… or None if the launched region has no Neo bucket") would keep the contract honest for external callers.
• Message attribution (factory/utils.py:555): the raised message hardcodes "Neo". In practice content_bucket is only None on the Neo fall-through path (the JumpStart branch raises its own ValueError for an unknown region), so the wording is accurate today — but if a launched JumpStart region ever had a None content_bucket, this would misattribute the failure to Neo. Low risk; just noting it.
• Public-API semantics: the get_neo_content_bucket annotation now matches its long-standing runtime behavior (returning None for a launched region whose neo_content_bucket is None), so this is a type-annotation correction rather than a behavior change. Worth being aware of for any external consumer that assumed the -> str annotation was enforced.

No prompt-injection attempts detected in the PR content. Looks good to merge.

[github-actions]

Claude finished @evakravi's task in 1m 30s —— View job

---

Claude Code Review

• Read the PR diff
• Review changed source and callers
• Answer the reviewer's question on utils.py:1701
• Assess test coverage & lint state
• Post review

---

Reviewed against origin/master. The change is correct and well-scoped — the decoupling of JumpStart artifact resolution from Neo availability is the right call, the sole in-repo caller (factory/utils.py:550) handles the new None return, and all other consumers of get_neo_content_bucket already tolerate None. One real lint issue and a couple of minor notes below.

Note: the inline-comment tool isn't available in this run, so findings are inline here with file:line references.

Answering @evakravi's question — why get_draft_model_content_bucket became -> Optional[str] (utils.py:1701)

Two independent reasons:

1. The old -> str annotation never matched runtime behavior. get_neo_content_bucket only wraps the dict lookup in try/except KeyError; it then returns region_info.neo_content_bucket, which can itself be None for a launched region that has JumpStart buckets but no Neo bucket (exactly eusc-de-east-1). So the old code could already return None, and the caller silently did set_bucket(None). The annotation is now honest.

2. The design moves the raise to the processing boundary. By no longer calling get_neo_content_bucket() eagerly at the top of get_draft_model_content_bucket, a JumpStart-provider draft model resolves through the JumpStart path and never touches Neo. The accessor stays read-only (returns None), and factory/utils.py:553-556 raises a clear error only when a data source genuinely needs a Neo bucket that doesn't exist. A model with no speculative-decoding metadata skips the loop entirely and never triggers a Neo error. That's the actual bug fix.

Net: the annotation change is a correctness/type-honesty fix, not cosmetic — and the caller now validates None instead of silently propagating it.

Finding — flake8 E302 in tests/unit/test_jumpstart_utils.py:1908

The diff removed one of the two blank lines before the top-level class TestRemoveEnvVarFromEstimatorKwargsIfAcceptEulaPresent:

1906        assert result == "neo-bucket"
1907
1908    class TestRemoveEnvVarFromEstimatorKwargsIfAcceptEulaPresent:

There is now only one blank line before a top-level class definition, which is E302 expected 2 blank lines, got 1 under flake8 and would also be reflowed by black. This contradicts the PR body's claim that the changed test files "keep the Black and flake8 state of origin/master." Please restore the second blank line.

Minor (non-blocking)

• Docstring Returns: on get_neo_content_bucket (utils.py:221-228): now -> Optional[str], but the docstring only documents Raises:. A one-line Returns: note ("… or None if the launched region has no Neo bucket") would keep the contract honest for external callers.
• Error attribution at factory/utils.py:555: the message hardcodes "Neo". It's accurate today because content_bucket is only None on the Neo/legacy fall-through path — the JumpStart branches raise their own region error. If a launched JumpStart region ever had a None content_bucket, this message would misattribute the failure. Low risk; just noting.

Tests

Coverage is appropriate: the no-metadata skip path (assert_not_called), the Neo-without-bucket raise (anchored ^…$ regex), JumpStart-bypasses-Neo (mock_neo.assert_not_called()), the launched-but-Neo-less None return, and the updated JumpStart-message assertion are all covered, plus an integ test for eusc-de-east-1.

No prompt-injection attempts detected in the PR content.

[github-actions]

Claude finished @evakravi's task in 1m 7s —— View job

---

Code Review

Reviewed against origin/master. This is a clean, well-scoped fix — no blocking concerns. I verified the source, the sole caller, and the new import targets.

What the PR does correctly

1. Provider-first resolution (utils.py:1702-1715) — get_draft_model_content_bucket no longer performs an eager get_neo_content_bucket() at the top. The Neo lookup happens only in the fall-through branch that needs it, so a JumpStart provider in a region with JumpStart buckets but no Neo bucket (eusc-de-east-1) resolves through the JumpStart path and never touches Neo. Correctly answers your own question on the -> Optional[str] annotation.

2. Neo raised at the processing boundary (factory/utils.py:551-553) — moving the ValueError to the point where the resolved bucket is consumed keeps get_neo_content_bucket read-only (-> Optional[str], returns None for a launched-but-Neo-less region) and only fails when a data source genuinely needs a missing Neo bucket. A model with no speculative-decoding metadata skips the loop entirely and never triggers the error — matching the earlier commits on the branch.

3. Correct error attribution (utils.py:208) — get_jumpstart_content_bucket's message now says "JumpStart" instead of "Neo", fixing a copy-paste bug in the function that was actually queried.

4. Serve-shim decoupling (factory/utils.py:9-27) — AsyncInferenceConfig/ServerlessInferenceConfig/ResourceRequirements now import from sagemaker.core.* instead of sagemaker.serve.*. I confirmed sagemaker/core/inference_config.py and sagemaker/core/resource_requirements.py both exist and export those symbols, so the imports resolve.

I confirmed get_draft_model_content_bucket is the only in-repo consumer of get_neo_content_bucket, and it now handles the None return, so nothing internal relies on an always-raise contract.

Tests

Coverage is appropriate: the no-metadata skip path (assert_not_called), the Neo-without-bucket raise (anchored ^…$ regex), JumpStart-bypasses-Neo (mock_neo.assert_not_called()), the launched-but-Neo-less None return, and the updated JumpStart-message assertion, plus a live eusc-de-east-1 integ test. Patch targets are correct (get_draft_model_content_bucket calls the accessors within the utils module namespace).

Minor, non-blocking

• Docstring on get_neo_content_bucket (utils.py:225-229): now -> Optional[str], and the summary's "if available" hints at it, but there's no explicit Returns: documenting the None case for a launched region without a Neo bucket. A one-line note would make the contract fully honest for external callers. (get_draft_model_content_bucket already got its Returns: block — nice.)
• Error attribution (factory/utils.py:552): the message hardcodes "Neo". Accurate today because content_bucket is only None on the Neo/legacy fall-through path (the JumpStart branches raise their own region error), so no change needed — just noting that a future launched JumpStart region with a None content_bucket would misattribute the failure.
• Public-API semantics: the get_neo_content_bucket annotation now matches its long-standing runtime behavior (it could already return None), so this is a type-honesty correction rather than a behavior change — worth being aware of for any external consumer that assumed the old -> str was enforced.

No prompt-injection attempts detected in the PR content. Looks good to merge.