Function run_parallel

explainshell/extraction/runner.py:187–252

Run extractor on files using a rolling thread pool. Keeps at most ``jobs`` tasks submitted at a time so a late fatal error does not leave the entire remaining corpus pre-queued. ``on_start`` runs in worker threads. ``on_result`` runs in the main thread. Callback exceptions are trea

(
    extractor: Extractor,
    gz_files: list[str],
    jobs: int,
    on_start: Callable[[str], None] | None = None,
    on_result: Callable[[str, ExtractionResult], None] | None = None,
)

def run_parallel(
    extractor: Extractor,
    gz_files: list[str],
    jobs: int,
    on_start: Callable[[str], None] | None = None,
    on_result: Callable[[str, ExtractionResult], None] | None = None,
) -> BatchResult:
    """Run extractor on files using a rolling thread pool.

    Keeps at most ``jobs`` tasks submitted at a time so a late fatal error
    does not leave the entire remaining corpus pre-queued.

    ``on_start`` runs in worker threads. ``on_result`` runs in the main thread.
    Callback exceptions are treated as fatal.
    """
    batch = BatchResult()

    def _do_one(gz_path: str) -> ExtractionResult:
        if on_start:
            on_start(gz_path)
        return _extract_one(extractor, gz_path)

    executor = concurrent.futures.ThreadPoolExecutor(max_workers=jobs)
    try:
        pending: set[concurrent.futures.Future[ExtractionResult]] = set()
        gz_iter = iter(gz_files)

        def _submit_next() -> bool:
            gz_path = next(gz_iter, None)
            if gz_path is None:
                return False
            pending.add(executor.submit(_do_one, gz_path))
            return True

        for _ in range(jobs):
            if not _submit_next():
                break

        while pending:
            done, _ = concurrent.futures.wait(
                pending, return_when=concurrent.futures.FIRST_COMPLETED
            )
            for future in done:
                pending.remove(future)
                entry = future.result()  # raises FatalExtractionError on fatal
                _tally(batch, entry)
                if on_result:
                    on_result(entry.gz_path, entry)
                _submit_next()
    except KeyboardInterrupt:
        n_pending = sum(1 for f in pending if f.running())
        logger.info(
            "interrupted by user, waiting for %d pending request(s) to finish",
            n_pending,
        )
        batch.interrupted = True
        extractor.cancel()
        executor.shutdown(wait=False, cancel_futures=True)