From 40f504314aba67c0c1c1d074c6f042589b1c43b0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 9 Jun 2026 22:21:11 +0000 Subject: [PATCH] fix(pickled-core): make mine features stage --max-parallel actually parallelize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Calling FeatureDrafter.draft_from_story (sync HTTP) from inside an `async def` blocked the event loop for the entire LLM round-trip, so the asyncio.Semaphore never released to a waiting coroutine. The advertised `--max-parallel N` flag silently degraded to 1-way serial execution despite the loop and semaphore looking correct on paper. Concrete blast radius: a workspace with 50 surfaces and 5s LLM calls and `--max-parallel 8` should mine in ~63s but actually took ~250s — a 4× regression on a CLI flag advertised to do parallel work, paid in real LLM-latency wallclock. The stories stage was already correct because it scheduled the same kind of sync work via asyncio.to_thread; only the features stage missed that handoff. Fix: extract `_draft_one_sync` as a synchronous helper and dispatch it through `asyncio.to_thread` from the parallel runner. The interactive (non-quick) path now also calls `_draft_one_sync` directly instead of round-tripping through `asyncio.run` — a small cleanup that lets a test reliably observe the LLM call without dragging in event-loop machinery. Adds a regression test that monkeypatches FeatureDrafter to track peak concurrent in-flight calls; with the fix peak >= 2 and 8 jobs of 0.2s finish in well under 1.5s, while the previous code pinned peak at 1 and took ~1.6s on the same hardware. Co-authored-by: Bartłomiej Rosa --- .../src/pickled_core/mine/features_stage.py | 35 ++++++---- .../tests/test_mine_features_stage.py | 69 +++++++++++++++++++ 2 files changed, 90 insertions(+), 14 deletions(-) diff --git a/packages/pickled-core/src/pickled_core/mine/features_stage.py b/packages/pickled-core/src/pickled_core/mine/features_stage.py index 49f1ab1..1354622 100644 --- a/packages/pickled-core/src/pickled_core/mine/features_stage.py +++ b/packages/pickled-core/src/pickled_core/mine/features_stage.py @@ -47,7 +47,7 @@ def _filter_story_paths( return filtered -async def _draft_one( +def _draft_one_sync( story_path: Path, feature_path: Path, llm: LLMClient, @@ -55,6 +55,15 @@ async def _draft_one( overwrite: bool, interactive: bool, ) -> FeatureResult: + """Draft one feature synchronously (used by both serial and parallel paths). + + This function performs blocking I/O (LLM HTTP call + file writes). The + parallel path schedules it on a worker thread via :func:`asyncio.to_thread` + so that ``--max-parallel`` actually overlaps LLM calls. Calling the LLM + directly inside an ``async def`` would block the event loop and serialise + every draft, defeating the flag (and producing N×slower mining runs for + workspaces with many surfaces). + """ surface_id = _story_surface_id(story_path) if feature_path.is_file() and not overwrite: return FeatureResult( @@ -122,7 +131,8 @@ async def _run_parallel( async def _one(story_path: Path, feature_path: Path) -> None: async with sem: results.append( - await _draft_one( + await asyncio.to_thread( + _draft_one_sync, story_path, feature_path, llm, @@ -178,19 +188,16 @@ def run_features( _run_parallel(jobs, llm, overwrite=overwrite, max_parallel=max_parallel) ) else: - results = [] - for story_path, feature_path in jobs: - results.append( - asyncio.run( - _draft_one( - story_path, - feature_path, - llm, - overwrite=overwrite, - interactive=True, - ) - ) + results = [ + _draft_one_sync( + story_path, + feature_path, + llm, + overwrite=overwrite, + interactive=True, ) + for story_path, feature_path in jobs + ] return FeaturesStageResult(output_dir=paths.root, results=results, skipped_entire_stage=False) diff --git a/packages/pickled-core/tests/test_mine_features_stage.py b/packages/pickled-core/tests/test_mine_features_stage.py index f2a2f7a..6ad0f12 100644 --- a/packages/pickled-core/tests/test_mine_features_stage.py +++ b/packages/pickled-core/tests/test_mine_features_stage.py @@ -2,6 +2,8 @@ from __future__ import annotations +import threading +import time from pathlib import Path from typing import Any @@ -92,3 +94,70 @@ class _Client: assert not result.skipped_entire_stage assert len(result.results) == 1 assert result.results[0].surface_id == "tiny_target_greet" + + +def test_max_parallel_actually_overlaps_llm_calls( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """`--max-parallel N` must let up to N LLM drafts run concurrently. + + Regression for a bug where ``_draft_one`` was an ``async def`` that + invoked the synchronous ``FeatureDrafter.draft_from_story`` directly: + every coroutine blocked the event loop for the entire LLM call, so + the asyncio.Semaphore never released to a waiting task and the + advertised ``--max-parallel`` flag silently degraded to 1-way + serial execution. With 50 surfaces × 5s LLM calls, that turns a + ~63 s mining run with ``--max-parallel 8`` into ~250 s. + """ + stories = tmp_path / "stories" + stories.mkdir() + for index in range(8): + (stories / f"surf_{index}.story.md").write_text("# story\n", encoding="utf-8") + + state: dict[str, int] = {"current": 0, "peak": 0} + lock = threading.Lock() + + class _SlowDrafter: + def __init__(self, llm: Any) -> None: + _ = llm + + def draft_from_story(self, story: str) -> DraftResult: + _ = story + with lock: + state["current"] += 1 + if state["current"] > state["peak"]: + state["peak"] = state["current"] + try: + time.sleep(0.2) + finally: + with lock: + state["current"] -= 1 + return DraftResult( + text="Feature: x\n\n Scenario: y\n Given z\n", + rationale="", + warnings=(), + ) + + monkeypatch.setattr("pickled_core.mine.features_stage.FeatureDrafter", _SlowDrafter) + + class _Client: + pass + + started = time.time() + result = run_features( + tmp_path, + llm=_Client(), # type: ignore[arg-type] + quick=True, + overwrite=True, + max_parallel=4, + ) + elapsed = time.time() - started + + assert len(result.results) == 8 + # 4-way parallelism over 8 jobs of 0.2s each: ideal ~0.4s, accept up to 1.5s + # for slow CI jitter; serial would take ~1.6s and we want a clear gap. + assert state["peak"] >= 2, ( + f"max-parallel=4 should overlap LLM calls; observed peak={state['peak']}" + ) + assert elapsed < 1.5, f"max-parallel=4 should run in well under serial time; got {elapsed:.2f}s"