[Data][1/2] - Dynamic work queue for traversals

[goutamvenkat-anyscale]

Description

Adds parallel_process_work_stealing, a thread-pool-backed generator for parallelizing workloads where the full set of work items is discovered at runtime (e.g., recursive directory listing). Unlike make_async_gen which maps over a static input iterator, this utility lets workers dynamically enqueue new work items back into a shared queue, enabling work-stealing-style load balancing across threads.

1. Introduces dynamic_work_queue.py with parallel_process_work_stealing
2. _WorkerPool (thread lifecycle management)
3. _raise_if_error (cross-thread error propagation).

Supports optional deterministic ordering via preserve_order + order_key, grouping results by their originating seed item.
Adds comprehensive tests covering flat processing, recursive tree traversal, deep chains, error propagation, early stopping, empty inputs, and ordering guarantees.

Diagram:

seed_items ──────▶ [ work_queue ] ◀──────────┐
                        │                     │
            ┌───────────┼───────────┐         │
            ▼           ▼           ▼         │
        worker_0    worker_1    worker_N       │
            │           │           │         │
            │   process_fn(item, add_work, add_result)
            │         /            \          │
            │   add_work(new)   add_result(r) │
            │       │                │        │
            │       └────────────────┼────────┘
            │                        ▼
            └──────────────▶ [ output_queue ]
                                     │
                          completion signaler
                         (join work_queue, then
                          send SENTINEL)
                                     │
                                     ▼
                            main thread yields
                            results to caller

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request introduces a new dynamic work queue implementation (parallel_process_work_stealing) and corresponding unit tests to support parallel processing with dynamic work generation and load balancing. The review feedback highlights several critical robustness improvements: catching BaseException instead of Exception in worker threads to prevent silent failures, updating the type annotation of _WorkerError.exception to Optional[BaseException] to satisfy type checkers, implementing a global timeout pattern when joining threads during shutdown, and wrapping the completion signaler in a try-except block to avoid indefinite hangs on unexpected errors.

[gemini-code-assist]

In _worker, catching only Exception means that any BaseException (such as KeyboardInterrupt or SystemExit) raised inside process_fn will not be caught. This causes the worker thread to die silently while the task is still marked as done in the finally block, leading to silent data loss and incomplete processing without the main thread knowing. Catching BaseException instead ensures all critical errors are propagated to the main thread.

            except BaseException as e:\n                pool.output_queue.put(_WorkerError(e))

[gemini-code-assist]

The exception field in _WorkerError is typed as BaseException, but in _raise_if_error it is set to None to avoid reference cycles. This will cause type checkers like mypy to raise an error. The type should be updated to Optional[BaseException].

Suggested change

	exception: BaseException
	exception: Optional[BaseException]

[gemini-code-assist]

In _WorkerPool.shutdown, joining threads sequentially with a fixed timeout can block the main thread for up to num_workers * join_timeout if multiple threads hang. It is highly recommended to use a global timeout pattern to ensure the total shutdown time does not exceed join_timeout.

        if join_timeout is not None:\n            import time\n            end_time = time.monotonic() + join_timeout\n            for t in self._threads:\n                remaining = end_time - time.monotonic()\n                if remaining <= 0:\n                    break\n                t.join(timeout=remaining)

[gemini-code-assist]

In _signal_completion, if any unexpected exception occurs, the signaler thread will die silently, causing the main thread to block indefinitely on pool.output_queue.get(). Wrapping the signaler logic in a try-except block to propagate unexpected exceptions to the output_queue improves robustness and prevents hangs.

    try:\n        if not _interruptible_join(pool.work_queue, pool.interrupted):\n            return\n        # Stop consumers\n        pool.output_queue.put(SENTINEL)\n        for _ in range(num_workers):\n            # Stop producers\n            pool.work_queue.put(SENTINEL)\n    except InterruptedError:\n        pass\n    except Exception as e:\n        pool.output_queue.put(_WorkerError(e))

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit c843a27. Configure here.}

[cursor]

Sentinel equality stops output early

Medium Severity

The main thread drains output_queue with iter(pool.output_queue.get, SENTINEL), which ends when a dequeued value compares equal to the sentinel via ==, not identity. A legitimate add_result value whose __eq__ returns true for the sentinel can terminate draining early, yielding incomplete results without error while workers may still be running.

 

^{Reviewed by Cursor Bugbot for commit c843a27. Configure here.}