fix: Made Endpoint.raw_predict thread-safe by setting authorized_session last

calvingiles · copybara-github · commit c72c1ef4c58c · 2024-03-15T11:12:39.000-07:00
Fixes: #2068 Closes: #3385 COPYBARA_INTEGRATE_REVIEW=#3385 from calvingiles:make_raw_predict_threadsafe fa2c719 PiperOrigin-RevId: 616187206
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -1687,10 +1687,10 @@ def raw_predict(
         """
         if not self.authorized_session:
             self.credentials._scopes = constants.base.DEFAULT_AUTHED_SCOPES
+            self.raw_predict_request_url = f"https://{self.location}-{constants.base.API_BASE_PATH}/v1/projects/{self.project}/locations/{self.location}/endpoints/{self.name}:rawPredict"
             self.authorized_session = google_auth_requests.AuthorizedSession(
                 self.credentials
             )
-            self.raw_predict_request_url = f"https://{self.location}-{constants.base.API_BASE_PATH}/v1/projects/{self.project}/locations/{self.location}/endpoints/{self.name}:rawPredict"
 
         return self.authorized_session.post(
             url=self.raw_predict_request_url, data=body, headers=headers
