1717from botocore .exceptions import ClientError
1818from sagemaker .s3 import parse_s3_url
1919from sagemaker .async_inference import WaiterConfig
20- from sagemaker .exceptions import ObjectNotExistedError , UnexpectedClientError
20+ from sagemaker .exceptions import (
21+ ObjectNotExistedError ,
22+ UnexpectedClientError ,
23+ AsyncInferenceModelError ,
24+ )
2125
2226
2327class AsyncInferenceResponse (object ):
@@ -32,6 +36,7 @@ def __init__(
3236 self ,
3337 predictor_async ,
3438 output_path ,
39+ failure_path ,
3540 ):
3641 """Initialize an AsyncInferenceResponse object.
3742
@@ -43,10 +48,13 @@ def __init__(
4348 that return this response.
4449 output_path (str): The Amazon S3 location that endpoints upload inference responses
4550 to.
51+ failure_path (str): The Amazon S3 location that endpoints upload model errors
52+ for failed requests.
4653 """
4754 self .predictor_async = predictor_async
4855 self .output_path = output_path
4956 self ._result = None
57+ self .failure_path = failure_path
5058
5159 def get_result (
5260 self ,
@@ -71,28 +79,34 @@ def get_result(
7179
7280 if self ._result is None :
7381 if waiter_config is None :
74- self ._result = self ._get_result_from_s3 (self .output_path )
82+ self ._result = self ._get_result_from_s3 (self .output_path , self . failure_path )
7583 else :
7684 self ._result = self .predictor_async ._wait_for_output (
77- self .output_path , waiter_config
85+ self .output_path , self . failure_path , waiter_config
7886 )
7987 return self ._result
8088
81- def _get_result_from_s3 (
82- self ,
83- output_path ,
84- ):
89+ def _get_result_from_s3 (self , output_path , failure_path ):
8590 """Get inference result from the output Amazon S3 path"""
8691 bucket , key = parse_s3_url (output_path )
8792 try :
8893 response = self .predictor_async .s3_client .get_object (Bucket = bucket , Key = key )
8994 return self .predictor_async .predictor ._handle_response (response )
90- except ClientError as ex :
91- if ex .response ["Error" ]["Code" ] == "NoSuchKey" :
92- raise ObjectNotExistedError (
93- message = "Inference could still be running" ,
94- output_path = output_path ,
95- )
96- raise UnexpectedClientError (
97- message = ex .response ["Error" ]["Message" ],
98- )
95+ except ClientError as e :
96+ if e .response ["Error" ]["Code" ] == "NoSuchKey" :
97+ try :
98+ failure_bucket , failure_key = parse_s3_url (failure_path )
99+ failure_response = self .predictor_async .s3_client .get_object (
100+ Bucket = failure_bucket , Key = failure_key
101+ )
102+ failure_response = self .predictor_async .predictor ._handle_response (
103+ failure_response
104+ )
105+ raise AsyncInferenceModelError (message = failure_response )
106+ except ClientError as ex :
107+ if ex .response ["Error" ]["Code" ] == "NoSuchKey" :
108+ raise ObjectNotExistedError (
109+ message = "Inference could still be running" , output_path = output_path
110+ )
111+ raise UnexpectedClientError (message = ex .response ["Error" ]["Message" ])
112+ raise UnexpectedClientError (message = e .response ["Error" ]["Message" ])
0 commit comments