#!/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": ""} {"ok": false, "refused": true, "reason": ""} plus whatever extra keys that task's input object carried. Printed to stdout by default, or written to `--out `. 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())