will
25592fa6c2
ci / test (push) Successful in 6m14s
History-free snapshot of the private development repository at 77a82eb. Code, docs and the public catalogue only: no plant or job knowledge layers, no client, job or producer names (tools/scrub.py clean against the widened term list; a second word-level grep clean). The intro deck is left out because its slide images cannot be checked for job drawings.
214 lines
9.1 KiB
Python
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())
|