chore(weave): add rich progress bar to client flush

tests/trace/test_weave_client.py

@@ -7,6 +7,7 @@
 import pydantic
 import pytest
 import requests
+from PIL import Image
 
 import weave
 import weave.trace_server.trace_server_interface as tsi
@@ -1886,3 +1887,61 @@ def my_op(a: int) -> int:
 
         # Local attributes override global ones
         assert call.attributes["env"] == "override"
+
+
+def test_flush_progress_bar(client):
+    # expensive op
+    @weave.op
+    def make_image():
+        img = Image.new("RGB", (300, 300))
+        pixels = img.load()
+        for x in range(300):
+            for y in range(300):
+                pixels[x, y] = (x, y, 0)
+        return img
+
+    ds2 = [make_image() for i in range(10)]
+
+    # really expensive op
+    @weave.op
+    def make_dataset(ds2):
+        new_ds = [make_image() for i in range(10)]
+        return ds2 + new_ds
+
+    make_dataset(ds2)
+
+    # flush with progress bar
+    client.flush(use_progress_bar=True)
+
+    # make sure there are no pending jobs
+    assert client._get_pending_jobs() == (0, 0, 0, 0)
+    assert client._has_pending_jobs() == False
+
+
+def test_flush_status_callback(client, capsys):
+    @weave.op
+    def make_image():
+        img = Image.new("RGB", (300, 300))
+        pixels = img.load()
+        for x in range(300):
+            for y in range(300):
+                pixels[x, y] = (x, y, 0)
+        return img
+
+    ds = [make_image() for i in range(10)]
+
+    @weave.op
+    def use_dataset(ds):
+        return ds
+
+    use_dataset(ds)
+
+    def logger(status):
+        assert status["main_jobs"] == 0
+
+    # use custom logger callback
+    client.flush(callback=logger)
+
+    # make sure there are no pending jobs
+    assert client._get_pending_jobs() == (0, 0, 0, 0)
+    assert client._has_pending_jobs() == False

weave/trace/client_progress_bar.py

@@ -0,0 +1,122 @@
+"""Progress bar utilities for WeaveClient.
+
+This module provides functionality for displaying progress bars when flushing
+tasks in the WeaveClient.
+"""
+
+from typing import Callable, TypedDict
+
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+
+
+class WeaveFlushStatus(TypedDict):
+    """Status information about the current flush operation."""
+
+    # Current job counts
+    main_jobs: int
+    fastlane_jobs: int
+    call_processor_jobs: int
+    total_jobs: int
+
+    # Tracking of completed jobs
+    completed_since_last_update: int
+    total_completed: int
+
+    # Maximum number of jobs seen during this flush operation
+    max_total_jobs: int
+
+    # Whether there are any pending jobs
+    has_pending_jobs: bool
+
+
+def create_progress_bar_callback() -> Callable[[WeaveFlushStatus], None]:
+    """Create a callback function that displays a progress bar for flush status.
+
+    Returns:
+        A callback function that can be passed to WeaveClient._flush.
+    """
+    console = Console()
+
+    # Create a progress bar instance
+    progress = Progress(
+        SpinnerColumn(),
+        TextColumn("[bold blue]{task.description}"),
+        BarColumn(bar_width=None, complete_style="magenta"),
+        TaskProgressColumn(),
+        TimeElapsedColumn(),
+        console=console,
+        refresh_per_second=10,
+        expand=True,
+        transient=False,
+    )
+
+    # Start the progress display
+    progress.start()
+
+    # Create a task for tracking progress
+    task_id = None
+
+    def progress_callback(status: WeaveFlushStatus) -> None:
+        """Update the progress bar based on the flush status.
+
+        Args:
+            status: The current flush status.
+        """
+        nonlocal task_id
+
+        # If this is the first update, create the task
+        if task_id is None:
+            if status["max_total_jobs"] == 0:
+                # No jobs to track, just return
+                progress.stop()
+                return
+
+            # Print initial message
+            if not progress.live.is_started:
+                print(f"Flushing {status['max_total_jobs']} pending tasks...")
+
+            # Create the task
+            task_id = progress.add_task(
+                "Flushing tasks", total=status["max_total_jobs"]
+            )
+
+        # If there are no more pending jobs, complete the progress bar
+        if not status["has_pending_jobs"]:
+            progress.update(task_id, completed=status["max_total_jobs"])
+            progress.stop()
+            return
+
+        # If new jobs were added, update the total
+        if status["max_total_jobs"] > progress.tasks[task_id].total:
+            progress.update(task_id, total=status["max_total_jobs"])
+
+        # Update progress bar with completed jobs
+        if status["completed_since_last_update"] > 0:
+            progress.update(task_id, advance=status["completed_since_last_update"])
+
+        # Format job details for description
+        job_details = []
+        if status["main_jobs"] > 0:
+            job_details.append(f"{status['main_jobs']} main")
+        if status["fastlane_jobs"] > 0:
+            job_details.append(f"{status['fastlane_jobs']} file-upload")
+        if status["call_processor_jobs"] > 0:
+            job_details.append(f"{status['call_processor_jobs']} call-batch")
+
+        job_details_str = ", ".join(job_details) if job_details else "none"
+
+        # Update progress bar description
+        progress.update(
+            task_id,
+            description=f"Flushing tasks ({status['total_jobs']} remaining: {job_details_str})",
+        )
+
+    return progress_callback

weave/trace/concurrent/futures.py

@@ -79,6 +79,11 @@ def __init__(
         self._in_thread_context = ContextVar("in_deferred_context", default=False)
         atexit.register(self._shutdown)
 
+    @property
+    def num_outstanding_futures(self) -> int:
+        with self._active_futures_lock:
+            return len(self._active_futures)
+
     def defer(self, f: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
         """
         Defer a function to be executed in a thread pool.