Files
willtools-demo/tools/dsh_pool.py
T
will 8e6ee31596
ci / test (push) Successful in 6m18s
WillTools 2026-09-09 -- clean snapshot of the development repo for review
History-free snapshot of the private development repository at cbe45c5
(2026-09-09). Code, docs and public catalogue only; no plant or job
knowledge layers, no client or job names (tools/scrub.py clean). The
intro deck is left out of this snapshot because its slide images could
not be checked for job drawings.
2026-09-09 08:32:15 -05:00

214 lines
9.1 KiB
Python

#!/usr/bin/env python3
"""dsh_pool - fan N tasks out to dsh_worker concurrently, bounded by a worker
cap, and hand back one result per task, in input order, each independently
marked ok-or-refused.
python tools/dsh_pool.py --tasks tasks.json --out results.json
python tools/dsh_pool.py --tasks tasks.json --max-workers 8 --task-timeout 120
`tasks.json` is a JSON array. Each element is either a plain string (the
task text) or an object with a `"task"` key (any other keys on that object
-- an id, a label -- are echoed back unchanged alongside that task's
result, so a caller can match results back to whatever it cared about
without relying on array position). Example:
["summarize file_a.py", "summarize file_b.py"]
[{"id": "a", "task": "summarize file_a.py"},
{"id": "b", "task": "summarize file_b.py"}]
Output is a JSON array, same length and same order as the input, one
object per task:
{"ok": true, "reply": "<the agent's final message>"}
{"ok": false, "refused": true, "reason": "<what went wrong>"}
plus whatever extra keys that task's input object carried. Printed to
stdout by default, or written to `--out <path>`.
WHY A POOL, NOT JUST N CALLS TO dsh_worker.py
-----------------------------------------------
`tools/dsh_worker.py`'s own docstring says a cheap backend's value is in
wide, cheap fan-out -- "N similar small tasks where a wrong one in isolation
is not costly" (docs/dsh-worker.md's routing table). Running dsh_worker.py N
times in a row gets none of that: it is still one task at a time, at
dsh_worker's own per-call timeout, serialized. This file is what actually
makes it a *pool*: a bounded ThreadPoolExecutor runs up to `max_workers`
dsh CLI processes at once (dsh_worker.run_task shells out to a real
subprocess per task, so there is no GIL contention worth worrying about --
the threads are waiting on I/O, not competing for the interpreter), and
every task gets its own DshRefusal-or-reply outcome independently of every
other task's.
THE RULE THIS FILE MUST NOT BREAK: one task failing must never take the
batch down, and a refused task must be distinguishable from a task that
succeeded with a short answer. Both are enforced the same way dsh_worker.py
enforces "could not judge" vs "judged, found nothing" (CLAUDE.md SS3/SS5):
every task's outcome is captured into its own result dict inside the worker
function passed to the executor, so an exception in task 3 becomes task 3's
`{"ok": false, "refused": true, ...}" entry, not an exception that aborts
tasks 4..N or corrupts the ordering promise.
Standard library only (concurrent.futures). CPython 3.8+, host-only --
same exemption from the IronPython 2.7 rules as tools/dsh_worker.py, which
this module imports as a sibling in tools/.
"""
from __future__ import annotations
import argparse
import json
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
_TOOLS_DIR = os.path.dirname(os.path.abspath(__file__))
if _TOOLS_DIR not in sys.path:
sys.path.insert(0, _TOOLS_DIR)
import dsh_worker # noqa: E402 (sibling module, path set up above)
DEFAULT_MAX_WORKERS = int(os.environ.get('DSH_POOL_MAX_WORKERS', '4'))
DEFAULT_TASK_TIMEOUT_S = int(os.environ.get(
'DSH_POOL_TASK_TIMEOUT_S', str(dsh_worker.DEFAULT_TIMEOUT_S)))
def _run_one(task_text, profile, timeout_s, run_task):
"""Run one task and turn any DshRefusal (or any other exception -- a
task must never be allowed to kill the batch just because it raised
something dsh_worker did not anticipate) into a result dict shaped
exactly like dsh_worker's own CLI output.
"""
try:
reply = run_task(task_text, profile=profile, timeout_s=timeout_s)
return {'ok': True, 'reply': reply}
except dsh_worker.DshRefusal as exc:
return {'ok': False, 'refused': True, 'reason': str(exc)}
except Exception as exc: # noqa: BLE001 - deliberately broad: see
# module docstring, "must never take the
# batch down"
return {'ok': False, 'refused': True,
'reason': 'unexpected error running task ({0}): {1}'
.format(type(exc).__name__, exc)}
def run_pool(tasks, profile=None, max_workers=DEFAULT_MAX_WORKERS,
timeout_s=DEFAULT_TASK_TIMEOUT_S, run_task=None):
"""Run `tasks` (a list of task strings) concurrently, bounded by
`max_workers` at once. Returns a list the same length as `tasks`, in
the same order as `tasks`, each entry either
`{"ok": True, "reply": ...}` or
`{"ok": False, "refused": True, "reason": ...}`.
`run_task` is injected the same way dsh_worker.run_task injects `which`
/ `runner` -- so a test can fake the whole dsh round trip per task
without a real dsh install. Left as `None` it is resolved to
`dsh_worker.run_task` at call time (not bound as a default at import
time), so a caller/test that monkeypatches `dsh_worker.run_task` still
reaches its fake through `tools/dsh_pool.py main()`, which never passes
`run_task` explicitly.
"""
if run_task is None:
run_task = dsh_worker.run_task
profile = profile or dsh_worker.DEFAULT_PROFILE
max_workers = max(1, int(max_workers))
results = [None] * len(tasks)
if not tasks:
return results
with ThreadPoolExecutor(max_workers=max_workers) as pool:
future_to_index = {
pool.submit(_run_one, task_text, profile, timeout_s, run_task): i
for i, task_text in enumerate(tasks)
}
for future in as_completed(future_to_index):
i = future_to_index[future]
# _run_one already catches everything it can reach; .result()
# cannot raise here except for a framework-level cancellation,
# which this code never issues.
results[i] = future.result()
return results
def _load_tasks(path):
"""Read the tasks file. Returns (task_texts, extras) where `extras[i]`
is the dict of any non-"task" keys the i'th input object carried (empty
for a plain-string task), to be merged back into that task's result.
"""
with open(path, 'r', encoding='utf-8') as fh:
data = json.load(fh)
if not isinstance(data, list):
raise SystemExit('tasks file must contain a JSON array, got {0}'
.format(type(data).__name__))
task_texts = []
extras = []
for i, item in enumerate(data):
if isinstance(item, str):
task_texts.append(item)
extras.append({})
elif isinstance(item, dict) and 'task' in item:
task_texts.append(item['task'])
extras.append({k: v for k, v in item.items() if k != 'task'})
else:
raise SystemExit(
'tasks[{0}] must be a string or an object with a "task" '
'key, got: {1!r}'.format(i, item))
return task_texts, extras
def main(argv=None):
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
ap.add_argument('--tasks', required=True,
help='JSON file: array of task strings, or array of '
'{"task": ...} objects')
ap.add_argument('--out',
help='write the JSON results array here (default: '
'stdout)')
ap.add_argument('--profile', default=dsh_worker.DEFAULT_PROFILE,
help='dsh profile every task boots (default: {0}; also '
'settable with $DSH_WORKER_PROFILE)'
.format(dsh_worker.DEFAULT_PROFILE))
ap.add_argument('--max-workers', type=int, default=DEFAULT_MAX_WORKERS,
help='how many dsh calls run at once (default: {0}; '
'also settable with $DSH_POOL_MAX_WORKERS)'
.format(DEFAULT_MAX_WORKERS))
ap.add_argument('--task-timeout', type=int,
default=DEFAULT_TASK_TIMEOUT_S,
help='seconds to wait before refusing any one task '
'(default: {0}; also settable with '
'$DSH_POOL_TASK_TIMEOUT_S)'
.format(DEFAULT_TASK_TIMEOUT_S))
args = ap.parse_args(argv)
task_texts, extras = _load_tasks(args.tasks)
results = run_pool(task_texts, profile=args.profile,
max_workers=args.max_workers,
timeout_s=args.task_timeout)
merged = []
for extra, result in zip(extras, results):
row = dict(extra)
row.update(result)
merged.append(row)
payload = json.dumps(merged, ensure_ascii=False, indent=2)
if args.out:
with open(args.out, 'w', encoding='utf-8') as fh:
fh.write(payload)
fh.write('\n')
else:
print(payload)
# The batch itself ran to completion -- that is what this exit code
# reports. Whether any individual task was refused is in the JSON,
# per-task, by design (see module docstring): one refusal must never
# look like the whole run failing, or a caller scripting around this
# would stop trusting the N-1 tasks that actually succeeded.
return 0
if __name__ == '__main__':
raise SystemExit(main())