diff --git a/sagemaker-core/src/sagemaker/core/helper/session_helper.py b/sagemaker-core/src/sagemaker/core/helper/session_helper.py index d3c7096030..e8eab15519 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -1274,9 +1274,15 @@ def wait_for_endpoint(self, endpoint, poll=DEFAULT_EP_POLL, live_logging=False): cloudwatch_client = self.boto_session.client("logs") paginator = cloudwatch_client.get_paginator("filter_log_events") paginator_config = create_paginator_config() + not_found_budget = _EndpointNotFoundBudget() desc = _wait_until( lambda: _live_logging_deploy_done( - self.sagemaker_client, endpoint, paginator, paginator_config, EP_LOGGER_POLL + self.sagemaker_client, + endpoint, + paginator, + paginator_config, + EP_LOGGER_POLL, + not_found_budget=not_found_budget, ), poll=EP_LOGGER_POLL, ) @@ -2979,8 +2985,49 @@ def _deploy_done(sagemaker_client, endpoint_name): return None if status in in_progress_statuses else desc -def _live_logging_deploy_done(sagemaker_client, endpoint_name, paginator, paginator_config, poll): - """Placeholder docstring""" +# DescribeEndpoint can briefly report "Could not find endpoint" right after +# CreateEndpoint returns. An endpoint still missing after this many consecutive +# polls was deleted out of band, and waiting on it can never succeed. +_MAX_ENDPOINT_NOT_FOUND_POLLS = 10 + + +class _EndpointNotFoundBudget(object): + """Bounds how many consecutive polls a deploy wait tolerates a missing endpoint.""" + + def __init__(self, max_polls=_MAX_ENDPOINT_NOT_FOUND_POLLS): + """Allow up to ``max_polls`` consecutive "endpoint not found" polls.""" + self.max_polls = max_polls + self.misses = 0 + + def miss(self, error): + """Record one "not found" poll, re-raising ``error`` once the budget is spent.""" + self.misses += 1 + if self.misses > self.max_polls: + raise error + + def reset(self): + """Clear the count once the endpoint is visible.""" + self.misses = 0 + + +def _live_logging_deploy_done( + sagemaker_client, endpoint_name, paginator, paginator_config, poll, not_found_budget=None +): + """Return the ``DescribeEndpoint`` response once the endpoint is no longer in progress. + + Streams the endpoint's CloudWatch logs on every poll and returns ``None`` while + the endpoint is still ``Creating`` or ``Updating``, the same in-progress statuses + as ``_deploy_done``. The log group only exists once a container + has started, so an endpoint that fails before any instance is provisioned (for + example on InsufficientInstanceCapacity) never gets one: a missing log group + must not keep a finished deployment waiting. + + Args: + not_found_budget (_EndpointNotFoundBudget): Optional bound on how many + consecutive polls a missing endpoint is tolerated. Without it a + missing endpoint is waited on indefinitely. + """ + in_progress_statuses = ["Creating", "Updating"] stop = False endpoint_status = None try: @@ -2988,21 +3035,25 @@ def _live_logging_deploy_done(sagemaker_client, endpoint_name, paginator, pagina endpoint_status = desc["EndpointStatus"] except ClientError as e: if e.response["Error"]["Code"] == "ValidationException": + if not_found_budget is not None: + not_found_budget.miss(e) LOGGER.debug("Waiting for endpoint to become visible") return None raise e + if not_found_budget is not None: + not_found_budget.reset() - try: - # if endpoint is in an invalid state -> set stop to true, sleep, and flush the logs - if endpoint_status != "Creating": - stop = True - if endpoint_status == "InService": - LOGGER.info( - "Created endpoint with name %s. Waiting for it to be InService", endpoint_name - ) - else: - time.sleep(poll) + # if endpoint is in an invalid state -> set stop to true, sleep, and flush the logs + if endpoint_status not in in_progress_statuses: + stop = True + if endpoint_status == "InService": + LOGGER.info( + "Created endpoint with name %s. Waiting for it to be InService", endpoint_name + ) + else: + time.sleep(poll) + try: pages = paginator.paginate( logGroupName=f"/aws/sagemaker/Endpoints/{endpoint_name}", logStreamNamePrefix="AllTraffic/", @@ -3016,17 +3067,13 @@ def _live_logging_deploy_done(sagemaker_client, endpoint_name, paginator, pagina LOGGER.info(event["message"]) else: LOGGER.debug("No log events available") - - # if stop is true -> return the describe response and stop polling - if stop: - return desc except ClientError as e: - if e.response["Error"]["Code"] == "ResourceNotFoundException": - LOGGER.debug("Waiting for endpoint log group to appear") - return None - raise e + if e.response["Error"]["Code"] != "ResourceNotFoundException": + raise e + LOGGER.debug("Waiting for endpoint log group to appear") - return None + # if stop is true -> return the describe response and stop polling + return desc if stop else None def _deployment_entity_exists(describe_fn): diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index d807523935..9ebfe7748e 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -505,22 +505,29 @@ def __next__(self) -> T: # Otherwise, get the next page of summaries by calling the list method with the next token if available else: - if self.next_token: - response = getattr(self.client, self.list_method)( - NextToken=self.next_token, **self.list_method_kwargs - ) - else: - response = getattr(self.client, self.list_method)(**self.list_method_kwargs) - - self.summary_list = response.get(self.summaries_key, []) - self.next_token = response.get("NextToken", None) - self.index = 0 - - # If list_method returned an empty list, raise StopIteration - if len(self.summary_list) == 0: - raise StopIteration - - return self.__next__() + while True: + previous_token = self.next_token + if self.next_token: + response = getattr(self.client, self.list_method)( + NextToken=self.next_token, **self.list_method_kwargs + ) + else: + response = getattr(self.client, self.list_method)(**self.list_method_kwargs) + + self.summary_list = response.get(self.summaries_key, []) + self.next_token = response.get("NextToken", None) + self.index = 0 + + if len(self.summary_list) > 0: + return self.__next__() + + # A page can be empty and still carry a NextToken (the service + # filters results after paginating), so an empty page only ends + # the iteration when there is no further page. A NextToken that + # comes back unchanged would page forever, so it also ends the + # iteration (the same consecutive-repeat check botocore paginators do). + if not self.next_token or self.next_token == previous_token: + raise StopIteration def serialize(value: Any) -> Any: diff --git a/sagemaker-core/tests/unit/generated/test_utils.py b/sagemaker-core/tests/unit/generated/test_utils.py index 20a83d3306..4b4847eba9 100644 --- a/sagemaker-core/tests/unit/generated/test_utils.py +++ b/sagemaker-core/tests/unit/generated/test_utils.py @@ -174,6 +174,41 @@ def test_next_client_returns_empty_list(resource_iterator): next(iterator) +def test_next_follows_next_token_past_empty_pages(resource_iterator): + iterator, client, _ = resource_iterator + client.list_training_jobs.side_effect = [ + {"TrainingJobSummaries": [], "NextToken": "token-1"}, + {"TrainingJobSummaries": [], "NextToken": "token-2"}, + LIST_TRAINING_JOB_RESPONSE_WITHOUT_NEXT_TOKEN, + ] + + with patch.object(TrainingJob, "refresh"): + names = [job.training_job_name for job in iterator] + + assert names == [ + summary["TrainingJobName"] + for summary in LIST_TRAINING_JOB_RESPONSE_WITHOUT_NEXT_TOKEN["TrainingJobSummaries"] + ] + assert client.list_training_jobs.call_args_list == [ + call(), + call(NextToken="token-1"), + call(NextToken="token-2"), + ] + + +def test_next_stops_when_empty_pages_repeat_the_next_token(resource_iterator): + iterator, client, _ = resource_iterator + client.list_training_jobs.side_effect = [ + {"TrainingJobSummaries": [], "NextToken": "token-1"}, + {"TrainingJobSummaries": [], "NextToken": "token-1"}, + ] + + with pytest.raises(StopIteration): + next(iterator) + + assert client.list_training_jobs.call_count == 2 + + def test_next_without_next_token(resource_iterator): iterator, client, _ = resource_iterator client.list_training_jobs.return_value = LIST_TRAINING_JOB_RESPONSE_WITHOUT_NEXT_TOKEN diff --git a/sagemaker-core/tests/unit/helper/test_session_helper.py b/sagemaker-core/tests/unit/helper/test_session_helper.py index e95f251285..ce3dfedff7 100644 --- a/sagemaker-core/tests/unit/helper/test_session_helper.py +++ b/sagemaker-core/tests/unit/helper/test_session_helper.py @@ -19,7 +19,12 @@ from unittest.mock import Mock, patch from botocore.exceptions import ClientError -from sagemaker.core.helper.session_helper import Session +from sagemaker.core.exceptions import UnexpectedStatusException +from sagemaker.core.helper.session_helper import ( + Session, + _EndpointNotFoundBudget, + _live_logging_deploy_done, +) from sagemaker.core.session_settings import SessionSettings @@ -375,6 +380,124 @@ def test_wait_for_endpoint_failure(self, mock_boto_session, mock_sagemaker_clien with pytest.raises(Exception, match="Error hosting endpoint"): session.wait_for_endpoint("my-endpoint") + @patch("sagemaker.core.helper.session_helper._has_permission_for_live_logging") + @patch("time.sleep") + def test_live_logging_wait_raises_when_endpoint_fails_without_log_group( + self, mock_sleep, mock_permission, mock_boto_session, mock_sagemaker_client + ): + """A Failed endpoint that never got a log group ends the wait instead of hanging.""" + mock_permission.return_value = True + failed = { + "EndpointStatus": "Failed", + "FailureReason": "Unable to provision requested ML compute capacity due to " + "InsufficientInstanceCapacity error.", + } + # A finite side_effect so a regression surfaces as a failure, not a hang. + mock_sagemaker_client.describe_endpoint.side_effect = [failed, failed] + paginator = mock_boto_session.client.return_value.get_paginator.return_value + paginator.paginate.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException"}}, "FilterLogEvents" + ) + session = Session(boto_session=mock_boto_session, sagemaker_client=mock_sagemaker_client) + + with pytest.raises(UnexpectedStatusException, match="InsufficientInstanceCapacity"): + session.wait_for_endpoint("my-endpoint", live_logging=True) + + assert mock_sagemaker_client.describe_endpoint.call_count == 1 + + +class TestLiveLoggingDeployDone: + """Test _live_logging_deploy_done.""" + + RESOURCE_NOT_FOUND = ClientError( + {"Error": {"Code": "ResourceNotFoundException"}}, "FilterLogEvents" + ) + ENDPOINT_NOT_FOUND = ClientError( + {"Error": {"Code": "ValidationException", "Message": "Could not find endpoint"}}, + "DescribeEndpoint", + ) + + @patch("time.sleep") + def test_finished_endpoint_without_log_group_returns_desc(self, mock_sleep): + """InService or Failed with no log group is finished, not "still waiting".""" + for status in ("InService", "Failed"): + client = Mock() + desc = {"EndpointStatus": status} + client.describe_endpoint.return_value = desc + paginator = Mock() + paginator.paginate.side_effect = self.RESOURCE_NOT_FOUND + + assert _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) == desc + + @patch("time.sleep") + def test_in_progress_endpoint_keeps_waiting(self, mock_sleep): + """Creating or Updating is still in progress, with or without a log group.""" + for status in ("Creating", "Updating"): + for log_error in (self.RESOURCE_NOT_FOUND, None): + client = Mock() + client.describe_endpoint.return_value = {"EndpointStatus": status} + paginator = Mock() + paginator.paginate.side_effect = log_error + paginator.paginate.return_value = [] + + assert _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) is None + mock_sleep.assert_not_called() + + def test_other_log_errors_are_raised(self): + """Log-fetch errors other than a missing log group still propagate.""" + client = Mock() + client.describe_endpoint.return_value = {"EndpointStatus": "InService"} + paginator = Mock() + paginator.paginate.side_effect = ClientError( + {"Error": {"Code": "ThrottlingException"}}, "FilterLogEvents" + ) + + with pytest.raises(ClientError): + _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) + + def test_missing_endpoint_is_waited_on_without_budget(self): + """Without a budget a missing endpoint keeps the legacy "keep waiting" result.""" + client = Mock() + client.describe_endpoint.side_effect = self.ENDPOINT_NOT_FOUND + + assert _live_logging_deploy_done(client, "my-endpoint", Mock(), {}, 5) is None + + def test_missing_endpoint_raises_once_budget_is_spent(self): + """A missing endpoint is tolerated only for the budgeted number of polls.""" + client = Mock() + client.describe_endpoint.side_effect = self.ENDPOINT_NOT_FOUND + budget = _EndpointNotFoundBudget(max_polls=2) + + for _ in range(2): + assert ( + _live_logging_deploy_done( + client, "my-endpoint", Mock(), {}, 5, not_found_budget=budget + ) + is None + ) + with pytest.raises(ClientError, match="Could not find endpoint"): + _live_logging_deploy_done(client, "my-endpoint", Mock(), {}, 5, not_found_budget=budget) + + def test_budget_resets_once_endpoint_is_found(self): + """Only consecutive "not found" polls count against the budget.""" + client = Mock() + client.describe_endpoint.side_effect = [ + self.ENDPOINT_NOT_FOUND, + {"EndpointStatus": "Creating"}, + self.ENDPOINT_NOT_FOUND, + ] + paginator = Mock() + paginator.paginate.return_value = [] + budget = _EndpointNotFoundBudget(max_polls=1) + + for _ in range(3): + assert ( + _live_logging_deploy_done( + client, "my-endpoint", paginator, {}, 5, not_found_budget=budget + ) + is None + ) + class TestUpdateEndpoint: """Test update_endpoint method.""" diff --git a/sagemaker-serve/src/sagemaker/serve/deployment_progress.py b/sagemaker-serve/src/sagemaker/serve/deployment_progress.py index 4178d8a852..70d42a36bc 100644 --- a/sagemaker-serve/src/sagemaker/serve/deployment_progress.py +++ b/sagemaker-serve/src/sagemaker/serve/deployment_progress.py @@ -74,31 +74,49 @@ def _deploy_done_with_progress(sagemaker_client, endpoint_name, progress_tracker def _live_logging_deploy_done_with_progress( - sagemaker_client, endpoint_name, paginator, paginator_config, poll, progress_tracker=None + sagemaker_client, + endpoint_name, + paginator, + paginator_config, + poll, + progress_tracker=None, + not_found_budget=None, ): - """Live logging deployment checker that routes logs to Rich progress tracker""" + """Live logging deployment checker that routes logs to Rich progress tracker. + + Returns the ``DescribeEndpoint`` response once the endpoint is no longer + ``Creating``/``Updating``, and ``None`` while it still is. The endpoint log + group only exists once a container has started, so an endpoint that fails + before any instance is provisioned (for example on InsufficientInstanceCapacity) + never gets one: a missing log group must not keep a finished deployment waiting. + ``not_found_budget`` optionally bounds how many consecutive polls a missing + endpoint is tolerated. + """ import time from botocore.exceptions import ClientError - stop = False - endpoint_status = None + in_progress_statuses = ["Creating", "Updating"] try: desc = sagemaker_client.describe_endpoint(EndpointName=endpoint_name) endpoint_status = desc["EndpointStatus"] except ClientError as e: if e.response["Error"]["Code"] == "ValidationException": + if not_found_budget is not None: + not_found_budget.miss(e) return None raise e + if not_found_budget is not None: + not_found_budget.reset() - try: - # Update status and check if we should stop - if endpoint_status != "Creating": - stop = True - if endpoint_status == "InService" and progress_tracker: - progress_tracker.log(f"✅ Created endpoint with name {endpoint_name}") - elif endpoint_status != "InService": - time.sleep(poll) + # Update status and check if we should stop + stop = endpoint_status not in in_progress_statuses + if stop: + if endpoint_status == "InService" and progress_tracker: + progress_tracker.log(f"Created endpoint with name {endpoint_name}") + elif endpoint_status != "InService": + time.sleep(poll) + try: # Fetch and route CloudWatch logs to progress tracker pages = paginator.paginate( logGroupName=f"/aws/sagemaker/Endpoints/{endpoint_name}", @@ -112,17 +130,13 @@ def _live_logging_deploy_done_with_progress( for event in page["events"]: if progress_tracker: progress_tracker.log(event["message"]) - - # Update progress tracker status - if progress_tracker: - progress_tracker.update_status(endpoint_status) - - # Return desc if we should stop polling - if stop: - return desc except ClientError as e: - if e.response["Error"]["Code"] == "ResourceNotFoundException": - return None - raise e + if e.response["Error"]["Code"] != "ResourceNotFoundException": + raise e + + # Update progress tracker status + if progress_tracker: + progress_tracker.update_status(endpoint_status) - return None + # Return desc if we should stop polling + return desc if stop else None diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 4195ec862e..10a29fa5f8 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -58,6 +58,7 @@ ModelPackageModelCard, ) from sagemaker.core.utils.utils import logger +from sagemaker.core.exceptions import CapacityError, UnexpectedStatusException from sagemaker.core.helper import session_helper from sagemaker.core.helper.iam_role_resolver import resolve_and_validate_role from sagemaker.core.helper.session_helper import ( @@ -3767,9 +3768,27 @@ def _deploy_local_endpoint(self, **kwargs): return LocalEndpoint.get(endpoint_name=endpoint_name, local_session=local_session) def _wait_for_endpoint( - self, endpoint, poll=30, live_logging=False, show_progress=True, wait=True + self, + endpoint, + poll=30, + live_logging=False, + show_progress=True, + wait=True, + stream_endpoint_logs=True, ): - """Enhanced wait with rich progress bar and status logging""" + """Wait for an endpoint deployment to finish, with a rich progress bar. + + Args: + stream_endpoint_logs (bool): Stream the endpoint's CloudWatch logs while + waiting (default: True). Inference-component based endpoints host no + model themselves and never get an endpoint log group, so they wait on + status only. + + Raises: + CapacityError: If the endpoint fails with a CapacityError. + UnexpectedStatusException: If the endpoint finishes in any state other + than ``InService``. + """ if not wait: logger.info( "🚀 Deployment started: Endpoint '%s' using %s in %s mode (deployment in progress)", @@ -3793,16 +3812,20 @@ def _wait_for_endpoint( # Check if we have permission for live logging from sagemaker.core.helper.session_helper import _has_permission_for_live_logging - if _has_permission_for_live_logging(self.sagemaker_session.boto_session, endpoint): + if stream_endpoint_logs and _has_permission_for_live_logging( + self.sagemaker_session.boto_session, endpoint + ): # Use live logging with Rich progress tracker cloudwatch_client = self.sagemaker_session.boto_session.client("logs") paginator = cloudwatch_client.get_paginator("filter_log_events") from sagemaker.core.helper.session_helper import ( create_paginator_config, EP_LOGGER_POLL, + _EndpointNotFoundBudget, ) paginator_config = create_paginator_config() + not_found_budget = _EndpointNotFoundBudget() desc = _wait_until( lambda: _live_logging_deploy_done_with_progress( sagemaker_client, @@ -3811,6 +3834,7 @@ def _wait_for_endpoint( paginator_config, EP_LOGGER_POLL, progress, + not_found_budget=not_found_budget, ), poll=EP_LOGGER_POLL, ) @@ -3824,29 +3848,37 @@ def _wait_for_endpoint( # Existing implementation desc = _wait_until(lambda: _deploy_done(sagemaker_client, endpoint), poll) - # Check final endpoint status and log accordingly - try: - endpoint_desc = sagemaker_client.describe_endpoint(EndpointName=endpoint) - endpoint_status = endpoint_desc["EndpointStatus"] - if endpoint_status == "InService": - endpoint_arn_info = ( - f" (ARN: {endpoint_desc['EndpointArn']})" - if self.mode == Mode.SAGEMAKER_ENDPOINT - else "" - ) - logger.info( - "✅ Deployment successful: Endpoint '%s' using %s in %s mode%s", - endpoint, - self.model_server, - self.mode, - endpoint_arn_info, - ) - else: - logger.error( - "❌ Deployment failed: Endpoint '%s' status is '%s'", endpoint, endpoint_status - ) - except Exception as e: - logger.error("❌ Deployment failed: Unable to verify endpoint status - %s", str(e)) + # _wait_until returns the DescribeEndpoint response of the final status. + endpoint_status = desc.get("EndpointStatus") + if endpoint_status != "InService": + reason = desc.get("FailureReason") + logger.error( + "Deployment failed: Endpoint '%s' status is '%s'. Reason: %s", + endpoint, + endpoint_status, + reason, + ) + error_class = ( + CapacityError if "CapacityError" in str(reason) else UnexpectedStatusException + ) + raise error_class( + message=f"Error hosting endpoint {endpoint}: {endpoint_status}. Reason: {reason}.", + allowed_statuses=["InService"], + actual_status=endpoint_status, + ) + + endpoint_arn_info = ( + f" (ARN: {desc['EndpointArn']})" + if self.mode == Mode.SAGEMAKER_ENDPOINT and desc.get("EndpointArn") + else "" + ) + logger.info( + "Deployment successful: Endpoint '%s' using %s in %s mode%s", + endpoint, + self.model_server, + self.mode, + endpoint_arn_info, + ) return desc @@ -4110,7 +4142,14 @@ def _deploy_core_endpoint(self, **kwargs): live_logging=False, # TODO: enable when IC supports this wait=False, ) - self._wait_for_endpoint(endpoint=self.endpoint_name, show_progress=True, wait=wait) + # An IC-based endpoint hosts no model itself, so it never gets an + # endpoint log group to stream: wait on its status only. + self._wait_for_endpoint( + endpoint=self.endpoint_name, + show_progress=True, + wait=wait, + stream_endpoint_logs=False, + ) core_endpoint = Endpoint.get( endpoint_name=self.endpoint_name, @@ -4151,7 +4190,12 @@ def _deploy_core_endpoint(self, **kwargs): tags=tags, wait=False, ) - self._wait_for_endpoint(endpoint=self.endpoint_name, show_progress=True, wait=wait) + self._wait_for_endpoint( + endpoint=self.endpoint_name, + show_progress=True, + wait=wait, + stream_endpoint_logs=False, + ) return core_endpoint @@ -6060,6 +6104,12 @@ def deploy( resource representing the deployed endpoint, a ``LocalEndpoint`` for local mode, or a ``Transformer`` for batch transform inference. + Raises: + UnexpectedStatusException: If ``wait`` is True and the endpoint finishes in a + state other than ``InService`` (``CapacityError`` for capacity failures). + Model customization and recommendation deployments raise + ``FailedStatusError`` instead. + Example: >>> model_builder = ModelBuilder(model=my_model, role_arn=role, instance_type="ml.m5.xlarge") >>> model = model_builder.build() # Creates Model resource diff --git a/sagemaker-serve/tests/integ/conftest.py b/sagemaker-serve/tests/integ/conftest.py index 8c237eb059..7bdfe4ee74 100644 --- a/sagemaker-serve/tests/integ/conftest.py +++ b/sagemaker-serve/tests/integ/conftest.py @@ -34,11 +34,21 @@ fail the test loudly (rather than being converted to a skip), so a persistent rate-limit regression stays visible instead of silently disappearing from the results. + +GPU capacity is a separate environmental failure: in us-west-2 SageMaker +regularly cannot provision the requested GPU instances and fails the endpoint / +job with an InsufficientInstanceCapacity reason after tens of minutes. Tests +marked ``xfail_on_insufficient_capacity`` report XFAIL instead of FAIL when they +fail for exactly that reason. The failure reason set by SageMaker is the signal: +the test still ends when SageMaker gives up (no client-side timeout that could +misclassify a slow but healthy deployment), and any other failure still fails +the test. """ from __future__ import absolute_import import os +import re import pytest @@ -66,3 +76,54 @@ def _configure_boto_adaptive_retries(): os.environ.pop(key, None) else: os.environ[key] = value + + +# Failure reasons SageMaker reports when it cannot provision the requested +# instance capacity, as seen in this account: +# Endpoint: "Unable to provision requested ML compute capacity due to +# InsufficientInstanceCapacity error ..." +# AIRecommendationJob: "... Could not deploy an endpoint for instance type +# 'ml.g5.2xlarge': all reservation and on-demand capacity +# attempts were exhausted." +# OptimizationJob: "EC2InsufficientCapacityException: ... not available due +# to insufficient capacity." +_INSUFFICIENT_CAPACITY_REASON = re.compile( + r"insufficient\s*(?:instance\s*)?capacity|capacity attempts were exhausted", + re.IGNORECASE, +) +_XFAIL_REASON_MAX_CHARS = 500 + + +def pytest_configure(config): + """Register the ``xfail_on_insufficient_capacity`` marker.""" + config.addinivalue_line( + "markers", + "xfail_on_insufficient_capacity: report the test as XFAIL instead of FAIL when it " + "fails because SageMaker could not provision the requested instance capacity.", + ) + + +def _is_insufficient_capacity(exc): + """True if ``exc``, or an exception it was raised from, reports a capacity shortage.""" + seen = set() + while exc is not None and id(exc) not in seen: + seen.add(id(exc)) + if _INSUFFICIENT_CAPACITY_REASON.search(str(exc)): + return True + exc = exc.__cause__ or exc.__context__ + return False + + +@pytest.hookimpl(wrapper=True) +def pytest_runtest_call(item): + """XFAIL a marked test whose failure is an insufficient-capacity error.""" + try: + return (yield) + except Exception as exc: + marked = item.get_closest_marker("xfail_on_insufficient_capacity") is not None + if marked and _is_insufficient_capacity(exc): + pytest.xfail( + "SageMaker could not provision the requested instance capacity: " + f"{str(exc)[:_XFAIL_REASON_MAX_CHARS]}" + ) + raise diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py index d8215c811c..2385581983 100644 --- a/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py @@ -96,6 +96,7 @@ def test_list_benchmarks_and_recommendations_plumbing(): @pytest.mark.gpu_intensive +@pytest.mark.xfail_on_insufficient_capacity def test_recommendation_deploy_best_and_compare_e2e(): """Full flow across all three enhancements, sharing one rec job + endpoint: diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py index 2e4d5f513c..ee2b4d4f59 100644 --- a/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py @@ -71,6 +71,7 @@ def _build_jumpstart_model_builder(role_arn): @pytest.mark.slow_test +@pytest.mark.xfail_on_insufficient_capacity def test_benchmark_workflow_end_to_end(): """Deploy a JumpStart endpoint, run a benchmark against it, parse the result.""" logger.info("Starting AI inference recommender benchmark integration test...") @@ -127,6 +128,7 @@ def test_benchmark_workflow_end_to_end(): @pytest.mark.gpu_intensive +@pytest.mark.xfail_on_insufficient_capacity def test_recommendation_workflow_end_to_end(): """Run an AI recommendation via generate_deployment_recommendations and deploy the top recommendation.""" logger.info("Starting AI inference recommender recommendation integration test...") diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py index 1582c718d4..df3a1d6c21 100644 --- a/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_sdkt_ic_integration.py @@ -64,6 +64,7 @@ def _src(uri): @pytest.mark.gpu_intensive +@pytest.mark.xfail_on_insufficient_capacity def test_deploy_sdkt_model_as_inference_component(): """A model carrying SD/KT AdditionalModelDataSources deploys as an Inference Component and reaches InService via ``ModelBuilder.deploy``.""" diff --git a/sagemaker-serve/tests/integ/test_huggingface_integration.py b/sagemaker-serve/tests/integ/test_huggingface_integration.py index acb64d266e..bcdd616927 100644 --- a/sagemaker-serve/tests/integ/test_huggingface_integration.py +++ b/sagemaker-serve/tests/integ/test_huggingface_integration.py @@ -31,6 +31,7 @@ @pytest.mark.slow_test +@pytest.mark.xfail_on_insufficient_capacity def test_huggingface_build_deploy_invoke_cleanup(): """Integration test for HuggingFace model build, deploy, invoke, and cleanup workflow""" logger.info("Starting HuggingFace integration test...") diff --git a/sagemaker-serve/tests/integ/test_jumpstart_integration.py b/sagemaker-serve/tests/integ/test_jumpstart_integration.py index 3fdf21e52b..49ad53a027 100644 --- a/sagemaker-serve/tests/integ/test_jumpstart_integration.py +++ b/sagemaker-serve/tests/integ/test_jumpstart_integration.py @@ -33,6 +33,7 @@ @pytest.mark.gpu_intensive +@pytest.mark.xfail_on_insufficient_capacity def test_jumpstart_build_deploy_invoke_cleanup(): """Integration test for JumpStart model build, deploy, invoke, and cleanup workflow""" logger.info("Starting JumpStart integration test...") diff --git a/sagemaker-serve/tests/integ/test_model_customization_deployment.py b/sagemaker-serve/tests/integ/test_model_customization_deployment.py index f8c5b02a2e..3aaedd4999 100644 --- a/sagemaker-serve/tests/integ/test_model_customization_deployment.py +++ b/sagemaker-serve/tests/integ/test_model_customization_deployment.py @@ -326,6 +326,7 @@ def test_build_from_model_package(self, model_package_arn, sagemaker_session): assert model.model_arn is not None @pytest.mark.slow_test + @pytest.mark.xfail_on_insufficient_capacity def test_deploy_from_model_package( self, model_package_arn, cleanup_endpoints, sagemaker_session ): diff --git a/sagemaker-serve/tests/integ/test_optimize_integration.py b/sagemaker-serve/tests/integ/test_optimize_integration.py index bd40616b3e..8242e9bb49 100644 --- a/sagemaker-serve/tests/integ/test_optimize_integration.py +++ b/sagemaker-serve/tests/integ/test_optimize_integration.py @@ -40,6 +40,7 @@ @pytest.mark.slow_test +@pytest.mark.xfail_on_insufficient_capacity def test_optimize_build_deploy_invoke_cleanup(): """Integration test for Optimize workflow""" logger.info("Starting Optimize integration test...") diff --git a/sagemaker-serve/tests/integ/test_tei_integration.py b/sagemaker-serve/tests/integ/test_tei_integration.py index 9c85b2f02c..33a920849e 100644 --- a/sagemaker-serve/tests/integ/test_tei_integration.py +++ b/sagemaker-serve/tests/integ/test_tei_integration.py @@ -31,6 +31,7 @@ @pytest.mark.slow_test +@pytest.mark.xfail_on_insufficient_capacity def test_tei_build_deploy_invoke_cleanup(): """Integration test for TEI model build, deploy, invoke, and cleanup workflow""" logger.info("Starting TEI integration test...") diff --git a/sagemaker-serve/tests/integ/test_tgi_integration.py b/sagemaker-serve/tests/integ/test_tgi_integration.py index b5774a96c0..eac40f48f8 100644 --- a/sagemaker-serve/tests/integ/test_tgi_integration.py +++ b/sagemaker-serve/tests/integ/test_tgi_integration.py @@ -31,6 +31,7 @@ @pytest.mark.slow_test +@pytest.mark.xfail_on_insufficient_capacity def test_tgi_build_deploy_invoke_cleanup(): """Integration test for TGI model build, deploy, invoke, and cleanup workflow""" logger.info("Starting TGI integration test...") diff --git a/sagemaker-serve/tests/unit/test_deployment_progress.py b/sagemaker-serve/tests/unit/test_deployment_progress.py index 54f6389c66..bd41bcb79d 100644 --- a/sagemaker-serve/tests/unit/test_deployment_progress.py +++ b/sagemaker-serve/tests/unit/test_deployment_progress.py @@ -3,6 +3,7 @@ import unittest from unittest.mock import Mock, patch from botocore.exceptions import ClientError +from sagemaker.core.helper.session_helper import _EndpointNotFoundBudget from sagemaker.serve.deployment_progress import ( EndpointDeploymentProgress, _deploy_done_with_progress, @@ -197,6 +198,93 @@ def test_pagination_with_next_token(self): self.assertEqual(paginator_config.get("StartingToken"), "token123") + def test_inservice_endpoint_without_log_group_returns_desc(self): + """An InService endpoint with no log group (IC-based endpoint) is finished.""" + mock_client = Mock() + expected_desc = {"EndpointStatus": "InService"} + mock_client.describe_endpoint.return_value = expected_desc + mock_paginator = Mock() + mock_paginator.paginate.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException"}}, "FilterLogEvents" + ) + mock_tracker = Mock() + + result = _live_logging_deploy_done_with_progress( + mock_client, "test-endpoint", mock_paginator, {}, 5, mock_tracker + ) + + self.assertEqual(result, expected_desc) + mock_tracker.update_status.assert_called_once_with("InService") + + @patch("time.sleep") + def test_failed_endpoint_without_log_group_returns_desc(self, mock_sleep): + """An endpoint that failed before any container started (no log group) is finished.""" + mock_client = Mock() + expected_desc = { + "EndpointStatus": "Failed", + "FailureReason": "Unable to provision requested ML compute capacity due to " + "InsufficientInstanceCapacity error.", + } + mock_client.describe_endpoint.return_value = expected_desc + mock_paginator = Mock() + mock_paginator.paginate.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException"}}, "FilterLogEvents" + ) + + result = _live_logging_deploy_done_with_progress( + mock_client, "test-endpoint", mock_paginator, {}, 5 + ) + + self.assertEqual(result, expected_desc) + mock_sleep.assert_called_once_with(5) + + def test_updating_endpoint_keeps_waiting(self): + """Updating is in progress, matching _deploy_done_with_progress.""" + mock_client = Mock() + mock_client.describe_endpoint.return_value = {"EndpointStatus": "Updating"} + mock_paginator = Mock() + mock_paginator.paginate.return_value = [] + + result = _live_logging_deploy_done_with_progress( + mock_client, "test-endpoint", mock_paginator, {}, 5 + ) + + self.assertIsNone(result) + + def test_other_log_errors_are_raised(self): + """Log-fetch errors other than a missing log group still propagate.""" + mock_client = Mock() + mock_client.describe_endpoint.return_value = {"EndpointStatus": "InService"} + mock_paginator = Mock() + mock_paginator.paginate.side_effect = ClientError( + {"Error": {"Code": "ThrottlingException"}}, "FilterLogEvents" + ) + + with self.assertRaises(ClientError): + _live_logging_deploy_done_with_progress( + mock_client, "test-endpoint", mock_paginator, {}, 5 + ) + + def test_missing_endpoint_raises_once_budget_is_spent(self): + """A missing endpoint is tolerated only for the budgeted number of polls.""" + mock_client = Mock() + mock_client.describe_endpoint.side_effect = ClientError( + {"Error": {"Code": "ValidationException", "Message": "Could not find endpoint"}}, + "DescribeEndpoint", + ) + budget = _EndpointNotFoundBudget(max_polls=2) + + for _ in range(2): + self.assertIsNone( + _live_logging_deploy_done_with_progress( + mock_client, "test-endpoint", Mock(), {}, 5, not_found_budget=budget + ) + ) + with self.assertRaises(ClientError): + _live_logging_deploy_done_with_progress( + mock_client, "test-endpoint", Mock(), {}, 5, not_found_budget=budget + ) + if __name__ == "__main__": unittest.main() diff --git a/sagemaker-serve/tests/unit/test_model_builder_build.py b/sagemaker-serve/tests/unit/test_model_builder_build.py index 83a839996e..d0d5eebc28 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_build.py +++ b/sagemaker-serve/tests/unit/test_model_builder_build.py @@ -8,6 +8,7 @@ import tempfile import os +from sagemaker.core.exceptions import UnexpectedStatusException from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.utils.types import ModelServer from sagemaker.serve.mode.function_pointers import Mode @@ -520,8 +521,9 @@ def test_wait_for_endpoint_failure(self, mock_wait): sagemaker_session=self.mock_session, ) - # _wait_for_endpoint doesn't raise, just waits - builder._wait_for_endpoint("test-endpoint", wait=True, show_progress=False) + # A deployment that does not reach InService raises + with self.assertRaises(UnexpectedStatusException): + builder._wait_for_endpoint("test-endpoint", wait=True, show_progress=False) # Verify wait was called mock_wait.assert_called_once() diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy.py b/sagemaker-serve/tests/unit/test_model_builder_deploy.py index 2bb381b274..d255f42b1e 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_deploy.py +++ b/sagemaker-serve/tests/unit/test_model_builder_deploy.py @@ -313,6 +313,38 @@ def test_deploy_core_endpoint_invalid_serverless_config_raises_error(self): self.assertIn("ServerlessInferenceConfig object", str(context.exception)) + @patch("sagemaker.serve.model_builder.Endpoint.get") + @patch("sagemaker.serve.model_builder.session_helper.production_variant") + def test_deploy_core_endpoint_ic_based_waits_on_status_only( + self, mock_production_variant, mock_endpoint_get + ): + """IC-based endpoints never get an endpoint log group, so both waits skip log streaming.""" + mock_production_variant.return_value = {"VariantName": "AllTraffic"} + mock_endpoint_get.return_value = Mock(spec=Endpoint) + builder = ModelBuilder( + model=Mock(), + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + builder.built_model = Mock() + builder.built_model.model_name = "test-model" + builder.model_name = "test-model" + + with patch.object(builder, "_wait_for_endpoint") as mock_wait: + builder._deploy_core_endpoint( + instance_type="ml.g5.xlarge", + initial_instance_count=1, + endpoint_name="test-endpoint", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=ResourceRequirements(requests={"memory": 1024, "copies": 1}), + inference_component_name="test-ic", + ) + + self.mock_session.create_inference_component.assert_called_once() + self.assertEqual(mock_wait.call_count, 2) + for wait_call in mock_wait.call_args_list: + self.assertIs(wait_call.kwargs["stream_endpoint_logs"], False) + @unittest.skip("Missing inference_component_name attribute - complex deployment flow") def test_deploy_core_endpoint_sharded_model_forces_ic_based(self): """Test _deploy_core_endpoint forces INFERENCE_COMPONENT_BASED for sharded models.""" diff --git a/sagemaker-serve/tests/unit/test_model_builder_v3.py b/sagemaker-serve/tests/unit/test_model_builder_v3.py index 76a4db8ebb..64fb842181 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_v3.py +++ b/sagemaker-serve/tests/unit/test_model_builder_v3.py @@ -10,6 +10,9 @@ import unittest from unittest.mock import Mock, patch +from botocore.exceptions import ClientError + +from sagemaker.core.exceptions import CapacityError, UnexpectedStatusException from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.utils.types import ModelServer from sagemaker.serve.mode.function_pointers import Mode @@ -734,6 +737,80 @@ def test_wait_for_endpoint_with_wait_false(self): # Should log deployment started message mock_logger.assert_called() + def _builder(self): + return ModelBuilder( + model=Mock(), + role_arn=self.mock_role_arn, + sagemaker_session=self.mock_session, + model_server=ModelServer.TORCHSERVE, + ) + + @patch("sagemaker.serve.model_builder._wait_until") + def test_wait_for_endpoint_raises_when_endpoint_fails(self, mock_wait_until): + """A failed endpoint raises with its failure reason instead of returning quietly.""" + mock_wait_until.return_value = { + "EndpointStatus": "Failed", + "FailureReason": "Unable to provision requested ML compute capacity due to " + "InsufficientInstanceCapacity error.", + } + + with self.assertRaises(UnexpectedStatusException) as context: + self._builder()._wait_for_endpoint("test-endpoint", wait=True, show_progress=False) + + self.assertNotIsInstance(context.exception, CapacityError) + self.assertEqual(context.exception.actual_status, "Failed") + self.assertIn("InsufficientInstanceCapacity", str(context.exception)) + + @patch("sagemaker.serve.model_builder._wait_until") + def test_wait_for_endpoint_raises_capacity_error(self, mock_wait_until): + """A CapacityError failure reason raises CapacityError.""" + mock_wait_until.return_value = { + "EndpointStatus": "Failed", + "FailureReason": "CapacityError: not enough capacity", + } + + with self.assertRaises(CapacityError): + self._builder()._wait_for_endpoint("test-endpoint", wait=True, show_progress=False) + + @patch("sagemaker.serve.deployment_progress.Live") + @patch("sagemaker.core.helper.session_helper._has_permission_for_live_logging") + @patch("sagemaker.serve.model_builder._wait_until") + def test_wait_for_endpoint_status_only_skips_log_streaming( + self, mock_wait_until, mock_permission, mock_live + ): + """stream_endpoint_logs=False never looks for the endpoint log group.""" + mock_wait_until.return_value = {"EndpointStatus": "InService"} + + self._builder()._wait_for_endpoint("test-endpoint", wait=True, stream_endpoint_logs=False) + + mock_permission.assert_not_called() + mock_wait_until.assert_called_once() + + @patch("time.sleep") + @patch("sagemaker.serve.deployment_progress.Live") + @patch("sagemaker.core.helper.session_helper._has_permission_for_live_logging") + def test_wait_for_endpoint_live_logging_ends_when_endpoint_fails_without_log_group( + self, mock_permission, mock_live, mock_sleep + ): + """The live-logging wait ends (and raises) for a Failed endpoint with no log group.""" + mock_permission.return_value = True + failed = { + "EndpointStatus": "Failed", + "FailureReason": "Unable to provision requested ML compute capacity due to " + "InsufficientInstanceCapacity error.", + } + # A finite side_effect so a regression surfaces as a failure, not a hang. + self.mock_client.describe_endpoint.side_effect = [failed, failed] + paginator = self.mock_session.boto_session.client.return_value.get_paginator.return_value + paginator.paginate.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException"}}, "FilterLogEvents" + ) + + with self.assertRaises(UnexpectedStatusException): + self._builder()._wait_for_endpoint("test-endpoint", wait=True) + + self.assertEqual(self.mock_client.describe_endpoint.call_count, 1) + class TestModelBuilderV3EndToEnd(unittest.TestCase): """Test V3 ModelBuilder end-to-end workflow."""