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:
+33
-3
@@ -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)
|
||||||
rewritten = [fmt(entry) for entry in parsed.entries]
|
|
||||||
|
if concurrency <= 1 or len(parsed.entries) <= 1:
|
||||||
|
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
|
||||||
|
|||||||
@@ -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)",
|
||||||
|
]
|
||||||
|
|||||||
Reference in New Issue
Block a user