fix: Raise exception for RoV BQ Write for too many rate limit exceeded

PiperOrigin-RevId: 577271978
googleapis · Oct 27, 2023 · 7e09529 · 7e09529
1 parent dc1b82a
commit 7e09529
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/google/cloud/aiplatform/preview/vertex_ray/bigquery_datasource.py b/google/cloud/aiplatform/preview/vertex_ray/bigquery_datasource.py
@@ -50,6 +50,8 @@
     gapic_version=_BQS_GAPIC_VERSION, user_agent=f"ray-on-vertex/{_BQS_GAPIC_VERSION}"
 )
 
+MAX_RETRY_CNT = 20
+
 
 class _BigQueryDatasourceReader(Reader):
     def __init__(
@@ -192,7 +194,7 @@ def _write_single_block(
                 pq.write_table(block, fp, compression="SNAPPY")
 
                 retry_cnt = 0
-                while retry_cnt < 10:
+                while retry_cnt < MAX_RETRY_CNT:
                     with open(fp, "rb") as source_file:
                         job = client.load_table_from_file(
                             source_file, dataset, job_config=job_config
@@ -207,6 +209,14 @@ def _write_single_block(
                         )
                         logging.debug(e)
                         time.sleep(11)
+
+                # Raise exception if retry_cnt hits MAX_RETRY_CNT
+                if retry_cnt >= MAX_RETRY_CNT:
+                    raise RuntimeError(
+                        f"[Ray on Vertex AI]: Write failed due to {MAX_RETRY_CNT}"
+                        + " repeated API rate limit exceeded responses"
+                    )
+
             print("[Ray on Vertex AI]: Finished writing", metadata.num_rows, "rows")
 
         project_id = project_id or initializer.global_config.project