Limit Deep CFR traversal futures in flight

src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py

@@ -4,7 +4,7 @@ import copy
 import json
 import multiprocessing as mp
 import time
-from concurrent.futures import ProcessPoolExecutor, as_completed
+from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, as_completed, wait
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -415,10 +415,19 @@ class DeepCFRTrainer:
             max_workers=max_workers,
             mp_context=mp.get_context("spawn"),
         ) as executor:
-            futures = [executor.submit(run_traversal_worker_batch, batch) for batch in batches]
-            total_batches = len(futures)
-            for completed_batches, future in enumerate(as_completed(futures), start=1):
+            total_batches = len(batches)
+            in_flight_limit = min(total_batches, max(1, max_workers * 2))
+            batch_iter = iter(batches)
+            futures = {
+                executor.submit(run_traversal_worker_batch, batch)
+                for _, batch in zip(range(in_flight_limit), batch_iter, strict=False)
+            }
+            completed_batches = 0
+            while futures:
+                done, futures = wait(futures, return_when=FIRST_COMPLETED)
+                for future in done:
                     result = future.result()
+                    completed_batches += 1
                     total_stats.accumulate(result.stats)
                     memory_add_started = time.perf_counter()
                     self._add_advantage_samples(result.advantage_samples)
@@ -435,12 +444,18 @@ class DeepCFRTrainer:
                         self.tracker.log_event(
                             f"Traversal multiprocessing progress iteration={iteration} "
                             f"completed_batches={completed_batches}/{total_batches} "
-                        f"completed_traversals={progress_traversals} elapsed_seconds={elapsed:.2f} "
+                            f"completed_traversals={progress_traversals} "
+                            f"elapsed_seconds={elapsed:.2f} "
                             f"total_nodes={progress_nodes} "
                             f"nodes_per_second={progress_nodes / max(elapsed, 1.0e-12):.1f}"
                         )
-                    while next_progress_at is not None and next_progress_at <= progress_traversals:
+                        while (
+                            next_progress_at is not None and next_progress_at <= progress_traversals
+                        ):
                             next_progress_at += progress_every
+                    next_batch = next(batch_iter, None)
+                    if next_batch is not None:
+                        futures.add(executor.submit(run_traversal_worker_batch, next_batch))
         return total_stats
 
     def _worker_batches(self, iteration: int) -> list[TraversalWorkerBatch]: