cli: parallelize formatter calls with thread pool

reformat_draft now uses concurrent.futures.ThreadPoolExecutor with a
default of 8 workers. OpenAI SDK calls are synchronous but network-
bound, so threads release the GIL during I/O and give real speedup.
ThreadPoolExecutor.map preserves input order regardless of completion
order, so output is deterministic.

Empirical: HML draft (134 entries) went from ~25 min serial to 1m55s
with concurrency=12. ~12x speedup; further increases hit OpenAI rate
limits.

CLI gains a --concurrency flag (default 8) for tuning per draft size /
rate limit headroom. concurrency=1 forces serial execution for
debugging. New test asserts that order is preserved when concurrent
calls finish out of input order (uses a sleep-by-index fake formatter).
This commit is contained in:
cmos dev
2026-04-11 00:42:56 -04:00
parent 1b11a03a0b
commit 2bd1bdf0d4
2 changed files with 68 additions and 3 deletions
+32 -2
View File
@@ -9,6 +9,7 @@ from __future__ import annotations
import argparse import argparse
import sys import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path from pathlib import Path
from typing import Callable from typing import Callable
@@ -17,16 +18,36 @@ from cmos.parser import split_bibliography
FormatterFn = Callable[[str], str] FormatterFn = Callable[[str], str]
DEFAULT_CONCURRENCY = 8
def reformat_draft(text: str, formatter: FormatterFn | None = None) -> str:
def reformat_draft(
text: str,
formatter: FormatterFn | None = None,
concurrency: int = DEFAULT_CONCURRENCY,
) -> str:
"""Rewrite the bibliography section of ``text`` using ``formatter``. """Rewrite the bibliography section of ``text`` using ``formatter``.
The formatter takes a single messy entry and returns the CMOS 18 form. The formatter takes a single messy entry and returns the CMOS 18 form.
If omitted, the real OpenAI-backed formatter is used. If omitted, the real OpenAI-backed formatter is used.
Entries are reformatted concurrently using a thread pool of size
``concurrency``. OpenAI SDK calls are synchronous but network-bound, so
threads give real speedup. ``ThreadPoolExecutor.map`` preserves input
order regardless of which calls finish first, so the output is
deterministic.
Set ``concurrency=1`` for sequential execution (useful for
deterministic debugging or when hitting rate limits).
""" """
fmt = formatter or format_bibliography_entry fmt = formatter or format_bibliography_entry
parsed = split_bibliography(text) parsed = split_bibliography(text)
if concurrency <= 1 or len(parsed.entries) <= 1:
rewritten = [fmt(entry) for entry in parsed.entries] rewritten = [fmt(entry) for entry in parsed.entries]
else:
with ThreadPoolExecutor(max_workers=concurrency) as pool:
rewritten = list(pool.map(fmt, parsed.entries))
pieces: list[str] = [] pieces: list[str] = []
if parsed.before: if parsed.before:
@@ -46,11 +67,20 @@ def main(argv: list[str] | None = None) -> int:
fmt = sub.add_parser("format", help="Reformat a draft's bibliography to CMOS 18.") fmt = sub.add_parser("format", help="Reformat a draft's bibliography to CMOS 18.")
fmt.add_argument("path", type=Path, help="Path to a markdown draft.") fmt.add_argument("path", type=Path, help="Path to a markdown draft.")
fmt.add_argument(
"--concurrency",
type=int,
default=DEFAULT_CONCURRENCY,
help=(
"Number of parallel formatter calls (default: "
f"{DEFAULT_CONCURRENCY}). Set to 1 for serial execution."
),
)
args = ap.parse_args(argv) args = ap.parse_args(argv)
if args.command == "format": if args.command == "format":
text = args.path.read_text(encoding="utf-8") text = args.path.read_text(encoding="utf-8")
sys.stdout.write(reformat_draft(text)) sys.stdout.write(reformat_draft(text, concurrency=args.concurrency))
return 0 return 0
return 2 return 2
+35
View File
@@ -66,3 +66,38 @@ def test_reformat_draft_bibliography_heading_preserved():
draft = "## Bibliography\n\nentry.\n" draft = "## Bibliography\n\nentry.\n"
output = reformat_draft(draft, formatter=_fake_formatter) output = reformat_draft(draft, formatter=_fake_formatter)
assert "## Bibliography" in output assert "## Bibliography" in output
def test_reformat_draft_preserves_order_under_concurrency():
# With concurrent execution the formatter is called on all entries in
# parallel; the CLI must reassemble them in input order regardless of
# which call finishes first. This test uses a fake formatter that
# sleeps based on the entry content so earlier entries finish AFTER
# later ones if order were naively tied to completion.
import time
def slow_fake(messy: str) -> str:
# Entries with lower index sleep longer so they finish last.
n = int(messy.split()[-1])
time.sleep(0.05 * (5 - n))
return f"FORMATTED({messy})"
draft = """\
## Bibliography
entry 0
entry 1
entry 2
entry 3
entry 4
"""
output = reformat_draft(draft, formatter=slow_fake, concurrency=4)
# Entries must appear in input order.
lines = [l for l in output.splitlines() if l.startswith("FORMATTED(")]
assert lines == [
"FORMATTED(entry 0)",
"FORMATTED(entry 1)",
"FORMATTED(entry 2)",
"FORMATTED(entry 3)",
"FORMATTED(entry 4)",
]