From aed8f9a940d81544be04ae01503dae126b15a892 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Fri, 9 Oct 2026 13:09:44 -0500 Subject: [PATCH] fix(serve): raise PollingTimeoutError instead of hanging in AsyncPredictor _check_output_and_failure_paths waited in a loop until either the output or the failure file was found. When both S3 waiters gave up (WaiterError after max_attempts * delay), neither event was set and the loop spun forever, so predict()/the wait on an async endpoint never returned, and the PollingTimeoutError branch was unreachable. That branch also fetched the failure object before checking whether it existed. Exit the loop once both waiter threads have finished, raise AsyncInferenceModelError only when the failure file was found, and PollingTimeoutError otherwise. v2 handles this case (#4951); unlike v2, this keeps waiting for the failure waiter if only the output waiter has timed out, so a model error that lands slightly later is still reported as such. --- .../src/sagemaker/serve/predictor_async.py | 26 ++++---- .../tests/unit/test_predictor_async.py | 62 +++++++++++++++++++ 2 files changed, 77 insertions(+), 11 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/predictor_async.py b/sagemaker-serve/src/sagemaker/serve/predictor_async.py index 986d51c3d7..acaa87b824 100644 --- a/sagemaker-serve/src/sagemaker/serve/predictor_async.py +++ b/sagemaker-serve/src/sagemaker/serve/predictor_async.py @@ -335,7 +335,13 @@ def check_failure_file(): output_thread.start() failure_thread.start() - while not output_file_found.is_set() and not failure_file_found.is_set(): + # Stop once either file is found, or once both waiters have given up (WaiterError); + # otherwise a timeout on both would leave this loop spinning forever. + while ( + not output_file_found.is_set() + and not failure_file_found.is_set() + and (output_thread.is_alive() or failure_thread.is_alive()) + ): time.sleep(1) if output_file_found.is_set(): @@ -343,17 +349,15 @@ def check_failure_file(): result = self.predictor._handle_response(response=s3_object) return result - failure_object = self.s3_client.get_object(Bucket=failure_bucket, Key=failure_key) - failure_response = self.predictor._handle_response(response=failure_object) + if failure_file_found.is_set(): + failure_object = self.s3_client.get_object(Bucket=failure_bucket, Key=failure_key) + failure_response = self.predictor._handle_response(response=failure_object) + raise AsyncInferenceModelError(message=failure_response) - raise ( - AsyncInferenceModelError(message=failure_response) - if failure_file_found.is_set() - else PollingTimeoutError( - message="Inference could still be running", - output_path=output_path, - seconds=waiter_config.delay * waiter_config.max_attempts, - ) + raise PollingTimeoutError( + message="Inference could still be running", + output_path=output_path, + seconds=waiter_config.delay * waiter_config.max_attempts, ) def update_endpoint( diff --git a/sagemaker-serve/tests/unit/test_predictor_async.py b/sagemaker-serve/tests/unit/test_predictor_async.py index aaf77a9c4a..a8ab24c8ab 100644 --- a/sagemaker-serve/tests/unit/test_predictor_async.py +++ b/sagemaker-serve/tests/unit/test_predictor_async.py @@ -1,9 +1,15 @@ import io +import threading +import time import unittest from unittest.mock import Mock, patch from sagemaker.core.deserializers import JSONDeserializer from sagemaker.core.serializers import JSONSerializer +from botocore.exceptions import WaiterError + +from sagemaker.core.exceptions import AsyncInferenceModelError, PollingTimeoutError +from sagemaker.serve.async_inference import WaiterConfig from sagemaker.serve.predictor_async import AsyncPredictor @@ -231,3 +237,59 @@ def test_reflects_serializers_set_on_wrapped_predictor(self): if __name__ == "__main__": unittest.main() + + +class TestCheckOutputAndFailurePaths(unittest.TestCase): + """_check_output_and_failure_paths must return or raise once both waiters are done.""" + + def setUp(self): + self.mock_predictor = Mock() + self.mock_predictor.endpoint_name = "test-endpoint" + self.async_predictor = AsyncPredictor(self.mock_predictor) + self.async_predictor.s3_client = Mock() + self.waiter_config = WaiterConfig(max_attempts=1, delay=1) + + def _run_with_timeout(self, timeout=10): + outcome = {} + + def run(): + try: + outcome["result"] = self.async_predictor._check_output_and_failure_paths( + "s3://bucket/output", "s3://bucket/failure", self.waiter_config + ) + except Exception as e: # pylint: disable=broad-except + outcome["error"] = e + + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(timeout=timeout) + self.assertFalse(thread.is_alive(), "_check_output_and_failure_paths did not return") + return outcome + + def test_raises_polling_timeout_when_both_waiters_time_out(self): + waiter = Mock() + waiter.wait.side_effect = WaiterError("ObjectExists", "Max attempts exceeded", {}) + self.async_predictor.s3_client.get_waiter.return_value = waiter + + outcome = self._run_with_timeout() + + self.assertIsInstance(outcome.get("error"), PollingTimeoutError) + self.async_predictor.s3_client.get_object.assert_not_called() + + def test_raises_model_error_when_failure_file_appears_after_output_waiter_times_out(self): + def wait(Bucket, Key, WaiterConfig): + if Key == "output": + raise WaiterError("ObjectExists", "Max attempts exceeded", {}) + time.sleep(0.2) + + waiter = Mock() + waiter.wait.side_effect = wait + self.async_predictor.s3_client.get_waiter.return_value = waiter + self.mock_predictor._handle_response.return_value = "model failure" + + outcome = self._run_with_timeout() + + self.assertIsInstance(outcome.get("error"), AsyncInferenceModelError) + self.async_predictor.s3_client.get_object.assert_called_once_with( + Bucket="bucket", Key="failure" + )