diff --git a/summarize/runner.py b/summarize/runner.py index e3a3801..fc08627 100644 --- a/summarize/runner.py +++ b/summarize/runner.py @@ -437,21 +437,38 @@ def _progress(i: int, n: int, t0: float, c: dict) -> None: print(f" [{i}/{n}] {state} {elapsed:.0f}s 경과, 약 {eta:.0f}s 남음", file=sys.stderr, flush=True) -def _run_calls_parallel(llm, calls: list[dict], concurrency: int) -> None: - """calls 각 원소의 'prompt' 를 호출하고 결과를 그 자리에 채운다 (in-place).""" +def _run_calls_parallel(llm, calls: list[dict], concurrency: int, on_done=None) -> None: + """calls 각 원소의 'prompt' 를 호출하고 결과를 그 자리에 채운다 (in-place). + + on_done(call) 은 호출 하나가 끝날 때마다 **메인 스레드에서** 불린다 — 여기서 DB 에 저장하면 + 중간에 끊어도 그때까지 끝난 unit 은 남는다 (고객사 요청 2026-09-21). Ctrl+C 는 남은 호출을 + 취소하고 바로 올라간다 (진행 중인 호출은 응답 상한까지 기다릴 수 있다). + """ if not calls: return t0 = time.time() - if concurrency <= 1 or len(calls) == 1: + n = len(calls) + if concurrency <= 1 or n == 1: for i, c in enumerate(calls, 1): c.update(_llm_call(llm, SYSTEM_PROMPT, c["prompt"])) - _progress(i, len(calls), t0, c) + _progress(i, n, t0, c) + if on_done: + on_done(c) return - with ThreadPoolExecutor(max_workers=min(concurrency, len(calls))) as ex: + ex = ThreadPoolExecutor(max_workers=min(concurrency, n)) + try: futures = {ex.submit(_llm_call, llm, SYSTEM_PROMPT, c["prompt"]): c for c in calls} for i, fut in enumerate(as_completed(futures), 1): - futures[fut].update(fut.result()) - _progress(i, len(futures), t0, futures[fut]) + c = futures[fut] + c.update(fut.result()) + _progress(i, n, t0, c) + if on_done: + on_done(c) + except KeyboardInterrupt: + print("\n중단 요청 — 남은 호출을 취소한다 (이미 끝난 unit 은 저장됨)", file=sys.stderr, flush=True) + ex.shutdown(wait=False, cancel_futures=True) + raise + ex.shutdown(wait=True) def _dedupe_plan(con: sqlite3.Connection, pending: list[sqlite3.Row]) -> tuple[list[sqlite3.Row], dict]: @@ -581,13 +598,12 @@ def summarize_units(program: str | None, limit: int | None, fake: bool = False, jobs.append((u, [dict(x) for x in prompts])) # 2) LLM 호출 (병렬, DB 접근 없음) — 수정사항 10번 + # + 3) 검증·저장: unit 의 호출이 모두 끝나는 즉시 메인 스레드에서 저장 (중간에 끊어도 보존) all_calls = [c for _, calls in jobs for c in calls] print(f"[{p['name']}] unit {len(units)}건 → LLM 호출 {len(all_calls)}건 (동시 {conc})", file=sys.stderr, flush=True) - _run_calls_parallel(llm, all_calls, conc) - # 3) 검증·저장 (DB 쓰기, 메인 스레드) - for u, calls in jobs: + def _save_unit(u: sqlite3.Row, calls: list[dict]) -> None: lines = lines_by_inc.get(u["include"], []) try: r = finish_unit(con, u, calls, lines, text_symbols) @@ -598,15 +614,26 @@ def summarize_units(program: str | None, limit: int | None, fake: bool = False, con.commit() stats["failed"] += 1 print(f"[FAIL] {u['unit_id']}: {err}") - continue + return if r["status"] == "pending": stats["awaiting_response"] += 1 - continue + return stats["done"] += 1 stats["chunks"] += r["chunks"] stats["dropped"] += r["dropped"] con.commit() + owner = {id(c): (u, calls) for u, calls in jobs for c in calls} + remaining = {u["unit_id"]: len(calls) for u, calls in jobs} + + def _on_call_done(c: dict) -> None: + u, calls = owner[id(c)] + remaining[u["unit_id"]] -= 1 + if remaining[u["unit_id"]] == 0: # 창이 여러 개인 unit 은 마지막 창까지 기다린다 + _save_unit(u, calls) + + _run_calls_parallel(llm, all_calls, conc, on_done=_on_call_done) + if run_id is not None: _update_run(con, run_id, stats, llm, t0, status="running") diff --git a/tests/test_summarize_runner.py b/tests/test_summarize_runner.py index 1fc512e..aa5daf8 100644 --- a/tests/test_summarize_runner.py +++ b/tests/test_summarize_runner.py @@ -204,3 +204,26 @@ def test_reload_keeps_chunks_for_unchanged_units(ingested, monkeypatch): assert con.execute("SELECT summary_status FROM program WHERE name='ZRUNNER_T1'").fetchone()[0] == "stale" finally: con.close() + + + +def test_run_calls_parallel_saves_each_call_on_main_thread_as_it_finishes(): + """on_done 은 호출이 끝나는 대로, 메인 스레드에서 불린다 (중간 저장의 전제).""" + import threading + import time as _t + from summarize import runner + + class SlowLLM: + usage = {} + def complete_json(self, system, user): + _t.sleep(0.3 if user == "slow" else 0.01) + return {"unit_purpose_ko": user, "chunks": []} + + calls = [{"prompt": "slow"}, {"prompt": "a"}, {"prompt": "b"}] + seen = [] + main = threading.get_ident() + runner._run_calls_parallel(SlowLLM(), calls, 3, on_done=lambda c: seen.append((c["prompt"], threading.get_ident()))) + assert [p for p, _ in seen][-1] == "slow" # 느린 호출이 끝나기 전에 빠른 것들이 먼저 저장된다 + assert {p for p, _ in seen} == {"slow", "a", "b"} + assert all(tid == main for _, tid in seen) + assert all("raw" in c for c in calls)