diff --git a/.gitignore b/.gitignore index 6c23839..7f6221a 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,6 @@ notifiers/* # This lab's own settings: what it runs, its models, its channels. lab.yaml litellm/config.yaml +# pixi environments +.pixi/* +!.pixi/config.toml diff --git a/campaigns/example-local-compression/task.py b/campaigns/example-local-compression/task.py index 0fdcf01..33c2ce8 100644 --- a/campaigns/example-local-compression/task.py +++ b/campaigns/example-local-compression/task.py @@ -12,7 +12,8 @@ BANDWIDTH_MB_S = 5.0 CORPUS_MB = 8 -LOCAL_DESC = """ +LOCAL_DESC = ( + """ Compress the benchmark corpus with one zlib configuration and return its timings. Returns `total_seconds` -- encode time plus transmit time at @@ -30,21 +31,45 @@ Each run rebuilds the same corpus from a fixed seed, so results are comparable across jobs. A job takes a few seconds. -""" % BANDWIDTH_MB_S +""" + % BANDWIDTH_MB_S +) LOCAL_SCHEMA = {"level": int, "strategy": str, "memlevel": int, "windowlog": int} -_STRATEGIES = {"default": "Z_DEFAULT_STRATEGY", "filtered": "Z_FILTERED", - "huffman_only": "Z_HUFFMAN_ONLY", "rle": "Z_RLE", "fixed": "Z_FIXED"} +_STRATEGIES = { + "default": "Z_DEFAULT_STRATEGY", + "filtered": "Z_FILTERED", + "huffman_only": "Z_HUFFMAN_ONLY", + "rle": "Z_RLE", + "fixed": "Z_FIXED", +} def _corpus(mb): """Deterministic mixed-entropy corpus: mostly repeated vocabulary, some noise, so that both match-finding and entropy coding have something to do.""" import random + rnd = random.Random(1234) - words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta", - "run", "job", "node", "rank", "queue", "kernel", "buffer", "stride"] + words = [ + "alpha", + "beta", + "gamma", + "delta", + "epsilon", + "zeta", + "eta", + "theta", + "run", + "job", + "node", + "rank", + "queue", + "kernel", + "buffer", + "stride", + ] out, size, target = [], 0, mb * 1024 * 1024 while size < target: if rnd.random() < 0.75: @@ -68,8 +93,10 @@ def local_fn(args): if not 1 <= level <= 9: return {"error": f"level must be 1-9, got {level}", "args": args} if strategy not in _STRATEGIES: - return {"error": f"strategy must be one of {sorted(_STRATEGIES)}, got {strategy!r}", - "args": args} + return { + "error": f"strategy must be one of {sorted(_STRATEGIES)}, got {strategy!r}", + "args": args, + } if not 1 <= memlevel <= 9: return {"error": f"memlevel must be 1-9, got {memlevel}", "args": args} if not 9 <= windowlog <= 15: @@ -85,8 +112,12 @@ def local_fn(args): transmit_seconds = len(blob) / (BANDWIDTH_MB_S * 1024 * 1024) return { - "args": {"level": level, "strategy": strategy, - "memlevel": memlevel, "windowlog": windowlog}, + "args": { + "level": level, + "strategy": strategy, + "memlevel": memlevel, + "windowlog": windowlog, + }, "total_seconds": round(encode_seconds + transmit_seconds, 4), "encode_seconds": round(encode_seconds, 4), "transmit_seconds": round(transmit_seconds, 4), @@ -94,6 +125,9 @@ def local_fn(args): "encode_mb_s": round(len(data) / (1024 * 1024) / encode_seconds, 2), "raw_bytes": len(data), "compressed_bytes": len(blob), - "diagnostics": {"bandwidth_mb_s": BANDWIDTH_MB_S, "corpus_mb": CORPUS_MB, - "zlib_version": zlib.ZLIB_VERSION}, + "diagnostics": { + "bandwidth_mb_s": BANDWIDTH_MB_S, + "corpus_mb": CORPUS_MB, + "zlib_version": zlib.ZLIB_VERSION, + }, } diff --git a/campaigns/example-quick-optimum/task.py b/campaigns/example-quick-optimum/task.py index 2d2ada7..fd67b8b 100644 --- a/campaigns/example-quick-optimum/task.py +++ b/campaigns/example-quick-optimum/task.py @@ -16,11 +16,11 @@ # What the agent is looking for, and never sees. Kept here rather than hidden away so # whoever runs this can check the agent's conclusion against the truth afterwards. -TRUE_OPTIMUM = 3.4 # where the response is genuinely lowest -TRUE_FLOOR = 12.0 # the response there -CURVATURE = 1.7 # how sharply it rises either side -ASYMMETRY = 0.35 # rises faster above the optimum than below -NOISE_SD = 0.6 # spread of one reading, so nearby settings overlap +TRUE_OPTIMUM = 3.4 # where the response is genuinely lowest +TRUE_FLOOR = 12.0 # the response there +CURVATURE = 1.7 # how sharply it rises either side +ASYMMETRY = 0.35 # rises faster above the optimum than below +NOISE_SD = 0.6 # spread of one reading, so nearby settings overlap LOCAL_DESC = """ Measure the response of the system at one setting. @@ -46,7 +46,7 @@ def _response(setting, rnd): offset = setting - TRUE_OPTIMUM rise = CURVATURE * offset * offset if offset > 0: - rise += ASYMMETRY * offset * offset * offset # steeper above the optimum + rise += ASYMMETRY * offset * offset * offset # steeper above the optimum return TRUE_FLOOR + rise + rnd.gauss(0.0, NOISE_SD) @@ -57,12 +57,17 @@ def local_fn(args): try: setting = float(args.get("setting")) except (TypeError, ValueError): - return {"error": f"setting must be a number, got {args.get('setting')!r}", - "args": args} + return { + "error": f"setting must be a number, got {args.get('setting')!r}", + "args": args, + } replicates = int(args.get("replicates", 1) or 1) if not 0.0 <= setting <= 10.0: - return {"error": f"setting must be between 0 and 10, got {setting}", "args": args} + return { + "error": f"setting must be between 0 and 10, got {setting}", + "args": args, + } if not 1 <= replicates <= 9: return {"error": f"replicates must be 1-9, got {replicates}", "args": args} @@ -71,7 +76,7 @@ def local_fn(args): rnd = random.Random() readings = [] for _ in range(replicates): - time.sleep(0.8) # a job is work, not an instant lookup + time.sleep(0.8) # a job is work, not an instant lookup readings.append(round(_response(setting, rnd), 4)) mean = sum(readings) / len(readings) @@ -81,6 +86,5 @@ def local_fn(args): "response": round(mean, 4), "noise_sd": round(spread, 4), "readings": readings, - "diagnostics": {"replicates": replicates, - "single_reading_sd": NOISE_SD}, + "diagnostics": {"replicates": replicates, "single_reading_sd": NOISE_SD}, } diff --git a/campaigns/example-vllm-inference-opt/task.py b/campaigns/example-vllm-inference-opt/task.py index a3a21f1..adce5c2 100644 --- a/campaigns/example-vllm-inference-opt/task.py +++ b/campaigns/example-vllm-inference-opt/task.py @@ -112,26 +112,39 @@ def remote_fn(args, target): timeout = int(target.get("timeout", 5400)) os.makedirs(work_dir, exist_ok=True) - tag = (f"{bench_mode}_in{input_len}_out{output_len}_tp{tp}_{dtype}" - f"_eager{int(enforce_eager)}_seqs{max_num_seqs}_{int(time.time())}") + tag = ( + f"{bench_mode}_in{input_len}_out{output_len}_tp{tp}_{dtype}" + f"_eager{int(enforce_eager)}_seqs{max_num_seqs}_{int(time.time())}" + ) log_path = os.path.join(work_dir, f"{tag}.log") # --- environment --- env = dict(os.environ) for k, v in (target.get("env") or {}).items(): env[str(k)] = str(v) - for k, v in env_extra.items(): # agent's overrides win + for k, v in env_extra.items(): # agent's overrides win env[str(k)] = str(v) # --- build the vLLM command --- - cmd = ["vllm", "bench", bench_mode, - "--model", model, - "--input-len", str(input_len), - "--output-len", str(output_len), - "--dtype", dtype, - "--tensor-parallel-size", str(tp), - "--max-model-len", str(max_model_len), - "--max-num-seqs", str(max_num_seqs)] + cmd = [ + "vllm", + "bench", + bench_mode, + "--model", + model, + "--input-len", + str(input_len), + "--output-len", + str(output_len), + "--dtype", + dtype, + "--tensor-parallel-size", + str(tp), + "--max-model-len", + str(max_model_len), + "--max-num-seqs", + str(max_num_seqs), + ] if bench_mode == "latency": cmd += ["--batch-size", "1", "--num-iters-warmup", "2", "--num-iters", "2"] else: @@ -142,14 +155,24 @@ def remote_fn(args, target): # The module load has to happen in the same shell as vllm, so go through bash. setup = target.get("worker_setup", "") shell_cmd = (setup + " && " if setup else "") + " ".join( - "'" + c + "'" if " " in c else c for c in cmd) + "'" + c + "'" if " " in c else c for c in cmd + ) started = time.time() try: - proc = subprocess.run(["bash", "-lc", shell_cmd], capture_output=True, - text=True, timeout=timeout, env=env) + proc = subprocess.run( + ["bash", "-lc", shell_cmd], + capture_output=True, + text=True, + timeout=timeout, + env=env, + ) except Exception as e: - return {"error": f"{type(e).__name__}: {e}", "cmd": shell_cmd, "key_env": env_extra} + return { + "error": f"{type(e).__name__}: {e}", + "cmd": shell_cmd, + "key_env": env_extra, + } out = (proc.stdout or "") + "\n" + (proc.stderr or "") @@ -157,42 +180,68 @@ def remote_fn(args, target): # startup banner is the most informative artefact available. try: with open(log_path, "w") as f: - f.write("$ " + shell_cmd + "\n\nENV_EXTRA: " + json.dumps(env_extra) + "\n\n" + out) + f.write( + "$ " + + shell_cmd + + "\n\nENV_EXTRA: " + + json.dumps(env_extra) + + "\n\n" + + out + ) except Exception: pass # --- startup diagnostics: what did vLLM actually do? --- diagnostics = {} patterns = { - "platform": r"[Pp]latform[:\s]+(\S+)", + "platform": r"[Pp]latform[:\s]+(\S+)", "attention_backend": r"[Uu]sing (\S+) backend", - "graph_capture": r"(?i)(graph capturing finished|CUDA graphs|Capturing.*graph|enforce_eager)", - "kv_cache_blocks": r"(?i)GPU KV cache size[:\s]+([\d,]+)", - "num_devices": r"(?i)(?:world_size|tensor.parallel.size)[=:\s]+(\d+)", - "vllm_version": r"(?i)vLLM (?:API server )?version[:\s]+(\S+)", + "graph_capture": r"(?i)(graph capturing finished|CUDA graphs|Capturing.*graph|enforce_eager)", + "kv_cache_blocks": r"(?i)GPU KV cache size[:\s]+([\d,]+)", + "num_devices": r"(?i)(?:world_size|tensor.parallel.size)[=:\s]+(\d+)", + "vllm_version": r"(?i)vLLM (?:API server )?version[:\s]+(\S+)", } for name, pat in patterns.items(): m = re.search(pat, out) if m: diagnostics[name] = m.group(1) if m.groups() else m.group(0) - warn = [l.strip() for l in out.splitlines() - if re.search(r"(?i)\b(warning|fallback|not supported|disabled)\b", l)] + warn = [ + l.strip() + for l in out.splitlines() + if re.search(r"(?i)\b(warning|fallback|not supported|disabled)\b", l) + ] if warn: diagnostics["warnings"] = warn[:15] if proc.returncode != 0: - return {"error": out[-3000:], "cmd": shell_cmd, "log": log_path, - "diagnostics": diagnostics, "key_env": env_extra} + return { + "error": out[-3000:], + "cmd": shell_cmd, + "log": log_path, + "diagnostics": diagnostics, + "key_env": env_extra, + } # --- metrics --- - result = {"bench_mode": bench_mode, "model": model, "input_len": input_len, - "output_len": output_len, "tensor_parallel_size": tp, "dtype": dtype, - "enforce_eager": enforce_eager, "max_num_seqs": max_num_seqs, - "max_model_len": max_model_len, "env_extra": env_extra, - "wall_seconds": round(time.time() - started, 1), - "log": log_path, "diagnostics": diagnostics} - - m = re.search(r"Throughput:\s*([\d.]+)\s*requests/s,\s*([\d.]+)\s*(?:total\s+)?tokens/s", out) + result = { + "bench_mode": bench_mode, + "model": model, + "input_len": input_len, + "output_len": output_len, + "tensor_parallel_size": tp, + "dtype": dtype, + "enforce_eager": enforce_eager, + "max_num_seqs": max_num_seqs, + "max_model_len": max_model_len, + "env_extra": env_extra, + "wall_seconds": round(time.time() - started, 1), + "log": log_path, + "diagnostics": diagnostics, + } + + m = re.search( + r"Throughput:\s*([\d.]+)\s*requests/s,\s*([\d.]+)\s*(?:total\s+)?tokens/s", out + ) if m: result["requests_per_sec"] = float(m.group(1)) result["throughput_tokens_per_sec"] = float(m.group(2)) diff --git a/framework/agent.py b/framework/agent.py index e765b2f..4e05285 100644 --- a/framework/agent.py +++ b/framework/agent.py @@ -12,8 +12,8 @@ import asyncio import glob import json -import re import os +import re import shutil import signal import socket @@ -23,31 +23,42 @@ from datetime import datetime from claude_agent_sdk import ( - ClaudeSDKClient, - ClaudeAgentOptions, AgentDefinition, AssistantMessage, + ClaudeAgentOptions, + ClaudeSDKClient, ResultMessage, ) + SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) # Resolved and exported BEFORE importing tools: tools.py reads WORKSPACE_DIR at import # time and falls back to its own directory, which puts claims.jsonl, jobs.jsonl and # ANNOUNCEMENTS.md in framework/ instead of the campaign's workspace. WORKSPACE_DIR = os.environ.get("WORKSPACE_DIR") or ( - os.path.join(os.path.abspath(os.environ.get("LAB_DIR", - os.path.join(SCRIPT_DIR, ".."))), "workspace", os.environ["CAMPAIGN"]) - if os.environ.get("CAMPAIGN") else SCRIPT_DIR) + os.path.join( + os.path.abspath(os.environ.get("LAB_DIR", os.path.join(SCRIPT_DIR, ".."))), + "workspace", + os.environ["CAMPAIGN"], + ) + if os.environ.get("CAMPAIGN") + else SCRIPT_DIR +) os.environ["WORKSPACE_DIR"] = WORKSPACE_DIR -import critic # noqa: E402 -import tools # noqa: E402 -from tools import create_server, shutdown_executor # noqa: E402 +import critic +import tools +from tools import create_server, shutdown_executor + # Campaign files (prompt.md, user prompt) live with the campaign, not the framework. LAB_DIR = os.path.abspath(os.environ.get("LAB_DIR", os.path.join(SCRIPT_DIR, ".."))) CAMPAIGN = os.environ.get("CAMPAIGN", "") -CAMPAIGN_DIR = os.path.abspath(os.environ.get( - "CAMPAIGN_DIR", os.path.join(LAB_DIR, "campaigns", CAMPAIGN) if CAMPAIGN else SCRIPT_DIR)) -SYSTEM = tools.SYSTEM # from the campaign's campaign.json +CAMPAIGN_DIR = os.path.abspath( + os.environ.get( + "CAMPAIGN_DIR", + os.path.join(LAB_DIR, "campaigns", CAMPAIGN) if CAMPAIGN else SCRIPT_DIR, + ) +) +SYSTEM = tools.SYSTEM # from the campaign's campaign.json ROLE = os.environ.get("ROLE", "both") # Roles only mean something when a campaign splits work between agents. Left unset, # they are noise in anything a person reads, so they are shown only when set. @@ -61,18 +72,21 @@ RUN_STAMP = datetime.now().strftime("%Y%m%d_%H%M%S") RUN_ID = f"{SYSTEM}_{ROLE}_{RUN_STAMP}" if ROLE_SET else f"{SYSTEM}_{RUN_STAMP}" RUN_DIR = os.path.join(WORKSPACE_DIR, "runs", RUN_ID) -os.environ["RUN_ID"] = RUN_ID # tools stamps the job log with it +os.environ["RUN_ID"] = RUN_ID # tools stamps the job log with it LOG_PATH = os.path.join(LOG_DIR, f"run_{SYSTEM}_{RUN_STAMP}.log") -HEARTBEAT_INTERVAL = 30 # s; minimum gap between heartbeat writes during a wait -CRITIC_MODEL = None # resolved in preflight() +HEARTBEAT_INTERVAL = 30 # s; minimum gap between heartbeat writes during a wait +CRITIC_MODEL = None # resolved in preflight() CRITIC_LABEL = "no critic" # What a cycle gets written up in depends on the method: the standard one records to # LOGBOOK.md and writes JOURNAL.md only at the end, the research one keeps both as it # goes. The critic watches both and reviews whichever grew. -CYCLE_RECORDS = [os.path.join(WORKSPACE_DIR, name) - for name in ("LOGBOOK.md", "JOURNAL.md")] +CYCLE_RECORDS = [ + os.path.join(WORKSPACE_DIR, name) for name in ("LOGBOOK.md", "JOURNAL.md") +] -AGENT_ALIVE_WITHIN = int(os.environ.get("AGENT_ALIVE_WITHIN", "300")) # s; fresher heartbeat = agent is up +AGENT_ALIVE_WITHIN = int( + os.environ.get("AGENT_ALIVE_WITHIN", "300") +) # s; fresher heartbeat = agent is up def _live_handles(): @@ -81,7 +95,7 @@ def _live_handles(): numbers are reused once a run ends.""" out = [] now = time.time() - root = os.path.dirname(WORKSPACE_DIR) # workspace/, one dir per campaign + root = os.path.dirname(WORKSPACE_DIR) # workspace/, one dir per campaign for hb in glob.glob(os.path.join(root, "*", "runs", "*", "heartbeat")): try: with open(hb) as f: @@ -90,7 +104,7 @@ def _live_handles(): with open(os.path.join(os.path.dirname(hb), "meta.json")) as f: meta = json.load(f) except Exception: - continue # unreadable run: treat its handle as free + continue # unreadable run: treat its handle as free if meta.get("handle"): out.append((meta["handle"], meta.get("campaign", ""))) return out @@ -116,7 +130,7 @@ def _allocate_handle(): slugs.append(name.replace("-", "").replace("_", "")[:12]) for slug in slugs: if any(h.rstrip("0123456789") == slug and c != CAMPAIGN for h, c in live): - continue # another campaign already answers to this slug + continue # another campaign already answers to this slug for n in range(1, 100): cand = f"{slug}{n}" if cand not in {h for h, _ in live}: @@ -148,39 +162,52 @@ def _prompt(name, default): return v if v else default -CONTINUE_PROMPT = _prompt("CONTINUE_PROMPT", +CONTINUE_PROMPT = _prompt( + "CONTINUE_PROMPT", f"One or more jobs have finished. Collect them with {_COMPLETED_TOOL}, " - "fit and log each, then continue: submit new jobs as needed." + "fit and log each, then continue: submit new jobs as needed.", ) # What "more work" means belongs to the method, so this says only that the run has # capacity and asks for the next step from the records the method already keeps. -EXPLORE_PROMPT = _prompt("EXPLORE_PROMPT", +EXPLORE_PROMPT = _prompt( + "EXPLORE_PROMPT", "No jobs are running and there is budget left. From your own records, decide what " "the next step is and submit it. If the goal is met, or nothing further is worth " - "running, say so and stop rather than filling the budget." + "running, say so and stop rather than filling the budget.", ) -WINDDOWN_PROMPT = _prompt("WINDDOWN_PROMPT", +WINDDOWN_PROMPT = _prompt( + "WINDDOWN_PROMPT", "Wind-down requested: this run is ending. Submit no new work -- the submit tools " "will refuse it. Collect and log the jobs already in flight as they finish. Once " - "everything is collected you get a final turn to write up the cycle." + "everything is collected you get a final turn to write up the cycle.", ) -FINALIZE_PROMPT = _prompt("FINALIZE_PROMPT", +FINALIZE_PROMPT = _prompt( + "FINALIZE_PROMPT", # Which records a cycle is written up in is the method's business, not the # runner's: naming a file here produces one that the method never asked for. "All outstanding work is collected and this run is now ending. Close out the " "current cycle: write it up in the records your method keeps, and note anything a " - "later run needs to pick up where you left off. Submit no new work." + "later run needs to pick up where you left off. Submit no new work.", ) # A session id to start from. Its whole conversation becomes this run's context, which # costs what it costs and brings any stale conclusions with it, so it is off by default. # The session must belong to this user on this machine. RESUME_SESSION = (os.environ.get("RESUME_SESSION") or "").strip() -MAX_ROUNDS = 500 # backstop against a runaway loop -MAX_EMPTY_ROUNDS = 3 # consecutive idle rounds (no work proposed) before giving up -MAX_RUNTIME = int(os.environ["MAX_RUNTIME"]) if os.environ.get("MAX_RUNTIME") else None # total agent wallclock (s); None = no time limit -WAIT_TIMEOUT = 1800 # s between "still-alive" logs / backend-health checks during a wait -ANNOUNCE_POLL = int(os.environ.get("ANNOUNCE_POLL", "2")) # s between announcement-board checks during a job wait -STALL_LIMIT = int(os.environ["STALL_LIMIT"]) if os.environ.get("STALL_LIMIT") else None # None = wait indefinitely (HPC queues can take many hours); set seconds to cap (tests do) +MAX_ROUNDS = 500 # backstop against a runaway loop +MAX_EMPTY_ROUNDS = 3 # consecutive idle rounds (no work proposed) before giving up +MAX_RUNTIME = ( + int(os.environ["MAX_RUNTIME"]) if os.environ.get("MAX_RUNTIME") else None +) # total agent wallclock (s); None = no time limit +WAIT_TIMEOUT = ( + 1800 # s between "still-alive" logs / backend-health checks during a wait +) +ANNOUNCE_POLL = int( + os.environ.get("ANNOUNCE_POLL", "2") +) # s between announcement-board checks during a job wait +STALL_LIMIT = ( + int(os.environ["STALL_LIMIT"]) if os.environ.get("STALL_LIMIT") else None +) # None = wait indefinitely (HPC queues can take many hours); set seconds to cap (tests do) + # --- Slack notifications (optional; see SLACK_NOTIFY.md). Missing webhook/script # or a failed post is ignored so a run is never affected. --- @@ -188,6 +215,7 @@ def _bool_env(name, default=False): v = os.environ.get(name) return default if v is None else v.strip().lower() in ("1", "true", "yes", "on") + # A browser view of this run: the log as it is written and the files it writes. Off # unless asked for, and it never affects the run -- it only reads the workspace. WATCH = _bool_env("WATCH", False) @@ -195,21 +223,28 @@ def _bool_env(name, default=False): # The viewer outlives the run -- the end of a run is when its records are worth reading # -- and stops itself once nobody has looked for this long. WATCH_IDLE = int(os.environ.get("WATCH_IDLE", "600")) -_watcher = None # the viewer process, stopped when the run ends +_watcher = None # the viewer process, stopped when the run ends NOTIFY_START = _bool_env("NOTIFY_START", False) NOTIFY_DAILY = _bool_env("NOTIFY_DAILY", True) NOTIFY_FINISH = _bool_env("NOTIFY_FINISH", True) -DAILY_INTERVAL = int(os.environ.get("NOTIFY_DAILY_INTERVAL", "86400")) # seconds between periodic summaries -PROBLEM_GRACE = int(os.environ.get("NOTIFY_PROBLEM_GRACE", "1800")) # shut down this long (s) after the agent flags an unresolved blocking problem -NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join(SCRIPT_DIR, "slack_notify.sh") +DAILY_INTERVAL = int( + os.environ.get("NOTIFY_DAILY_INTERVAL", "86400") +) # seconds between periodic summaries +PROBLEM_GRACE = int( + os.environ.get("NOTIFY_PROBLEM_GRACE", "1800") +) # shut down this long (s) after the agent flags an unresolved blocking problem +NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join( + SCRIPT_DIR, "slack_notify.sh" +) # When a periodic summary is due, the runner asks the agent to write it (its own # words) via the notify tool, instead of a fixed harness string. -REPORT_PROMPT = _prompt("REPORT_PROMPT", +REPORT_PROMPT = _prompt( + "REPORT_PROMPT", "Before anything else this turn, post a brief (1-2 line) status summary to Slack " "with the notify tool: what you are currently working on, recent progress, and any " - "concern. Then continue as normal." + "concern. Then continue as normal.", ) @@ -217,8 +252,12 @@ def slack_notify(msg): if not os.path.isfile(NOTIFY_SCRIPT): return try: - subprocess.run(["bash", NOTIFY_SCRIPT, msg], timeout=30, - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + subprocess.run( + ["bash", NOTIFY_SCRIPT, msg], + timeout=30, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) except Exception as e: print(f"[slack_notify] failed (ignored): {e}", flush=True) @@ -252,7 +291,7 @@ async def _context_usage(client): return None -_context_task = None # the in-flight context lookup, if any +_context_task = None # the in-flight context lookup, if any def _refresh_context(client): @@ -270,19 +309,27 @@ async def _record_context(client): if not ctx: return _last_context.update(ctx) - _write_meta(context_tokens=ctx["tokens"], context_window=ctx["window"], - context_pct=ctx["pct"], **({"model": ctx["model"]} if ctx["model"] else {})) + _write_meta( + context_tokens=ctx["tokens"], + context_window=ctx["window"], + context_pct=ctx["pct"], + **({"model": ctx["model"]} if ctx["model"] else {}), + ) def _as_context(usage): """The /context answer in the shape the rest of the run records.""" if not usage or usage.get("totalTokens") is None: return None - return {"tokens": usage.get("totalTokens"), "window": usage.get("rawMaxTokens"), - "pct": usage.get("percentage"), "model": usage.get("model")} + return { + "tokens": usage.get("totalTokens"), + "window": usage.get("rawMaxTokens"), + "pct": usage.get("percentage"), + "model": usage.get("model"), + } -_last_context = {} # tokens/window/pct/model, from the most recent turn +_last_context = {} # tokens/window/pct/model, from the most recent turn async def _post_scheduled_status(client, round_num, start_time): @@ -290,12 +337,17 @@ async def _post_scheduled_status(client, round_num, start_time): u = _last_context model = u.get("model") or "?" tok, win, pct = u.get("tokens"), u.get("window"), u.get("pct") - ctx = (f"ctx ~{tok}/{win} (~{pct:.0f}%)" - if tok is not None and win and pct is not None else "ctx n/a") - slack_notify(f":calendar: Scheduled Status — {model}, " - f"round {round_num} · {tools.submit_count()} remote / {tools.local_submit_count()} local " - f"this session · {tools.jobs_in_flight()} in-flight · {ctx} · " - f"uptime {_fmt_uptime(time.time() - start_time)}") + ctx = ( + f"ctx ~{tok}/{win} (~{pct:.0f}%)" + if tok is not None and win and pct is not None + else "ctx n/a" + ) + slack_notify( + f":calendar: Scheduled Status — {model}, " + f"round {round_num} · {tools.submit_count()} remote / {tools.local_submit_count()} local " + f"this session · {tools.jobs_in_flight()} in-flight · {ctx} · " + f"uptime {_fmt_uptime(time.time() - start_time)}" + ) def _new_board_lines(seen, current): @@ -304,8 +356,8 @@ def _new_board_lines(seen, current): part' then. Without this, any change re-sends every old message and the agent re-acts on things it already handled.""" old, new = seen.splitlines(), current.splitlines() - if new[:len(old)] == old: - return "\n".join(new[len(old):]).strip() + if new[: len(old)] == old: + return "\n".join(new[len(old) :]).strip() return current @@ -327,7 +379,7 @@ def _new_record_text(before, after): for path, text in after.items(): old = before.get(path, "") if len(text) > len(old) and text.startswith(old): - chunk = text[len(old):].strip() + chunk = text[len(old) :].strip() if chunk: added.append(f"--- new in {os.path.basename(path)} ---\n{chunk}") return "\n\n".join(added) @@ -349,8 +401,11 @@ def _recent_results(budget=120000): break kept.append(row) kept.reverse() - head = (f"({len(rows)} rows recorded; all supplied)\n" if len(kept) == len(rows) - else f"({len(rows)} rows recorded, the {len(kept)} most recent supplied)\n") + head = ( + f"({len(rows)} rows recorded; all supplied)\n" + if len(kept) == len(rows) + else f"({len(rows)} rows recorded, the {len(kept)} most recent supplied)\n" + ) return head + "".join(kept) @@ -359,8 +414,10 @@ def _append_review(reply): belong in the record, and a later reader can see what was checked.""" try: with open(os.path.join(WORKSPACE_DIR, "REVIEWS.md"), "a") as f: - f.write(f"\n## {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} " - f"-- {CRITIC_LABEL}, run {RUN_ID}\n\n{reply}\n") + f.write( + f"\n## {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} " + f"-- {CRITIC_LABEL}, run {RUN_ID}\n\n{reply}\n" + ) except Exception as e: print(f"[critic] could not record the review (ignored): {e}", flush=True) @@ -370,26 +427,32 @@ def _critic_prompt(findings, reply, tail=""): a critic reading only the rows can be wrong about what the rows mean, and saying so with evidence is a legitimate answer.""" listed = "\n".join(f"- {claim} ({verdict})" for claim, verdict in findings) - return (f"The critic ({CRITIC_LABEL}) reviewed your latest journal section and " - f"found claims it says the recorded results do not support:\n\n" - f"{listed}\n\nIts full review:\n\n{reply}\n\n" - "Deal with each one before continuing: correct the write-up, run what " - "would settle it, or answer the objection in the journal citing the rows " - "that support you.\n\n" + tail) + return ( + f"The critic ({CRITIC_LABEL}) reviewed your latest journal section and " + f"found claims it says the recorded results do not support:\n\n" + f"{listed}\n\nIts full review:\n\n{reply}\n\n" + "Deal with each one before continuing: correct the write-up, run what " + "would settle it, or answer the objection in the journal citing the rows " + "that support you.\n\n" + tail + ) def _announcements_prompt(text, tail=""): """Wrap NEW announcements-board lines as the next turn's prompt.""" - body = ("New on the shared announcements board:\n" + text + - "\nAct on anything here that concerns you. Anything marked as already " - "answered by the secretary needs no reply from you. If it needs immediate " - "action, take it now; otherwise acknowledge it briefly and continue. " - "Pending jobs remain tracked.") + body = ( + "New on the shared announcements board:\n" + + text + + "\nAct on anything here that concerns you. Anything marked as already " + "answered by the secretary needs no reply from you. If it needs immediate " + "action, take it now; otherwise acknowledge it briefly and continue. " + "Pending jobs remain tracked." + ) return body + ("\n\n" + tail if tail else "") class Tee: """Write to both a file and the original stream.""" + def __init__(self, log_file, stream): self.log_file = log_file self.stream = stream @@ -413,8 +476,11 @@ def method_path(): into the campaign, so each campaign owns its own and can change it. The library default applies to a campaign created before this, or one whose copy is missing.""" campaign_copy = os.path.join(CAMPAIGN_DIR, "method.md") - return (campaign_copy if os.path.isfile(campaign_copy) - else os.path.join(LAB_DIR, "methods", "standard.md")) + return ( + campaign_copy + if os.path.isfile(campaign_copy) + else os.path.join(LAB_DIR, "methods", "standard.md") + ) def load_method(): @@ -449,8 +515,17 @@ def load_user_prompt(): # framework's bookkeeping. WebSearch is left out too: it runs on the API server rather # than here, so a gateway that does not carry server tools refuses it, and some sites # are behind one. AGENT_TOOLS changes this set. -BASE_TOOLS = ["Read", "Write", "Edit", "Glob", "Grep", "Bash", "Skill", - "WebFetch", "Agent"] +BASE_TOOLS = [ + "Read", + "Write", + "Edit", + "Glob", + "Grep", + "Bash", + "Skill", + "WebFetch", + "Agent", +] AGENT_TOOLS = os.environ.get("AGENT_TOOLS", "").strip() @@ -467,8 +542,10 @@ def claude_tools(): return list(BASE_TOOLS) signed = [e for e in entries if e[0] in "+-"] if signed and len(signed) != len(entries): - raise ValueError("AGENT_TOOLS mixes +/- adjustments with plain tool names; " - "use one form or the other") + raise ValueError( + "AGENT_TOOLS mixes +/- adjustments with plain tool names; " + "use one form or the other" + ) if not signed: return list(dict.fromkeys(entries)) out = list(BASE_TOOLS) @@ -497,7 +574,7 @@ def agent_tools(): # campaign on a cheaper model before giving a machine to a long run. An alias # ('sonnet', 'opus') or a full model name; which ones work depends on the lab's gateway. AGENT_MODEL = os.environ.get("AGENT_MODEL", "").strip() -RESOLVED_MODEL = "" # what the agent will actually run as; filled in by preflight +RESOLVED_MODEL = "" # what the agent will actually run as; filled in by preflight # One worked example of delegating: a subagent that reads a long record and returns @@ -535,8 +612,9 @@ def _parse_subagent(path): # can act on it. A background call returns a launch stub and the answer lands # some turns later. fields["background"] = meta["background"].lower() == "true" - return meta["name"], AgentDefinition(description=meta["description"], - prompt=body.strip(), **fields) + return meta["name"], AgentDefinition( + description=meta["description"], prompt=body.strip(), **fields + ) def subagent_defs(): @@ -562,12 +640,20 @@ def subagent_defs(): # What a tool call means for someone watching. The tool name says which function was # called; a phase says what the run is doing. _PHASES = { - "submit_job": "submitting jobs", "submit_local": "submitting jobs", - "get_completed_jobs": "collecting results", "get_local_completed": "collecting results", - "check_backend": "checking the backend", "release_claim": "releasing a claim", - "notify": "posting to Slack", "cycle_done": "closing the cycle", - "Read": "reading records", "Grep": "reading records", "Glob": "reading records", - "Write": "writing up", "Edit": "writing up", "NotebookEdit": "writing up", + "submit_job": "submitting jobs", + "submit_local": "submitting jobs", + "get_completed_jobs": "collecting results", + "get_local_completed": "collecting results", + "check_backend": "checking the backend", + "release_claim": "releasing a claim", + "notify": "posting to Slack", + "cycle_done": "closing the cycle", + "Read": "reading records", + "Grep": "reading records", + "Glob": "reading records", + "Write": "writing up", + "Edit": "writing up", + "NotebookEdit": "writing up", "Bash": "running analysis", } # Which subagent each Agent call started, keyed by the call's id, so a turn arriving @@ -608,9 +694,18 @@ def _start_watcher(): os.makedirs(RUN_DIR, exist_ok=True) watch_log = open(os.path.join(RUN_DIR, "watch.log"), "w") _watcher = subprocess.Popen( - [sys.executable, os.path.join(SCRIPT_DIR, "watch.py"), CAMPAIGN, - "--port", str(port), "--no-open", f"--exit-when-idle={WATCH_IDLE}"], - stdout=watch_log, stderr=subprocess.STDOUT) + [ + sys.executable, + os.path.join(SCRIPT_DIR, "watch.py"), + CAMPAIGN, + "--port", + str(port), + "--no-open", + f"--exit-when-idle={WATCH_IDLE}", + ], + stdout=watch_log, + stderr=subprocess.STDOUT, + ) print(f"watch: http://127.0.0.1:{port}/", flush=True) except Exception as e: print(f"[watch] could not start the viewer (ignored): {e}", flush=True) @@ -662,8 +757,11 @@ def _heartbeat(force=True): def _start_run_dir(): os.makedirs(RUN_DIR, exist_ok=True) - for path in (os.path.join(CAMPAIGN_DIR, "prompt.md"), method_path(), - os.path.join(CAMPAIGN_DIR, USER_PROMPT_FILE)): + for path in ( + os.path.join(CAMPAIGN_DIR, "prompt.md"), + method_path(), + os.path.join(CAMPAIGN_DIR, USER_PROMPT_FILE), + ): name = os.path.basename(path) try: shutil.copy2(path, os.path.join(RUN_DIR, name)) @@ -672,15 +770,25 @@ def _start_run_dir(): # The budgets this run stops at, recorded so anything reading the run -- a watcher, # a later reader -- can say how far through it is without knowing the environment # it was launched in. - _write_meta(max_submits=tools.MAX_SUBMITS, max_runtime_s=MAX_RUNTIME, - max_rounds=MAX_ROUNDS, critic=CRITIC_LABEL, - run_id=RUN_ID, handle=HANDLE, system=SYSTEM, role=ROLE, - started_by=os.environ.get("STARTED_BY", ""), - host=socket.gethostname(), pid=os.getpid(), - started_at=datetime.now().isoformat(timespec="seconds"), - user_prompt_file=USER_PROMPT_FILE, - campaign=CAMPAIGN, - shared_dir=WORKSPACE_DIR, log=LOG_PATH, status="running") + _write_meta( + max_submits=tools.MAX_SUBMITS, + max_runtime_s=MAX_RUNTIME, + max_rounds=MAX_ROUNDS, + critic=CRITIC_LABEL, + run_id=RUN_ID, + handle=HANDLE, + system=SYSTEM, + role=ROLE, + started_by=os.environ.get("STARTED_BY", ""), + host=socket.gethostname(), + pid=os.getpid(), + started_at=datetime.now().isoformat(timespec="seconds"), + user_prompt_file=USER_PROMPT_FILE, + campaign=CAMPAIGN, + shared_dir=WORKSPACE_DIR, + log=LOG_PATH, + status="running", + ) _heartbeat() print(f"Run dir: {RUN_DIR}", flush=True) @@ -704,11 +812,13 @@ def _stop_file_present(): # not opened, so the last run's log survives, and no watcher or Slack post is made. # The environment variable is the form campaigns use, since it reaches here whatever a # run.sh looks like; --preflight is accepted too, for running this file directly. -CHECK_ONLY = (os.environ.get("PREFLIGHT", "").lower() in ("1", "true", "yes") - or "--preflight" in sys.argv) +CHECK_ONLY = ( + os.environ.get("PREFLIGHT", "").lower() in ("1", "true", "yes") + or "--preflight" in sys.argv +) -GATEWAY_URL = None # set when this run is routed through the lab's gateway +GATEWAY_URL = None # set when this run is routed through the lab's gateway def _on_lab_gateway(): @@ -737,7 +847,7 @@ def _route_to_gateway(): return try: with open(conf) as f: - served = re.findall(r"^\s*-?\s*model_name:\s*(\S+)", f.read(), re.M) + served = re.findall(r"^\s*-?\s*model_name:\s*(\S+)", f.read(), re.MULTILINE) except OSError: return if AGENT_MODEL in served: @@ -748,8 +858,10 @@ def _route_to_gateway(): # The critic and the framework's own checks read this; the CLI does not, which # _gateway_settings_file handles. os.environ["ANTHROPIC_BASE_URL"] = url - print(f"gateway: {AGENT_MODEL} is served by {url} — routing the agent there", - flush=True) + print( + f"gateway: {AGENT_MODEL} is served by {url} — routing the agent there", + flush=True, + ) def _gateway_settings_file(): @@ -779,7 +891,8 @@ def _probe_model(): async def _ask(): opts = ClaudeAgentOptions( cwd=SCRIPT_DIR, - **({"settings": _gateway_settings_file()} if GATEWAY_URL else {})) + **({"settings": _gateway_settings_file()} if GATEWAY_URL else {}), + ) async with ClaudeSDKClient(options=opts) as c: return (await _context_usage(c) or {}).get("model") @@ -803,12 +916,15 @@ def preflight(): if tools.HAS_LOCAL: required += ["LOCAL_DESC", "LOCAL_SCHEMA", "local_fn"] if not required: - problems.append(f"task {tools.TASK_DIR} defines neither 'remote_fn' nor " - f"'local_fn' (see AGENTS.md)") + problems.append( + f"task {tools.TASK_DIR} defines neither 'remote_fn' nor " + f"'local_fn' (see AGENTS.md)" + ) for attr in required: if not hasattr(tools.task, attr): - problems.append(f"task {tools.TASK_DIR} is missing '{attr}' " - f"(see AGENTS.md)") + problems.append( + f"task {tools.TASK_DIR} is missing '{attr}' (see AGENTS.md)" + ) # A task may declare its own checks -- e.g. that its binary is where it expects. if hasattr(tools.task, "preflight"): try: @@ -820,7 +936,9 @@ def preflight(): except ValueError as e: problems.append(str(e)) if not os.path.isfile(method_path()): - problems.append(f"method.md missing: {method_path()} (how-to-work prompt loaded into the agent)") + problems.append( + f"method.md missing: {method_path()} (how-to-work prompt loaded into the agent)" + ) _fw = os.path.join(SCRIPT_DIR, "framework_prompt.md") if not os.path.isfile(_fw): problems.append(f"framework_prompt.md missing: {_fw}") @@ -830,21 +948,32 @@ def preflight(): if tools.HAS_REMOTE: try: import globus_compute_sdk as _gc + _c = _gc.Client() _st = _c.get_endpoint_status(tools.ENDPOINT_ID).get("status") if _st != "online": try: - _nm = _c.get_endpoint_metadata(tools.ENDPOINT_ID).get("name") or tools.ENDPOINT_ID + _nm = ( + _c.get_endpoint_metadata(tools.ENDPOINT_ID).get("name") + or tools.ENDPOINT_ID + ) except Exception: _nm = tools.ENDPOINT_ID if not CHECK_ONLY: - slack_notify(f":rotating_light: Agent exiting -- Globus Compute " - f"endpoint '{_nm}' is not online (status={_st}). Start it: " - f"globus-compute-endpoint start {_nm} --detach") - problems.append(f"Globus Compute endpoint '{_nm}' ({tools.ENDPOINT_ID}) is not online " - f"(status={_st}); start it: globus-compute-endpoint start {_nm} --detach") + slack_notify( + f":rotating_light: Agent exiting -- Globus Compute " + f"endpoint '{_nm}' is not online (status={_st}). Start it: " + f"globus-compute-endpoint start {_nm} --detach" + ) + problems.append( + f"Globus Compute endpoint '{_nm}' ({tools.ENDPOINT_ID}) is not online " + f"(status={_st}); start it: globus-compute-endpoint start {_nm} --detach" + ) except Exception as e: - print(f"[preflight] WARNING: could not query endpoint status ({tools.ENDPOINT_ID}): {e}", flush=True) + print( + f"[preflight] WARNING: could not query endpoint status ({tools.ENDPOINT_ID}): {e}", + flush=True, + ) try: os.makedirs(WORKSPACE_DIR, exist_ok=True) _t = os.path.join(WORKSPACE_DIR, ".preflight_write_test") @@ -859,7 +988,10 @@ def preflight(): print(f" - {pr}", flush=True) sys.exit(1) backend = "endpoint online" if tools.HAS_REMOTE else "local execution only" - print(f"preflight OK: task_dir={tools.TASK_DIR}, method.md, WORKSPACE_DIR, {backend}.", flush=True) + print( + f"preflight OK: task_dir={tools.TASK_DIR}, method.md, WORKSPACE_DIR, {backend}.", + flush=True, + ) # The gateway converts between the Messages API and a backend that does not speak # it. The agent needs it whenever it is pointed at one, whether or not there is a # critic: a run on a non-Anthropic model goes through the same proxy. @@ -878,7 +1010,8 @@ def preflight(): RESOLVED_MODEL = _probe_model() try: CRITIC_MODEL, CRITIC_LABEL = critic.resolve( - RESOLVED_MODEL or os.environ.get("ANTHROPIC_MODEL", "")) + RESOLVED_MODEL or os.environ.get("ANTHROPIC_MODEL", "") + ) except critic.CriticUnavailable as e: print(f"preflight FAILED: {e}", flush=True) sys.exit(1) @@ -892,15 +1025,19 @@ def preflight(): print(f"job tools: {' '.join(job)}", flush=True) print(f"claude tools: {' '.join(claude)}", flush=True) if RESOLVED_MODEL or CHECK_ONLY: - print(f"model: {RESOLVED_MODEL or '(could not be determined)'}", flush=True) + print( + f"model: {RESOLVED_MODEL or '(could not be determined)'}", flush=True + ) print(f"critic: {CRITIC_LABEL}", flush=True) # The budget and the resources a job asks for. They come from three files and the # environment, so the resolved values are the only honest way to show them -- and # they are what a run gets wrong most often. # Named as the environment variables that set them, so a value that looks wrong # can be searched for in run.sh without a translation step. - limits = [f"MAX_SUBMITS={tools.MAX_SUBMITS}", - f"MAX_CONCURRENT={tools.MAX_CONCURRENT}"] + limits = [ + f"MAX_SUBMITS={tools.MAX_SUBMITS}", + f"MAX_CONCURRENT={tools.MAX_CONCURRENT}", + ] if MAX_RUNTIME: limits.append(f"MAX_RUNTIME={MAX_RUNTIME}s") print(f"budget: {', '.join(limits)}", flush=True) @@ -910,18 +1047,24 @@ def preflight(): buckets = tools._SYS["buckets"] for name, b in buckets.items(): res = b["user_config"] - shown = ", ".join(f"{k}={v}" for k, v in sorted(res.items()) - if v != "" and v is not None - and k not in ("init_blocks", "min_blocks")) + shown = ", ".join( + f"{k}={v}" + for k, v in sorted(res.items()) + if v != "" and v is not None and k not in ("init_blocks", "min_blocks") + ) label = f"resources[{name}]" if len(buckets) > 1 else "resources" - marker = " (default)" if name == tools._default_bucket and len(buckets) > 1 else "" + marker = ( + " (default)" + if name == tools._default_bucket and len(buckets) > 1 + else "" + ) print(f"{label}: {shown}{marker}", flush=True) print(f"job timeout: {tools.TARGET.get('timeout', '(unset)')}s", flush=True) if not CHECK_ONLY: _start_watcher() -_session_id = None # this run's Claude session, for reopening it later +_session_id = None # this run's Claude session, for reopening it later async def drain_turn(client, round_num): @@ -946,8 +1089,9 @@ async def drain_turn(client, round_num): sub = (getattr(block, "input", None) or {}).get("subagent_type") if sub: _DELEGATES[getattr(block, "id", None)] = sub - _set_phase(f"round {round_num}: waiting on subagent " - f"{sub or '(unnamed)'}") + _set_phase( + f"round {round_num}: waiting on subagent {sub or '(unnamed)'}" + ) elif who: _set_phase(f"round {round_num}: subagent {who} ({bare})") else: @@ -977,10 +1121,13 @@ async def main(): at_once.append(f"local jobs running at once: {tools.LOCAL_MAX_CONCURRENT}") if MAX_RUNTIME: at_once.append(f"wall clock for this run: {MAX_RUNTIME}s") - system_prompt += ("\n\n# This run\n" + "\n".join(at_once) - + "\n\nSubmitting more at once than that queues the rest, which " - "tells you nothing sooner. Each submit answers with how much of " - "the run's job budget it has used.") + system_prompt += ( + "\n\n# This run\n" + + "\n".join(at_once) + + "\n\nSubmitting more at once than that queues the rest, which " + "tells you nothing sooner. Each submit answers with how much of " + "the run's job budget it has used." + ) system_prompt += f"\n\n# This agent\nSYSTEM={SYSTEM}.{f' ROLE={ROLE}.' if ROLE_SET else ''}\nThe shared files (results.jsonl, LOGBOOK.md, JOURNAL.md, claims.jsonl) live in {WORKSPACE_DIR} \u2014 always read and write them by full path there (e.g. {WORKSPACE_DIR}/results.jsonl). Follow the role rules in the Collaboration section of the prompt." server = create_server() @@ -1020,8 +1167,9 @@ async def main(): # several run at once. Removed on clean exit. run_dir = os.path.join(WORKSPACE_DIR, "run") os.makedirs(run_dir, exist_ok=True) - pid_file = os.path.join(run_dir, f"agent_{SYSTEM}_{ROLE}.pid" if ROLE_SET - else f"agent_{SYSTEM}.pid") + pid_file = os.path.join( + run_dir, f"agent_{SYSTEM}_{ROLE}.pid" if ROLE_SET else f"agent_{SYSTEM}.pid" + ) with open(pid_file, "w") as f: f.write(str(os.getpid())) _start_run_dir() @@ -1030,11 +1178,13 @@ async def main(): # and no finish ping fired. Cancel the main task instead so shutdown runs. SIGINT # (kill -INT / Ctrl-C) already unwinds via KeyboardInterrupt. main_task = asyncio.current_task() + def _on_sigterm(): nonlocal stop_reason stop_reason = "signal (SIGTERM)" print("SIGTERM received -- shutting down gracefully.", flush=True) main_task.cancel() + try: loop.add_signal_handler(signal.SIGTERM, _on_sigterm) except (NotImplementedError, RuntimeError): @@ -1050,15 +1200,19 @@ def _on_sigterm(): model = (start or {}).get("model") or RESOLVED_MODEL or AGENT_MODEL or "?" if start: _last_context.update(start) - _write_meta(context_tokens=start["tokens"], - context_window=start["window"], - context_pct=start["pct"]) + _write_meta( + context_tokens=start["tokens"], + context_window=start["window"], + context_pct=start["pct"], + ) print(f"Agent started -- {SYSTEM}{ROLE_NOTE} · model {model}", flush=True) _write_meta(model=model) if NOTIFY_START: - slack_notify(f":rocket: Agent {HANDLE} started — " - f"{CAMPAIGN or 'no campaign'} on {SYSTEM}{ROLE_NOTE} · {model}" - f" · critic {CRITIC_LABEL}.") + slack_notify( + f":rocket: Agent {HANDLE} started — " + f"{CAMPAIGN or 'no campaign'} on {SYSTEM}{ROLE_NOTE} · {model}" + f" · critic {CRITIC_LABEL}." + ) prompt = load_user_prompt() empty_rounds = 0 last_daily = start_time @@ -1074,7 +1228,7 @@ def _on_sigterm(): # Set while the agent is dealing with findings, so its answer to them is # not itself put up for review. answering_critic = False - stopping = None # set to the reason once the run starts winding down + stopping = None # set to the reason once the run starts winding down finalize_rounds = 0 for round_num in range(1, MAX_ROUNDS + 1): print(f"\n===== ROUND {round_num} =====", flush=True) @@ -1084,7 +1238,10 @@ def _on_sigterm(): # around and it has stayed unresolved past the grace period. ps = tools.problem_since() if ps is not None and time.time() - ps >= PROBLEM_GRACE: - print(f"Agent-flagged problem unresolved for >{PROBLEM_GRACE}s -- stopping.", flush=True) + print( + f"Agent-flagged problem unresolved for >{PROBLEM_GRACE}s -- stopping.", + flush=True, + ) stop_reason = "agent-flagged problem unresolved past grace" break # Scheduled status: post the fixed-metrics line, then ask the agent to @@ -1097,8 +1254,11 @@ def _on_sigterm(): if stopping is None and _stop_file_present(): stopping = "stop requested" tools.request_stop() - print("Stop requested -- winding down: no new work, finishing " - "what is in flight.", flush=True) + print( + "Stop requested -- winding down: no new work, finishing " + "what is in flight.", + flush=True, + ) # Replace, not prepend: whatever was queued (CONTINUE/EXPLORE) tells # the agent to submit the next region, which contradicts winding down. prompt = WINDDOWN_PROMPT @@ -1108,7 +1268,9 @@ def _on_sigterm(): stopping = "goal met" tools.request_stop() _write_meta(goal_met=tools.goal_is_met()) - print(f"Goal met -- winding down: {tools.goal_is_met()}", flush=True) + print( + f"Goal met -- winding down: {tools.goal_is_met()}", flush=True + ) prompt = WINDDOWN_PROMPT # A cycle write-up is the trigger: the journal gains a section, so a # longer journal than last round means there is something to review. @@ -1125,23 +1287,35 @@ def _on_sigterm(): new_section = _new_record_text(last_records, records) if conclusion: last_records = records - new_section = (f"The agent's stated conclusion for this cycle:\n" - f"{conclusion}\n\n{new_section}") - print(f"[critic] reviewing {len(new_section)} new chars " - f"with {CRITIC_LABEL}", flush=True) + new_section = ( + f"The agent's stated conclusion for this cycle:\n" + f"{conclusion}\n\n{new_section}" + ) + print( + f"[critic] reviewing {len(new_section)} new chars " + f"with {CRITIC_LABEL}", + flush=True, + ) if stopping: # A review takes a couple of minutes. During a wind-down # that silence looks like a hang, so say what it is waiting # for. - slack_notify(f":mag: Reviewing the last cycle with " - f"{CRITIC_LABEL} before exit.") - _set_phase(f"round {round_num}: critic reviewing ({CRITIC_LABEL})") - reply = critic.review(CRITIC_MODEL, new_section, - _recent_results()) + slack_notify( + f":mag: Reviewing the last cycle with " + f"{CRITIC_LABEL} before exit." + ) + _set_phase( + f"round {round_num}: critic reviewing ({CRITIC_LABEL})" + ) + reply = critic.review( + CRITIC_MODEL, new_section, _recent_results() + ) if reply: _append_review(reply) found = critic.blocking(reply) - print(f"[critic] {len(found)} blocking finding(s)", flush=True) + print( + f"[critic] {len(found)} blocking finding(s)", flush=True + ) answering_critic = bool(found) if found: prompt = _critic_prompt(found, reply, tail=prompt) @@ -1162,20 +1336,34 @@ def _on_sigterm(): tools.cycle_done_pending() answering_critic = False new_submits = tools.submit_count() - submits_before - print(f"[round {round_num}] new_submits={new_submits} " - f"in_flight={tools.jobs_in_flight()} pending={tools.pending_count()}", - flush=True) + print( + f"[round {round_num}] new_submits={new_submits} " + f"in_flight={tools.jobs_in_flight()} pending={tools.pending_count()}", + flush=True, + ) # Start winding down when the budget is spent or time is up. Only the # DECISION happens here -- outstanding work still drains below, so a # run never orphans in-flight jobs. Finding a good result is NOT a # stop condition; keep exploring new regions. - over_time = MAX_RUNTIME is not None and (time.time() - start_time) >= MAX_RUNTIME - if stopping is None and (tools.submit_count() >= tools.MAX_SUBMITS or over_time): - stopping = "time limit" if over_time else f"budget ({tools.MAX_SUBMITS} submits)" + over_time = ( + MAX_RUNTIME is not None + and (time.time() - start_time) >= MAX_RUNTIME + ) + if stopping is None and ( + tools.submit_count() >= tools.MAX_SUBMITS or over_time + ): + stopping = ( + "time limit" + if over_time + else f"budget ({tools.MAX_SUBMITS} submits)" + ) tools.request_stop() - print(f"{stopping} reached -- winding down: no new work, finishing " - f"what is in flight.", flush=True) + print( + f"{stopping} reached -- winding down: no new work, finishing " + f"what is in flight.", + flush=True, + ) # Everything collected and the run is ending: give the agent a turn (or # two) to write the journal and LOGBOOK BEFORE exiting. Without this the # loop would break the moment the last job landed and the write-up would @@ -1183,11 +1371,17 @@ def _on_sigterm(): if stopping is not None and tools.jobs_in_flight() == 0: if finalize_rounds < MAX_FINALIZE_ROUNDS: finalize_rounds += 1 - print(f"Drained -- finalize turn {finalize_rounds}/{MAX_FINALIZE_ROUNDS} " - f"(write-up).", flush=True) + print( + f"Drained -- finalize turn {finalize_rounds}/{MAX_FINALIZE_ROUNDS} " + f"(write-up).", + flush=True, + ) prompt = FINALIZE_PROMPT continue - print(f"{stopping}: drained and written up — run complete.", flush=True) + print( + f"{stopping}: drained and written up — run complete.", + flush=True, + ) stop_reason = stopping break # Idle round with budget left: re-prompt the agent to propose a NEW @@ -1196,8 +1390,11 @@ def _on_sigterm(): if tools.jobs_in_flight() == 0 and new_submits == 0: empty_rounds += 1 if empty_rounds >= MAX_EMPTY_ROUNDS: - print(f"No new work proposed for {MAX_EMPTY_ROUNDS} rounds — " - f"stopping.", flush=True) + print( + f"No new work proposed for {MAX_EMPTY_ROUNDS} rounds — " + f"stopping.", + flush=True, + ) stop_reason = f"no new work for {MAX_EMPTY_ROUNDS} rounds" break prompt = EXPLORE_PROMPT @@ -1214,9 +1411,12 @@ def _on_sigterm(): backend_problem = None board_update = None while tools.pending_count() > 0: - _set_phase(f"round {round_num}: waiting for " - f"{tools.pending_count()} job(s)") - done = await loop.run_in_executor(None, tools.wait_for_any, ANNOUNCE_POLL) + _set_phase( + f"round {round_num}: waiting for {tools.pending_count()} job(s)" + ) + done = await loop.run_in_executor( + None, tools.wait_for_any, ANNOUNCE_POLL + ) _heartbeat(force=False) if done > 0: break @@ -1237,13 +1437,19 @@ def _on_sigterm(): continue since_tick = 0 stalled += WAIT_TIMEOUT - print(f"[waiting] {tools.pending_count()} job(s) still queued/running " - f"after {stalled // 60} min -- still alive.", flush=True) + print( + f"[waiting] {tools.pending_count()} job(s) still queued/running " + f"after {stalled // 60} min -- still alive.", + flush=True, + ) # Catch a stuck/dead backend AT THIS TICK (not on the summary interval): # if nothing is really running, break now and let the agent recover or alert. backend_problem = tools.backend_trouble() if backend_problem: - print(f"[backend check] problem detected: {backend_problem}", flush=True) + print( + f"[backend check] problem detected: {backend_problem}", + flush=True, + ) break # Periodic summary during a long queue wait: break to give the agent # a turn to report in its own words, then resume waiting next round. @@ -1251,24 +1457,31 @@ def _on_sigterm(): report_due = True break if STALL_LIMIT is not None and stalled >= STALL_LIMIT: - print(f"Stall cap ({STALL_LIMIT}s) reached; stopping. Pending jobs will " - f"need recovery on restart.", flush=True) + print( + f"Stall cap ({STALL_LIMIT}s) reached; stopping. Pending jobs will " + f"need recovery on restart.", + flush=True, + ) stop_reason = f"stall cap ({STALL_LIMIT}s)" return if board_update: prompt = _announcements_prompt(board_update) elif backend_problem: - prompt = ("A backend health check during the wait found a problem: " - f"{backend_problem}. Nothing is completing. Act now: if it is " - "recoverable, resubmit the affected config(s); if not, call " - "notify(blocking=true) with a clear one-line message so the run stops.") + prompt = ( + "A backend health check during the wait found a problem: " + f"{backend_problem}. Nothing is completing. Act now: if it is " + "recoverable, resubmit the affected config(s); if not, call " + "notify(blocking=true) with a clear one-line message so the run stops." + ) backend_problem = None elif report_due: report_due = False last_daily = time.time() await _post_scheduled_status(client, round_num, start_time) - prompt = (REPORT_PROMPT + " Nothing new has completed; just post the " - "summary and take no other action.") + prompt = ( + REPORT_PROMPT + " Nothing new has completed; just post the " + "summary and take no other action." + ) elif stopping is not None: # CONTINUE asks for the next region, which is the one thing a run # that is winding down must not do. @@ -1286,21 +1499,26 @@ def _on_sigterm(): # down can block (a local job runs for as long as it runs), and if that # happens the run must still end up correctly recorded rather than frozen as # "running" forever. The run dir itself STAYS -- it is the run history. - _write_meta(status="stopped", stop_reason=stop_reason, - ended_at=datetime.now().isoformat(timespec="seconds"), - remote_submitted=tools.submit_count(), - local_submitted=tools.local_submit_count(), - uptime_s=int(time.time() - start_time)) + _write_meta( + status="stopped", + stop_reason=stop_reason, + ended_at=datetime.now().isoformat(timespec="seconds"), + remote_submitted=tools.submit_count(), + local_submitted=tools.local_submit_count(), + uptime_s=int(time.time() - start_time), + ) for path in (pid_file, os.path.join(RUN_DIR, "heartbeat")): try: os.remove(path) except OSError: pass if NOTIFY_FINISH: - slack_notify(f":checkered_flag: Agent stopped — " - f"reason: {stop_reason} · {tools.submit_count()} remote / " - f"{tools.local_submit_count()} local submitted · " - f"uptime {_fmt_uptime(time.time() - start_time)}.") + slack_notify( + f":checkered_flag: Agent stopped — " + f"reason: {stop_reason} · {tools.submit_count()} remote / " + f"{tools.local_submit_count()} local submitted · " + f"uptime {_fmt_uptime(time.time() - start_time)}." + ) shutdown_executor() print("Executor shut down.", flush=True) _stop_watcher() diff --git a/framework/critic.py b/framework/critic.py index 36ae069..c7b5ef5 100644 --- a/framework/critic.py +++ b/framework/critic.py @@ -40,22 +40,32 @@ SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) MODEL_SETTING = (os.environ.get("CRITIC_MODEL") or "").strip() -REQUIRED = (os.environ.get("CRITIC_REQUIRED", "").strip().lower() - in ("1", "true", "yes", "on")) +REQUIRED = os.environ.get("CRITIC_REQUIRED", "").strip().lower() in ( + "1", + "true", + "yes", + "on", +) MAX_TOKENS = int(os.environ.get("CRITIC_MAX_TOKENS", "8000")) # Usually the agent's own gateway, with a different model on it. Separable because the # common arrangement is the reverse of that: the agent on Claude directly, and a proxy # standing alongside purely to reach a second family for review. -BASE_URL = (os.environ.get("CRITIC_BASE_URL") - or os.environ.get("ANTHROPIC_BASE_URL") - or "https://api.anthropic.com").rstrip("/") +BASE_URL = ( + os.environ.get("CRITIC_BASE_URL") + or os.environ.get("ANTHROPIC_BASE_URL") + or "https://api.anthropic.com" +).rstrip("/") + + def _claude_key(): """Whatever Claude Code itself authenticates with. A lab reaching a second model usually reaches it through the same gateway the agent uses, with the same credential, and that credential is already configured -- asking a person to write it out again invites a second, staler copy of it.""" - cfg = os.path.join(os.path.expanduser(os.environ.get("CLAUDE_CONFIG_DIR", "~/.claude")), - "settings.json") + cfg = os.path.join( + os.path.expanduser(os.environ.get("CLAUDE_CONFIG_DIR", "~/.claude")), + "settings.json", + ) try: with open(cfg) as f: settings = json.load(f) @@ -67,8 +77,9 @@ def _claude_key(): helper = settings.get("apiKeyHelper") if helper: try: - out = subprocess.run(helper, shell=True, capture_output=True, text=True, - timeout=15) + out = subprocess.run( + helper, shell=True, capture_output=True, text=True, timeout=15 + ) return out.stdout.strip() except Exception as e: print(f"[critic] apiKeyHelper failed (ignored): {e}", flush=True) @@ -85,18 +96,24 @@ def _key(): with open(os.path.expanduser(path)) as f: return f.read().strip() except OSError as e: - print(f"[critic] cannot read CRITIC_API_KEY_FILE ({e}); " - "falling back", flush=True) - return (os.environ.get("CRITIC_API_KEY") - or os.environ.get("ANTHROPIC_API_KEY") - or _claude_key()) + print( + f"[critic] cannot read CRITIC_API_KEY_FILE ({e}); falling back", + flush=True, + ) + return ( + os.environ.get("CRITIC_API_KEY") + or os.environ.get("ANTHROPIC_API_KEY") + or _claude_key() + ) API_KEY = _key() BLOCK_RE = re.compile( r"CLAIM:\s*(?P.+?)\n\s*VERDICT:\s*(?P\w+).*?" - r"SEVERITY:\s*(?P\w+)", re.S | re.I) + r"SEVERITY:\s*(?P\w+)", + re.DOTALL | re.IGNORECASE, +) class CriticUnavailable(Exception): @@ -107,14 +124,16 @@ def _prompt_text(): """What the critic is asked to do. The level picks how much of a write-up is in scope: everything it asserts, or only what a recorded number can settle.""" level = (os.environ.get("CRITIC_LEVEL") or "full").strip().lower() - path = (os.environ.get("CRITIC_PROMPT_FILE") - or os.path.join(SCRIPT_DIR, f"critic_prompt_{level}.md")) + path = os.environ.get("CRITIC_PROMPT_FILE") or os.path.join( + SCRIPT_DIR, f"critic_prompt_{level}.md" + ) try: with open(path) as f: return f.read() except OSError as e: raise CriticUnavailable( - f"no critic prompt at {path} (CRITIC_LEVEL={level}): {e}") + f"no critic prompt at {path} (CRITIC_LEVEL={level}): {e}" + ) def gateway_up(): @@ -142,24 +161,31 @@ def ensure_gateway(needed=False): return None wait = int(os.environ.get("CRITIC_GATEWAY_WAIT", "60")) try: - subprocess.Popen(shlex.split(start), start_new_session=True, - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + subprocess.Popen( + shlex.split(start), + start_new_session=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) except Exception as e: return f"could not start the model gateway (ignored): {e}" for _ in range(wait): time.sleep(1) if gateway_up(): return f"started the model gateway at {BASE_URL}" - return (f"started the model gateway but {BASE_URL} did not answer within " - f"{wait}s; continuing without it") + return ( + f"started the model gateway but {BASE_URL} did not answer within " + f"{wait}s; continuing without it" + ) def _served_models(): """Model names the gateway offers, with what each resolves to upstream. Empty when there is no gateway -- talking to Anthropic directly means Claude only.""" try: - req = urllib.request.Request(BASE_URL + "/model/info", - headers={"x-api-key": API_KEY}) + req = urllib.request.Request( + BASE_URL + "/model/info", headers={"x-api-key": API_KEY} + ) with urllib.request.urlopen(req, timeout=10) as r: data = json.load(r).get("data", []) except Exception: @@ -168,7 +194,9 @@ def _served_models(): for m in data: name = m.get("model_name") if name: - out[name] = (m.get("litellm_params", {}).get("model", "") or "").split("/")[-1] + out[name] = (m.get("litellm_params", {}).get("model", "") or "").split("/")[ + -1 + ] return out @@ -181,13 +209,15 @@ def resolve(agent_model=""): if served and MODEL_SETTING not in served: raise CriticUnavailable( f"CRITIC_MODEL={MODEL_SETTING} is not served by {BASE_URL} " - f"(it has: {', '.join(sorted(served)) or 'nothing'})") + f"(it has: {', '.join(sorted(served)) or 'nothing'})" + ) return MODEL_SETTING, served.get(MODEL_SETTING) or MODEL_SETTING if not served: if REQUIRED: raise CriticUnavailable( f"CRITIC_MODEL=auto but {BASE_URL} lists no models, so there is nothing " - "to choose from") + "to choose from" + ) return None, "none (asked for one, but no second model is reachable)" # A different family than the agent's own, where there is one. agent_family = (agent_model or "").split("-")[0].lower() @@ -207,15 +237,29 @@ def review(model, write_up, evidence, prompt=None): if prompt is not None: prompt = prompt + "\n\n" + write_up else: - prompt = (_prompt_text() - + "\n\n# The write-up\n\n" + write_up - + "\n\n# The recorded results\n\n" + (evidence or "(no rows recorded)")) - body = json.dumps({"model": model, "max_tokens": MAX_TOKENS, - "messages": [{"role": "user", "content": prompt}]}).encode() + prompt = ( + _prompt_text() + + "\n\n# The write-up\n\n" + + write_up + + "\n\n# The recorded results\n\n" + + (evidence or "(no rows recorded)") + ) + body = json.dumps( + { + "model": model, + "max_tokens": MAX_TOKENS, + "messages": [{"role": "user", "content": prompt}], + } + ).encode() req = urllib.request.Request( - BASE_URL + "/v1/messages", data=body, - headers={"content-type": "application/json", "x-api-key": API_KEY, - "anthropic-version": "2023-06-01"}) + BASE_URL + "/v1/messages", + data=body, + headers={ + "content-type": "application/json", + "x-api-key": API_KEY, + "anthropic-version": "2023-06-01", + }, + ) try: with urllib.request.urlopen(req, timeout=300) as r: reply = json.load(r) @@ -223,8 +267,11 @@ def review(model, write_up, evidence, prompt=None): # A reasoning model can spend the whole budget thinking and return nothing. # Silence and truncation look identical from here, so say which it was. if not text and reply.get("stop_reason") == "max_tokens": - print(f"[critic] no review: the reply hit CRITIC_MAX_TOKENS " - f"({MAX_TOKENS}) before writing anything", flush=True) + print( + f"[critic] no review: the reply hit CRITIC_MAX_TOKENS " + f"({MAX_TOKENS}) before writing anything", + flush=True, + ) return text except Exception as e: print(f"[critic] review failed (ignored): {e}", flush=True) diff --git a/framework/engineer.py b/framework/engineer.py index 0cf2142..dc69640 100644 --- a/framework/engineer.py +++ b/framework/engineer.py @@ -28,8 +28,6 @@ """ import asyncio -import glob -import json import os import subprocess import sys @@ -44,18 +42,21 @@ SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) LAB_DIR = os.path.abspath(os.environ.get("LAB_DIR", os.path.join(SCRIPT_DIR, ".."))) -INBOX = os.environ.get("SLACK_INBOX") or os.path.join(LAB_DIR, "workspace", "run", - "engineer_inbox.md") +INBOX = os.environ.get("SLACK_INBOX") or os.path.join( + LAB_DIR, "workspace", "run", "engineer_inbox.md" +) STATE = os.path.join(os.path.dirname(INBOX), "engineer_seen.txt") -HEARTBEAT = (os.environ.get("ENGINEER_HEARTBEAT") - or os.path.join(os.path.dirname(INBOX), "engineer_heartbeat")) +HEARTBEAT = os.environ.get("ENGINEER_HEARTBEAT") or os.path.join( + os.path.dirname(INBOX), "engineer_heartbeat" +) SESSION_FILE = os.path.join(os.path.dirname(INBOX), "engineer_session") POLL = int(os.environ.get("ENGINEER_POLL", "5")) BRANCH = (os.environ.get("ENGINEER_BRANCH") or "").strip() RESUME_SESSION = (os.environ.get("RESUME_SESSION") or "").strip() COMPACT_FIRST = RESUME_SESSION.lower() == "compact" -NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join(SCRIPT_DIR, - "slack_notify.sh") +NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join( + SCRIPT_DIR, "slack_notify.sh" +) SYSTEM_PROMPT = f"""You are the engineer for AgentLab, the framework in {LAB_DIR}, and you work on it from a Slack channel. Someone types there; you answer, and change the @@ -125,8 +126,8 @@ def beat(): def new_lines(seen, current): old, new = seen.splitlines(), current.splitlines() - if new[:len(old)] == old: - return "\n".join(new[len(old):]).strip() + if new[: len(old)] == old: + return "\n".join(new[len(old) :]).strip() return current @@ -135,7 +136,9 @@ def last_session(): The record carries the checkout it belonged to, so a second lab on the same machine does not pick up this one's conversation.""" lines = _read(SESSION_FILE).splitlines() - if len(lines) >= 2 and os.path.abspath(lines[1].strip()) == os.path.abspath(LAB_DIR): + if len(lines) >= 2 and os.path.abspath(lines[1].strip()) == os.path.abspath( + LAB_DIR + ): return lines[0].strip() return "" @@ -176,13 +179,20 @@ def on_branch(): if not BRANCH: return "wherever the repository already is" try: - current = subprocess.run(["git", "-C", LAB_DIR, "rev-parse", "--abbrev-ref", - "HEAD"], capture_output=True, text=True, - timeout=15).stdout.strip() + current = subprocess.run( + ["git", "-C", LAB_DIR, "rev-parse", "--abbrev-ref", "HEAD"], + capture_output=True, + text=True, + timeout=15, + ).stdout.strip() if current == BRANCH: return f"on {BRANCH}" - made = subprocess.run(["git", "-C", LAB_DIR, "checkout", "-B", BRANCH], - capture_output=True, text=True, timeout=30) + made = subprocess.run( + ["git", "-C", LAB_DIR, "checkout", "-B", BRANCH], + capture_output=True, + text=True, + timeout=30, + ) if made.returncode != 0: return None return f"switched from {current} to {BRANCH}" @@ -207,18 +217,28 @@ async def answer(client, text): async def main(): once = "--once" in sys.argv branch = on_branch() - print(f"Engineer watching {INBOX} (poll {POLL}s)" - f"{' [once]' if once else ''}", flush=True) - print(f"branch: {branch or 'not a git repository -- commits will fail'}", flush=True) + print( + f"Engineer watching {INBOX} (poll {POLL}s){' [once]' if once else ''}", + flush=True, + ) + print( + f"branch: {branch or 'not a git repository -- commits will fail'}", flush=True + ) resume = resume_session() if RESUME_SESSION and not resume: - print(f"no earlier session recorded for {LAB_DIR} -- starting fresh", flush=True) + print( + f"no earlier session recorded for {LAB_DIR} -- starting fresh", flush=True + ) elif resume: - print(f"resuming session {resume}" - + (", compacting first" if COMPACT_FIRST else ""), flush=True) + print( + f"resuming session {resume}" + + (", compacting first" if COMPACT_FIRST else ""), + flush=True, + ) if not os.path.isfile(NOTIFY_SCRIPT): - print(f"[engineer] WARNING: {NOTIFY_SCRIPT} missing -- cannot reply.", - flush=True) + print( + f"[engineer] WARNING: {NOTIFY_SCRIPT} missing -- cannot reply.", flush=True + ) options = ClaudeAgentOptions( system_prompt=SYSTEM_PROMPT, diff --git a/framework/secretary.py b/framework/secretary.py index 9b2e572..a8937ef 100644 --- a/framework/secretary.py +++ b/framework/secretary.py @@ -52,8 +52,9 @@ # One secretary serves the whole lab. WORKSPACE_ROOT holds one directory per campaign, # each with its own ANNOUNCEMENTS.md, results.jsonl, LOGBOOK.md and JOURNAL.md, plus # run/ for state that belongs to the lab rather than to any one campaign. -WORKSPACE_ROOT = os.path.abspath(os.environ.get( - "WORKSPACE_ROOT", os.path.join(SCRIPT_DIR, "..", "workspace"))) +WORKSPACE_ROOT = os.path.abspath( + os.environ.get("WORKSPACE_ROOT", os.path.join(SCRIPT_DIR, "..", "workspace")) +) # Slack questions land here while this process is alive. Separate from the boards on # purpose: a board is broadcast to a campaign's agents, this is a queue for one reader. INBOX = os.path.join(WORKSPACE_ROOT, "run", "slack_inbox.md") @@ -61,9 +62,13 @@ # Liveness, read by slack_to_board.py to decide where to deliver. Same convention as # the agents' runs//heartbeat: a recent timestamp means alive. HEARTBEAT = os.path.join(WORKSPACE_ROOT, "run", "secretary_heartbeat") -POLL = int(os.environ.get("SECRETARY_POLL", "5")) # s between inbox checks -AGENT_ALIVE_WITHIN = int(os.environ.get("AGENT_ALIVE_WITHIN", "300")) # s; fresher heartbeat = agent is up -NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join(SCRIPT_DIR, "slack_notify.sh") +POLL = int(os.environ.get("SECRETARY_POLL", "5")) # s between inbox checks +AGENT_ALIVE_WITHIN = int( + os.environ.get("AGENT_ALIVE_WITHIN", "300") +) # s; fresher heartbeat = agent is up +NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join( + SCRIPT_DIR, "slack_notify.sh" +) SYSTEM_PROMPT = f"""You are the secretary for a collaborative agentic search workflow. Research agents run on compute nodes and coordinate through shared files. @@ -257,8 +262,8 @@ def new_lines(seen, current): """Lines added since the last look. Falls back to the whole inbox if it was edited rather than appended to, since there is no clean 'new part' then.""" old, new = seen.splitlines(), current.splitlines() - if new[:len(old)] == old: - return "\n".join(new[len(old):]).strip() + if new[: len(old)] == old: + return "\n".join(new[len(old) :]).strip() return current @@ -286,7 +291,7 @@ def live_agents(): return sorted(out) -_session_id = None # this secretary's Claude session, for reopening it later +_session_id = None # this secretary's Claude session, for reopening it later async def answer(client, text, agents): @@ -295,10 +300,14 @@ async def answer(client, text, agents): # here is re-sent with every question, so anything actionable becomes a standing # order -- which is how "relay if it needs live reasoning" turned into relaying # answers it had already given. - status = ("Research agents running:\n" + "\n".join(agents) - if agents else "No research agent is running.") - await client.query(status + "\n\nNew from Slack:\n\n" + text + - "\n\nAnswer it, then stop.") + status = ( + "Research agents running:\n" + "\n".join(agents) + if agents + else "No research agent is running." + ) + await client.query( + status + "\n\nNew from Slack:\n\n" + text + "\n\nAnswer it, then stop." + ) async for message in client.receive_response(): if isinstance(message, AssistantMessage): for block in message.content: @@ -310,20 +319,29 @@ async def answer(client, text, agents): if sid and sid != _session_id: _session_id = sid try: - with open(os.path.join(WORKSPACE_ROOT, "run", "secretary_session"), "w") as f: + with open( + os.path.join(WORKSPACE_ROOT, "run", "secretary_session"), "w" + ) as f: f.write(f"{sid}\n{WORKSPACE_ROOT}\n") except Exception as e: - print(f"[secretary] session id not recorded (ignored): {e}", flush=True) + print( + f"[secretary] session id not recorded (ignored): {e}", + flush=True, + ) print(f"[turn end] {message.subtype}", flush=True) async def main(): once = "--once" in sys.argv - print(f"Secretary watching {INBOX} (poll {POLL}s)" - f"{' [once]' if once else ''}", flush=True) + print( + f"Secretary watching {INBOX} (poll {POLL}s){' [once]' if once else ''}", + flush=True, + ) if not os.path.isfile(NOTIFY_SCRIPT): - print(f"[secretary] WARNING: {NOTIFY_SCRIPT} missing -- cannot post replies.", - flush=True) + print( + f"[secretary] WARNING: {NOTIFY_SCRIPT} missing -- cannot post replies.", + flush=True, + ) options = ClaudeAgentOptions( system_prompt=SYSTEM_PROMPT, allowed_tools=["Read", "Grep", "Glob", "Bash"], @@ -339,8 +357,11 @@ async def main(): fresh = new_lines(read_seen(), inbox) if inbox else "" if fresh: agents = live_agents() - print(f"\n--- answering ({', '.join(agents) or 'no agent running'}) " - f"---\n{fresh}\n-------------------------", flush=True) + print( + f"\n--- answering ({', '.join(agents) or 'no agent running'}) " + f"---\n{fresh}\n-------------------------", + flush=True, + ) # Record BEFORE answering, so a failure cannot loop on the same message. write_seen(inbox) try: diff --git a/framework/slack_to_board.py b/framework/slack_to_board.py index 49091c5..cde0029 100644 --- a/framework/slack_to_board.py +++ b/framework/slack_to_board.py @@ -41,8 +41,8 @@ SLACK_READER_HEARTBEAT that reader's liveness file """ -import json import glob +import json import os import sys import time @@ -52,24 +52,34 @@ SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) # One bridge serves the whole lab. WORKSPACE_ROOT holds one directory per campaign; # each campaign has its own ANNOUNCEMENTS.md that its agent reads between rounds. -WORKSPACE_ROOT = os.path.abspath(os.environ.get( - "WORKSPACE_ROOT", os.path.join(SCRIPT_DIR, "..", "workspace"))) -STATE = (os.environ.get("SLACK_STATE") - or os.path.join(WORKSPACE_ROOT, "run", "slack_last_ts")) -INBOX = os.environ.get("SLACK_INBOX") or os.path.join(WORKSPACE_ROOT, "run", - "slack_inbox.md") +WORKSPACE_ROOT = os.path.abspath( + os.environ.get("WORKSPACE_ROOT", os.path.join(SCRIPT_DIR, "..", "workspace")) +) +STATE = os.environ.get("SLACK_STATE") or os.path.join( + WORKSPACE_ROOT, "run", "slack_last_ts" +) +INBOX = os.environ.get("SLACK_INBOX") or os.path.join( + WORKSPACE_ROOT, "run", "slack_inbox.md" +) HEARTBEAT = os.environ.get("SLACK_READER_HEARTBEAT") or os.path.join( - WORKSPACE_ROOT, "run", "secretary_heartbeat") + WORKSPACE_ROOT, "run", "secretary_heartbeat" +) # s; a secretary heartbeat fresher than this means it is up and owns Slack questions. # It rewrites the file every poll (default 5s), so this tolerates many missed beats. SECRETARY_ALIVE_WITHIN = int(os.environ.get("SECRETARY_ALIVE_WITHIN", "60")) CHANNEL = os.environ.get("SLACK_CHANNEL", "") -TOKEN_FILE = os.environ.get("SLACK_BOT_TOKEN_FILE", - os.path.expanduser("~/.slack_bot_token")) +TOKEN_FILE = os.environ.get( + "SLACK_BOT_TOKEN_FILE", os.path.expanduser("~/.slack_bot_token") +) POLL = int(os.environ.get("SLACK_FETCH_POLL", "5")) # Plain-text fallback for a mention typed without Slack autocomplete. BOT_NAME = os.environ.get("SLACK_BOT_NAME", "@cas_agent") -READ_ALL = os.environ.get("SLACK_READ_ALL", "").strip().lower() in ("1", "true", "yes", "on") +READ_ALL = os.environ.get("SLACK_READ_ALL", "").strip().lower() in ( + "1", + "true", + "yes", + "on", +) def slack_get(method, token, **params): @@ -114,7 +124,9 @@ def secretary_up(): except FileNotFoundError: return False except Exception as e: - print(f"[slack] heartbeat read failed, assuming down (ignored): {e}", flush=True) + print( + f"[slack] heartbeat read failed, assuming down (ignored): {e}", flush=True + ) return False @@ -123,28 +135,34 @@ def forward(messages, me): up, to every campaign board if it is not.""" lines = [] read_all = READ_ALL and (DEDICATED or secretary_up()) - for m in reversed(messages): # Slack returns newest first + for m in reversed(messages): # Slack returns newest first if m.get("bot_id") or m.get("subtype"): - continue # never echo bot posts back at the agents + continue # never echo bot posts back at the agents text = m.get("text", "").strip() if not text: continue if ALLOW and m.get("user") not in ALLOW: - print(f"ignored, not on SLACK_ALLOW: <@{m.get('user')}>: {text}", flush=True) + print( + f"ignored, not on SLACK_ALLOW: <@{m.get('user')}>: {text}", flush=True + ) continue addressed = f"<@{me}>" in text or BOT_NAME in text if not addressed and not read_all: - continue # not addressed to the agents + continue # not addressed to the agents text = text.replace(f"<@{me}>", "").strip() # The author's Slack id travels with the message: it is what records who asked # for a run, and Slack renders it as their name when it is quoted back. who = f"<@{m['user']}>" if m.get("user") else "someone" - tag = ("reply with the notify tool" if addressed - else "overheard, not addressed to you") + tag = ( + "reply with the notify tool" + if addressed + else "overheard, not addressed to you" + ) # One reader, one channel: it is talking to whoever is there, so the author's # id is noise. The shared bridge needs it to say who asked for what. - lines.append(f"[from Slack -- {tag}] " - + (text if DEDICATED else f"{who}: {text}")) + lines.append( + f"[from Slack -- {tag}] " + (text if DEDICATED else f"{who}: {text}") + ) if not lines: return 0 if DEDICATED or secretary_up(): @@ -155,23 +173,31 @@ def forward(messages, me): for line in lines: print(f"forwarded to {where}:", line, flush=True) return len(lines) - campaigns = sorted(d for d in glob.glob(os.path.join(WORKSPACE_ROOT, "*")) - if os.path.isdir(d) and os.path.basename(d) != "run") + campaigns = sorted( + d + for d in glob.glob(os.path.join(WORKSPACE_ROOT, "*")) + if os.path.isdir(d) and os.path.basename(d) != "run" + ) if not campaigns: - print(f"no campaigns under {WORKSPACE_ROOT}; dropping {len(lines)} message(s)", - flush=True) + print( + f"no campaigns under {WORKSPACE_ROOT}; dropping {len(lines)} message(s)", + flush=True, + ) return 0 delivered = 0 for line in lines: # A message naming a campaign goes to that one; otherwise to all of them. named = [d for d in campaigns if os.path.basename(d).lower() in line.lower()] - for d in (named or campaigns): + for d in named or campaigns: os.makedirs(d, exist_ok=True) with open(os.path.join(d, "ANNOUNCEMENTS.md"), "a") as f: f.write(line + "\n") delivered += 1 - print(f"forwarded to {len(named) or len(campaigns)} campaign(s):", line, - flush=True) + print( + f"forwarded to {len(named) or len(campaigns)} campaign(s):", + line, + flush=True, + ) return delivered @@ -183,8 +209,9 @@ def check(token, me): write_state(now) print(f"first run -- starting from now ({now}); nothing forwarded", flush=True) return - resp = slack_get("conversations.history", token, channel=CHANNEL, - oldest=oldest, limit=50) + resp = slack_get( + "conversations.history", token, channel=CHANNEL, oldest=oldest, limit=50 + ) if not resp.get("ok"): print(f"[slack] read failed (ignored): {resp.get('error')}", flush=True) return @@ -209,8 +236,11 @@ def main(): if not who.get("ok"): sys.exit(f"token rejected by Slack: {who.get('error')}") me = who["user_id"] - print(f"Watching Slack channel {CHANNEL} as {who.get('user')} ({me}); " - f"campaigns under {WORKSPACE_ROOT}", flush=True) + print( + f"Watching Slack channel {CHANNEL} as {who.get('user')} ({me}); " + f"campaigns under {WORKSPACE_ROOT}", + flush=True, + ) once = "--once" in sys.argv while True: diff --git a/framework/tools.py b/framework/tools.py index 40cbdf4..277fc79 100644 --- a/framework/tools.py +++ b/framework/tools.py @@ -10,11 +10,11 @@ config.json. """ +import fcntl import importlib import itertools import json import os -import fcntl import subprocess import sys import time @@ -22,13 +22,12 @@ from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor from concurrent.futures import wait as _futures_wait -from claude_agent_sdk import tool, create_sdk_mcp_server - import transfer as _transfer +from claude_agent_sdk import create_sdk_mcp_server, tool from globus_compute_sdk import Executor -from globus_compute_sdk.serialize import ComputeSerializer, AllCodeStrategies +from globus_compute_sdk.serialize import AllCodeStrategies, ComputeSerializer -ROLE = os.environ.get("ROLE", "both") # free-form; the prompt defines what roles mean +ROLE = os.environ.get("ROLE", "both") # free-form; the prompt defines what roles mean SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) LAB_DIR = os.path.abspath(os.environ.get("LAB_DIR", os.path.join(SCRIPT_DIR, ".."))) @@ -59,12 +58,16 @@ def _read_json(path, what, needs=()): sys.exit("CAMPAIGN is not set (the directory name under campaigns/).") _CAMPAIGN_DIR = os.path.join(LAB_DIR, "campaigns", CAMPAIGN) -_cam = _read_json(os.path.join(_CAMPAIGN_DIR, "campaign.json"), - f"campaign '{CAMPAIGN}'", needs=("system",)) +_cam = _read_json( + os.path.join(_CAMPAIGN_DIR, "campaign.json"), + f"campaign '{CAMPAIGN}'", + needs=("system",), +) SYSTEM = _cam["system"] -_sys_cfg = _read_json(os.path.join(LAB_DIR, "systems", f"{SYSTEM}.json"), - f"system '{SYSTEM}'") +_sys_cfg = _read_json( + os.path.join(LAB_DIR, "systems", f"{SYSTEM}.json"), f"system '{SYSTEM}'" +) # --- the task plug-in ---------------------------------------------------------- # Supplies what a job IS: how it is described to the agent, what arguments it takes, @@ -86,35 +89,50 @@ def _read_json(path, what, needs=()): # work_dir then defaults to the campaign workspace. _user_path = os.path.join(LAB_DIR, "users", USER_NAME, f"{SYSTEM}.json") if HAS_REMOTE: - _usr = _read_json(_user_path, f"your access to '{SYSTEM}'", - needs=("endpoint", "account", "work_dir")) + _usr = _read_json( + _user_path, + f"your access to '{SYSTEM}'", + needs=("endpoint", "account", "work_dir"), + ) else: - _usr = (_read_json(_user_path, f"your access to '{SYSTEM}'") - if os.path.isfile(_user_path) else {}) + _usr = ( + _read_json(_user_path, f"your access to '{SYSTEM}'") + if os.path.isfile(_user_path) + else {} + ) _usr.setdefault("work_dir", os.path.join(LAB_DIR, "workspace", CAMPAIGN)) ENDPOINT_ID = _usr.get("endpoint", "") # Globus Transfer is optional: configured per user, and simply absent otherwise. It is # how the agent reads files on the compute system when the two do not share a filesystem. -_transfer.CFG = _transfer.configure(_usr, os.path.join(LAB_DIR, "workspace", CAMPAIGN), - _CAMPAIGN_DIR, _sys_cfg) +_transfer.CFG = _transfer.configure( + _usr, os.path.join(LAB_DIR, "workspace", CAMPAIGN), _CAMPAIGN_DIR, _sys_cfg +) HAS_TRANSFER = _transfer.CFG is not None # How many jobs may be in flight at once. The system file holds a site default, bounded # by queue policy and allocation rather than by the size of the machine; a campaign # overrides it, because what is sensible depends on what one job does. -MAX_CONCURRENT = int(os.environ.get("MAX_CONCURRENT", - _cam.get("max_concurrent", - _sys_cfg.get("max_concurrent", 1)))) +MAX_CONCURRENT = int( + os.environ.get( + "MAX_CONCURRENT", _cam.get("max_concurrent", _sys_cfg.get("max_concurrent", 1)) + ) +) # Named resource shapes on one system (e.g. a small quick queue and a large long one). # A task may route a job to one; otherwise the default is used. _bucket_defaults = dict(_sys_cfg.get("bucket_defaults", {})) -_bucket_defaults.update(_cam.get("resources", {})) # campaign: queue, walltime, nodes -_bucket_defaults.update(_usr.get("resources", {})) # user: anything they must override +_bucket_defaults.update(_cam.get("resources", {})) # campaign: queue, walltime, nodes +_bucket_defaults.update(_usr.get("resources", {})) # user: anything they must override _bucket_defaults["account"] = _usr.get("account", "") -_SYS = {"buckets": {"default": {"num_nodes": _bucket_defaults.get("num_nodes", 1), - "user_config": _bucket_defaults}}} +_SYS = { + "buckets": { + "default": { + "num_nodes": _bucket_defaults.get("num_nodes", 1), + "user_config": _bucket_defaults, + } + } +} _default_bucket = "default" # TARGET is handed to the task's remote_fn. Everything the remote side needs must be @@ -129,25 +147,28 @@ def _read_json(path, what, needs=()): TARGET.setdefault("ppn", _sys_cfg.get("ppn", 1)) TARGET["nranks"] = _SYS["buckets"][_default_bucket].get("num_nodes", 1) * TARGET["ppn"] -REMOTE_TIMEOUT = int(os.environ.get("JOB_TIMEOUT", "43200")) # 12h client-side wait -LOCAL_TIMEOUT = int(os.environ.get("LOCAL_JOB_TIMEOUT", "14400")) # 4h +REMOTE_TIMEOUT = int(os.environ.get("JOB_TIMEOUT", "43200")) # 12h client-side wait +LOCAL_TIMEOUT = int(os.environ.get("LOCAL_JOB_TIMEOUT", "14400")) # 4h # The same, for jobs run on this machine. One by default: a local job is assumed to use # the whole thing, and a task whose jobs are small enough to share it says so. -LOCAL_MAX_CONCURRENT = int(os.environ.get("LOCAL_MAX_CONCURRENT", - _cam.get("local_max_concurrent", - _sys_cfg.get("local_max_concurrent", 1)))) +LOCAL_MAX_CONCURRENT = int( + os.environ.get( + "LOCAL_MAX_CONCURRENT", + _cam.get("local_max_concurrent", _sys_cfg.get("local_max_concurrent", 1)), + ) +) # One Executor per bucket, created lazily and reused. Each distinct user_endpoint_config # gets its own block pool on the endpoint, so buckets can run concurrently. _executors = {} -_sa_executor = None # local backend, created lazily +_sa_executor = None # local backend, created lazily -_jobs = {} # remote: job_id -> {"future", "args", "key", "bucket"} +_jobs = {} # remote: job_id -> {"future", "args", "key", "bucket"} _job_counter = itertools.count(1) _submit_count = 0 -MAX_SUBMITS = int(os.environ.get("MAX_SUBMITS", "60")) # backstop on total jobs per run +MAX_SUBMITS = int(os.environ.get("MAX_SUBMITS", "60")) # backstop on total jobs per run -_local_jobs = {} # local: job_id -> {"future", "args"} +_local_jobs = {} # local: job_id -> {"future", "args"} _local_counter = itertools.count(1) _local_submit_count = 0 @@ -155,16 +176,20 @@ def _read_json(path, what, needs=()): # drain and the run can end cleanly. Collecting finished work is unaffected. _stop_requested = False -JOBS_LOG = os.path.join(WORKSPACE_DIR, "jobs.jsonl") # durable record of every job fired +JOBS_LOG = os.path.join( + WORKSPACE_DIR, "jobs.jsonl" +) # durable record of every job fired # --- Shared announcements board ------------------------------------------------- # One plain text file. Anyone (an operator, a Slack bridge) appends lines; every # agent reads it and decides what applies to it. No routing, no per-agent state. ANNOUNCEMENTS_FILE = os.path.join(WORKSPACE_DIR, "ANNOUNCEMENTS.md") -NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join(SCRIPT_DIR, "slack_notify.sh") -_last_success_time = None # last non-error completion (real progress) -_problem_since = None # when the agent flagged a blocking problem (None = none) +NOTIFY_SCRIPT = os.environ.get("NOTIFY_SCRIPT") or os.path.join( + SCRIPT_DIR, "slack_notify.sh" +) +_last_success_time = None # last non-error completion (real progress) +_problem_since = None # when the agent flagged a blocking problem (None = none) def _slack_post(msg): @@ -172,8 +197,12 @@ def _slack_post(msg): if not os.path.isfile(NOTIFY_SCRIPT): return try: - subprocess.run(["bash", NOTIFY_SCRIPT, msg], timeout=30, - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + subprocess.run( + ["bash", NOTIFY_SCRIPT, msg], + timeout=30, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) except Exception: pass @@ -191,6 +220,7 @@ def backend_trouble(): if not HAS_REMOTE: return None import globus_compute_sdk as _gc + try: c = _gc.Client() st = (c.get_endpoint_status(ENDPOINT_ID) or {}).get("status") @@ -205,7 +235,7 @@ def backend_trouble(): for jid, j in pend: tid = getattr(j["future"], "task_id", None) if not tid: - return None # not populated yet (just submitted) -> not stuck + return None # not populated yet (just submitted) -> not stuck try: t = c.get_task(tid) or {} except Exception: @@ -273,8 +303,18 @@ def _try_claim(key, stage=""): if held and held.get("agent") != AGENT_ID: return False, held.get("agent") with open(CLAIMS_FILE, "a") as f: - f.write(json.dumps({"key": key, "stage": stage, "agent": AGENT_ID, - "ts": time.time(), "state": "claimed"}) + "\n") + f.write( + json.dumps( + { + "key": key, + "stage": stage, + "agent": AGENT_ID, + "ts": time.time(), + "state": "claimed", + } + ) + + "\n" + ) return True, None finally: fcntl.flock(lk, fcntl.LOCK_UN) @@ -286,8 +326,17 @@ def _release_claim(key): fcntl.flock(lk, fcntl.LOCK_EX) try: with open(CLAIMS_FILE, "a") as f: - f.write(json.dumps({"key": key, "agent": AGENT_ID, - "ts": time.time(), "state": "done"}) + "\n") + f.write( + json.dumps( + { + "key": key, + "agent": AGENT_ID, + "ts": time.time(), + "state": "done", + } + ) + + "\n" + ) finally: fcntl.flock(lk, fcntl.LOCK_UN) except Exception: @@ -407,8 +456,10 @@ def shutdown_executor(): # --- tools the agent calls ------------------------------------------------------ -_WINDDOWN_REFUSAL = ("submit refused: this run is winding down. Collect and log the work " - "already in flight, but do not submit anything new.") +_WINDDOWN_REFUSAL = ( + "submit refused: this run is winding down. Collect and log the work " + "already in flight, but do not submit anything new." +) @tool("submit_job", getattr(task, "JOB_DESC", ""), getattr(task, "JOB_SCHEMA", {})) @@ -416,29 +467,58 @@ async def submit_job(args): """Fire one remote job and return immediately with a job_id.""" global _submit_count if _stop_requested: - return {"content": [{"type": "text", "text": _WINDDOWN_REFUSAL}], "is_error": True} + return { + "content": [{"type": "text", "text": _WINDDOWN_REFUSAL}], + "is_error": True, + } if _submit_count >= MAX_SUBMITS: - return {"content": [{"type": "text", "text": - f"submit refused: hit MAX_SUBMITS={MAX_SUBMITS} total-jobs cap for this run"}], - "is_error": True} + return { + "content": [ + { + "type": "text", + "text": f"submit refused: hit MAX_SUBMITS={MAX_SUBMITS} total-jobs cap for this run", + } + ], + "is_error": True, + } if _remote_pending_count() >= MAX_CONCURRENT: - return {"content": [{"type": "text", "text": - f"submit refused: at capacity ({_remote_pending_count()} running/queued, " - f"max_concurrent={MAX_CONCURRENT}). Collect a finished job " - f"with get_completed_jobs before submitting more."}], "is_error": True} + return { + "content": [ + { + "type": "text", + "text": f"submit refused: at capacity ({_remote_pending_count()} running/queued, " + f"max_concurrent={MAX_CONCURRENT}). Collect a finished job " + f"with get_completed_jobs before submitting more.", + } + ], + "is_error": True, + } key = task.job_key(args) # A fresh-context agent cannot remember what it already fired, so refuse a repeat. for info in _jobs.values(): if info["key"] == key and not info["future"].done(): - return {"content": [{"type": "text", "text": - f"submit refused: a job for {key} is already in flight. Collect it " - f"with get_completed_jobs before re-submitting."}], "is_error": True} + return { + "content": [ + { + "type": "text", + "text": f"submit refused: a job for {key} is already in flight. Collect it " + f"with get_completed_jobs before re-submitting.", + } + ], + "is_error": True, + } ok, holder = _try_claim(key, stage=ROLE) if not ok: - return {"content": [{"type": "text", "text": - f"submit skipped: {key} is already claimed by {holder}. Pick different work."}], - "is_error": True} + return { + "content": [ + { + "type": "text", + "text": f"submit skipped: {key} is already claimed by {holder}. Pick different work.", + } + ], + "is_error": True, + } bucket = task.bucket_for(args) if hasattr(task, "bucket_for") else _default_bucket target = dict(TARGET) @@ -448,20 +528,44 @@ async def submit_job(args): except Exception as e: _release_claim(key) traceback.print_exc(file=sys.stderr) - return {"content": [{"type": "text", "text": f"submit failed: {e}"}], "is_error": True} + return { + "content": [{"type": "text", "text": f"submit failed: {e}"}], + "is_error": True, + } job_id = next(_job_counter) _submit_count += 1 _jobs[job_id] = {"future": fut, "args": args, "key": key, "bucket": bucket} - _append_jobs_log({"event": "submit", "job_id": job_id, "key": key, "args": args, - "bucket": bucket, "task_id": getattr(fut, "task_id", None)}) + _append_jobs_log( + { + "event": "submit", + "job_id": job_id, + "key": key, + "args": args, + "bucket": bucket, + "task_id": getattr(fut, "task_id", None), + } + ) # The budget left, with the job that was just accepted counted. It changes as the # run goes, so it belongs in what a submit answers rather than in a prompt written # once at the start -- and it counts this run's submits, not the rows in a record # that outlives the run. - return {"content": [{"type": "text", "text": json.dumps( - {"job_id": job_id, "key": key, "bucket": bucket, - "submits_used": _submit_count, "submits_allowed": MAX_SUBMITS})}]} + return { + "content": [ + { + "type": "text", + "text": json.dumps( + { + "job_id": job_id, + "key": key, + "bucket": bucket, + "submits_used": _submit_count, + "submits_allowed": MAX_SUBMITS, + } + ), + } + ] + } GET_COMPLETED_DESC = ( @@ -490,33 +594,67 @@ async def get_completed_jobs(args): res.setdefault("key", info["key"]) if "error" not in res: _last_success_time = time.time() - _problem_since = None # real progress clears any flagged problem + _problem_since = None # real progress clears any flagged problem completed.append(res) - _append_jobs_log({"event": "completed", "job_id": job_id, "key": info["key"], - "error": "error" in res}) + _append_jobs_log( + { + "event": "completed", + "job_id": job_id, + "key": info["key"], + "error": "error" in res, + } + ) _release_claim(info["key"]) del _jobs[job_id] - pending = [{"job_id": jid, "key": i["key"], "args": i["args"], "bucket": i["bucket"]} - for jid, i in _jobs.items()] - return {"content": [{"type": "text", "text": - json.dumps({"completed": completed, "pending": pending}, indent=2, default=str)}]} - - -@tool("submit_local", getattr(task, "LOCAL_DESC", ""), getattr(task, "LOCAL_SCHEMA", {})) + pending = [ + {"job_id": jid, "key": i["key"], "args": i["args"], "bucket": i["bucket"]} + for jid, i in _jobs.items() + ] + return { + "content": [ + { + "type": "text", + "text": json.dumps( + {"completed": completed, "pending": pending}, indent=2, default=str + ), + } + ] + } + + +@tool( + "submit_local", getattr(task, "LOCAL_DESC", ""), getattr(task, "LOCAL_SCHEMA", {}) +) async def submit_local(args): """Fire the local comparator and return immediately with a job_id.""" global _local_submit_count if _stop_requested: - return {"content": [{"type": "text", "text": _WINDDOWN_REFUSAL}], "is_error": True} + return { + "content": [{"type": "text", "text": _WINDDOWN_REFUSAL}], + "is_error": True, + } if not HAS_REMOTE and _local_submit_count >= MAX_SUBMITS: - return {"content": [{"type": "text", "text": - f"submit refused: hit MAX_SUBMITS={MAX_SUBMITS} total-jobs cap for this run"}], - "is_error": True} + return { + "content": [ + { + "type": "text", + "text": f"submit refused: hit MAX_SUBMITS={MAX_SUBMITS} total-jobs cap for this run", + } + ], + "is_error": True, + } if _local_pending_count() >= LOCAL_MAX_CONCURRENT: - return {"content": [{"type": "text", "text": - f"submit refused: a local job is already running (max {LOCAL_MAX_CONCURRENT} " - f"at a time -- it uses the whole node). Collect it with get_local_completed " - f"before submitting more."}], "is_error": True} + return { + "content": [ + { + "type": "text", + "text": f"submit refused: a local job is already running (max {LOCAL_MAX_CONCURRENT} " + f"at a time -- it uses the whole node). Collect it with get_local_completed " + f"before submitting more.", + } + ], + "is_error": True, + } fut = get_local_executor().submit(task.local_fn, args) job_id = next(_local_counter) _local_submit_count += 1 @@ -552,8 +690,16 @@ async def get_local_completed(args): _append_jobs_log({"event": "local_completed", "job_id": job_id}) del _local_jobs[job_id] pending = [{"job_id": jid, "args": i["args"]} for jid, i in _local_jobs.items()] - return {"content": [{"type": "text", "text": - json.dumps({"completed": completed, "pending": pending}, indent=2, default=str)}]} + return { + "content": [ + { + "type": "text", + "text": json.dumps( + {"completed": completed, "pending": pending}, indent=2, default=str + ), + } + ] + } RELEASE_CLAIM_DESC = ( @@ -599,7 +745,7 @@ async def notify(args): "anything they do not support comes back to you next turn. One call per cycle." ) -_cycle_mark = None # set by cycle_done, read and cleared by the runner +_cycle_mark = None # set by cycle_done, read and cleared by the runner def cycle_done_pending(): @@ -614,8 +760,11 @@ def cycle_done_pending(): async def cycle_done(args): global _cycle_mark _cycle_mark = (args.get("conclusion") or "").strip() or "(no conclusion given)" - return {"content": [{"type": "text", "text": - "cycle recorded; its write-up will be reviewed"}]} + return { + "content": [ + {"type": "text", "text": "cycle recorded; its write-up will be reviewed"} + ] + } GOAL_MET_DESC = ( @@ -626,7 +775,7 @@ async def cycle_done(args): "taking new work, finishes what is in flight, and gives you a turn to write up." ) -_goal_met = None # what the agent said settles the goal, or None +_goal_met = None # what the agent said settles the goal, or None def goal_is_met(): @@ -639,9 +788,15 @@ async def goal_met(args): global _goal_met _goal_met = (args.get("reason") or "").strip() or "(no reason given)" request_stop() - return {"content": [{"type": "text", "text": - "goal recorded; this run is winding down -- collect what is in " - "flight and write the cycle up"}]} + return { + "content": [ + { + "type": "text", + "text": "goal recorded; this run is winding down -- collect what is in " + "flight and write the cycle up", + } + ] + } CHECK_BACKEND_DESC = ( @@ -656,12 +811,17 @@ async def goal_met(args): @tool("check_backend", CHECK_BACKEND_DESC, {}) async def check_backend(args): import globus_compute_sdk as _gc + info = {"endpoint_id": ENDPOINT_ID} try: c = _gc.Client() except Exception as e: - return {"content": [{"type": "text", "text": json.dumps({"error": f"client init: {e}"})}], - "is_error": True} + return { + "content": [ + {"type": "text", "text": json.dumps({"error": f"client init: {e}"})} + ], + "is_error": True, + } try: info["endpoint_status"] = c.get_endpoint_status(ENDPOINT_ID) except Exception as e: @@ -675,10 +835,19 @@ async def check_backend(args): st = c.get_task(tid) except Exception as e: st = {"error": str(e)} - tasks.append({"job_id": jid, "key": j["key"], "task_id": tid, - "future_done": j["future"].done(), "task_status": st}) + tasks.append( + { + "job_id": jid, + "key": j["key"], + "task_id": tid, + "future_done": j["future"].done(), + "task_status": st, + } + ) info["in_flight_tasks"] = tasks - return {"content": [{"type": "text", "text": json.dumps(info, indent=2, default=str)}]} + return { + "content": [{"type": "text", "text": json.dumps(info, indent=2, default=str)}] + } def create_server(): diff --git a/framework/transfer.py b/framework/transfer.py index 103bb8e..991e2e1 100644 --- a/framework/transfer.py +++ b/framework/transfer.py @@ -76,8 +76,9 @@ def configure(user_cfg, workspace_dir, campaign_dir=None, sys_cfg=None): # Writes to the compute system are confined to one subtree. Defaulting it to # work_dir means the tool cannot scribble outside the campaign's own directory # unless someone widens it deliberately. - "remote_write_root": str(g.get("remote_write_root") - or user_cfg.get("work_dir", "")).rstrip("/"), + "remote_write_root": str( + g.get("remote_write_root") or user_cfg.get("work_dir", "") + ).rstrip("/"), # Reads are unbounded by default: the usual job is fetching a log from a path # the campaign did not choose. Set it to confine reads to one subtree. "remote_read_root": str(g.get("remote_read_root") or "").rstrip("/"), @@ -88,9 +89,15 @@ def configure(user_cfg, workspace_dir, campaign_dir=None, sys_cfg=None): # One or several prefixes: the same filesystem is often reachable by a short # mount path and a long one (/flare and /lus/flare/projects), and a work_dir may # be written either way. The first that matches is stripped. - "collection_root": [str(r).rstrip("/") for r in _as_list( - g.get("collection_root") if g.get("collection_root") is not None - else sys_cfg.get("globus_collection_root", "")) if str(r).strip()], + "collection_root": [ + str(r).rstrip("/") + for r in _as_list( + g.get("collection_root") + if g.get("collection_root") is not None + else sys_cfg.get("globus_collection_root", "") + ) + if str(r).strip() + ], "workspace_dir": workspace_dir, # The agent may send from, and fetch into, either the campaign's own directory # (task.py and the scripts a job runs) or its workspace (results and artefacts). @@ -100,15 +107,16 @@ def configure(user_cfg, workspace_dir, campaign_dir=None, sys_cfg=None): } -CFG = None # set by tools.py at import; None disables the tools +CFG = None # set by tools.py at import; None disables the tools def _globus(*argv, timeout=_WAIT_SECONDS): """Run the globus CLI. It already holds the user's login, so this module never implements an auth flow of its own.""" try: - p = subprocess.run(["globus", *argv], capture_output=True, text=True, - timeout=timeout) + p = subprocess.run( + ["globus", *argv], capture_output=True, text=True, timeout=timeout + ) except FileNotFoundError: return 127, "", "the 'globus' CLI is not on PATH (pip install globus-cli)" except subprocess.TimeoutExpired: @@ -164,11 +172,11 @@ def _local_dest(rel, roots, must_exist=False): def _cpath(posix_path): """POSIX path -> the path this collection understands.""" p = os.path.normpath(posix_path) - for root in (CFG.get("collection_root") or []): + for root in CFG.get("collection_root") or []: if p == root: return "/" if p.startswith(root + "/"): - return p[len(root):] + return p[len(root) :] # Already collection-relative, or outside the collection: pass it through and let # Globus reject it, rather than silently rewriting into the wrong place. return posix_path @@ -192,8 +200,15 @@ def _under(path, root): def _wait(task_id): - rc, out, err = _globus("task", "wait", task_id, "--timeout", str(_WAIT_SECONDS), - "--polling-interval", "2") + rc, out, err = _globus( + "task", + "wait", + task_id, + "--timeout", + str(_WAIT_SECONDS), + "--polling-interval", + "2", + ) if rc != 0: return False, f"transfer {task_id} did not complete: {err or out}" return True, task_id @@ -263,24 +278,38 @@ async def transfer(args): # and a staging hop only adds a second place for permissions to be wrong. if CFG["remote_read_root"] and not _under(path, CFG["remote_read_root"]): return _err(f"refusing to read outside {CFG['remote_read_root']}: {path}") - rel = local_path or os.path.join("scratch", "transfers", - os.path.basename(path.rstrip("/"))) + rel = local_path or os.path.join( + "scratch", "transfers", os.path.basename(path.rstrip("/")) + ) dest = _local_dest(rel, CFG["local_roots"]) if dest is None: - return _err("refusing to write outside the campaign and workspace " - f"directories: {local_path}") + return _err( + "refusing to write outside the campaign and workspace " + f"directories: {local_path}" + ) recursive = _remote_is_dir(rc_coll, _cpath(path)) os.makedirs(dest if recursive else os.path.dirname(dest), exist_ok=True) - cmd = ["transfer", f"{rc_coll}:{_cpath(path)}", f"{lc_coll}:{dest}", - "--label", "agentlab-get", "--notify", "off", "--format", "json"] + cmd = [ + "transfer", + f"{rc_coll}:{_cpath(path)}", + f"{lc_coll}:{dest}", + "--label", + "agentlab-get", + "--notify", + "off", + "--format", + "json", + ] if recursive: cmd.insert(1, "--recursive") rc, out, err = _globus(*cmd, timeout=120) if rc != 0: - return _err(f"transfer submit failed: {err or out}\n" - "If the local collection is not connected, start Globus Connect " - "Personal. If the destination is refused, the workspace path is " - "not writable in ~/.globusonline/lta/config-paths.") + return _err( + f"transfer submit failed: {err or out}\n" + "If the local collection is not connected, start Globus Connect " + "Personal. If the destination is refused, the workspace path is " + "not writable in ~/.globusonline/lta/config-paths." + ) try: task_id = json.loads(out)["task_id"] except Exception: @@ -300,10 +329,11 @@ async def transfer(args): if size > _HEAD_BYTES + _TAIL_BYTES: f.seek(-_TAIL_BYTES, os.SEEK_END) tail = f.read() - body = (head.decode("utf-8", "replace") - + f"\n\n... [{size - _HEAD_BYTES - _TAIL_BYTES} bytes omitted;" - f" full file at {dest}] ...\n\n" - + tail.decode("utf-8", "replace")) + body = ( + head.decode("utf-8", "replace") + + f"\n\n... [{size - _HEAD_BYTES - _TAIL_BYTES} bytes omitted;" + f" full file at {dest}] ...\n\n" + tail.decode("utf-8", "replace") + ) else: body = (head + f.read()).decode("utf-8", "replace") return _ok(f"{path}\n -> {dest} ({size} bytes)\n\n{body}") @@ -312,17 +342,30 @@ async def transfer(args): if not path or not local_path: return _err("put needs both `local_path` and `path`.") if not _under(path, CFG["remote_write_root"]): - return _err(f"refusing to write outside {CFG['remote_write_root']}: {path}\n" - "Widen remote_write_root in your user file if that is intended.") + return _err( + f"refusing to write outside {CFG['remote_write_root']}: {path}\n" + "Widen remote_write_root in your user file if that is intended." + ) src = _local_dest(local_path, CFG["local_roots"], must_exist=True) if src is None: - return _err("refusing to send from outside the campaign and workspace " - f"directories: {local_path}") + return _err( + "refusing to send from outside the campaign and workspace " + f"directories: {local_path}" + ) if not os.path.exists(src): return _err(f"no such local path: {src}") recursive = os.path.isdir(src) - cmd = ["transfer", f"{lc_coll}:{src}", f"{rc_coll}:{_cpath(path)}", - "--label", "agentlab-put", "--notify", "off", "--format", "json"] + cmd = [ + "transfer", + f"{lc_coll}:{src}", + f"{rc_coll}:{_cpath(path)}", + "--label", + "agentlab-put", + "--notify", + "off", + "--format", + "json", + ] if recursive: cmd.insert(1, "--recursive") rc, out, err = _globus(*cmd, timeout=120) diff --git a/framework/watch.py b/framework/watch.py index e581c61..289641c 100644 --- a/framework/watch.py +++ b/framework/watch.py @@ -21,9 +21,9 @@ import json import os import re +import sys import threading import time -import sys import urllib.parse import webbrowser from datetime import datetime @@ -32,11 +32,17 @@ LAB_DIR = os.path.abspath(os.environ.get("LAB_DIR", os.path.join(SCRIPT_DIR, ".."))) # Files worth opening while a run is in flight. Anything else in the workspace is # listed but not offered as a tab: run directories, caches, figures. -READABLE = ("LOGBOOK.md", "JOURNAL.md", "REVIEWS.md", "results.jsonl", - "ANNOUNCEMENTS.md", "jobs.jsonl") -TAIL_BYTES = 400_000 # of a file view; the log is followed from an offset instead - -try: # in requirements.txt; without it records are plain text +READABLE = ( + "LOGBOOK.md", + "JOURNAL.md", + "REVIEWS.md", + "results.jsonl", + "ANNOUNCEMENTS.md", + "jobs.jsonl", +) +TAIL_BYTES = 400_000 # of a file view; the log is followed from an offset instead + +try: # in requirements.txt; without it records are plain text import markdown as _markdown except Exception: _markdown = None @@ -55,8 +61,11 @@ def newest_run(campaign): metas = glob.glob(os.path.join(workspace(campaign), "runs", "*", "meta.json")) if not metas: return None - live = [m for m in metas - if os.path.isfile(os.path.join(os.path.dirname(m), "heartbeat"))] + live = [ + m + for m in metas + if os.path.isfile(os.path.join(os.path.dirname(m), "heartbeat")) + ] return os.path.dirname(max(live or metas, key=os.path.getmtime)) @@ -113,8 +122,9 @@ def status(campaign): age = int(time.time() - float(f.read().strip())) except Exception: age = None - submits_total, submits_run, done_run = _submits(os.path.join(ws, "jobs.jsonl"), - meta.get("run_id")) + submits_total, submits_run, done_run = _submits( + os.path.join(ws, "jobs.jsonl"), meta.get("run_id") + ) # How long the run took, not how long ago it began: once it has ended, the clock # stops where it stopped. phase, phase_age = None, None @@ -133,20 +143,31 @@ def status(campaign): except Exception: elapsed = None return { - "run": meta.get("run_id"), "handle": meta.get("handle"), - "campaign": campaign, "status": meta.get("status"), - "stop_reason": meta.get("stop_reason"), "model": meta.get("model"), - "critic": meta.get("critic"), "host": meta.get("host"), + "run": meta.get("run_id"), + "handle": meta.get("handle"), + "campaign": campaign, + "status": meta.get("status"), + "stop_reason": meta.get("stop_reason"), + "model": meta.get("model"), + "critic": meta.get("critic"), + "host": meta.get("host"), "context_tokens": meta.get("context_tokens"), "context_window": meta.get("context_window"), "context_pct": meta.get("context_pct"), - "started_at": started, "ended_at": meta.get("ended_at"), - "elapsed_s": elapsed, "heartbeat_age_s": age, - "phase": phase, "phase_age_s": phase_age, + "started_at": started, + "ended_at": meta.get("ended_at"), + "elapsed_s": elapsed, + "heartbeat_age_s": age, + "phase": phase, + "phase_age_s": phase_age, "results": _count_lines(os.path.join(ws, "results.jsonl")), - "jobs": submits_total, "jobs_run": submits_run, "jobs_done": done_run, - "reviews": _count_lines(os.path.join(ws, "REVIEWS.md")) and - open(os.path.join(ws, "REVIEWS.md"), errors="replace").read().count("\n## "), + "jobs": submits_total, + "jobs_run": submits_run, + "jobs_done": done_run, + "reviews": _count_lines(os.path.join(ws, "REVIEWS.md")) + and open(os.path.join(ws, "REVIEWS.md"), errors="replace") + .read() + .count("\n## "), "max_submits": meta.get("max_submits"), "max_runtime_s": meta.get("max_runtime_s"), "max_rounds": meta.get("max_rounds"), @@ -361,31 +382,40 @@ def _render(text): html_out = _markdown.markdown(text, extensions=["tables", "fenced_code"]) except Exception: return "
" + html.escape(text) + "
" + # Figures are referenced relative to the workspace, which only this server can read. def _img(m): src = urllib.parse.quote(m.group("src")) alt = m.group(0) alt = re.search(r'alt="([^"]*)"', alt) - return (f'' - f'{alt.group(1) if alt else ') + return ( + f'' + f'{alt.group(1) if alt else ' + ) - html_out = re.sub(r']*?src="(?!https?:|/)(?P[^"]+)"[^>]*/?>', _img, html_out) + html_out = re.sub( + r']*?src="(?!https?:|/)(?P[^"]+)"[^>]*/?>', _img, html_out + ) # A record may link a figure rather than embed it. The link is relative to the # workspace, which only this server can read, so point it at the same route. return re.sub( r' TAIL_BYTES: f.seek(size - TAIL_BYTES) - return f"[showing the last {TAIL_BYTES // 1000} KB of {size // 1000} KB]\n\n" + f.read() + return ( + f"[showing the last {TAIL_BYTES // 1000} KB of {size // 1000} KB]\n\n" + + f.read() + ) return f.read() except OSError as e: return f"cannot read {name}: {e}" @@ -510,8 +564,10 @@ def main(): else: sys.exit(f"no free port between {port} and {port + 19}") url = f"http://127.0.0.1:{port}/" - print(f"watching {campaign} at {url} (Ctrl-C to stop; the run is unaffected)", - flush=True) + print( + f"watching {campaign} at {url} (Ctrl-C to stop; the run is unaffected)", + flush=True, + ) if "--no-open" not in sys.argv: webbrowser.open(url) # An open page polls every second or so, so a long silence means nobody is looking. diff --git a/systems/endpoints/README.md b/systems/endpoints/README.md index ce1766a..73aed83 100644 --- a/systems/endpoints/README.md +++ b/systems/endpoints/README.md @@ -71,13 +71,23 @@ so confirm it independently first: ```python from globus_compute_sdk import Executor + + def hello(): import socket + return f"ran on {socket.gethostname()}" -ex = Executor(endpoint_id="your-uuid", - user_endpoint_config={"account": "MYPROJECT", "queue": "debug", - "num_nodes": 1, "walltime": "00:10:00"}) + +ex = Executor( + endpoint_id="your-uuid", + user_endpoint_config={ + "account": "MYPROJECT", + "queue": "debug", + "num_nodes": 1, + "walltime": "00:10:00", + }, +) print(ex.submit(hello).result()) ```