Debugging Thread Pools

# before
for doc in docs:
    ex.submit(upsert, doc)
# after
futures = {ex.submit(upsert, doc): doc for doc in docs}
for fut in as_completed(futures):
    try:
        fut.result()
    except Exception as exc:
        failed[futures[fut]] = exc

# before
def embed_all(batches):
    ex = ThreadPoolExecutor(16)
    return list(ex.map(embed, batches))
# after
def embed_all(batches):
    with ThreadPoolExecutor(16, thread_name_prefix="embed") as ex:
        return list(ex.map(embed, batches))

# before
futures = [ex.submit(score, line) for line in open("prompts.jsonl")]
results = ex.map(score, (json.loads(l) for l in open("prompts.jsonl")))
# after: sliding window (any version) or buffersize (3.14+)
for i, res in window_map(score, open("prompts.jsonl"), ex, max_in_flight=64):
    write(i, res)
results = ex.map(score, (json.loads(l) for l in open("prompts.jsonl")), buffersize=64)

# before
with ThreadPoolExecutor(8) as ex:
    ex.map(upsert, batches)
    it = ex.map(fetch, urls, timeout=30)
pages = list(it)                       # after the with: everything already finished, timeout meaningless
# after
with ThreadPoolExecutor(8) as ex:
    for _ in ex.map(upsert, batches):
        pass
    pages = list(ex.map(fetch, urls, timeout=30))

# before
def work(doc):
    global TOKENS
    TOKENS += count(doc)
    RESULTS.append(parse(doc))
# after: return values and aggregate in the caller
def work(doc):
    return count(doc), parse(doc)
pairs = list(ex.map(work, docs))
tokens = sum(n for n, _ in pairs)

# before
SESSION = requests.Session()
def fetch(url):
    return SESSION.get(url, timeout=10).text
# after
_local = threading.local()
def fetch(url):
    if not hasattr(_local, "s"):
        _local.s = requests.Session()
    return _local.s.get(url, timeout=10).text

# before
def summarize_doc(doc):
    parts = [POOL.submit(summarize_chunk, c) for c in chunks(doc)]
    return combine([p.result() for p in parts])        # holds a POOL worker while waiting on POOL
results = [POOL.submit(summarize_doc, d) for d in docs]
# after: flatten (submit leaves from the caller), or a separate pool per level
chunk_futs = {POOL.submit(summarize_chunk, c): d for d in docs for c in chunks(d)}

# before
ex = ThreadPoolExecutor(max_workers=200)
# after
ex = ThreadPoolExecutor(max_workers=workers_for(rpm / 60, p95_latency_s), thread_name_prefix="llm")

# before
def search(queries):
    futures = [IO_POOL.submit(call_llm, q) for q in queries]
    return [f.result() for f in futures]
# after
def search(queries):
    futures = [IO_POOL.submit(call_llm, q) for q in queries]
    deadline = time.monotonic() + 10
    return [f.result(timeout=max(0.0, deadline - time.monotonic())) for f in futures]

# before
try:
    run_all(ex)
except KeyboardInterrupt:
    ex.shutdown(wait=False)
# after
try:
    run_all(ex)
except KeyboardInterrupt:
    stop.set()
    ex.shutdown(wait=True, cancel_futures=True)
    raise

# before
threading.Thread(target=flush_to_s3, args=(q,), daemon=True).start()
# after
t = threading.Thread(target=flush_to_s3, args=(q,), name="s3-flusher")
t.start()
q.put(None)
t.join()

# before
while not all(f.done() for f in futures):
    time.sleep(1)
# after
for f in as_completed(futures, timeout=3600):
    handle(f.result())

# before
try:
    rows.append(fut.result())
except Exception:
    pass
# after
try:
    rows.append(fut.result())
except Exception as exc:
    log.warning("task failed", exc_info=exc)
    errors.append(exc)

# before
@app.post("/answer")
def answer(q: str):
    with ThreadPoolExecutor(4) as ex:
        return list(ex.map(call_tool, plan(q)))
# after: one executor for the app, created in the lifespan, shut down on stop
def answer(q: str):
    return [f.result(timeout=20) for f in [IO_POOL.submit(call_tool, t) for t in plan(q)]]

# before
threading.Thread(target=warm_cache).start()
# after
t = threading.Thread(target=warm_cache, name="warm-cache")
t.start()
t.join(timeout=60)

# before
with ThreadPoolExecutor(8) as ex:
    scores = list(ex.map(edit_distance_matrix, pairs))
# after
with ProcessPoolExecutor() as ex:
    scores = list(ex.map(edit_distance_matrix, pairs, chunksize=64))

# before
fut = ex.submit(call_llm, prompt)                  # call_llm logs request_id.get() -> "-"
# after
fut = ex.submit(contextvars.copy_context().run, call_llm, prompt)

import sys
import threading
import time
import traceback
from concurrent.futures import ThreadPoolExecutor


def dump_threads() -> str:
    names = {t.ident: t.name for t in threading.enumerate()}
    out = []
    for ident, frame in sys._current_frames().items():
        top = traceback.extract_stack(frame)[-1]
        out.append(f"{names.get(ident, ident)}: {top.name} line {top.lineno}")
    return "\n".join(sorted(out))


def stuck_call() -> None:
    time.sleep(0.2)                                     # e.g. a request without a timeout


with ThreadPoolExecutor(max_workers=2, thread_name_prefix="llm") as ex:
    ex.submit(stuck_call).add_done_callback(lambda f: f.result())
    ex.submit(stuck_call).add_done_callback(lambda f: f.result())
    time.sleep(0.05)
    print("\n".join(line.split(" line ")[0] for line in dump_threads().splitlines() if line.startswith("llm")))
# llm_0: stuck_call
# llm_1: stuck_call

import threading


def report_live_threads() -> list:
    """Call at the end of main(): anything listed here will keep the process alive."""
    return sorted(t.name for t in threading.enumerate()
                  if t is not threading.main_thread() and not t.daemon and t.is_alive())


print(report_live_threads())                            # [] when everything was joined/shut down

import threading
import time
from concurrent.futures import Future, ThreadPoolExecutor
from typing import Any, Callable, Dict, List


class InstrumentedExecutor(ThreadPoolExecutor):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._lock = threading.Lock()
        self.in_flight = 0
        self.queue_wait: List[float] = []
        self.run_time: List[float] = []
        self.errors: Dict[str, int] = {}

    def submit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> Future:
        submitted = time.monotonic()

        def timed() -> Any:
            started = time.monotonic()
            with self._lock:
                self.in_flight += 1
                self.queue_wait.append(started - submitted)
            try:
                return fn(*args, **kwargs)
            except Exception as exc:
                with self._lock:
                    self.errors[type(exc).__name__] = self.errors.get(type(exc).__name__, 0) + 1
                raise
            finally:
                with self._lock:
                    self.in_flight -= 1
                    self.run_time.append(time.monotonic() - started)

        return super().submit(timed)


def call(i: int) -> int:
    time.sleep(0.01)
    if i % 10 == 0:
        raise TimeoutError("read timeout")
    return i


with InstrumentedExecutor(max_workers=4, thread_name_prefix="api") as ex:
    futures = [ex.submit(call, i) for i in range(40)]
    outcomes = [f.exception() is None for f in futures]
p95_wait = sorted(ex.queue_wait)[int(0.95 * len(ex.queue_wait))]
print(sum(outcomes), ex.errors, f"p95 queue wait > run time: {p95_wait > 0.01}")
# 36 {'TimeoutError': 4} p95 queue wait > run time: True