"""Worker processes on one store, faults, and the run's outcome. The test drives noeta only from outside: worker processes started as ``noeta worker --lease``, a :class:`~noeta.core.client.Client` on the same store to create tasks, send messages and read logs, and signals. """ from __future__ import annotations import asyncio import contextlib import json import os import random import signal import subprocess import sys import time from collections import Counter from collections.abc import Iterable from dataclasses import dataclass, field from pathlib import Path from typing import Any from noeta.core import registry from noeta.core.client import Client from noeta.core.events import InputTaken, Recorded, TimerFired, ToolFinished, ToolStarted from noeta.core.store import TaskRow from noeta.core.tool import MAX_ATTEMPTS from tests.cluster.invariants import ( EpochWatch, check_journal, check_log, child_ends, job_notices, jobs_of, jobs_running, jobs_untold, last_reply, taken_texts, timers_taken, ) #: The ``noeta`` command of the environment the test runs in. NOETA = str(Path(sys.executable).parent / "noeta") MODEL = "cluster-scripted" ROOT_AGENT = "cluster" CHILD_AGENT = "cluster-child" SPILL_AGENT = "cluster-spill" INTERRUPTED = "outcome is unknown" #: Workers write fold snapshots this often, so claims under faults start #: from them (the runs' logs are a few dozen events). SNAPSHOT_EVERY = 5 ROOT_FILE = """--- name: cluster description: The cluster test's root agent. model: inherit plugins: - cluster_crash: arm: {arm} - shell - cluster_effect: journal: {journal} - cluster_schedule: fire_after: {fire_after} - subtasks: agents: [cluster-child] inherit_config: [cwd] --- The cluster test's scripted root agent. """ CHILD_FILE = """--- name: cluster-child description: The cluster test's sub-agent; calls Effect once. model: inherit plugins: - shell - cluster_effect: journal: {journal} - subtasks: agents: [] --- The cluster test's scripted sub-agent. """ SPILL_FILE = """--- name: cluster-spill description: The cluster test's agent whose long output is saved and read back. model: inherit plugins: - shell - fs - permissions: mode: default --- The cluster test's scripted spill agent. """ def seed_from_env() -> int: raw = os.environ.get("NOETA_CLUSTER_SEED", "") return int(raw) if raw else random.SystemRandom().randrange(1, 2**31) @dataclass class Spec: """One root task's script (see ``noeta_cluster_scripted``).""" id: str effects: int = 1 bg: float = 0.0 agent: bool = False wake: bool = False pings: int = 0 lookup: float = 0.0 child_bg: float = 0.0 cron: int = 0 agent_first: bool = False def prompt(self, journal: str) -> str: return "CLUSTER-ROOT " + json.dumps({ "id": self.id, "journal": journal, "effects": self.effects, "bg": self.bg, "agent": self.agent, "wake": self.wake, "lookup": self.lookup, "child_bg": self.child_bg, "cron": self.cron, "agent_first": self.agent_first}) @dataclass class Faults: kills: int = 0 stops: int = 0 targeted: int = 0 log: list[str] = field(default_factory=list) class Cluster: """Worker processes serving one store from a shared directory.""" def __init__(self, base: Path, url: str, *, workers: int = 3, lease: float = 1.0, fire_after: float = 2.0, restart_limit: int | None = None) -> None: self.base = base #: ``noeta worker --restart-limit``; ``None`` keeps the default. self.restart_limit = restart_limit self.url = url self.lease = lease self.home = base / "home" self.work = base / "work" self.logs = base / "logs" self.journal = str(base / "journal") #: Write an ``External`` name here to have the root's worker die #: right after taking one (``cluster_crash``). self.arm = str(base / "crash-arm") for d in (self.home / "agents", self.work, self.logs): d.mkdir(parents=True, exist_ok=True) (self.home / "agents" / "cluster.md").write_text( ROOT_FILE.format(journal=self.journal, fire_after=fire_after, arm=self.arm)) (self.home / "agents" / "cluster-child.md").write_text( CHILD_FILE.format(journal=self.journal)) (self.home / "agents" / "cluster-spill.md").write_text(SPILL_FILE) models = self.home / "models.json" models.write_text(json.dumps({"models": [ {"id": MODEL, "context_window": 200_000, "max_output_tokens": 8_000}]})) self.env = {**os.environ, "NOETA_HOME": str(self.home), "NOETA_MODEL": MODEL, "NOETA_PROVIDER": "cluster_scripted", "NOETA_MODELS_FILE": str(models), "NOETA_STORE": url} for key in ("NOETA_BASE_URL", "NOETA_API_KEY", "NOETA_EFFORT"): self.env.pop(key, None) self.size = workers self.procs: list[subprocess.Popen[bytes]] = [] self.stopped: set[int] = set() self.spawned = 0 self.store: Any = None self.client: Client # ------------------------------------------------------------ processes def _spawn(self) -> subprocess.Popen[bytes]: self.spawned += 1 log = open(self.logs / f"worker-{self.spawned}.log", "wb") with log: limit = ([] if self.restart_limit is None else ["--restart-limit", str(self.restart_limit)]) return subprocess.Popen([NOETA, "worker", "--lease", str(self.lease), "--snapshot-every", str(SNAPSHOT_EVERY), *limit], cwd=self.work, env=self.env, stdout=log, stderr=log) async def start(self) -> None: self.store = registry.store(self.url) self.client = Client(self.store, poll_interval=0.2) await self.store.tasks(limit=1) # creates the schema before the workers race to self.procs = [self._spawn() for _ in range(self.size)] def kill(self, i: int, *, respawn: bool = True) -> None: p = self.procs[i] p.send_signal(signal.SIGKILL) p.wait() self.stopped.discard(p.pid) if respawn: self.procs[i] = self._spawn() def kill_pid(self, pid: int, *, respawn: bool = True) -> bool: for i, p in enumerate(self.procs): if p.pid == pid and p.poll() is None: self.kill(i, respawn=respawn) return True return False def index_of(self, pid: int) -> int | None: return next((i for i, p in enumerate(self.procs) if p.pid == pid), None) def pause(self, i: int) -> int: pid = self.procs[i].pid os.kill(pid, signal.SIGSTOP) self.stopped.add(pid) return pid def resume(self, pid: int) -> None: if pid in self.stopped: self.stopped.discard(pid) with contextlib.suppress(ProcessLookupError): os.kill(pid, signal.SIGCONT) async def close(self) -> None: self.stop_workers() if self.store is not None: await self.store.close() self.store = None def stop_workers(self) -> None: for pid in list(self.stopped): self.resume(pid) for p in self.procs: if p.poll() is None: p.send_signal(signal.SIGTERM) deadline = time.monotonic() + 10 for p in self.procs: try: p.wait(max(0.1, deadline - time.monotonic())) except subprocess.TimeoutExpired: p.kill() p.wait() # ------------------------------------------------------------ tasks async def create(self, spec: Spec) -> str: return await self.client.create(ROOT_AGENT, spec.prompt(self.journal), config={"cwd": str(self.work)}, dedupe_key=f"root-{spec.id}") async def ping(self, task_id: str, spec: Spec, j: int, *, rng: random.Random) -> None: when = "now" if rng.random() < 0.3 else "next_turn" text = f"CLUSTER-PING {spec.id} {j}" await self.client.send(task_id, text, when=when, dedupe_key=f"ping-{spec.id}-{j}") if rng.random() < 0.3: # a retried send: the dedupe key makes it one message await self.client.send(task_id, text, when=when, dedupe_key=f"ping-{spec.id}-{j}") async def rows(self) -> list[TaskRow]: return await self.client.tasks(limit=10_000) async def tree(self, root: str) -> list[TaskRow]: return [r for r in await self.rows() if r.root_id == root] def journal_lines(self) -> list[str]: path = Path(self.journal) return path.read_text().splitlines() if path.exists() else [] async def running_job_pids(self) -> list[tuple[str, int]]: """``(root id, worker pid)`` of every background job some task's log shows running, from the claim its task keeps meanwhile (the row's ``owner``, ``host:pid:random``).""" out = [] for row in await self.rows(): if row.owner is None: continue records = await self.client.events(row.task_id) if any([k for _s, k in marks] == ["started"] for marks in jobs_of(records).values()): out.append((row.root_id, int(str(row.owner).split(":")[1]))) return out # ------------------------------------------------------------------ outcome @dataclass class Outcome: """What one run left: rows, logs, journal; and the checks over them.""" seed: int specs: dict[str, Spec] roots: dict[str, str] pings: dict[str, int] rows: list[TaskRow] logs: dict[str, list[Recorded]] journal: list[str] watch: EpochWatch faults: Faults worker_logs: Path #: Tasks the workers left a fold snapshot of. snapshots: int = 0 def settled(self, row: TaskRow) -> bool: return row.status in ("waiting", "ended") and row.inbox_read_seq == row.inbox_seq def root_done(self, sid: str) -> list[str]: """Why root ``sid`` has not finished its script (empty: it has).""" tid = self.roots[sid] row = next(r for r in self.rows if r.task_id == tid) records = self.logs[tid] out = [] if not self.settled(row): out.append(f"root {sid} is {row.status} (inbox {row.inbox_read_seq}/{row.inbox_seq})") reply = last_reply(records) if reply != f"DONE {sid}": out.append(f"root {sid} last replied {reply!r}, not 'DONE {sid}'") texts = taken_texts(records) for j in range(self.pings.get(sid, 0)): n = sum(t == f"CLUSTER-PING {sid} {j}" for t in texts) if n != 1: out.append(f"root {sid}: ping {j} taken {n} times") for child in (r for r in self.rows if r.parent_id == tid): if child.status != "ended": out.append(f"child {child.task_id} of {sid} is {child.status}") elif child_ends(records)[child.task_id] != 1: out.append(f"child {child.task_id} of {sid}: its end taken " f"{child_ends(records)[child.task_id]} times") tree = [r for r in self.rows if r.root_id == tid] told = job_notices(records, [j for r in tree for j in jobs_of(self.logs[r.task_id])]) for row in tree: if row.task_id != tid and row.status == "ended": # Its jobs died with it, untold; the root never hears of them. out += [f"{row.task_id}: job {job} of an ended sub-agent: the root was told " f"{told[job]} times" for job in jobs_of(self.logs[row.task_id]) if told[job]] continue out += jobs_running(row.task_id, self.logs[row.task_id]) out += jobs_untold(row.task_id, self.logs[row.task_id]) spec = self.specs[sid] if spec.wake and "schedule.wakeup" not in timers_taken(records) \ and not self.interrupted(tid, f"wake-{sid}"): out.append(f"root {sid}: its wake-up never fired") if spec.cron: out += self.cron_problems(sid) return out def cron_problems(self, sid: str) -> list[str]: """The cron job fired at least ``cron`` times, and never once its deletion is in the log (unless creating it had an unknown outcome, which leaves no job).""" tid, spec = self.roots[sid], self.specs[sid] records = self.logs[tid] if self.interrupted(tid, f"cron-{sid}"): return [] deleted = next((r.seq for r in records if isinstance(r.event, ToolFinished) and r.event.result.call_id.startswith(f"cdel-{sid}-") and not r.event.result.is_error), None) fires = [r.seq for r in records if isinstance(r.event, InputTaken) for i in r.event.items if isinstance(i.body, TimerFired) and i.body.name.startswith("schedule.job:")] out = [] if deleted is None: out.append(f"root {sid}: its cron job was never deleted") elif any(seq > deleted for seq in fires): out.append(f"root {sid}: its cron job fired after it was deleted (seq {deleted})") if len(fires) < spec.cron: out.append(f"root {sid}: its cron job fired {len(fires)} times, not {spec.cron}") return out def interrupted(self, tid: str, call_id: str) -> bool: return any(isinstance(r.event, ToolFinished) and r.event.result.call_id == call_id and INTERRUPTED in r.event.result.text for r in self.logs[tid]) def finished_ok(self) -> dict[str, tuple[str, str]]: """Calls the logs show finished with a real result: token -> (task, call).""" out = {} for tid, records in self.logs.items(): for r in records: ev = r.event if isinstance(ev, ToolFinished) and not ev.result.is_error \ and ev.result.text.startswith("did "): out[ev.result.text[4:]] = (tid, ev.result.call_id) return out def log_violations(self) -> list[str]: out = [] for row in self.rows: out += check_log(row, self.logs[row.task_id], settled=self.settled(row)) return out def effect_violations(self) -> list[str]: """Each token at most once; exactly once when its call finished.""" out = check_journal(self.journal) counts = Counter(line.split(" ", 1)[0] for line in self.journal if line.strip()) for token, (tid, call) in sorted(self.finished_ok().items()): if counts[token] != 1: out.append(f"{tid} call {call}: finished as done but journalled " f"{counts[token]} times") return out def lookups(self) -> dict[str, tuple[int, str]]: """Each ``Lookup`` call (task:call) -> (its starts, its result).""" out = {} for tid, records in self.logs.items(): starts: Counter[str] = Counter(r.event.call.id for r in records if isinstance(r.event, ToolStarted) and r.event.call.name == "Lookup") for r in records: ev = r.event if isinstance(ev, ToolFinished) and ev.result.call_id in starts: out[f"{tid}:{ev.result.call_id}"] = (starts[ev.result.call_id], ev.result.text) return out def lookup_violations(self) -> list[str]: """A retry-safe call a crash interrupted ran again and finished with its real result, unless every attempt it had was cut short.""" return [f"{key}: {starts} start(s) ended {text!r}" for key, (starts, text) in sorted(self.lookups().items()) if not text.startswith("found ") and starts < MAX_ATTEMPTS] def interrupted_calls(self) -> int: return sum(isinstance(r.event, ToolFinished) and INTERRUPTED in r.event.result.text for records in self.logs.values() for r in records) def lost_jobs(self) -> int: return sum(1 for records in self.logs.values() for marks in jobs_of(records).values() for _s, k in marks if k == "lost") def describe(self, task_ids: Iterable[str]) -> str: """The seed and the event logs of ``task_ids`` (and their trees).""" lines = [f"NOETA_CLUSTER_SEED={self.seed} faults: kills={self.faults.kills} " f"stops={self.faults.stops} targeted={self.faults.targeted}", f"worker logs: {self.worker_logs}"] wanted = set(task_ids) rows = {r.task_id: r for r in self.rows} for tid in list(wanted): row = rows.get(tid) if row is not None: wanted |= {r.task_id for r in self.rows if r.root_id == row.root_id} for tid in sorted(wanted)[:6]: row = rows.get(tid) lines.append(f"--- task {tid} {row.status if row else '?'} " f"epoch={row.owner_epoch if row else '?'} " f"inbox={row.inbox_read_seq if row else '?'}/" f"{row.inbox_seq if row else '?'}") for r in self.logs.get(tid, []): lines.append(f" {r.seq:3d} {_brief(r)}") return "\n".join(lines) def _brief(r: Recorded) -> str: ev = r.event name = type(ev).__name__ text = repr(ev) if hasattr(ev, "response"): msg = ev.response.message text = msg.text or ", ".join(f"{u.name}#{u.id}" for u in msg.tool_uses) elif hasattr(ev, "result"): text = f"{ev.result.call_id}: {ev.result.text}" elif hasattr(ev, "call"): text = f"{ev.call.name}#{ev.call.id}" text = " ".join(text.split()) return f"{name} {text[:200]}" async def collect(cluster: Cluster, *, seed: int, specs: dict[str, Spec], roots: dict[str, str], pings: dict[str, int], watch: EpochWatch, faults: Faults) -> Outcome: rows = await cluster.rows() logs = {r.task_id: await cluster.client.events(r.task_id) for r in rows} snapshots = [await cluster.store.snapshot(r.task_id) for r in rows] return Outcome(seed=seed, specs=specs, roots=roots, pings=pings, rows=rows, logs=logs, journal=cluster.journal_lines(), watch=watch, faults=faults, worker_logs=cluster.logs, snapshots=sum(s is not None for s in snapshots)) async def watch_rows(cluster: Cluster, watch: EpochWatch, stop: asyncio.Event) -> None: while not stop.is_set(): with contextlib.suppress(Exception): watch.observe(await cluster.rows()) with contextlib.suppress(asyncio.TimeoutError): await asyncio.wait_for(stop.wait(), 0.15) async def wait_until(predicate: Any, timeout: float, interval: float = 0.25) -> bool: deadline = time.monotonic() + timeout while time.monotonic() < deadline: if await predicate(): return True await asyncio.sleep(interval) return False # ------------------------------------------------------------------ faults async def chaos(cluster: Cluster, rng: random.Random, seconds: float, faults: Faults, *, targets: set[str] | None = None) -> None: """For ``seconds``: SIGKILL a random worker (respawned at once), or SIGSTOP one for 1.5-3 leases (the zombie whose claims run out while it still thinks it owns them), at random moments. With ``targets`` (root ids) some faults are aimed: kill the worker running a background job of a target root that is waiting (the first such fault as soon as there is one).""" end = time.monotonic() + seconds paused: dict[int, float] = {} while time.monotonic() < end: await asyncio.sleep(rng.uniform(0.1, 0.6)) now = time.monotonic() for pid, until in list(paused.items()): if now >= until: cluster.resume(pid) paused.pop(pid) roll = rng.random() if targets and (roll < 0.3 or not faults.targeted): if await kill_job_worker(cluster, targets, faults): continue roll = rng.random() if roll < 0.6: i = rng.randrange(len(cluster.procs)) paused.pop(cluster.procs[i].pid, None) cluster.kill(i) faults.kills += 1 faults.log.append(f"{now:.1f} kill worker {i}") elif roll < 0.85 and not paused: i = rng.randrange(len(cluster.procs)) pid = cluster.pause(i) paused[pid] = now + rng.uniform(1.5, 3.0) * cluster.lease faults.stops += 1 faults.log.append(f"{now:.1f} stop worker {i}") for pid in paused: cluster.resume(pid) async def kill_all(cluster: Cluster, rng: random.Random, seconds: float, faults: Faults, *, targets: set[str] | None = None) -> None: """For ``seconds``: every 1.5-3 s SIGKILL every worker at once (each respawned), so no worker lives to finish what another started.""" end = time.monotonic() + seconds while time.monotonic() < end: await asyncio.sleep(rng.uniform(1.5, 3.0)) for i in range(len(cluster.procs)): cluster.kill(i) faults.kills += 1 faults.log.append(f"{time.monotonic():.1f} kill every worker") async def kill_job_worker(cluster: Cluster, targets: set[str], faults: Faults) -> bool: """Kill the worker running a background job of a waiting target root.""" for root, pid in await cluster.running_job_pids(): if root not in targets: continue row = await cluster.client.get(root) if row.status == "waiting" and cluster.kill_pid(pid): faults.targeted += 1 faults.kills += 1 faults.log.append(f"{time.monotonic():.1f} kill job worker {pid} of {root}") return True return False async def send_pings(cluster: Cluster, rng: random.Random, seconds: float, plan: list[tuple[float, str, Spec, int]]) -> None: """Send each planned ping ``(at, task id, spec, j)`` at its time.""" start = time.monotonic() for at, tid, spec, j in sorted(plan, key=lambda p: p[0]): delay = start + min(at, seconds) - time.monotonic() if delay > 0: await asyncio.sleep(delay) await cluster.ping(tid, spec, j, rng=rng) async def root_finished(cluster: Cluster, tid: str, spec: Spec) -> bool: """Has root ``tid`` finished its script (as far as rows and its log tell; :meth:`Outcome.root_done` gives the reasons when not)?""" row = await cluster.client.get(tid) if row.status != "waiting" or row.inbox_read_seq != row.inbox_seq: return False records = await cluster.client.events(tid) if last_reply(records) != f"DONE {spec.id}": return False texts = taken_texts(records) if any(f"CLUSTER-PING {spec.id} {j}" not in texts for j in range(spec.pings)): return False tree = await cluster.tree(tid) return all(r.status == "ended" for r in tree if r.parent_id is not None)