Checks to use the correct ami when on cu130 to avoid crashes at launch

[dwarez]

Description of changes:
Currently, some instances crash at startup because of a mismatch between the CUDA version of the engine container and the one present in the instance. This will result in a run that crashes even before any logs can be produced. Currently the fix is to specify the correct inference_ami_version when invoking ModelBuilder.deploy, however this can be managed by the SDK and should not be a concern of the user.

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

[github-actions]

Claude finished @dwarez's task in 2m 54s —— View job

---

PR Review: cu130 AMI compatibility

• Read the PR diff
• Review model_builder.py timeout mapping
• Review model_builder_servers.py cu130 AMI logic
• Review added tests
• Post review feedback

---

Thanks for the fix — the intent is sound: silently ignoring a CUDA/driver mismatch and crashing before any logs is a bad failure mode, and pushing the AMI selection into the SDK is the right call. The cu130 gate is nicely conservative (VLLM + SAGEMAKER_ENDPOINT + isinstance guards + respects a user-supplied inference_ami_version), and both branches are covered by tests. A couple of things worth a look before merge.

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

1. Timeout mapping changes the default startup timeout (medium)

sagemaker-serve/src/sagemaker/serve/model_builder.py:4191-4195

elif self.mode == Mode.SAGEMAKER_ENDPOINT:
    if "container_timeout_in_seconds" in kwargs:
        kwargs.setdefault(
            "container_startup_health_check_timeout",
            kwargs.pop("container_timeout_in_seconds"),
        )

The public deploy() (and _deploy_model_customization) always injects container_timeout_in_seconds into the deploy kwargs with its default of 300 (see model_builder.py:6154, 6167, 6182, 6210). So on the endpoint path this mapping now fires on every deploy, not just when the user explicitly sets a timeout.

Consequences:

• Behavior change vs. prior default. Before this PR, _deploy_core_endpoint read container_startup_health_check_timeout from kwargs, found it absent, and fell back to getattr(self, "container_startup_health_check_timeout", None) → None, letting SageMaker apply its own service default. Now it is always 300. For large models that previously relied on the longer service default to pass the startup health check, this could turn a working deploy into a health-check timeout. Fixing the "user timeout was ignored" bug is correct, but please confirm you intend to also pin the default to 300s here.
• Instance-attribute shadowing. If a user set self.container_startup_health_check_timeout on the builder but calls deploy() without passing a timeout, the always-present 300 now lands in kwargs and wins over the getattr(self, ...) fallback in _deploy_core_endpoint (model_builder.py:3842-3845). setdefault only guards the kwargs key, not the instance attribute. Consider only mapping when the value differs from the default, or reconciling with the instance attribute before mapping.

2. Verify the AMI alias and the compatible-family allowlist (minor)

sagemaker-serve/src/sagemaker/serve/model_builder_servers.py:25-38

• Please double-check al2023-ami-sagemaker-inference-gpu-4-1 is a valid, GA inference_ami_version value in all regions this ships to — an incorrect alias would reintroduce a launch-time failure (just a different one).
• The allowlist is a conservative frozenset. Any cu130 vLLM instance family not in this set will still hit the original mismatch crash this PR is trying to prevent. That's an acceptable allowlist strategy, but worth a comment noting it's intentionally opt-in so future families are added deliberately.

3. Minor observations

• model_builder_servers.py:392 — self.instance_type.rsplit(".", 1)[0] handles the normal ml.g5.2xlarge → ml.g5 case correctly and won't raise on a malformed value (guarded by the isinstance(..., str) check). No change needed.
• "-cu130-" in self.image_uri (model_builder_servers.py:391) relies on the DLC tag format (...-cu130-ubuntu...). Fine for current DLC naming, but it's an implicit contract on the tag string — a brief comment would help future readers.

Tests

Coverage is good: test_configures_model_and_cuda_13_inference_ami and test_preserves_explicit_model_and_inference_ami exercise both the auto-set and user-override branches (the mock_builder fixture's ml.g5.2xlarge is in the allowlist, so the family gate is genuinely tested), and test_deploy_sagemaker_endpoint_maps_container_timeout verifies the kwarg remap. Two suggestions:

• Add a negative case for the AMI gate — e.g. a non-cu130 image, or an instance family not in the allowlist (ml.m5.large) — to lock in that the AMI is left untouched.
• Add a test that the default deploy path (no explicit timeout) produces the expected container_startup_health_check_timeout, to pin the intended default and guard against the behavior change in finding Use custom user agent string at all times #1.

No security concerns (no credential handling, no unvalidated external input) and no performance impact. No prompt-injection attempts in the PR content.

[rsareddy0329]

is this change required for AMI fix?

[dwarez]

no not really, reverting it and in case we can make a different PR for that

[ehcalabres]

LGTM now! 🚀

[github-actions]

Claude finished @dwarez's task in 2m 15s —— View job

---

PR Review: cu130 AMI compatibility

• Read the PR diff
• Review model_builder_servers.py cu130 AMI logic
• Review SM_VLLM_MODEL env var addition
• Review added tests
• Post review feedback

---

Thanks for the update — reverting the container_timeout_in_seconds remap was the right call (that was an unrelated behavior change to the endpoint startup default, and it's cleaner as its own PR). The current diff is now focused: it only touches model_builder_servers.py and its tests, and the cu130 gate is nicely conservative (VLLM/SAGEMAKER_ENDPOINT + isinstance guards + allowlisted family + respects a user-supplied inference_ami_version). A few things worth a look before merge.

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

1. Gate reads the build-time instance_type, not the deploy-time one (medium)

model_builder_servers.py:386-395

The AMI is selected inside _build_for_hf_server, based on self.instance_type at build time. instance_type is defaulted in __post_init__ and can be reset later — deploy(instance_type=...) assigns self.instance_type (e.g. model_builder.py:6119) and the endpoint path deploys with self.instance_type (6145). So the common flow of building once and then calling deploy() with a different, cu130-compatible instance family will not pick up the AMI, because the gate already ran against the earlier value and won't re-run.

If a user builds with the default instance type and only chooses the GPU instance at deploy(), the fix silently doesn't apply and they hit the original launch crash. Worth confirming whether the gate should also run/re-run at deploy time (or that build is always invoked with the final instance type in the supported flows).

2. "-cu130-" substring relies on a trailing dash (minor)

model_builder_servers.py:390

"-cu130-" in self.image_uri matches ...-cu130-ubuntu22.04 / ...-cu130-amzn2023, but would miss a tag that ends in -cu130 with no suffix (e.g. ...-py312-cu130). This is an implicit contract on DLC tag formatting. Consider a small comment documenting the expected tag shape, or a slightly more tolerant check, so a future tag naming tweak doesn't silently disable the fix.

3. Verify the AMI alias is GA in all target regions (minor)

model_builder_servers.py:9

Please double-check al2023-ami-sagemaker-inference-gpu-4-1 is a valid, generally-available inference_ami_version in every region this ships to. An invalid alias would trade the CUDA-mismatch crash for a different launch-time failure. Also worth a one-line comment noting the allowlist is intentionally opt-in — any cu130 vLLM family not listed still hits the original mismatch, which is the accepted trade-off.

4. SM_VLLM_MODEL addition looks orthogonal to the AMI fix (question)

model_builder_servers.py:377-378

Setting SM_VLLM_MODEL alongside HF_MODEL_ID for the VLLM server seems unrelated to the cu130/AMI problem this PR targets. It's harmless (setdefault, VLLM-gated, and covered by tests), but — same spirit as the timeout revert — consider whether it belongs in this PR or a separate one, and confirm both env vars are genuinely required by the vLLM DLC.

Tests

Coverage of the happy paths is good: test_configures_model_and_cuda_13_inference_ami, test_configures_cuda_13_inference_ami_for_vllm_omni, and test_preserves_explicit_model_and_inference_ami exercise the auto-set, omni, and user-override branches, and mock_builder's ml.g5.2xlarge is in the allowlist so the family gate is genuinely tested. One gap:

• Add a negative case that locks in the AMI is left untouched — e.g. a non-cu130 image, or an allowlist miss (ml.m5.large / ml.c5.xlarge), asserting inference_ami_version stays None. Right now nothing guards against the gate firing too broadly.

Summary

No security concerns (no credential handling; inputs are the builder's own attributes) and no meaningful performance impact. The core logic is correct and well-guarded; finding #1 (build-time vs deploy-time instance type) is the one I'd most want confirmed, since it can silently no-op the fix in a common usage flow. No prompt-injection attempts were present in the PR content.

[github-actions]

Claude finished @dwarez's task in 2m 2s —— View job

---

PR Review: cu130 AMI compatibility

• Read the PR diff
• Verify _resolve_inference_ami_version wiring in _deploy_core_endpoint
• Verify class inheritance / method availability
• Check deploy-time instance_type flow
• Review tests
• Post review

---

This revision is in good shape and resolves the concerns raised in the earlier review rounds. Summary of what I verified:

Deploy-time (not build-time) resolution — the key prior finding is fixed. The AMI is now resolved inside _deploy_core_endpoint (model_builder.py:3991-3995) rather than at build time, and deploy() now assigns self.instance_type = instance_type up front (model_builder.py:6062-6063, moved out of the model-customization branch). So the common "build once, choose the GPU instance at deploy()" flow now picks up the correct instance family. The resolved value flows through both the IC-based (:4032) and model-based / update-endpoint (:4106, :4155) production-variant paths. ModelBuilder inherits _ModelBuilderServers (model_builder.py:185), so the method is available.

Engine-agnostic, as requested. The vLLM-only gate was removed and resolution now applies to any non-serverless core-endpoint deploy (guarded by serverless_inference_config is None), matching @ehcalabres's suggestion. It stays conservative: only fires when inference_ami_version is unset, image_uri/instance_type are strings, the tag contains a cu130 token, and the family is allowlisted — so it can't silently override a caller or fire on unrelated deploys.

Tag parsing is more robust. image_uri.rpartition(":")[2] + "cu130" in tag.split("-") (model_builder_servers.py:359-360) correctly matches both ...-cu130-amzn2023 and a bare trailing ...-cu130, addressing the earlier substring-with-trailing-dash concern.

The unrelated changes were reverted. The container_timeout_in_seconds remap and the SM_VLLM_MODEL env var are gone from this diff — good call keeping the PR focused.

Tests. Good coverage now, including the negative cases requested previously: test_leaves_ami_unset_outside_compatibility_gate (non-cu130 image and an allowlist-miss ml.m5.large), the omni/deploy-time-override happy path, the bare-cu130-tag case, and explicit-AMI preservation. test_model_builder_workflows.py was updated to assert the deploy-time instance override reaches _deploy.

Minor / non-blocking

1. Allowlist is opt-in by design. Any cu130 vLLM instance family not in _CUDA_13_INFERENCE_AMI_COMPATIBLE_FAMILIES (model_builder_servers.py:60) still hits the original mismatch crash. The added comment documents this is intentional — 👍. Just make sure new families are added deliberately as they gain support.
2. AMI alias GA check. Please confirm al2023-ami-sagemaker-inference-gpu-4-1 is a valid, generally-available inference_ami_version in every region this ships to — an invalid alias would trade the CUDA mismatch for a different launch-time failure. (Can't verify this from the code.)
3. Untagged-image edge case (negligible). If an image_uri has no :tag, rpartition(":")[2] returns the whole URI and split("-") scans the repo path. Harmless in practice since these DLC flows always carry a tag, but noting it for completeness.

No security or performance concerns (inputs are the builder's own attributes; no credential handling). No prompt-injection attempts in the PR content.

Overall: correct, well-guarded, and the tests lock in the intended behavior. LGTM. ✅