Add status code to endpoint response (#172)

* add status code * filter out extraneous fields explicitly * black

Author: seanshi-scale

launch/model_endpoint.py

@@ -20,6 +20,9 @@
 TASK_SUCCESS_STATE="SUCCESS"
 TASK_FAILURE_STATE="FAILURE"
 
+# Echoes fields in EndpointResponse class
+ALLOWED_ENDPOINT_RESPONSE_FIELDS={"status", "result_url", "result", "traceback", "status_code"}
+
 
 @dataclass_json(undefined=Undefined.EXCLUDE)
 @dataclass
@@ -189,6 +192,7 @@ def __init__(
 result_url: Optional[str] =None,
 result: Optional[str] =None,
 traceback: Optional[str] =None,
+status_code: Optional[int] =None,
  ):
 """
  Parameters:
@@ -210,12 +214,15 @@ def __init__(
 
  traceback: The stack trace if the inference endpoint raised an error. Can be used for debugging
 
+ status_code: The underlying status code of the response, given from the inference endpoint itself.
+
  """
 self.client=client
 self.status=status
 self.result_url=result_url
 self.result=result
 self.traceback=traceback
+self.status_code=status_code
 
 def__str__(self) ->str:
 return (
@@ -271,6 +278,7 @@ def get(self, timeout: Optional[float] = None) -> EndpointResponse:
 result_url=async_response.get("result",{}).get("result_url", None),
 result=async_response.get("result",{}).get("result", None),
 traceback=None,
+status_code=async_response.get("status_code", None),
  )
 elifstatus=="FAILURE":
 returnEndpointResponse(
@@ -279,6 +287,7 @@ def get(self, timeout: Optional[float] = None) -> EndpointResponse:
 result_url=None,
 result=None,
 traceback=async_response.get("traceback", None),
+status_code=async_response.get("status_code", None),
  )
 else:
 raiseValueError(f"Unrecognized status: {async_response['status']}")
@@ -312,6 +321,7 @@ def __next__(self):
 result_url=result.get("result_url", None),
 result=result.get("result", None),
 traceback=data.get("traceback"),
+status_code=data.get("status_code", None),
  )
 
 
@@ -397,7 +407,10 @@ def predict(self, request: EndpointRequest) -> EndpointResponse:
 args=request.args,
 return_pickled=request.return_pickled,
  )
-raw_response={k: vfork, vinraw_response.items() ifvisnotNone}
+
+raw_response={
+k: vfork, vinraw_response.items() ifvisnotNoneandkinALLOWED_ENDPOINT_RESPONSE_FIELDS
+ }
 returnEndpointResponse(client=self.client, **raw_response)
 
 
@@ -632,6 +645,7 @@ def single_request(inner_url, inner_task_id):
 result_url=raw_response.get("result_url", None),
 result=raw_response.get("result", None),
 traceback=raw_response.get("traceback", None),
+status_code=raw_response.get("status_code", None),
  )
 self.responses[url] =response_object