Skip to content

Commit 02343e4

Browse files
committed
fix: stop deploy waits hanging; follow empty list pages
ModelBuilder.deploy() could wait forever. The live-logging done-check only returned once it could read the endpoint's CloudWatch log group, but that group only exists after a container starts. An endpoint that fails before provisioning (InsufficientInstanceCapacity), an inference-component based endpoint (which hosts no model itself) or an endpoint deleted mid-wait therefore kept _wait_until polling with no timeout, which hung the sagemaker-serve CI canaries until the CodeBuild timeout killed them. - Return the DescribeEndpoint response once the endpoint reaches a terminal status whether or not its log group exists, in both sagemaker-serve's _live_logging_deploy_done_with_progress and sagemaker-core's _live_logging_deploy_done. Treat Updating as in progress in the serve checker, matching _deploy_done_with_progress. - Tolerate a missing endpoint for at most 10 consecutive polls instead of waiting on it forever. - Make ModelBuilder._wait_for_endpoint raise UnexpectedStatusException (CapacityError for capacity failures) when the endpoint ends in any state other than InService, as Session.wait_for_endpoint does, instead of logging an error and returning a failed endpoint. - Wait on status only for inference-component based endpoints, which never get an endpoint log group. ResourceIterator stopped at the first empty page even when the response carried a NextToken. ListAIRecommendationJobs returns such pages, so AIRecommendationJob.get_all() and list_recommendations() always came back empty. Keep following NextToken past empty pages, and stop if a token repeats so a misbehaving list call cannot loop. Add an xfail_on_insufficient_capacity marker to the sagemaker-serve integ tests and apply it to the GPU tests that hit ICE: a test whose endpoint or job SageMaker fails for lack of capacity reports XFAIL instead of FAIL, while any other failure still fails the test.
1 parent 9a4a9a0 commit 02343e4

20 files changed

Lines changed: 631 additions & 93 deletions

‎sagemaker-core/src/sagemaker/core/helper/session_helper.py‎

Lines changed: 67 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1274,9 +1274,15 @@ def wait_for_endpoint(self, endpoint, poll=DEFAULT_EP_POLL, live_logging=False):
12741274
cloudwatch_client = self.boto_session.client("logs")
12751275
paginator = cloudwatch_client.get_paginator("filter_log_events")
12761276
paginator_config = create_paginator_config()
1277+
not_found_budget = _EndpointNotFoundBudget()
12771278
desc = _wait_until(
12781279
lambda: _live_logging_deploy_done(
1279-
self.sagemaker_client, endpoint, paginator, paginator_config, EP_LOGGER_POLL
1280+
self.sagemaker_client,
1281+
endpoint,
1282+
paginator,
1283+
paginator_config,
1284+
EP_LOGGER_POLL,
1285+
not_found_budget=not_found_budget,
12801286
),
12811287
poll=EP_LOGGER_POLL,
12821288
)
@@ -2979,30 +2985,73 @@ def _deploy_done(sagemaker_client, endpoint_name):
29792985
return None if status in in_progress_statuses else desc
29802986

29812987

2982-
def _live_logging_deploy_done(sagemaker_client, endpoint_name, paginator, paginator_config, poll):
2983-
"""Placeholder docstring"""
2988+
# DescribeEndpoint can briefly report "Could not find endpoint" right after
2989+
# CreateEndpoint returns. An endpoint still missing after this many consecutive
2990+
# polls was deleted out of band, and waiting on it can never succeed.
2991+
_MAX_ENDPOINT_NOT_FOUND_POLLS = 10
2992+
2993+
2994+
class _EndpointNotFoundBudget(object):
2995+
"""Bounds how many consecutive polls a deploy wait tolerates a missing endpoint."""
2996+
2997+
def __init__(self, max_polls=_MAX_ENDPOINT_NOT_FOUND_POLLS):
2998+
"""Allow up to ``max_polls`` consecutive "endpoint not found" polls."""
2999+
self.max_polls = max_polls
3000+
self.misses = 0
3001+
3002+
def miss(self, error):
3003+
"""Record one "not found" poll, re-raising ``error`` once the budget is spent."""
3004+
self.misses += 1
3005+
if self.misses > self.max_polls:
3006+
raise error
3007+
3008+
def reset(self):
3009+
"""Clear the count once the endpoint is visible."""
3010+
self.misses = 0
3011+
3012+
3013+
def _live_logging_deploy_done(
3014+
sagemaker_client, endpoint_name, paginator, paginator_config, poll, not_found_budget=None
3015+
):
3016+
"""Return the ``DescribeEndpoint`` response once the endpoint leaves ``Creating``.
3017+
3018+
Streams the endpoint's CloudWatch logs on every poll and returns ``None`` while
3019+
the endpoint is still being created. The log group only exists once a container
3020+
has started, so an endpoint that fails before any instance is provisioned (for
3021+
example on InsufficientInstanceCapacity) never gets one: a missing log group
3022+
must not keep a finished deployment waiting.
3023+
3024+
Args:
3025+
not_found_budget (_EndpointNotFoundBudget): Optional bound on how many
3026+
consecutive polls a missing endpoint is tolerated. Without it a
3027+
missing endpoint is waited on indefinitely.
3028+
"""
29843029
stop = False
29853030
endpoint_status = None
29863031
try:
29873032
desc = sagemaker_client.describe_endpoint(EndpointName=endpoint_name)
29883033
endpoint_status = desc["EndpointStatus"]
29893034
except ClientError as e:
29903035
if e.response["Error"]["Code"] == "ValidationException":
3036+
if not_found_budget is not None:
3037+
not_found_budget.miss(e)
29913038
LOGGER.debug("Waiting for endpoint to become visible")
29923039
return None
29933040
raise e
3041+
if not_found_budget is not None:
3042+
not_found_budget.reset()
29943043

2995-
try:
2996-
# if endpoint is in an invalid state -> set stop to true, sleep, and flush the logs
2997-
if endpoint_status != "Creating":
2998-
stop = True
2999-
if endpoint_status == "InService":
3000-
LOGGER.info(
3001-
"Created endpoint with name %s. Waiting for it to be InService", endpoint_name
3002-
)
3003-
else:
3004-
time.sleep(poll)
3044+
# if endpoint is in an invalid state -> set stop to true, sleep, and flush the logs
3045+
if endpoint_status != "Creating":
3046+
stop = True
3047+
if endpoint_status == "InService":
3048+
LOGGER.info(
3049+
"Created endpoint with name %s. Waiting for it to be InService", endpoint_name
3050+
)
3051+
else:
3052+
time.sleep(poll)
30053053

3054+
try:
30063055
pages = paginator.paginate(
30073056
logGroupName=f"/aws/sagemaker/Endpoints/{endpoint_name}",
30083057
logStreamNamePrefix="AllTraffic/",
@@ -3016,17 +3065,13 @@ def _live_logging_deploy_done(sagemaker_client, endpoint_name, paginator, pagina
30163065
LOGGER.info(event["message"])
30173066
else:
30183067
LOGGER.debug("No log events available")
3019-
3020-
# if stop is true -> return the describe response and stop polling
3021-
if stop:
3022-
return desc
30233068
except ClientError as e:
3024-
if e.response["Error"]["Code"] == "ResourceNotFoundException":
3025-
LOGGER.debug("Waiting for endpoint log group to appear")
3026-
return None
3027-
raise e
3069+
if e.response["Error"]["Code"] != "ResourceNotFoundException":
3070+
raise e
3071+
LOGGER.debug("Waiting for endpoint log group to appear")
30283072

3029-
return None
3073+
# if stop is true -> return the describe response and stop polling
3074+
return desc if stop else None
30303075

30313076

30323077
def _deployment_entity_exists(describe_fn):

‎sagemaker-core/src/sagemaker/core/utils/utils.py‎

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -505,22 +505,28 @@ def __next__(self) -> T:
505505

506506
# Otherwise, get the next page of summaries by calling the list method with the next token if available
507507
else:
508-
if self.next_token:
509-
response = getattr(self.client, self.list_method)(
510-
NextToken=self.next_token, **self.list_method_kwargs
511-
)
512-
else:
513-
response = getattr(self.client, self.list_method)(**self.list_method_kwargs)
514-
515-
self.summary_list = response.get(self.summaries_key, [])
516-
self.next_token = response.get("NextToken", None)
517-
self.index = 0
518-
519-
# If list_method returned an empty list, raise StopIteration
520-
if len(self.summary_list) == 0:
521-
raise StopIteration
522-
523-
return self.__next__()
508+
while True:
509+
previous_token = self.next_token
510+
if self.next_token:
511+
response = getattr(self.client, self.list_method)(
512+
NextToken=self.next_token, **self.list_method_kwargs
513+
)
514+
else:
515+
response = getattr(self.client, self.list_method)(**self.list_method_kwargs)
516+
517+
self.summary_list = response.get(self.summaries_key, [])
518+
self.next_token = response.get("NextToken", None)
519+
self.index = 0
520+
521+
if len(self.summary_list) > 0:
522+
return self.__next__()
523+
524+
# A page can be empty and still carry a NextToken (the service
525+
# filters results after paginating), so an empty page only ends
526+
# the iteration when there is no further page. A NextToken that
527+
# repeats would page forever, so it also ends the iteration.
528+
if not self.next_token or self.next_token == previous_token:
529+
raise StopIteration
524530

525531

526532
def serialize(value: Any) -> Any:

‎sagemaker-core/tests/unit/generated/test_utils.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,41 @@ def test_next_client_returns_empty_list(resource_iterator):
174174
next(iterator)
175175

176176

177+
def test_next_follows_next_token_past_empty_pages(resource_iterator):
178+
iterator, client, _ = resource_iterator
179+
client.list_training_jobs.side_effect = [
180+
{"TrainingJobSummaries": [], "NextToken": "token-1"},
181+
{"TrainingJobSummaries": [], "NextToken": "token-2"},
182+
LIST_TRAINING_JOB_RESPONSE_WITHOUT_NEXT_TOKEN,
183+
]
184+
185+
with patch.object(TrainingJob, "refresh"):
186+
names = [job.training_job_name for job in iterator]
187+
188+
assert names == [
189+
summary["TrainingJobName"]
190+
for summary in LIST_TRAINING_JOB_RESPONSE_WITHOUT_NEXT_TOKEN["TrainingJobSummaries"]
191+
]
192+
assert client.list_training_jobs.call_args_list == [
193+
call(),
194+
call(NextToken="token-1"),
195+
call(NextToken="token-2"),
196+
]
197+
198+
199+
def test_next_stops_when_empty_pages_repeat_the_next_token(resource_iterator):
200+
iterator, client, _ = resource_iterator
201+
client.list_training_jobs.side_effect = [
202+
{"TrainingJobSummaries": [], "NextToken": "token-1"},
203+
{"TrainingJobSummaries": [], "NextToken": "token-1"},
204+
]
205+
206+
with pytest.raises(StopIteration):
207+
next(iterator)
208+
209+
assert client.list_training_jobs.call_count == 2
210+
211+
177212
def test_next_without_next_token(resource_iterator):
178213
iterator, client, _ = resource_iterator
179214
client.list_training_jobs.return_value = LIST_TRAINING_JOB_RESPONSE_WITHOUT_NEXT_TOKEN

‎sagemaker-core/tests/unit/helper/test_session_helper.py‎

Lines changed: 119 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,12 @@
1919
from unittest.mock import Mock, patch
2020
from botocore.exceptions import ClientError
2121

22-
from sagemaker.core.helper.session_helper import Session
22+
from sagemaker.core.exceptions import UnexpectedStatusException
23+
from sagemaker.core.helper.session_helper import (
24+
Session,
25+
_EndpointNotFoundBudget,
26+
_live_logging_deploy_done,
27+
)
2328
from sagemaker.core.session_settings import SessionSettings
2429

2530

@@ -375,6 +380,119 @@ def test_wait_for_endpoint_failure(self, mock_boto_session, mock_sagemaker_clien
375380
with pytest.raises(Exception, match="Error hosting endpoint"):
376381
session.wait_for_endpoint("my-endpoint")
377382

383+
@patch("sagemaker.core.helper.session_helper._has_permission_for_live_logging")
384+
@patch("time.sleep")
385+
def test_live_logging_wait_raises_when_endpoint_fails_without_log_group(
386+
self, mock_sleep, mock_permission, mock_boto_session, mock_sagemaker_client
387+
):
388+
"""A Failed endpoint that never got a log group ends the wait instead of hanging."""
389+
mock_permission.return_value = True
390+
failed = {
391+
"EndpointStatus": "Failed",
392+
"FailureReason": "Unable to provision requested ML compute capacity due to "
393+
"InsufficientInstanceCapacity error.",
394+
}
395+
# A finite side_effect so a regression surfaces as a failure, not a hang.
396+
mock_sagemaker_client.describe_endpoint.side_effect = [failed, failed]
397+
paginator = mock_boto_session.client.return_value.get_paginator.return_value
398+
paginator.paginate.side_effect = ClientError(
399+
{"Error": {"Code": "ResourceNotFoundException"}}, "FilterLogEvents"
400+
)
401+
session = Session(boto_session=mock_boto_session, sagemaker_client=mock_sagemaker_client)
402+
403+
with pytest.raises(UnexpectedStatusException, match="InsufficientInstanceCapacity"):
404+
session.wait_for_endpoint("my-endpoint", live_logging=True)
405+
406+
assert mock_sagemaker_client.describe_endpoint.call_count == 1
407+
408+
409+
class TestLiveLoggingDeployDone:
410+
"""Test _live_logging_deploy_done."""
411+
412+
RESOURCE_NOT_FOUND = ClientError(
413+
{"Error": {"Code": "ResourceNotFoundException"}}, "FilterLogEvents"
414+
)
415+
ENDPOINT_NOT_FOUND = ClientError(
416+
{"Error": {"Code": "ValidationException", "Message": "Could not find endpoint"}},
417+
"DescribeEndpoint",
418+
)
419+
420+
@patch("time.sleep")
421+
def test_finished_endpoint_without_log_group_returns_desc(self, mock_sleep):
422+
"""InService or Failed with no log group is finished, not "still waiting"."""
423+
for status in ("InService", "Failed"):
424+
client = Mock()
425+
desc = {"EndpointStatus": status}
426+
client.describe_endpoint.return_value = desc
427+
paginator = Mock()
428+
paginator.paginate.side_effect = self.RESOURCE_NOT_FOUND
429+
430+
assert _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) == desc
431+
432+
def test_creating_endpoint_without_log_group_keeps_waiting(self):
433+
"""Creating with no log group yet is still in progress."""
434+
client = Mock()
435+
client.describe_endpoint.return_value = {"EndpointStatus": "Creating"}
436+
paginator = Mock()
437+
paginator.paginate.side_effect = self.RESOURCE_NOT_FOUND
438+
439+
assert _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) is None
440+
441+
def test_other_log_errors_are_raised(self):
442+
"""Log-fetch errors other than a missing log group still propagate."""
443+
client = Mock()
444+
client.describe_endpoint.return_value = {"EndpointStatus": "InService"}
445+
paginator = Mock()
446+
paginator.paginate.side_effect = ClientError(
447+
{"Error": {"Code": "ThrottlingException"}}, "FilterLogEvents"
448+
)
449+
450+
with pytest.raises(ClientError):
451+
_live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5)
452+
453+
def test_missing_endpoint_is_waited_on_without_budget(self):
454+
"""Without a budget a missing endpoint keeps the legacy "keep waiting" result."""
455+
client = Mock()
456+
client.describe_endpoint.side_effect = self.ENDPOINT_NOT_FOUND
457+
458+
assert _live_logging_deploy_done(client, "my-endpoint", Mock(), {}, 5) is None
459+
460+
def test_missing_endpoint_raises_once_budget_is_spent(self):
461+
"""A missing endpoint is tolerated only for the budgeted number of polls."""
462+
client = Mock()
463+
client.describe_endpoint.side_effect = self.ENDPOINT_NOT_FOUND
464+
budget = _EndpointNotFoundBudget(max_polls=2)
465+
466+
for _ in range(2):
467+
assert (
468+
_live_logging_deploy_done(
469+
client, "my-endpoint", Mock(), {}, 5, not_found_budget=budget
470+
)
471+
is None
472+
)
473+
with pytest.raises(ClientError, match="Could not find endpoint"):
474+
_live_logging_deploy_done(client, "my-endpoint", Mock(), {}, 5, not_found_budget=budget)
475+
476+
def test_budget_resets_once_endpoint_is_found(self):
477+
"""Only consecutive "not found" polls count against the budget."""
478+
client = Mock()
479+
client.describe_endpoint.side_effect = [
480+
self.ENDPOINT_NOT_FOUND,
481+
{"EndpointStatus": "Creating"},
482+
self.ENDPOINT_NOT_FOUND,
483+
]
484+
paginator = Mock()
485+
paginator.paginate.return_value = []
486+
budget = _EndpointNotFoundBudget(max_polls=1)
487+
488+
for _ in range(3):
489+
assert (
490+
_live_logging_deploy_done(
491+
client, "my-endpoint", paginator, {}, 5, not_found_budget=budget
492+
)
493+
is None
494+
)
495+
378496

379497
class TestUpdateEndpoint:
380498
"""Test update_endpoint method."""

0 commit comments

Comments
 (0)