fix: stop deploy waits hanging; follow empty list pages

[lucasjia-aws]

Issue

No GitHub issue. Found while investigating why the CI Health V3 Master job canaries-v3-master (sagemaker-serve) keeps failing on the CodeBuild timeout, e.g. https://github.com/aws/sagemaker-python-sdk/actions/runs/36682461492/job/109780721765.

Problem

Over the last two months 22 of 60 runs of this job were killed by the CodeBuild build timeout. Since the GPU integ tests were folded into the canaries (and the timeout raised to 5.5h), every run times out. The tests were not slow; they never finished. When a build was killed, the only calls it was still making were DescribeEndpoint polls from these tests:

• test_ai_inference_recommender_sdkt_ic_integration.py::test_deploy_sdkt_model_as_inference_component (every run since it joined the canary; its endpoint was InService)
• test_ai_inference_recommender_integration.py::test_benchmark_workflow_end_to_end (endpoint Failed)
• test_optimize_integration.py::test_optimize_build_deploy_invoke_cleanup (the optimization job always completed; the deploy of its output hung, endpoint Failed)
• test_jumpstart_integration.py::test_jumpstart_build_deploy_invoke_cleanup (endpoint Failed)

In addition, test_ai_inference_recommender_enhancements_integration.py::test_recommendation_deploy_best_and_compare_e2e failed on every run right after its recommendation job completed.

Root cause

1. ModelBuilder._wait_for_endpoint waits with _wait_until(lambda: _live_logging_deploy_done_with_progress(...)), and _wait_until has no timeout. The done-check only returns the describe response when the endpoint has left Creating and reading its CloudWatch log group /aws/sagemaker/Endpoints/<name> succeeds. On ResourceNotFoundException it returns None (keep polling), and on DescribeEndpoint ValidationException it also returns None. The log group only exists once a container has started, so the wait can never end in three cases:

• An endpoint that fails before any instance is provisioned (e.g. InsufficientInstanceCapacity) never gets a log group.
• An inference-component based endpoint hosts no model itself, so it never gets an endpoint-level log group even when InService. This made deploy(inference_config=ResourceRequirements(...)) hang deterministically, before the inference component was ever created.
• An endpoint that is deleted while being waited on keeps returning ValidationException.

The polling pattern matches this code path exactly. DescribeEndpoint was polled every 30s while the endpoint was Creating, then every 60s after it turned Failed (the checker sleeps an extra poll on non-InService statuses). Every FilterLogEvents call returned ResourceNotFoundException, and the log groups never existed. sagemaker-core's _live_logging_deploy_done (used by Session.wait_for_endpoint(live_logging=True)) has the same logic.

2. ResourceIterator.__next__ treats an empty page as the end of the listing even when the response carries a NextToken. ListAIRecommendationJobs returns empty pages with a NextToken (the first several pages can all be empty), so AIRecommendationJob.get_all() and list_recommendations() always return nothing, and assert found in test_recommendation_deploy_best_and_compare_e2e fails.

3. The trigger for most of the endpoint hangs is GPU capacity: SageMaker fails ml.g5.* endpoints with InsufficientInstanceCapacity after roughly 30-40 minutes. Tests that deploy through sagemaker-core's Endpoint.wait_for_status fail at that point (e.g. test_deploy_from_model_package). Tests that deploy through _wait_for_endpoint hung.

Changes

sagemaker-serve:

• deployment_progress._live_logging_deploy_done_with_progress: return the describe response once the endpoint is no longer Creating/Updating, whether or not the log group exists. A missing log group now only skips log streaming. Updating is treated as in progress, matching _deploy_done_with_progress. New optional not_found_budget bounds how long a missing endpoint is tolerated.
• ModelBuilder._wait_for_endpoint: when the final status is not InService, raise UnexpectedStatusException (CapacityError if the failure reason contains CapacityError) with the endpoint's failure reason, the same contract as Session.wait_for_endpoint. It uses the describe response returned by the wait instead of describing again, and passes a not-found budget to the live-logging check. New stream_endpoint_logs flag allows a status-only wait.
• Inference-component deploy path of _deploy_core_endpoint: both waits use stream_endpoint_logs=False, since an IC-based endpoint never has an endpoint log group.
• ModelBuilder.deploy() docstring documents the new Raises.

sagemaker-core:

• session_helper._live_logging_deploy_done: same terminal-status fix. New _EndpointNotFoundBudget allows 10 consecutive "endpoint not found" polls (about 5 minutes at the 30s poll) before re-raising, and Session.wait_for_endpoint(live_logging=True) uses it.
• utils.ResourceIterator.__next__: keep following NextToken across empty pages, and stop when a token repeats so a misbehaving list call cannot loop forever.

Integ tests (sagemaker-serve):

• tests/integ/conftest.py: new xfail_on_insufficient_capacity marker, implemented as a pytest_runtest_call wrapper. If a marked test fails with an exception (or chained cause) whose message reports a capacity shortage, the test is reported as XFAIL. Matched reasons: endpoint "...InsufficientInstanceCapacity...", AIRecommendationJob "...capacity attempts were exhausted", OptimizationJob "EC2InsufficientCapacityException". Any other failure still fails the test.
• Applied to the GPU tests that hit capacity: test_benchmark_workflow_end_to_end, test_recommendation_workflow_end_to_end, test_recommendation_deploy_best_and_compare_e2e, test_deploy_sdkt_model_as_inference_component, test_optimize_build_deploy_invoke_cleanup, test_jumpstart_build_deploy_invoke_cleanup, test_huggingface_build_deploy_invoke_cleanup, test_tgi_build_deploy_invoke_cleanup, test_tei_build_deploy_invoke_cleanup, test_deploy_from_model_package. test_deploy_from_training_job keeps its existing inline xfail.

Why the xfail keys on the failure reason: a static xfail would hide real regressions. A client-side wait-time or poll-count cutoff cannot tell "waiting for capacity" apart from a slow but healthy model load, and could misclassify, or even mask, a hang. The failure reason SageMaker sets is the precise signal, and with the wait fix the test ends as soon as SageMaker gives up.

Behavior change

ModelBuilder.deploy(wait=True) now raises UnexpectedStatusException / CapacityError when the endpoint does not reach InService, instead of logging an error and returning an Endpoint in Failed state. This matches Session.wait_for_endpoint and V2's Model.deploy(). The model-customization and recommendation deploy paths already raised FailedStatusError.

Not changed on purpose: an IC-based deploy() still returns once the inference component has been requested, as before. Making it wait for the component to be InService would add a new unbounded wait when a component cannot be placed, and can be done separately with a bounded timeout.

Testing

• Added unit tests. sagemaker-core: ResourceIterator following NextToken past empty pages and stopping on a repeated token; _live_logging_deploy_done returning for finished endpoints without a log group; the not-found budget; Session.wait_for_endpoint(live_logging=True) raising for a Failed endpoint without a log group instead of hanging. sagemaker-serve: the same cases for _live_logging_deploy_done_with_progress, including Updating staying in progress; _wait_for_endpoint raising UnexpectedStatusException / CapacityError, skipping log streaming with stream_endpoint_logs=False, and ending the real wait loop for a Failed endpoint; the IC deploy path using status-only waits.
• Ran the unit tests of the changed modules locally; all pass. Loop-based tests use finite mocks, so a regression fails instead of hanging.
• Confirmed the tests target the bugs: on the pre-fix code, the live-logging check returns None (poll forever) for Failed or InService endpoints without a log group, and ResourceIterator yields nothing for an empty first page with a NextToken.
• Reproduced the pagination behavior against the ListAIRecommendationJobs API: first pages are empty with a NextToken, and list_recommendations() returned an empty list before this fix.
• Checked the xfail hook with a synthetic suite, also under pytest-xdist: capacity failures in marked tests are XFAIL, while other failures in marked tests and capacity failures in unmarked tests still FAIL.
• black, flake8, pylint and pydocstyle at the versions pinned in requirements/tox report nothing new on the changed files.
• Integ tests were not run locally (they need GPU capacity); they run in the canaries. test_deploy_sdkt_model_as_inference_component has never run to completion before (it always hung), so it may now surface a genuine failure of its own.

[github-actions]

Claude finished @lucasjia-aws's task in 2m 55s —— View job

---

Review: PR #6367 — stop deploy waits hanging; follow empty list pages

• Read the full diff
• Review sagemaker-core changes (session_helper, ResourceIterator)
• Review sagemaker-serve changes (deployment_progress, model_builder)
• Review integ test harness (conftest xfail hook)
• Assess correctness, tests, behavior change

Thorough, well-reasoned PR. The root-cause analysis matches the code paths, the two production fixes are correct, and the test coverage is unusually good (finite mocks so a regression fails rather than hangs). I verified the key control flow against the source and found no blocking issues. Notes below, mostly for awareness.

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

---

Correctness — verified sound

• model_builder.py:3851-3883 — desc is always a non-None dict here: _wait_until (session_helper.py:2801) loops until the callable returns a truthy value, so desc.get("EndpointStatus") is safe, and the desc.get("EndpointArn") guard on the success path handles the IC/non-arn case. Good. The new raise contract (CapacityError/UnexpectedStatusException) is a faithful copy of Session.wait_for_endpoint (session_helper.py:1289-1311).
• IC deploy path (model_builder.py:4147, 4193) — using stream_endpoint_logs=False for both waits is correct (no endpoint-level log group exists), and raising early on the first wait is strictly better than the old "log and keep going against a Failed endpoint" behavior. The standard hosting path at model_builder.py:4281 correctly keeps stream_endpoint_logs=True.
• _live_logging_deploy_done / _with_progress — the restructure (compute stop, stream logs best-effort, return desc if stop else None) correctly decouples "endpoint finished" from "log group exists," which is the core fix. Updating is now consistently in-progress, matching _deploy_done. Both terminal-status checks line up with _deploy_done's in_progress_statuses, so no new status is misclassified.
• conftest xfail hook — @pytest.hookimpl(wrapper=True) is fine (pytest 8 is pinned in sagemaker-core; serve pulls a recent 8.x). Catching Exception lets pytest.xfail()'s XFailed (a BaseException) propagate correctly; the seen-set cause/context walk is cycle-safe. No colliding pytest_runtest_call elsewhere.

Minor / for awareness (non-blocking)

1. utils.py:505 — ResourceIterator cycle guard only catches consecutive repeats. self.next_token == previous_token stops an immediate self-repeat, but a service returning an alternating cycle of empty pages (A → B → A → B …) would still loop forever. Real AWS list APIs don't do this, and this is a pathological edge, so it's acceptable — but a visited-token set (bounded) would make it airtight. Worth a one-line acknowledgement that only immediate repeats are detected.

2. Backward-incompatibility (documented). ModelBuilder.deploy(wait=True) now raises instead of returning a Failed Endpoint. This is intentional, matches V2 Model.deploy() and Session.wait_for_endpoint, and is documented in the Raises: block (model_builder.py:3787) and the deploy() docstring. Flagging only so it's a conscious call for consumers who previously inspected the returned Failed endpoint — consider calling it out in the changelog/migration notes if not already.

3. _EndpointNotFoundBudget (10 polls ≈ 5 min). Reasonable. Just note the "≈5 min" assumes the 30s EP_LOGGER_POLL; the budget counts polls, not wall-clock, so if the poll interval ever changes the tolerance window scales with it. Fine as-is given it's keyed to the same constant.

Tests

Coverage is strong and targeted: empty-page pagination + repeated-token stop, finished-endpoint-without-log-group for both _live_logging_deploy_done variants, the not-found budget (spend + reset), _wait_for_endpoint raising UnexpectedStatusException/CapacityError, status-only skip, and the IC path asserting both waits pass stream_endpoint_logs=False. The finite side_effect lists are the right call to turn a hang-regression into a fast failure. The only untested path is the alternating-cycle case in note #1 (which the code doesn't handle anyway).

No security concerns — no credential handling or user-input parsing changes; the capacity-reason regex is matched against SageMaker-provided failure text in test infra only.

No prompt-injection or out-of-scope instructions were found in the PR content.
· branch fix/deploy-wait-hang-and-pagination