aiplatform.PrivateEndpoint.raw_predict crash (TypeError: Session.request() got an unexpected keyword argument 'body') #6101
Description
Environment details
google-cloud-aiplatformversion:pip show google-cloud-aiplatform: 1.126
Steps to reproduce
- Create VAI PrivateEndpoint
- call raw_predict
Code example
# example
endpoint = aiplatform.PrivateEndpoint.create()
response = endpoint.raw_predict()Stack trace
# example
ypeError Traceback (most recent call last)
Cell In[14], line 10
3 #ENDPOINT_ID = "5931165942456254464"
4 #endpoint = aiplatform.Endpoint(ENDPOINT_ID, location=LOCATION)
5 headers = {
6 "Content-Type": "application/json",
7 "x-request-id": "ebb94475-1ca2-4e4b-baa3-8d039c0e616e", #works when --enable-request-id-headers option enabled
8 "x-vertex-ai-timeout-ms": "60000"
9 }
---> 10 response = endpoint.raw_predict(headers=headers, body=json.dumps(prediction_input, indent=2).encode('utf-8'),
11 endpoint_override=IP_ADDRESS
12 )
13 print(response.headers)
14 print(response.json()['choices'][0]['message']['content'])
File ~/.local/lib/python3.10/site-packages/google/cloud/aiplatform/models.py:4151, in PrivateEndpoint.raw_predict(self, body, headers, endpoint_override)
4148 self._authorized_session.verify = False
4150 url = f"[https://{endpoint_override}/v1/projects/{self.project}/locations/{self.location}/endpoints/{self.name}:rawPredict](https://{endpoint_override}/v1/projects/%7Bself.project%7D/locations/%7Bself.location%7D/endpoints/%7Bself.name%7D:rawPredict)"
-> 4151 return self._authorized_session.post(
4152 url=url,
4153 body=body,
4154 headers=headers,
4155 )
File /opt/conda/lib/python3.10/site-packages/requests/sessions.py:637, in Session.post(self, url, data, json, **kwargs)
626 def post(self, url, data=None, json=None, **kwargs):
627 r"""Sends a POST request. Returns :class:`Response` object.
628
629 :param url: URL for the new :class:`Request` object.
(...)
634 :rtype: requests.Response
635 """
--> 637 return self.request("POST", url, data=data, json=json, **kwargs)
File /opt/conda/lib/python3.10/site-packages/google/auth/transport/requests.py:537, in AuthorizedSession.request(self, method, url, data, headers, max_allowed_time, timeout, **kwargs)
534 remaining_time = guard.remaining_timeout
536 with TimeoutGuard(remaining_time) as guard:
--> 537 response = super(AuthorizedSession, self).request(
538 method,
539 url,
540 data=data,
541 headers=request_headers,
542 timeout=timeout,
543 **kwargs
544 )
545 remaining_time = guard.remaining_timeout
547 # If the response indicated that the credentials needed to be
548 # refreshed, then refresh the credentials and re-attempt the
549 # request.
550 # A stored token may expire between the time it is retrieved and
551 # the time the request is made, so we may need to try twice.
TypeError: Session.request() got an unexpected keyword argument 'body'
The reason for error is changed parameter name of google-auth-transport requests from 'body' to 'data'
The fix already reflected to normal Endpoint class but PrivateEndpoint still remains.
Simply change required end of PrivateEndpoint.raw_predict like below
return self._authorized_session.post(
url=url,
data=body, #Currently body=body
headers=headers,
)