Thread-safe registry of in-flight provider batches. Used to cancel batches on KeyboardInterrupt so they don't linger on the provider side after the process exits.
| 50 | |
| 51 | |
| 52 | class _InflightBatches: |
| 53 | """Thread-safe registry of in-flight provider batches. |
| 54 | |
| 55 | Used to cancel batches on KeyboardInterrupt so they don't linger |
| 56 | on the provider side after the process exits. |
| 57 | """ |
| 58 | |
| 59 | def __init__(self) -> None: |
| 60 | self._lock = threading.Lock() |
| 61 | self._batches: list[tuple[Any, Any, str]] = [] # (bp, client, job_id) |
| 62 | self.stop_event = threading.Event() |
| 63 | |
| 64 | def register(self, bp: Any, client: Any, job_id: str) -> None: |
| 65 | with self._lock: |
| 66 | self._batches.append((bp, client, job_id)) |
| 67 | |
| 68 | def deregister(self, job_id: str) -> None: |
| 69 | with self._lock: |
| 70 | self._batches = [(b, c, j) for b, c, j in self._batches if j != job_id] |
| 71 | |
| 72 | def cancel_all(self) -> None: |
| 73 | self.stop_event.set() |
| 74 | with self._lock: |
| 75 | batches = list(self._batches) |
| 76 | self._batches.clear() |
| 77 | for bp, client, job_id in batches: |
| 78 | try: |
| 79 | bp.cancel_batch(client, job_id) |
| 80 | logger.info("cancelled in-flight batch %s", job_id) |
| 81 | except Exception as e: |
| 82 | logger.warning("failed to cancel batch %s: %s", job_id, e) |
| 83 | |
| 84 | |
| 85 | class WorkItem(NamedTuple): |