From 22a58ffc20d75fb8eebe22f4ddde6922c857ee41 Mon Sep 17 00:00:00 2001 From: safishamsi Date: Wed, 24 Jun 2026 13:53:33 +0100 Subject: [PATCH] feat: parallel community labeling via --max-concurrency / --batch-size (#1390) label_communities ran batches one LLM call at a time, so a large graph needed hundreds of sequential calls even on backends that allow heavy concurrency. It now fans batches out across a thread pool, mirroring extract_corpus_parallel: results are returned per batch and merged on the main thread (labels dict is never mutated concurrently, no lock), and workers==1 keeps the original sequential path verbatim. ollama and claude-cli are forced serial unless the matching GRAPHIFY_*_PARALLEL env opt-in is set (same guard as extract). generate_community_labels threads max_concurrency + batch_size through, and the cluster-only/label CLI parses --max-concurrency and --batch-size (both `--flag N` and `--flag=N` forms; the space form is parsed explicitly so the value is not mistaken for the positional scan path by the arg-walk's catch-all). Output is deterministic regardless of concurrency (keyed by community id). Tests: parallel == sequential result, batch-size controls batch count, batches actually run concurrently, ollama forced serial, and the CLI parses both new flags. Full suite 2393 passed; ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 1 + graphify/__main__.py | 17 +++++++- graphify/llm.py | 70 +++++++++++++++++++++++-------- tests/test_labeling.py | 94 +++++++++++++++++++++++++++++++++++++++++- 4 files changed, 162 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f5cdde72..d30cc251 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ Full release notes with details on each version: [GitHub Releases](https://githu ## Unreleased +- Feat: community labeling can now run in parallel (#1390). `graphify cluster-only` and `graphify label` accept `--max-concurrency N` (default 4) to fan labeling batches out across a thread pool, and `--batch-size N` (default 100) to tune communities per LLM call. A large graph that previously needed hundreds of sequential calls now runs them in rounds. Mirrors the existing `extract` parallelism, including the safety guards: `ollama` and `claude-cli` are forced serial (set `GRAPHIFY_OLLAMA_PARALLEL=1` / `GRAPHIFY_CLAUDE_CLI_PARALLEL=1` to override). Output is unchanged and deterministic regardless of concurrency, since results are keyed by community id and merged on the main thread. - Fix: `graphify reflect` no longer duplicates lines in the "known dead ends" and "corrections" sections when the same Q&A is saved more than once. Those lists were appended per memory doc with no key (node scoring already dedups by node, but these two did not); they now collapse by question, keeping the most recent entry — so a re-corrected question shows its latest correction. Output stays deterministic (ordered by date then question). - Fix: the work-memory loop no longer depends on the git hook. The skill now tells the agent to run `graphify reflect --if-stale` itself at the start of graph work (cheap, deterministic, a no-op when no outcomes have been saved), then read `LESSONS.md`. Previously a skill-only install (without `graphify hook install`) would keep recording outcomes via `save-result` but never regenerate `LESSONS.md`, so the lessons never surfaced. The post-commit hook is now an optimization for between-session freshness rather than a requirement. The new `--if-stale` flag skips the run when `LESSONS.md` is already newer than every input (the memory docs and the graph), so when the hook just refreshed it the agent's session-start run costs almost nothing. diff --git a/graphify/__main__.py b/graphify/__main__.py index 1a153d4a..0d620058 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -2227,9 +2227,13 @@ def main() -> None: print(" --no-label keep 'Community N' placeholders (skip LLM community naming)") print(" --backend= backend to use for community naming (default: auto-detect)") print(" --model= model to use for community naming") + print(" --max-concurrency=N parallel community-labeling LLM calls (default 4; forced to 1 for ollama/claude-cli)") + print(" --batch-size=N communities per labeling LLM call (default 100)") print(" label (re)name communities with the configured LLM backend, regenerate report") print(" --backend= backend to use (default: auto-detect from API keys)") print(" --model= model to use for community naming") + print(" --max-concurrency=N parallel labeling LLM calls (default 4; forced to 1 for ollama/claude-cli)") + print(" --batch-size=N communities per labeling LLM call (default 100)") print(" query \"\" BFS traversal of graph.json for a question") print(" --dfs use depth-first instead of breadth-first") print(" --context C explicit edge-context filter (repeatable)") @@ -3309,6 +3313,8 @@ def main() -> None: graph_override: Path | None = None co_resolution: float = 1.0 co_exclude_hubs: float | None = None + label_max_concurrency: int = 4 + label_batch_size: int = 100 i_arg = 0 while i_arg < len(args): a = args[i_arg] @@ -3330,6 +3336,14 @@ def main() -> None: co_exclude_hubs = float(args[i_arg + 1]); i_arg += 2 elif a.startswith("--exclude-hubs="): co_exclude_hubs = float(a.split("=", 1)[1]); i_arg += 1 + elif a == "--max-concurrency" and i_arg + 1 < len(args): + label_max_concurrency = int(args[i_arg + 1]); i_arg += 2 + elif a.startswith("--max-concurrency="): + label_max_concurrency = int(a.split("=", 1)[1]); i_arg += 1 + elif a == "--batch-size" and i_arg + 1 < len(args): + label_batch_size = int(args[i_arg + 1]); i_arg += 2 + elif a.startswith("--batch-size="): + label_batch_size = int(a.split("=", 1)[1]); i_arg += 1 elif a == "--no-viz" or a.startswith("--min-community-size="): i_arg += 1 elif a.startswith("--"): @@ -3419,7 +3433,8 @@ def main() -> None: # The final labels (LLM or placeholder fallback) are persisted to # .graphify_labels.json by the unconditional write below. labels, _ = generate_community_labels( - G, communities, backend=label_backend, model=label_model, gods=gods + G, communities, backend=label_backend, model=label_model, gods=gods, + max_concurrency=label_max_concurrency, batch_size=label_batch_size, ) questions = suggest_questions(G, communities, labels) tokens = {"input": 0, "output": 0} diff --git a/graphify/llm.py b/graphify/llm.py index f1d61478..9e4afcf6 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -2230,6 +2230,7 @@ def label_communities( max_communities: int | None = None, top_k: int = _LABEL_TOP_K, batch_size: int = _LABEL_BATCH_SIZE, + max_concurrency: int = 4, ) -> dict[int, str]: """Return a complete ``{cid: name}`` map using ``backend`` for naming. @@ -2257,32 +2258,62 @@ def label_communities( return labels n_batches = (len(labeled_cids) + batch_size - 1) // batch_size - written = 0 - first_error: Exception | None = None - for batch_idx in range(n_batches): + + # Mirror extract_corpus_parallel's backend guards: Ollama serves one request at + # a time per loaded model (parallel batches cause VRAM pressure and hollow + # replies, #798) and claude-cli shells out to a single Claude Code session that + # parallel subprocesses corrupt. Force serial for these unless the user opts in + # via the same env switches. + if backend == "ollama" and os.environ.get("GRAPHIFY_OLLAMA_PARALLEL", "").strip() != "1": + max_concurrency = 1 + if backend == "claude-cli" and os.environ.get("GRAPHIFY_CLAUDE_CLI_PARALLEL", "").strip() != "1": + max_concurrency = 1 + workers = max(1, min(max_concurrency, n_batches)) + + def _run_batch(batch_idx: int): start = batch_idx * batch_size end = min(start + batch_size, len(labeled_cids)) - batch_lines = lines[start:end] - batch_cids = labeled_cids[start:end] try: parsed = _label_batch_with_retry( - batch_cids, batch_lines, backend=backend, model=model, + labeled_cids[start:end], lines[start:end], backend=backend, model=model, ) - labels.update(parsed) - written += len(parsed) - except Exception as exc: - if first_error is None: - first_error = exc + return batch_idx, parsed, None + except Exception as exc: # noqa: BLE001 - reported per-batch; surfaced below + return batch_idx, None, exc + + written = 0 + errors: dict[int, Exception] = {} + + def _merge(batch_idx: int, parsed, exc) -> None: + nonlocal written + if exc is not None: + errors[batch_idx] = exc + start = batch_idx * batch_size + end = min(start + batch_size, len(labeled_cids)) print( f"[graphify label] batch {batch_idx + 1}/{n_batches} " - f"({len(batch_cids)} communities) failed: {exc}", + f"({end - start} communities) failed: {exc}", file=sys.stderr, ) - continue + return + labels.update(parsed) + written += len(parsed) - if written == 0 and first_error is not None: - # Every batch failed; propagate so generate_community_labels degrades cleanly. - raise first_error + # Fan out batches; merge on the main thread so `labels` is never mutated + # concurrently. workers == 1 keeps the original sequential path verbatim. + if workers == 1: + for batch_idx in range(n_batches): + _merge(*_run_batch(batch_idx)) + else: + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = [pool.submit(_run_batch, b) for b in range(n_batches)] + for future in as_completed(futures): + _merge(*future.result()) + + if written == 0 and errors: + # Every batch failed; propagate the lowest-index error so the message is + # deterministic and generate_community_labels degrades cleanly. + raise errors[min(errors)] return labels @@ -2294,6 +2325,8 @@ def generate_community_labels( model: str | None = None, gods=None, quiet: bool = False, + max_concurrency: int = 4, + batch_size: int = _LABEL_BATCH_SIZE, ) -> tuple[dict[int, str], str]: """CLI entry point: resolve a backend, name communities, and degrade to ``Community N`` placeholders on any failure (no backend, API error, malformed @@ -2313,7 +2346,10 @@ def generate_community_labels( ) return _placeholder_community_labels(communities), "placeholder" try: - labels = label_communities(G, communities, backend=backend, model=model, gods=gods) + labels = label_communities( + G, communities, backend=backend, model=model, gods=gods, + max_concurrency=max_concurrency, batch_size=batch_size, + ) return labels, "llm" except Exception as exc: if not quiet: diff --git a/tests/test_labeling.py b/tests/test_labeling.py index a228b6d7..b846a00a 100644 --- a/tests/test_labeling.py +++ b/tests/test_labeling.py @@ -81,9 +81,12 @@ def test_label_cli_passes_model_override(tmp_path, monkeypatch): captured = {} - def fake_generate(G, communities, *, backend=None, model=None, gods=None, quiet=False): + def fake_generate(G, communities, *, backend=None, model=None, gods=None, + quiet=False, max_concurrency=4, batch_size=100): captured["backend"] = backend captured["model"] = model + captured["max_concurrency"] = max_concurrency + captured["batch_size"] = batch_size return {0: "Orders"}, "llm" monkeypatch.setattr("graphify.llm.generate_community_labels", fake_generate) @@ -99,13 +102,22 @@ def test_label_cli_passes_model_override(tmp_path, monkeypatch): "gemini", "--model", "gemini-3.1-flash-lite", + "--max-concurrency", + "8", + "--batch-size", + "50", "--no-viz", ], ) cli.main() - assert captured == {"backend": "gemini", "model": "gemini-3.1-flash-lite"} + # Also verifies the space-separated forms parse (the value must not be mistaken + # for the positional path) and reach generate_community_labels. + assert captured == { + "backend": "gemini", "model": "gemini-3.1-flash-lite", + "max_concurrency": 8, "batch_size": 50, + } def test_label_communities_partial_reply_fills_placeholder(monkeypatch): @@ -278,3 +290,81 @@ def test_label_communities_max_communities_caps_total(monkeypatch): label_communities(G, communities, backend="gemini", max_communities=40, batch_size=100) # Only 40 communities should have been sent to the backend. assert len(captured_cids) == 40 + + +# --- #1390: parallel labeling (--max-concurrency) + --batch-size -------------- + +import threading +import time as _time + + +def _many_communities(n): + G = nx.Graph() + comms = {} + for i in range(n): + nid = f"n{i}" + G.add_node(nid, label=f"sym_{i}") + comms[i] = [nid] + return G, comms + + +def test_label_communities_parallel_matches_sequential(monkeypatch): + """Concurrency must not change the result: same cid->name map either way.""" + G, communities = _many_communities(6) + + def fake_batch(batch_cids, batch_lines, *, backend, model=None): + return {cid: f"name-{cid}" for cid in batch_cids} + + monkeypatch.setattr("graphify.llm._label_batch_with_retry", fake_batch) + seq = label_communities(G, communities, backend="gemini", batch_size=1, max_concurrency=1) + par = label_communities(G, communities, backend="gemini", batch_size=1, max_concurrency=4) + assert seq == par == {i: f"name-{i}" for i in range(6)} + + +def test_label_communities_batch_size_controls_batch_count(monkeypatch): + G, communities = _many_communities(5) + calls = [] + + def fake_batch(batch_cids, batch_lines, *, backend, model=None): + calls.append(list(batch_cids)) + return {cid: f"n-{cid}" for cid in batch_cids} + + monkeypatch.setattr("graphify.llm._label_batch_with_retry", fake_batch) + labels = label_communities(G, communities, backend="gemini", batch_size=2, max_concurrency=1) + assert len(calls) == 3 # 5 communities / batch 2 -> 3 batches + assert sum(len(c) for c in calls) == 5 + assert labels == {i: f"n-{i}" for i in range(5)} + + +def _peak_tracker(): + lock = threading.Lock() + state = {"now": 0, "peak": 0} + + def fake_batch(batch_cids, batch_lines, *, backend, model=None): + with lock: + state["now"] += 1 + state["peak"] = max(state["peak"], state["now"]) + _time.sleep(0.03) + with lock: + state["now"] -= 1 + return {cid: f"n-{cid}" for cid in batch_cids} + + return fake_batch, state + + +def test_label_communities_runs_batches_concurrently(monkeypatch): + G, communities = _many_communities(8) + fake_batch, state = _peak_tracker() + monkeypatch.setattr("graphify.llm._label_batch_with_retry", fake_batch) + label_communities(G, communities, backend="gemini", batch_size=1, max_concurrency=4) + assert state["peak"] > 1, "batches should run in parallel with max_concurrency>1" + + +def test_label_communities_forces_serial_for_ollama(monkeypatch): + """ollama/claude-cli must stay serial regardless of --max-concurrency.""" + G, communities = _many_communities(8) + fake_batch, state = _peak_tracker() + monkeypatch.setattr("graphify.llm._label_batch_with_retry", fake_batch) + monkeypatch.delenv("GRAPHIFY_OLLAMA_PARALLEL", raising=False) + label_communities(G, communities, backend="ollama", batch_size=1, max_concurrency=8) + assert state["peak"] == 1, "ollama must be forced serial"