feat(sdk): add wait_until_ready and wait_for_completion helpers

Move the dataset polling loop and training job polling loop out of every user's integration scripts and into the SDK. wait_until_ready raises on failed or timeout. wait_for_completion returns the terminal job on any of completed/failed/cancelled and accepts an optional on_progress callback for live metric streaming.

Author: Francisco

src/projectdavid/clients/datasets_client.py

@@ -97,6 +97,48 @@ def retrieve(self, dataset_id: str) -> validator.DatasetRead:
         response.raise_for_status()
         return validator.DatasetRead.model_validate(response.json())
 
+    # ------------------------------------------------------------------
+    # WAIT UNTIL READY
+    # ------------------------------------------------------------------
+    def wait_until_ready(
+        self,
+        dataset_id: str,
+        *,
+        timeout: float = 300.0,
+        poll_interval: float = 3.0,
+    ) -> validator.DatasetRead:
+        """
+        Block until the dataset reaches 'active' status.
+
+        Polls the dataset every `poll_interval` seconds. Returns the final
+        DatasetRead once active. Raises RuntimeError if the dataset enters
+        the 'failed' status, TimeoutError if it has not reached a terminal
+        state within `timeout` seconds.
+        """
+        import time  # local import keeps top-of-file imports minimal
+
+        deadline = time.monotonic() + timeout
+        last_status = None
+
+        while time.monotonic() < deadline:
+            ds = self.retrieve(dataset_id)
+            if ds.status != last_status:
+                logging_utility.info("Dataset %s status=%s", dataset_id, ds.status)
+                last_status = ds.status
+
+            if ds.status == "active":
+                return ds
+            if ds.status == "failed":
+                raise RuntimeError(
+                    f"Dataset {dataset_id} preparation failed (status=failed)"
+                )
+            time.sleep(poll_interval)
+
+        raise TimeoutError(
+            f"Dataset {dataset_id} did not reach 'active' within {timeout}s "
+            f"(last status={last_status})"
+        )
+
     # ------------------------------------------------------------------
     # LIST
     # ------------------------------------------------------------------

src/projectdavid/clients/training_client.py

@@ -85,6 +85,68 @@ def cancel(self, job_id: str) -> validator.TrainingJobCancelResponse:
         response.raise_for_status()
         return validator.TrainingJobCancelResponse.model_validate(response.json())
 
+    # ------------------------------------------------------------------
+    # WAIT FOR COMPLETION
+    # ------------------------------------------------------------------
+    def wait_for_completion(
+        self,
+        job_id: str,
+        *,
+        on_progress=None,
+        poll_interval: float = 10.0,
+        timeout: float = 7200.0,
+    ) -> validator.TrainingJobRead:
+        """
+        Block until the training job reaches a terminal state.
+
+        Terminal states: 'completed', 'failed', 'cancelled'.
+
+        An optional `on_progress` callback receives every distinct metrics
+        snapshot as it arrives (keyed by the 'step' field). Use it to drive
+        live progress display without coupling the SDK to any particular
+        output format.
+
+        Example:
+            def show(m):
+                print(f"step={m.get('step')} loss={m.get('loss')}")
+
+            job = client.training.wait_for_completion(job.id, on_progress=show)
+
+        Raises TimeoutError if the job has not reached a terminal state
+        within `timeout` seconds. Does NOT raise on a 'failed' or 'cancelled'
+        job — callers inspect the returned job.status themselves, since a
+        failed job is often still interesting (last_error, partial metrics).
+        """
+        import time  # local import keeps top-of-file imports minimal
+
+        TERMINAL_STATES = {"completed", "failed", "cancelled"}
+        deadline = time.monotonic() + timeout
+        last_step = -1
+
+        while time.monotonic() < deadline:
+            job = self.retrieve(job_id)
+
+            metrics = getattr(job, "metrics", None) or {}
+            step = metrics.get("step")
+            if on_progress is not None and step is not None and step != last_step:
+                last_step = step
+                try:
+                    on_progress(metrics)
+                except Exception as cb_err:
+                    # Callback failures must not break polling.
+                    logging_utility.warning(
+                        "on_progress callback raised %s: %s",
+                        type(cb_err).__name__,
+                        cb_err,
+                    )
+
+            if job.status in TERMINAL_STATES:
+                return job
+
+            time.sleep(poll_interval)
+
+        raise TimeoutError(f"Training job {job_id} did not complete within {timeout}s")
+
     # ------------------------------------------------------------------
     # DIAGNOSTIC PEEK (Secure Multi-tenant Gateway)
     # ------------------------------------------------------------------