From f6ce9a9a1171fde8e4ff4499e408dd9430c27397 Mon Sep 17 00:00:00 2001 From: MaxWang007 Date: Fri, 13 Nov 2026 02:43:25 +0000 Subject: [PATCH] fix(llm): cancel queued chunks when on_chunk_done raises extract_corpus_parallel submits every chunk before handling callbacks. The executor then waits for that whole queue, so a callback error still runs the model calls that have not started. Cancel those futures, let in-flight calls finish, and re-raise the original exception (#5184). Co-Authored-By: Cursor --- graphify/llm.py | 41 +++++++++++++++--------- tests/test_chunking.py | 73 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 15 deletions(-) diff --git a/graphify/llm.py b/graphify/llm.py index 21b5e9bfcf..9bb67d4f6c 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -3027,21 +3027,32 @@ def _save_checkpoint(result: dict, allowed: "list[Path]") -> None: results_by_idx: dict[int, dict] = {} with ThreadPoolExecutor(max_workers=workers) as pool: futures = [pool.submit(_run_one, idx, chunk) for idx, chunk in enumerate(chunks)] - for future in as_completed(futures): - idx, result, exc = future.result() - if exc is not None: - print( - f"[graphify] chunk {idx + 1}/{total} failed: {exc}", - file=sys.stderr, - ) - merged["failed_chunks"] += 1 - failed_files.update(unit_path(f) for f in chunks[idx]) - continue - assert result is not None - results_by_idx[idx] = result - _checkpoint_chunk(idx, result, chunks[idx]) - if callable(on_chunk_done): - on_chunk_done(idx, total, result) + try: + for future in as_completed(futures): + idx, result, exc = future.result() + if exc is not None: + print( + f"[graphify] chunk {idx + 1}/{total} failed: {exc}", + file=sys.stderr, + ) + merged["failed_chunks"] += 1 + failed_files.update(unit_path(f) for f in chunks[idx]) + continue + assert result is not None + results_by_idx[idx] = result + _checkpoint_chunk(idx, result, chunks[idx]) + if callable(on_chunk_done): + on_chunk_done(idx, total, result) + except Exception: + # Every chunk was submitted above, and the callback runs only + # after one returns. Cancel work that has not started so a + # callback failure does not keep queued model calls running + # (#5184). cancel() leaves anything already running alone; the + # pool then drains that in-flight work and this exception + # propagates unchanged. + for pending in futures: + pending.cancel() + raise for idx in sorted(results_by_idx): _merge_into(merged, results_by_idx[idx]) diff --git a/tests/test_chunking.py b/tests/test_chunking.py index a5a6e8820d..ed4618e1b5 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -1,4 +1,5 @@ """Tests for token-aware chunking and parallel chunk execution in graphify.llm.""" +import threading import time from pathlib import Path from unittest.mock import patch @@ -979,6 +980,78 @@ def stub(chunk, **kwargs): assert len(result["nodes"]) == 4 +def test_corpus_parallel_cancels_queued_chunks_when_callback_raises(tmp_path, monkeypatch): + """A raising on_chunk_done must not run chunks that are still queued (#5184). + + The pool submits every chunk before the first callback. Two workers take + the first pair and the callback rejects the first result. Work already + inside the model may finish; a chunk that has not started must not be + called, and the callback's own exception must propagate. + """ + from graphify.llm import extract_corpus_parallel + + monkeypatch.setenv("GRAPHIFY_NO_INCREMENTAL_CACHE", "1") + files = [tmp_path / f"d{index}.md" for index in range(6)] + for path in files: + path.write_text("Cedar requires Beacon approval.\n") + + pair = threading.Barrier(2) + failed = threading.Event() + lock = threading.Lock() + calls = [] + callbacks = [] + # Long enough that cancel wins over the worker picking up the rest of the + # queue, short enough that the pre-fix run still finishes in this test. + hold_s = 0.3 + + def _result(path: Path) -> dict: + return { + "nodes": [{"id": path.stem, "label": path.stem, "source_file": str(path)}], + "edges": [], + "hyperedges": [], + "finish_reason": "stop", + } + + def fake(chunk, **kwargs): + index = int(chunk[0].stem[1:]) + with lock: + calls.append(index) + if index >= 2: + time.sleep(hold_s) + return _result(chunk[0]) + pair.wait(5) + if index != 0: + time.sleep(hold_s) + return _result(chunk[0]) + + def reject(idx, total, result): + callbacks.append(idx) + failed.set() + raise RuntimeError("callback rejected chunk") + + with patch("graphify.llm.extract_files_direct", side_effect=fake): + with pytest.raises(RuntimeError, match="callback rejected chunk"): + extract_corpus_parallel( + files, + backend="kimi", + api_key="unused", + model="unused", + root=tmp_path, + chunk_size=1, + token_budget=None, + max_concurrency=2, + max_retry_depth=0, + on_chunk_done=reject, + ) + + assert callbacks == [0] + assert failed.is_set() + # 0 and 1 were in flight. The worker freed by chunk 0 may enter chunk 2 + # before the callback returns. 3, 4, and 5 are still queued. + assert 0 in calls and 1 in calls, calls + assert all(index < 3 for index in calls), calls + + # ---- #1685: special-token strings in docs must not crash token estimation ---- def test_estimate_file_tokens_handles_tiktoken_special_token(tmp_path):