Use exponential capped backoff for job waiting

no need to run in hot-loop

Author: FlorentinD

graphdatascience/arrow_client/v2/job_client.py

@@ -3,6 +3,7 @@
 
 from pandas import ArrowDtype, DataFrame
 from pyarrow._flight import Ticket
+from tenacity import Retrying, retry_always, wait_exponential
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
@@ -34,26 +35,28 @@ def run_job(client: AuthenticatedArrowClient, endpoint: str, config: dict[str, A
 
     def wait_for_job(self, client: AuthenticatedArrowClient, job_id: str, show_progress: bool) -> None:
         progress_bar: Optional[TqdmProgressBar] = None
-        while True:
-            arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
-            job_status = JobStatus(**deserialize_single(arrow_res))
-
-            if job_status.succeeded() or job_status.aborted():
-                if progress_bar:
-                    progress_bar.finish(success=job_status.succeeded())
-                break
-
-            if show_progress:
-                if progress_bar is None:
-                    base_task = job_status.base_task()
-                    if base_task:
-                        progress_bar = TqdmProgressBar(
-                            task_name=base_task,
-                            relative_progress=job_status.progress_percent(),
-                            bar_options=self._progress_bar_options,
-                        )
-                if progress_bar:
-                    progress_bar.update(job_status.status, job_status.progress_percent(), job_status.sub_tasks())
+
+        for attempt in Retrying(retry=retry_always, wait=wait_exponential(min=0.1, max=5)):
+            with attempt:
+                arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
+                job_status = JobStatus(**deserialize_single(arrow_res))
+
+                if job_status.succeeded() or job_status.aborted():
+                    if progress_bar:
+                        progress_bar.finish(success=job_status.succeeded())
+                    return
+
+                if show_progress:
+                    if progress_bar is None:
+                        base_task = job_status.base_task()
+                        if base_task:
+                            progress_bar = TqdmProgressBar(
+                                task_name=base_task,
+                                relative_progress=job_status.progress_percent(),
+                                bar_options=self._progress_bar_options,
+                            )
+                    if progress_bar:
+                        progress_bar.update(job_status.status, job_status.progress_percent(), job_status.sub_tasks())
 
     @staticmethod
     def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]:

graphdatascience/tests/unit/arrow_client/V2/test_job_client.py

@@ -173,8 +173,8 @@ def test_wait_for_job_progress_bar_qualitative(mocker: MockerFixture) -> None:
 
         progress_output = pbarOutputStream.getvalue().split("\r")
         assert "Algo [elapsed: 00:00 ]" in progress_output
-        assert "Algo [elapsed: 00:00 , status: RUNNING, task: Halfway there]" in progress_output
-        assert any("Algo [elapsed: 00:00 , status: FINISHED]" in line for line in progress_output)
+        assert "Algo [elapsed: 00:01 , status: RUNNING, task: Halfway there]" in progress_output
+        assert any("Algo [elapsed: 00:03 , status: FINISHED]" in line for line in progress_output)
 
 
 def test_get_summary(mocker: MockerFixture) -> None: