From 34f997884b747e54c7fcd4a5f18bf5fd7863125e Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Sat, 7 Nov 2026 20:10:40 +0800 Subject: [PATCH] bot: Telegram and Lark training-ops bot as a read-only dashboard client python -m yeto.bot polls the dashboard's /api/view over HTTP (never ssh, never the training processes), answers chat commands and pushes tiered alerts. Configuration is environment variables only. - Channels behind one interface (send / poll_commands / answer): Telegram Bot API with getUpdates long polling, the Lark self-built app bot (tenant_access_token cached and refreshed before expiry, im/v1 messages as interactive cards, im/v1/images for charts, commands and card.action.trigger over the lark-oapi long connection when installed) and the Lark group webhook (signed, send-only). Each channel is enabled when its configuration is complete; Markdown/card sends fall back to plain text; long messages are truncated. - Alerts: P0 (red, mentions the owners, re-notified every 30 min until acked: stale island, loss NaN/Inf, HBM peak at the refs limit, syncer outer step not advancing past median round + 2 x grace, OOM / checkpoint-failure / disk hooks), P1 (orange, 1 h cooldown per subject: loss > mean + 3 sd, grad_norm, missed rounds, HBM warn, host memory, syncer RSS, rejected / timed-out pushes, eval worse than the previous point) and P2 (grey: checkpoints, new eval point, joins, finished islands). Transitions only, one merged card per poll with Ack / Details / Open dashboard buttons; /ack, /mute, /unmute and quiet hours (P0 only). - Progress card every 2 h and a daily digest with loss / eval charts (matplotlib when installed, text sparklines otherwise). - Commands: /status /island /eval /merge /res /ckpt /alerts /ack /details /mute /unmute /help /start; only the registered report chat may command or press buttons. Users are authorized by numeric id (Telegram from.id, Lark open_id), never by display name. Registering the report chat with /start, and anything that changes alerting (/ack, /mute, /unmute, decision presses), needs a configured allowlist; read-only commands do not. Commands given through the local --command CLI are trusted. - Decision-request cards record the pressed option in an append-only JSONL audit log and execute nothing. Text that is not a command goes to an optional natural-language handler with read-only tools. - Alerts held back by mute, quiet hours or a cooldown are delivered when the hold ends if their condition is still active. - View rows are read by the column names the dashboard sends; eval references are every vs_ comparison; the HBM warning line derives from the HBM limit. - YETO_BOT_DRY_RUN=1 prints every request body instead of calling out; bot and webhook tokens are redacted from it and from HTTP errors. Co-Authored-By: Claude Opus 5.5 --- tests/bot_helpers.py | 137 +++++++++ tests/test_bot_app.py | 397 ++++++++++++++++++++++++ tests/test_bot_channels.py | 449 +++++++++++++++++++++++++++ tests/test_bot_monitor.py | 410 +++++++++++++++++++++++++ yeto/bot/__init__.py | 7 + yeto/bot/__main__.py | 60 ++++ yeto/bot/app.py | 419 +++++++++++++++++++++++++ yeto/bot/channels.py | 482 +++++++++++++++++++++++++++++ yeto/bot/charts.py | 99 ++++++ yeto/bot/config.py | 200 ++++++++++++ yeto/bot/decisions.py | 103 +++++++ yeto/bot/http.py | 88 ++++++ yeto/bot/lark_ws.py | 137 +++++++++ yeto/bot/monitor.py | 291 ++++++++++++++++++ yeto/bot/render.py | 607 +++++++++++++++++++++++++++++++++++++ yeto/bot/rules.py | 286 +++++++++++++++++ 16 files changed, 4172 insertions(+) create mode 100644 tests/bot_helpers.py create mode 100644 tests/test_bot_app.py create mode 100644 tests/test_bot_channels.py create mode 100644 tests/test_bot_monitor.py create mode 100644 yeto/bot/__init__.py create mode 100644 yeto/bot/__main__.py create mode 100644 yeto/bot/app.py create mode 100644 yeto/bot/channels.py create mode 100644 yeto/bot/charts.py create mode 100644 yeto/bot/config.py create mode 100644 yeto/bot/decisions.py create mode 100644 yeto/bot/http.py create mode 100644 yeto/bot/lark_ws.py create mode 100644 yeto/bot/monitor.py create mode 100644 yeto/bot/render.py create mode 100644 yeto/bot/rules.py diff --git a/tests/bot_helpers.py b/tests/bot_helpers.py new file mode 100644 index 00000000..5b52a9f8 --- /dev/null +++ b/tests/bot_helpers.py @@ -0,0 +1,137 @@ +"""Fake dashboard views (``/api/view`` shape, ``sft`` section) and fake HTTP for the bot tests.""" + +import copy + +from yeto.bot.channels import Channel + +T0 = 1_790_000_000.0 +REFS = {"sft_ref_step_s": 571.0, "sft_hbm_limit_gib": 53.9, "sft_syncer_rss_limit_bytes": 1.4e12, + "sft_slow_factor": 1.1} +# column names the view sends with the positional rows (sft.cols, sft.syncer.merge_cols) +STEP_COLS = ("step", "ts", "loss", "grad_norm", "lr", "step_s", "tps", "tokens", "target_tokens", + "sync_s", "export_s", "apply_s", "hbm_gib", "mfu", "pushes", "applies") +MERGE_COLS = ("ts", "step", "fragment", "merge_s", "gnorm", "round_s", "quorum_s", "grace_s", "wait_s") + + +def row(step, *, loss=0.7, gn=1.0, lr=1e-5, step_s=570.0, hbm=52.0, ts=None): + # in STEP_COLS order + return [step, T0 + 570 * step if ts is None else ts, loss, gn, lr, step_s, 1.75e4, 1e7, 4e6, 0.5, 0.0, 0.0, + hbm, 14.5, 0, 0] + + +def island(iid, steps=10, status="ok", total=100, **row_kw): + rows = [row(s, **row_kw) for s in range(1, steps + 1)] + return {"id": str(iid), "name": f"isl{iid}", "nodes": "a,b", "layout": "TP4/EP8/PP3", "cards": 24, + "dataset": "v2", "gbs": 384, "steps_total": total, "status": status, "last_ts": rows[-1][1], + "age_s": 30.0, "last_push": {"step": 9, "ts": T0 + 5130, "export_s": 3.0, "sync_s": 0.5}, + "last_apply": None, "rows": rows, "host": {}} + + +def view(n_islands=2, *, now=None, outer_steps=8, round_s=600.0, events=(), rss=9e11, eval_rows=(), alerts=(), + rounds=()): + outer = [[T0 + round_s * k, k, k % 2, 2.0, False, 2] for k in range(1, outer_steps + 1)] + merges = [[T0 + round_s * k, k, k % 2, 4.0, 1.3, round_s, 30.0, 10.0, round_s - 4.0] + for k in range(1, outer_steps + 1)] + last = outer[-1][0] if outer else T0 + return { + "overview": {"run": "sft-test", "alerts": list(alerts), "islands": [], "global_status": {"text": "ok"}}, + "sft": { + "now": last + 60 if now is None else now, "refs": dict(REFS), "cols": list(STEP_COLS), + "islands": [island(i) for i in range(n_islands)], + "syncer": {"meta": {}, "total_steps": 960, "outer_step": outer_steps, "eta_ts": None, + "outer": outer, "merges": merges, "merge_cols": list(MERGE_COLS), "rss": [[last, rss]], + "events": [list(e) for e in events], "log_seen": True}, + "eval": {"rows": [dict(r) for r in eval_rows]}, + }, + "rounds_syncer": [dict(r) for r in rounds], + } + + +def round_row(step, missed=(), resend=0): + # subset of yeto.dashboard.reducer.Reducer.rounds() rows + return {"round": step, "missed": [str(m) for m in missed], "resend": resend, "bad": bool(missed or resend)} + + +def eval_row(step, loss=0.70, d_base=-0.05, d_ck=-0.01, ref="ckpt100"): + # every vs_ key is a reference comparison; the bot labels it + return {"step": step, "loss": loss, "ppl": 2.01, + "vs_base": {"dloss": d_base, "mean": d_base, "se": 0.0015, "better": 0.9, "ref_loss": loss - d_base}, + f"vs_{ref}": {"dloss": d_ck, "mean": d_ck, "se": 0.001, "better": 0.6, "ref_loss": loss - d_ck}, + "vs_prev": None} + + +def clone(v): + return copy.deepcopy(v) + + +def relayout(v, cols=None, merge_cols=None): + """Copy of ``v`` whose island / merge rows use another column layout: names + reordered, dropped, or new ones added (filled with None).""" + v = clone(v) + sft = v["sft"] + + def remap(rows, old, new): + return [[r[old.index(n)] if n in old and old.index(n) < len(r) else None for n in new] for r in rows] + + if cols is not None: + for isl in sft["islands"]: + isl["rows"] = remap(isl["rows"], sft["cols"], cols) + sft["cols"] = list(cols) + if merge_cols is not None: + syn = sft["syncer"] + syn["merges"] = remap(syn["merges"], syn["merge_cols"], merge_cols) + syn["merge_cols"] = list(merge_cols) + return v + + +class FakeClock: + def __init__(self, t=T0): + self.t = t + + def __call__(self): + return self.t + + +class FakeHttp: + """Records requests; answers from a list of (predicate substring, response or exception).""" + + def __init__(self, routes=None): + self.calls = [] + self.routes = list(routes or []) + + def __call__(self, method, url, *, body=None, headers=None, timeout=0, files=None): + self.calls.append({"method": method, "url": url, "body": body, "headers": headers or {}, "files": files}) + for needle, resp in self.routes: + if needle in url: + if isinstance(resp, list): # a queue of responses for this route + item = resp.pop(0) if len(resp) > 1 else resp[0] + else: + item = resp + if isinstance(item, BaseException): + raise item + return copy.deepcopy(item(body) if callable(item) else item) + return {"ok": True, "code": 0, "result": []} + + def bodies(self, needle): + return [c["body"] for c in self.calls if needle in c["url"]] + + +class FakeChannel(Channel): + """Channel double that records sent messages and button acknowledgements.""" + + def __init__(self, name="fake", chat="", can_receive=True): + super().__init__(chat) + self.name = name + self.can_receive = can_receive + self.sent = [] + self.answers = [] + + def send(self, msg, chat=None): + target = chat or self.report_chat + if not target: + return False + self.sent.append((target, msg)) + return True + + def answer(self, cmd, text): + self.answers.append(text) diff --git a/tests/test_bot_app.py b/tests/test_bot_app.py new file mode 100644 index 00000000..b6d9e785 --- /dev/null +++ b/tests/test_bot_app.py @@ -0,0 +1,397 @@ +"""yeto.bot: message formatting, command routing, chat / user authorization, schedules, +decision audit log, config.""" + +import json +import time + +import pytest + +from yeto.bot.app import Bot, parse_command, parse_duration +from yeto.bot.channels import Command +from yeto.bot.config import Config, Policy +from yeto.bot.render import (ckpt_message, eval_message, find_island, island_message, merge_health, merge_message, + problem_counts, res_message, status_message, strip_bold) + +from bot_helpers import (MERGE_COLS, STEP_COLS, T0, FakeChannel, FakeClock, clone, eval_row, relayout, round_row, + row, view) + + +# -- formatting --------------------------------------------------------------------- + +def test_progress_card(): + v = view(events=[[T0, "rejected", 1, "log", None]], eval_rows=[eval_row(64), eval_row(96, loss=0.69)], + rounds=[round_row(1), round_row(2, missed=[1]), round_row(3)]) + v["sft"]["syncer"]["eta_ts"] = T0 + 86400 + v["sft"]["islands"][1]["status"] = "warn" + msg = status_message(v, alerts_line="Alerts active P0 0 | P1 1") + text = msg.plain() + assert msg.title == "yeto progress: sft-test" + assert text.splitlines()[1].startswith("Outer 8/960 (0.8%) | ETA ") + assert "isl0 10/100 | loss 0.7000 | 9.5min/step" in text + assert "isl1 10/100 | loss 0.7000 | 9.5min/step | warn" in text + assert ("Merge direct-in-round 2/3 (67%) | median round 10min (quorum 30s, grace 10s, merge 4s)" + " | missed 1 x1") in text + assert "Eval step 96 loss 0.6900 | vs base -0.0500 | vs ckpt100 -0.0100" in text + assert text.endswith("Alerts active P0 0 | P1 1") + assert [c.title for c in msg.charts] == ["loss, last 6 h"] + assert ("Merge", "/merge") in msg.buttons + + +def test_progress_uses_latest_logged_values_and_rl_fallback(): + v = view(1) + v["sft"]["islands"][0]["rows"].append(row(11, loss=None, gn=None, lr=None, hbm=None)) + assert "isl0 11/100 | loss 0.7000" in status_message(v).plain() + rl = {"overview": {"run": "rl", "global_status": {"text": "all healthy"}, + "islands": [{"id": "0", "status": "ok", "round": 3, "last_event_age_s": 5}], "alerts": []}} + assert "island 0 ok | round 3" in status_message(rl).plain() + + +def test_island_lookup_detail_and_charts(): + v = view(3) + sft = v["sft"] + assert find_island(sft, "2")["id"] == "2" and find_island(sft, "ISL1")["id"] == "1" + assert find_island(sft, "island 0")["id"] == "0" + msg = island_message(v, "isl1", hbm_limit=53.9) + text = msg.plain() + assert msg.title == "island 1: isl1" and "layout TP4/EP8/PP3" in text and "HBM 52/53.9 GiB" in text + assert [c.title for c in msg.charts] == ["isl1 loss", "isl1 grad_norm", "isl1 step time s"] + assert island_message(v, "nope").level == "warn" + + +def test_merge_res_ckpt_eval_messages(tmp_path): + v = view(events=[[T0 + 1, "rejected", 1, "log", None], [T0 + 2, "checkpoint", None, "log", 12.5]], + rounds=[round_row(1), round_row(2, missed=[0], resend=1)], eval_rows=[eval_row(64, loss=0.72), + eval_row(96)]) + syn = v["sft"]["syncer"] + syn["outer"][3][4] = True + syn["merges"][5][5] = 3000.0 + m = merge_message(v).plain() + assert "Timed-out rounds (1)" in m and "Slow rounds (> 2x median 10min): 1" in m + assert "rejected 1 | timeout 0 | not carried 0" in m and "round 2: missed 0 | resend 1" in m + assert merge_health(v)["direct"] == 1 + v["sft"]["islands"][0]["host"] = {"node-a": [[T0, 80e9, 1.5e12]]} + r = res_message(v, hbm_limit=53.9, hbm_warn=52.0, rss_warn=1.4e12, host_min=100e9).plain() + assert "isl0: last 52 | max of 20 steps 52 GiB" in r and "isl0 node-a: 80.00 GB of 1.50 TB" in r + assert "Syncer RSS 900.00 GB (warn > 1.40 TB)" in r + bot, ch, client = make_bot(tmp_path, chat="1", v=clone(v)) + assert bot.handle(ch, cmd("/res")).plain().splitlines()[1] == "HBM peak (warn > 52, limit 53.9 GiB)" + del client.v["sft"]["refs"]["sft_hbm_limit_gib"] + assert bot.handle(ch, cmd("/res")).plain().splitlines()[1] == "HBM peak (warn > -, limit - GiB)" + c = ckpt_message(v).plain() + assert "Syncer checkpoint writes (1)" in c and "written in 12s" in c + e = eval_message(v) + assert e.title == "yeto eval: step 96" + assert "vs base: dloss -0.0500 (ref 0.7500) | paired se 0.0015 | better on 90%" in e.plain() + assert "64 | 0.7200 | -0.0500 | -0.0100 | -" in e.plain() + assert [c.title for c in e.charts] == ["held-out eval loss"] + assert problem_counts(v["sft"])["rejected"] == 1 + + +def test_messages_resolve_columns_by_name(tmp_path): + v = view(2, eval_rows=[eval_row(96)], rounds=[round_row(1)]) + v["sft"]["islands"][0]["host"] = {"node-a": [[T0, 80e9, 1.5e12]]} + v["sft"]["syncer"]["merges"][5][5] = 3000.0 # one slow round + shuffled = relayout(v, cols=list(reversed(STEP_COLS)), merge_cols=list(reversed(MERGE_COLS))) + bot, ch, client = make_bot(tmp_path, chat="1", v=v) + + def texts(): + out = {t: bot.handle(ch, cmd(t)).plain() for t in ("/status", "/island 0", "/merge", "/res")} + out["digest"] = bot.digest(client.v).plain() + return out + + want = texts() + assert "HBM 52/53.9 GiB | tokens/s 1.75e+04 | MFU 14.5% | lr 1e-05" in want["/island 0"] + assert "Slow rounds (> 2x median 10min): 1" in want["/merge"] + client.v = shuffled + assert texts() == want + client.v = relayout(v, cols=[c for c in STEP_COLS if c not in ("hbm_gib", "mfu")], merge_cols=["ts", "step"]) + island = bot.handle(ch, cmd("/island 0")).plain() + assert "HBM -/53.9 GiB" in island and "MFU -%" in island and "| loss 0.7000 |" in island + assert "isl0: last - | max of 20 steps - GiB" in bot.handle(ch, cmd("/res")).plain() + assert "Slow rounds" not in bot.handle(ch, cmd("/merge")).plain() + + +def test_eval_references_are_every_vs_key(): + rows = [eval_row(64, ref="warmstart", d_ck=-0.02), eval_row(96, loss=0.69, ref="warmstart")] + rows[1]["vs_prev"] = {"dloss": -0.01, "se": 0.001, "better": 0.55, "ref_loss": 0.70} + rows[1]["vs_sft_v1"] = {"dloss": 0.03, "ref_loss": 0.66} # a reference that only newer points carry + v = view(eval_rows=rows) + assert ("Eval step 96 loss 0.6900 | vs base -0.0500 | vs warmstart -0.0100 | vs sft_v1 +0.0300" + " | vs prev -0.0100") in status_message(v).plain() + e = eval_message(v) + text = e.plain() + assert "step | loss | vs base | vs warmstart | vs sft_v1 | vs prev" in text + assert "64 | 0.7000 | -0.0500 | -0.0200 | - | -" in text and "vs sft_v1: dloss +0.0300 (ref 0.6600)" in text + assert "ckpt100" not in text + assert [s[0] for s in e.charts[0].series] == ["eval loss", "base ref", "warmstart ref", "sft_v1 ref"] + + +# -- commands and authorization ----------------------------------------------------- + +class FakeClient: + def __init__(self, v): + self.v = v + self.fail = None + + def view(self): + if self.fail: + raise self.fail + return self.v + + def bad_rounds(self): + return [] + + +def make_bot(tmp_path, chat="", v=None, allowed=None, nl=None, clock=None, **cfg_kw): + cfg_kw.setdefault("policy", Policy(digest_at=None)) + cfg = Config(state_file=str(tmp_path / "state.json"), audit_log=str(tmp_path / "audit.jsonl"), **cfg_kw) + ch = FakeChannel("telegram", chat) + ch.allowed = list(allowed or []) + client = FakeClient(v or view(eval_rows=[eval_row(96)])) + return Bot(cfg, client, [ch], clock=clock or FakeClock(), nl=nl), ch, client + + +def cmd(text, chat="1", user="al", button=False, user_id="1001"): + return Command("telegram", chat, text, user, button, None, user_id) + + +def test_parse_command_and_duration(): + assert parse_command("/Island@yeto_bot 2 ") == ("/island", "2") + assert parse_duration("2h") == 7200 and parse_duration("30m") == 1800 and parse_duration("45") == 2700 + assert parse_duration("soon") is None + + +def test_start_registers_then_only_report_chat(tmp_path): + bot, ch, _ = make_bot(tmp_path, allowed=["7"]) + assert bot.handle(ch, cmd("/status", "111")).title == "not registered" + assert bot.handle(ch, cmd("hello", "111")) is None + assert bot.handle(ch, cmd("/start", "111", user_id="8")) is None and ch.report_chat == "" # not allowlisted + assert "Registered" in bot.handle(ch, cmd("/start", "111", user_id="7")).plain() and ch.report_chat == "111" + assert bot.handle(ch, cmd("/start", "222", user_id="7")) is None # a second chat cannot take over + assert bot.handle(ch, cmd("/status", "222", button=True, user_id="7")) is None + assert ch.answers[-1] == "Not allowed" + assert all(target == "111" for target, _ in ch.sent) + assert bot.handle(ch, cmd("/status", "111", user_id="7")).title == "yeto progress: sft-test" + bot2, ch2, _ = make_bot(tmp_path) # registration survives a restart via the state file + assert ch2.report_chat == "111" + + +def test_start_without_allowlist_never_registers(tmp_path): + bot, ch, _ = make_bot(tmp_path) + ch.chat_env, ch.allow_env = "TELEGRAM_CHAT_ID", "TELEGRAM_ALLOWED_USERS" + reply = bot.handle(ch, cmd("/start", "111", user_id="7")) + assert reply.title == "registration disabled" and ch.report_chat == "" + assert "Set TELEGRAM_CHAT_ID or TELEGRAM_ALLOWED_USERS" in reply.plain() + assert not (tmp_path / "state.json").exists() + + +def test_user_allowlist(tmp_path): + bot, ch, _ = make_bot(tmp_path, chat="1", allowed=["ou_owner", "42"]) + assert bot.handle(ch, cmd("/status", user="ou_other", user_id="ou_other")) is None + assert bot.handle(ch, cmd("/status", user="ou_owner", user_id="ou_owner")).title.startswith("yeto progress") + assert bot.handle(ch, cmd("/status", user="tg_name", user_id="42")) is not None + assert bot.handle(ch, cmd("/status", user="42", user_id="7")) is None # display names never authorize + + +def test_cli_command_is_trusted_past_the_allowlist(tmp_path, monkeypatch): + from yeto.bot import __main__ as main_mod + + ch = FakeChannel("telegram", "1") + ch.allowed = ["42"] + monkeypatch.setattr(main_mod.Config, "from_env", classmethod( + lambda cls: Config(audit_log=str(tmp_path / "audit.jsonl"), policy=Policy(digest_at=None)))) + monkeypatch.setattr(main_mod, "build_channels", lambda cfg, receive: [ch]) + monkeypatch.setattr(main_mod, "DashboardClient", lambda *a, **kw: FakeClient(view())) + assert main_mod.main(["--command", "/status"]) == 0 + assert [(target, m.title) for target, m in ch.sent] == [("1", "yeto progress: sft-test")] + bot, ch2, _ = make_bot(tmp_path, chat="1", allowed=["42"]) + assert bot.handle(ch2, Command("telegram", "1", "/status", "cli")) is None # the flag, not the name + + +def test_commands_dispatch_and_errors(tmp_path): + bot, ch, client = make_bot(tmp_path, chat="1") + for text, title in (("/island 1", "island 1: isl1"), ("/eval", "yeto eval: step 96"), ("/merge", "yeto merge"), + ("/res", "yeto resources"), ("/ckpt", "yeto checkpoints"), ("/alerts", "yeto alerts"), + ("/help", "yeto training bot")): + assert bot.handle(ch, cmd(text)).title == title, text + assert "unknown command /foo" in bot.handle(ch, cmd("/foo")).plain() + client.fail = OSError("dashboard down") + reply = bot.handle(ch, cmd("/status")) + assert reply.title == "command failed" and "dashboard down" in reply.plain() + + +def test_ack_details_mute_flow(tmp_path): + clock = FakeClock() + bot, ch, client = make_bot(tmp_path, chat="1", clock=clock, allowed=["1001"], + policy=Policy(digest_at=None, progress_interval=0), dashboard_link="https://d/") + bot.tick() + st = clone(client.v) + st["sft"]["islands"][0]["status"] = "stale" + client.v = st + clock.t += 60 + [card] = bot.tick() + assert card.level == "p0" and ("Ack 2h", "/ack c1") in card.buttons + det = bot.handle(ch, cmd("/details c1", button=True)) + assert det.title == "details c1" and "isl0 (status stale" in det.plain() + assert det.plain().count("\n10 | ") == 1 # last 10 steps table + ack = bot.handle(ch, cmd("/ack c1", user="bob", button=True)) + assert ack.title == "acked 1 alert(s) for 2h" and "isl0 stale" in ack.plain() + clock.t += 3600 + assert bot.tick() == [] # re-notify held back by the ack + assert bot.handle(ch, cmd("/ack c99")).title == "unknown card" + assert bot.handle(ch, cmd("/mute 2h")).title == "muted" + assert "muted for 2h" in bot.alerts_line() + assert bot.handle(ch, cmd("/mute later")).title == "usage" + assert bot.handle(ch, cmd("/unmute")).title == "unmuted" + assert bot.monitor.mute_until == 0 + + +def test_alert_control_needs_an_allowlist(tmp_path): + bot, ch, _ = make_bot(tmp_path, chat="1") + ch.allow_env = "TELEGRAM_ALLOWED_USERS" + for text in ("/mute 1d all", "/unmute", "/ack all", "/ack c1"): + reply = bot.handle(ch, cmd(text, button=True)) + assert reply.title == "not allowed" and "TELEGRAM_ALLOWED_USERS" in reply.plain(), text + assert ch.answers[-1] == "Needs a user allowlist" + assert bot.monitor.mute_until == 0 + req = bot.request_decision("restart island 2?", [], ["yes", "no"]) + reply = bot.handle(ch, cmd(f"/decide {req.buttons[0][1].split()[1]} yes", button=True)) + assert reply.title == "not allowed" and "records a decision" in reply.plain() + assert bot.handle(ch, cmd("/alerts")).title == "yeto alerts" # read-only commands still work + assert bot.handle(ch, cmd("/status")).title == "yeto progress: sft-test" + trusted = Command("telegram", "1", "/mute 1h", "cli", trusted=True) + assert bot.handle(ch, trusted).title == "muted" # the local CLI is not a chat user + assert bot.allowlist_warnings() == [ + "telegram: no user allowlist (TELEGRAM_ALLOWED_USERS unset): anyone in the report chat can run read-only" + " commands; /ack, /mute, /unmute, /decide are refused"] + ch.allowed = ["1001"] + assert bot.allowlist_warnings() == [] + assert bot.handle(ch, cmd("/unmute")).title == "unmuted" + + +def test_progress_schedule_and_daily_digest(tmp_path): + start = time.mktime((2026, 11, 7, 8, 0, 0, 0, 0, -1)) + clock = FakeClock(start) + bot, ch, client = make_bot(tmp_path, chat="1", clock=clock, + policy=Policy(progress_interval=7200, digest_at=9 * 60)) + assert [m.title for m in bot.tick()] == ["yeto progress: sft-test"] # startup card + clock.t += 3600 - 60 + assert bot.tick() == [] + clock.t += 120 # 09:01: digest + msgs = bot.tick() + assert [m.title.split(":")[0] for m in msgs] == ["yeto daily digest"] + assert "Alerts, last 24 h P0 0 | P1 0 | P2 0" in msgs[0].plain() + clock.t += 3600 # 10:01: the 2 h progress card + assert [m.title for m in bot.tick()] == ["yeto progress: sft-test"] + clock.t += 600 + assert bot.tick() == [] # digest once a day + late = make_bot(tmp_path, chat="1", clock=FakeClock(start + 4 * 3600), + policy=Policy(progress_interval=0, digest_at=9 * 60))[0] + assert late.tick() == [] # started after 09:00: first digest tomorrow + + +def test_decision_request_and_audit_log(tmp_path): + bot, ch, _ = make_bot(tmp_path, chat="1", allowed=["1001"]) + msg = bot.request_decision("island 2 stale for 40 min", ["last step 812", "syncer still merging"], + ["restart", "wait", "ignore"], "wait") + req_id = msg.buttons[0][1].split()[1] + assert msg.level == "decision" and [b[0] for b in msg.buttons] == ["restart", "wait (suggested)", "ignore"] + assert ch.sent[-1][1] is msg + reply = bot.handle(ch, cmd(f"/decide {req_id} restart", user="bob", button=True)) + assert reply.title == "decision recorded" and "recorded; operator will act" in reply.plain() + again = bot.handle(ch, cmd(f"/decide {req_id} wait", user="eve", button=True)) + assert again.title == "decision already recorded" and "restart by bob" in again.plain() + assert bot.handle(ch, cmd(f"/decide {req_id} reboot")).title == "unknown option" + assert bot.handle(ch, cmd("/decide dnope wait")).title == "unknown decision" + recs = [json.loads(x) for x in (tmp_path / "audit.jsonl").read_text().splitlines()] + assert [r["type"] for r in recs] == ["request", "decision"] + assert recs[1]["option"] == "restart" and recs[1]["user"] == "bob" and recs[1]["chat"] == "1" + assert recs[1]["user_id"] == "1001" + with pytest.raises(ValueError): + bot.request_decision("x", [], ["two words"]) + # a request posted by another process (CLI) is answered by this bot through the shared log + other, _, _ = make_bot(tmp_path, chat="1") + req2 = other.decisions.request("lower lr?", [], ["yes", "no"]) + assert bot.handle(ch, cmd(f"/decide {req2.id} no")).title == "decision recorded" + + +def test_natural_language_hook(tmp_path): + seen = [] + + def nl(text, tools): + seen.append((text, tools.view()["overview"]["run"], tools.active_alerts())) + from yeto.bot.render import Message + return Message("answer", ["fine"]) + + bot, ch, _ = make_bot(tmp_path, chat="1", nl=nl) + assert bot.handle(ch, cmd("how is island 1?")).title == "answer" + assert seen == [("how is island 1?", "sft-test", [])] + assert bot.handle(ch, cmd("how is island 1?", chat="9")) is None + plain, ch2, _ = make_bot(tmp_path, chat="1") + assert plain.handle(ch2, cmd("hello")) is None # no handler installed: ignored + + +def test_no_report_chat_means_no_broadcast(tmp_path): + bot, ch, _ = make_bot(tmp_path) + bot.tick() + assert ch.sent == [] + + +# -- config ------------------------------------------------------------------------- + +def test_config_from_env_defaults_and_overrides(): + cfg = Config.from_env({}) + assert cfg.dashboard_url == "http://127.0.0.1:8787" and cfg.lark_domain == "https://open.larksuite.com" + assert cfg.policy.progress_interval == 7200 and cfg.policy.digest_at == 540 and cfg.policy.quiet_hours is None + assert cfg.policy.p0_renotify_s == 1800 and cfg.policy.p1_cooldown_s == 3600 and cfg.policy.ack_mute_s == 7200 + assert cfg.th.hbm_limit_gib is None and cfg.th.grad_norm_max == 1.0 and cfg.th.grace_s == 600 + assert cfg.th.hbm_warn_gib is None and cfg.th.hbm_warn_frac == 0.965 + assert cfg.lark_card_buttons and cfg.lark_mode == "auto" + cfg = Config.from_env({"YETO_DASHBOARD_URL": "http://dash:9/", "YETO_BOT_HBM_LIMIT_GIB": "60", + "YETO_BOT_QUIET_HOURS": "23-8", "YETO_BOT_DIGEST_AT": "off", "LARK_MODE": "webhook", + "LARK_OWNER_OPEN_IDS": "ou_a, ou_b", "TELEGRAM_OWNERS": "@ops_lead", + "YETO_BOT_DRY_RUN": "1", "LARK_CARD_BUTTONS": "0"}) + assert cfg.dashboard_url == "http://dash:9" and cfg.th.hbm_limit_gib == 60.0 + assert Config.from_env({"YETO_BOT_HBM_WARN_FRAC": "0.9"}).th.hbm_warn_frac == 0.9 + assert Config.from_env({"YETO_BOT_HBM_WARN_GIB": "70"}).th.hbm_warn_gib == 70.0 + assert cfg.policy.quiet_hours == (1380, 480) and cfg.policy.digest_at is None + assert cfg.lark_owners == ["ou_a", "ou_b"] and cfg.telegram_owners == ["ops_lead"] + assert cfg.dry_run and cfg.lark_mode == "webhook" and not cfg.lark_card_buttons + assert Config.from_env({"TELEGRAM_ALLOWED_USERS": "42, 1001"}).telegram_allowed == ["42", "1001"] + with pytest.raises(ValueError, match="numeric Telegram user ids"): + Config.from_env({"TELEGRAM_ALLOWED_USERS": "42,@ops_lead"}) + with pytest.raises(ValueError, match="LARK_MODE"): + Config.from_env({"LARK_MODE": "smtp"}) + with pytest.raises(ValueError, match="YETO_BOT_POLL_INTERVAL"): + Config.from_env({"YETO_BOT_POLL_INTERVAL": "soon"}) + + +def test_bold_markers_never_leak_into_plain_text(): + assert "\x02" not in status_message(view()).plain() + assert strip_bold("\x02a\x03") == "a" + + +def test_title_prefix_marks_every_outgoing_message(tmp_path): + bot, ch, _ = make_bot(tmp_path, chat="1", title_prefix="[TEST] ") + reply = bot.handle(ch, cmd("/eval")) + assert reply.title == "yeto eval: step 96" # the returned reply is unprefixed + assert ch.sent[-1][1].title == "[TEST] yeto eval: step 96" + bot.tick() + assert ch.sent[-1][1].title == "[TEST] yeto progress: sft-test" + + +def test_ack_mutes_event_subjects_too(tmp_path): + clock = FakeClock() + bot, ch, client = make_bot(tmp_path, chat="1", clock=clock, allowed=["1001"], + policy=Policy(digest_at=None, progress_interval=0, p1_cooldown_s=0)) + ev = [[T0, "rejected", 1, "log", None]] + client.v = view(events=ev) + bot.tick() + client.v = view(events=ev + [[T0 + 10, "rejected", 2, "log", None]]) + clock.t += 60 + [card] = bot.tick() + assert card.title == "P1 pushes rejected" and ("Ack 2h", "/ack c1") in card.buttons + assert bot.handle(ch, cmd("/ack c1")).title == "acked 1 alert(s) for 2h" + client.v = view(events=ev + [[T0 + 10, "rejected", 2, "log", None], [T0 + 20, "rejected", 2, "log", None]]) + clock.t += 60 + assert bot.tick() == [] # the subject is muted for 2 h diff --git a/tests/test_bot_channels.py b/tests/test_bot_channels.py new file mode 100644 index 00000000..cb0397c1 --- /dev/null +++ b/tests/test_bot_channels.py @@ -0,0 +1,449 @@ +"""yeto.bot channels: Telegram / Lark message bodies, token refresh, webhook signature, fallbacks, +and an end-to-end pass against a real loopback dashboard server.""" + +import base64 +import hashlib +import hmac +import json +import threading +from types import SimpleNamespace + +import pytest + +from yeto.bot import app as app_mod +from yeto.bot import channels as channels_mod +from yeto.bot import charts as charts_mod +from yeto.bot import lark_ws +from yeto.bot.charts import Chart, sparkline, text_lines +from yeto.bot.app import Bot, build_channels +from yeto.bot.channels import (LarkAppChannel, LarkWebhookChannel, TelegramChannel, lark_card, + tg_markdown, webhook_sign) +from yeto.bot.config import Config +from yeto.bot.http import DashboardClient, HttpError, multipart, redact_url, request_json +from yeto.bot.render import Message, b + +from bot_helpers import FakeClock, FakeHttp + +MSG = Message("isl_0 is stale", [f"{b('step')} 12 | loss 0.7 *x* [a]_b", "plain"], "p0", + [("Status", "/status"), ("Island 0", "/island 0")]) +CHART = Chart("loss", [("isl0", [1, 2, 3], [0.9, 0.8, 0.7])]) + + +@pytest.fixture(autouse=True) +def no_matplotlib(monkeypatch): + """Deterministic chart path: text sparklines unless a test installs a fake renderer.""" + monkeypatch.setattr(charts_mod, "available", lambda: False) + + +def fake_png(monkeypatch, png=b"\x89PNG fake"): + monkeypatch.setattr(channels_mod.charts_mod, "render_png", lambda chart: png) + + +# -- Telegram ----------------------------------------------------------------------- + +def test_telegram_markdown_escapes_and_bold(): + text = tg_markdown(MSG) + assert text.splitlines()[0] == "\U0001f6a8 *isl-0 is stale*" # no escapes possible inside bold + assert text.splitlines()[1] == r"*step* 12 | loss 0.7 \*x\* \[a]\_b" + assert "\x02" not in text + + +def test_telegram_send_body_and_plain_fallback(): + http = FakeHttp([("sendMessage", [HttpError(400, "can't parse entities", "u"), {"ok": True}])]) + ch = TelegramChannel("123:abc", "-100", http=http) + assert ch.send(MSG) + first, second = http.bodies("sendMessage") + assert http.calls[0]["url"] == "https://api.telegram.org/bot123:abc/sendMessage" + assert first["chat_id"] == "-100" and first["parse_mode"] == "Markdown" + assert first["reply_markup"] == {"inline_keyboard": [[{"text": "Status", "callback_data": "/status"}, + {"text": "Island 0", "callback_data": "/island 0"}]]} + assert "parse_mode" not in second and "*" not in second["text"].splitlines()[0] + assert second["text"].startswith("\U0001f6a8 isl_0 is stale") + + +def test_telegram_send_without_chat_is_skipped_and_long_text_truncated(): + http = FakeHttp() + ch = TelegramChannel("t", "", http=http) + assert not ch.send(MSG) and http.calls == [] + ch.configured_chat = "1" + ch.send(Message("big", ["x" * 5000])) + text = http.bodies("sendMessage")[0]["text"] + assert len(text) <= 4000 and text.endswith("(truncated)") + + +def test_telegram_parse_updates_messages_and_buttons(): + ch = TelegramChannel("t", "1", http=FakeHttp()) + cmds = ch.parse_updates({"result": [ + {"update_id": 10, "message": {"text": "/status", "chat": {"id": 1}, "from": {"username": "al"}}}, + {"update_id": 11, "message": {"text": "hello", "chat": {"id": 1}}}, + {"update_id": 12, "callback_query": {"id": "cq9", "data": "/island 0", "from": {"id": 77, "first_name": "Bo"}, + "message": {"chat": {"id": 1}, "message_id": 5}}}]}) + assert [(c.chat, c.text, c.user, c.user_id, c.button) for c in cmds] == [ + ("1", "/status", "al", "", False), ("1", "hello", "", "", False), ("1", "/island 0", "Bo", "77", True)] + assert ch.offset == 13 + ch.answer(cmds[2], "running /island") + assert ch.http.bodies("answerCallbackQuery") == [{"callback_query_id": "cq9", "text": "running /island"}] + + +def test_telegram_allowlist_checks_numeric_id_only(): + ch = TelegramChannel("t", "1", http=FakeHttp(), allowed=["42"]) + + def press(frm): + return ch.parse_updates({"result": [{"update_id": 1, "callback_query": { + "id": "cq", "data": "/mute 1d all", "from": frm, "message": {"chat": {"id": 1}}}}]})[0] + + assert ch.user_allowed(press({"id": 42, "first_name": "Eve"})) + assert not ch.user_allowed(press({"id": 7, "first_name": "42"})) # a first name is free text + assert not ch.user_allowed(press({"id": 7, "username": "42"})) + assert not ch.user_allowed(press({"first_name": "42"})) # no id at all + + +def test_telegram_poll_uses_offset_long_poll(): + http = FakeHttp([("getUpdates", {"ok": True, "result": [{"update_id": 4, "message": { + "text": "/eval", "chat": {"id": 7}}}]})]) + ch = TelegramChannel("t", "7", http=http) + assert [c.text for c in ch.poll_commands(25)] == ["/eval"] + ch.poll_commands(25) + bodies = http.bodies("getUpdates") + assert bodies[0]["offset"] == 0 and bodies[1]["offset"] == 5 and bodies[0]["timeout"] == 25 + assert bodies[0]["allowed_updates"] == ["message", "callback_query"] + + +# -- Lark card -------------------------------------------------------------------------- + +def test_lark_card_structure_and_escaping(): + card = lark_card(MSG, buttons=True) + assert card["header"]["template"] == "red" + assert card["header"]["title"] == {"tag": "plain_text", "content": "\U0001f6a8 isl_0 is stale"} + md = card["elements"][0]["text"] + assert md["tag"] == "lark_md" + assert md["content"].splitlines()[0] == "**step** 12 | loss 0.7 *x* [a]_b" + actions = card["elements"][1]["actions"] + assert [a["value"] for a in actions] == [{"cmd": "/status"}, {"cmd": "/island 0"}] + assert actions[0]["type"] == "primary" + assert len(lark_card(MSG, buttons=False)["elements"]) == 1 + + +# -- Lark app bot --------------------------------------------------------------------- + +def token_resp(tok, expire=7200): + return {"code": 0, "msg": "ok", "tenant_access_token": tok, "expire": expire} + + +def test_lark_app_send_body_and_token_cache_refresh(): + clock = FakeClock(1000.0) + http = FakeHttp([("tenant_access_token", [token_resp("t-1"), token_resp("t-2")]), + ("im/v1/messages", {"code": 0, "data": {}})]) + ch = LarkAppChannel("cli_a", "sec", "oc_1", http=http, clock=clock, refresh_margin=600, receive=True, + buttons=True) + assert ch.send(MSG) + tok_call = http.calls[0] + assert tok_call["url"] == "https://open.larksuite.com/open-apis/auth/v3/tenant_access_token/internal" + assert tok_call["body"] == {"app_id": "cli_a", "app_secret": "sec"} + send = http.calls[1] + assert send["url"].endswith("/open-apis/im/v1/messages?receive_id_type=chat_id") + assert send["headers"]["Authorization"] == "Bearer t-1" + assert send["body"]["receive_id"] == "oc_1" and send["body"]["msg_type"] == "interactive" + card = json.loads(send["body"]["content"]) # content is a JSON string + assert card["elements"][1]["tag"] == "action" + clock.t += 7200 - 601 # still valid outside the margin: cached + ch.send(MSG) + assert len(http.bodies("tenant_access_token")) == 1 + clock.t += 2 # inside the refresh margin: refreshed before expiry + ch.send(MSG) + assert len(http.bodies("tenant_access_token")) == 2 + assert http.calls[-1]["headers"]["Authorization"] == "Bearer t-2" + + +def test_lark_app_expired_token_retries_once_and_text_fallback(): + http = FakeHttp([("tenant_access_token", [token_resp("old"), token_resp("new")]), + ("im/v1/messages", [{"code": 99991663, "msg": "token expired"}, {"code": 0}])]) + ch = LarkAppChannel("a", "s", "oc", http=http) + assert ch.send(MSG) + sends = [c for c in http.calls if "messages" in c["url"]] + assert [c["headers"]["Authorization"] for c in sends] == ["Bearer old", "Bearer new"] + http2 = FakeHttp([("tenant_access_token", token_resp("t")), + ("im/v1/messages", [HttpError(400, json.dumps({"code": 230099, "msg": "bad card"}), "u"), + {"code": 0}])]) + ch2 = LarkAppChannel("a", "s", "oc", http=http2) + assert ch2.send(MSG) + last = http2.bodies("messages")[-1] + assert last["msg_type"] == "text" and json.loads(last["content"])["text"].startswith("\U0001f6a8 isl_0 is stale") + + +def test_lark_app_token_failure_and_no_chat(): + http = FakeHttp([("tenant_access_token", {"code": 10003, "msg": "invalid app_secret"})]) + ch = LarkAppChannel("a", "s", "oc", http=http) + assert not ch.send(MSG) + assert not LarkAppChannel("a", "s", "", http=http).send(MSG) + + +def test_lark_buttons_need_receive_and_opt_in(): + for receive, buttons, tags in ((False, True, ["div"]), (True, False, ["div"]), (True, True, ["div", "action"])): + http = FakeHttp([("tenant_access_token", token_resp("t"))]) + LarkAppChannel("a", "s", "oc", http=http, receive=receive, buttons=buttons).send(MSG) + card = json.loads(http.bodies("messages")[0]["content"]) + assert [e["tag"] for e in card["elements"]] == tags + + +# -- Lark webhook ------------------------------------------------------------------- + +def test_webhook_signature_matches_reference_algorithm(): + ts, secret = 1_700_000_000, "demo-secret" + expect = base64.b64encode(hmac.new(f"{ts}\n{secret}".encode(), b"", hashlib.sha256).digest()).decode() + assert webhook_sign(secret, ts) == expect + + +def test_webhook_payload_signed_and_fallback(): + http = FakeHttp([("hook", [{"code": 19021, "msg": "sign match fail"}, {"code": 0}])]) + ch = LarkWebhookChannel("https://open.feishu.cn/open-apis/bot/v2/hook/x", "sec", http=http, + clock=FakeClock(1_700_000_000.9)) + assert ch.report_chat and ch.send(MSG) + card_body, text_body = http.bodies("hook") + assert card_body["timestamp"] == "1700000000" and card_body["sign"] == webhook_sign("sec", 1_700_000_000) + assert card_body["msg_type"] == "interactive" and isinstance(card_body["card"], dict) + assert len(card_body["card"]["elements"]) == 1 # custom bots cannot take button callbacks + assert text_body["msg_type"] == "text" and text_body["content"]["text"].startswith("\U0001f6a8") + plain = FakeHttp([("hook", {"StatusCode": 0})]) + LarkWebhookChannel("https://x/hook/y", http=plain).send(MSG) + assert "sign" not in plain.bodies("hook")[0] + + +# -- Lark long-connection events ------------------------------------------------------ + +def ns(**kw): + return SimpleNamespace(**kw) + + +def test_lark_ws_message_and_card_parsing_with_dedupe(): + data = ns(event=ns(message=ns(message_type="text", content=json.dumps({"text": "@_user_1 /island 2"}), + chat_id="oc_9", message_id="om_1"), + sender=ns(sender_id=ns(open_id="ou_7")))) + cmd, mid = lark_ws.parse_message_event(data) + assert (cmd.chat, cmd.text, cmd.user, cmd.button, mid) == ("oc_9", "/island 2", "ou_7", False, "om_1") + data.event.message.content = json.dumps({"text": "@_user_1 how is island 2?"}) + assert lark_ws.parse_message_event(data)[0].text == "how is island 2?" # kept for the NL hook + data.event.message.message_type = "image" + assert lark_ws.parse_message_event(data) is None + card = ns(header=ns(event_id="ev1"), event=ns(action=ns(value={"cmd": "/status"}), + context=ns(open_chat_id="oc_9"), operator=ns(open_id="ou_7"))) + cmd2, eid = lark_ws.parse_card_event(card) + assert cmd2.button and cmd2.text == "/status" and eid == "ev1" + assert lark_ws.card_toast((cmd2, eid)) == "received" and lark_ws.card_toast(None) == "unknown button" + ch = LarkAppChannel("a", "s", "oc_9", http=FakeHttp(), receive=True) + ch.push(cmd, mid) + ch.push(cmd, mid) # redelivered + ch.push(cmd2, eid) + assert [c.text for c in ch.poll_commands(0.01)] == ["/island 2", "/status"] + assert ch.poll_commands(0.01) == [] + + +def test_lark_ws_callbacks_with_sdk(): + pytest.importorskip("lark_oapi") + ch = LarkAppChannel("a", "s", "oc_9", http=FakeHttp(), receive=True) + assert lark_ws.build_handler(ch) is not None + on_message, on_card = lark_ws.callbacks(ch) + on_message(ns(event=ns(message=ns(message_type="text", content=json.dumps({"text": "/eval"}), chat_id="oc_9", + message_id="om_2"), sender=ns(sender_id=ns(open_id="ou_1"))))) + press = ns(header=ns(event_id="ev2"), event=ns(action=ns(value={"cmd": "/ack c1"}), + context=ns(open_chat_id="oc_9"), operator=ns(open_id="ou_1"))) + resp = on_card(press) + assert resp.toast.content == "received" # the press is not authorized yet + assert [(c.text, c.button) for c in ch.poll_commands(0.01)] == [("/eval", False), ("/ack c1", True)] + + +# -- channel selection -------------------------------------------------------------- + +def test_build_channels_matrix(monkeypatch): + monkeypatch.setattr(lark_ws, "available", lambda: False) + http = FakeHttp() + assert build_channels(Config(), http) == [] + chans = build_channels(Config(telegram_token="t", lark_webhook_url="https://x/hook"), http) + assert [c.name for c in chans] == ["telegram", "lark-webhook"] + # app creds without the SDK and without a chat id: fall back to the webhook + chans = build_channels(Config(lark_app_id="a", lark_app_secret="s", lark_webhook_url="https://x/hook"), http) + assert [c.name for c in chans] == ["lark-webhook"] + chans = build_channels(Config(lark_app_id="a", lark_app_secret="s", lark_chat_id="oc"), http) + assert chans[0].name == "lark" and not chans[0].can_receive + assert build_channels(Config(lark_mode="ws", lark_webhook_url="https://x/hook"), http) == [] + started = [] + monkeypatch.setattr(lark_ws, "available", lambda: True) + monkeypatch.setattr(lark_ws, "start", lambda ch: started.append(ch)) + chans = build_channels(Config(lark_app_id="a", lark_app_secret="s"), http) + assert chans[0].can_receive and chans[0].buttons and started == chans + chans = build_channels(Config(lark_app_id="a", lark_app_secret="s", lark_card_buttons=False), http) + assert not chans[0].buttons + started.clear() + chans = build_channels(Config(lark_app_id="a", lark_app_secret="s", telegram_token="t"), http, receive=False) + assert started == [] and not any(c.can_receive for c in chans) # one-shot CLI runs never listen + started.clear() + chans = build_channels(Config(lark_app_id="a", lark_app_secret="s", dry_run=True), http) + assert not chans[0].can_receive and started == [] # dry run never opens the connection + + +def test_dry_run_prints_bodies_and_redacts(capsys, monkeypatch): + monkeypatch.setattr(lark_ws, "available", lambda: False) + cfg = Config(dry_run=True, telegram_token="123:secret", telegram_chat_id="-1", lark_app_id="cli", + lark_app_secret="topsecret", lark_chat_id="oc_1") + chans = build_channels(cfg) + for ch in chans: + assert ch.send(MSG) + out = capsys.readouterr().out + assert "POST https://api.telegram.org/bot/sendMessage" in out + assert "123:secret" not in out and "topsecret" not in out and "" in out + assert "/open-apis/im/v1/messages?receive_id_type=chat_id" in out + hook = build_channels(Config(dry_run=True, lark_webhook_url="https://open.feishu.cn/open-apis/bot/v2/hook/" + "a1b2-c3d4")) + assert hook[0].send(MSG) + out = capsys.readouterr().out + assert "POST https://open.feishu.cn/open-apis/bot/v2/hook/" in out and "a1b2-c3d4" not in out + + +def test_redact_url_hides_path_tokens(): + assert redact_url("https://api.telegram.org/bot123:AA-b_c/sendMessage") == \ + "https://api.telegram.org/bot/sendMessage" + assert redact_url("https://open.larksuite.com/open-apis/bot/v2/hook/0f9e-77aa") == \ + "https://open.larksuite.com/open-apis/bot/v2/hook/" + assert redact_url("https://open.feishu.cn/open-apis/bot/v2/hook/0f9e?x=1") == \ + "https://open.feishu.cn/open-apis/bot/v2/hook/?x=1" + plain = "https://open.feishu.cn/open-apis/im/v1/messages?receive_id_type=chat_id" + assert redact_url(plain) == plain + + +@pytest.mark.parametrize("url,secret", [("https://api.telegram.org/bot123:secret/sendMessage", "123:secret"), + ("https://open.feishu.cn/open-apis/bot/v2/hook/hook-secret", "hook-secret")]) +def test_http_error_never_carries_the_token(monkeypatch, url, secret): + import io + import urllib.error + + def boom(req, timeout): + raise urllib.error.HTTPError(req.full_url, 400, "Bad Request", {}, io.BytesIO(b'{"code": 19001}')) + + monkeypatch.setattr("urllib.request.urlopen", boom) + with pytest.raises(HttpError) as info: + request_json("POST", url, body={}) + assert secret not in str(info.value) and "" in str(info.value) + + +# -- end to end: real loopback dashboard -> bot -> fake channels ---------------------- + +@pytest.fixture +def dashboard(tmp_path): + from test_dashboard_sft import make_run + + from yeto.dashboard.reducer import Reducer + from yeto.dashboard.serve import DashboardState, make_server + + root = make_run(tmp_path / "run", steps=8) + state = DashboardState(Reducer(run="sft-e2e"), [str(root)]) + state.poll() + srv = make_server(state, "127.0.0.1", 0, live=False) + threading.Thread(target=srv.serve_forever, daemon=True).start() + yield f"http://127.0.0.1:{srv.server_address[1]}" + srv.shutdown() + srv.server_close() + + +def test_end_to_end_against_dashboard_server(dashboard, tmp_path): + from bot_helpers import FakeChannel + + client = DashboardClient(dashboard, http=request_json) + tg, lark = FakeChannel("telegram", "1"), FakeChannel("lark", "oc") + bot = Bot(Config(dashboard_url=dashboard, audit_log=str(tmp_path / "audit.jsonl")), client, [tg, lark], + clock=FakeClock()) + sent = bot.tick() + assert [m.title for m in sent] == ["yeto progress: sft-e2e"] + text = sent[0].plain() + assert "8/100" in text and "Outer 2/960" in text and "Merge direct-in-round" in text + assert [m.title for _, m in tg.sent] == [m.title for _, m in lark.sent] + assert bot.handle(tg, app_mod.Command("telegram", "1", "/island 0")).title.startswith("island 0") + merge = bot.handle(tg, app_mod.Command("telegram", "1", "/merge")).plain() + assert "Timed-out rounds (1)" in merge and "rejected 1" in merge + assert "Syncer checkpoint writes (1)" in bot.handle(tg, app_mod.Command("telegram", "1", "/ckpt")).plain() + + +# -- mentions, URL buttons, charts ------------------------------------------------------ + +ALERT = Message("P0 isl0 stale", ["no new step"], "p0", + [("Ack 2h", "/ack c1"), ("Open dashboard", "https://d.example/")], + [CHART], mention=True) + + +def test_telegram_mention_url_button_and_sparkline_fallback(): + http = FakeHttp() + TelegramChannel("t", "1", http=http, owners=["ops_lead"]).send(ALERT) + [body] = http.bodies("sendMessage") + assert body["text"].splitlines()[-1] == r"cc @ops\_lead" + assert "loss isl0: " in body["text"] and "0.7..0.9 (last 0.7)" in body["text"] + assert body["reply_markup"]["inline_keyboard"] == [[{"text": "Ack 2h", "callback_data": "/ack c1"}, + {"text": "Open dashboard", "url": "https://d.example/"}]] + assert http.bodies("sendPhoto") == [] + + +def test_telegram_send_photo_when_png_renders(monkeypatch): + fake_png(monkeypatch) + http = FakeHttp() + TelegramChannel("t", "1", http=http).send(ALERT) + [photo] = [c for c in http.calls if "sendPhoto" in c["url"]] + assert photo["body"] == {"chat_id": "1", "caption": "loss"} + assert photo["files"]["photo"] == ("chart.png", b"\x89PNG fake", "image/png") + assert "loss isl0" not in http.bodies("sendMessage")[0]["text"] # no sparkline when the image went out + + +def test_lark_image_upload_mention_and_fallback(monkeypatch): + fake_png(monkeypatch) + http = FakeHttp([("tenant_access_token", token_resp("t")), + ("im/v1/images", {"code": 0, "data": {"image_key": "img_v3_1"}})]) + ch = LarkAppChannel("a", "s", "oc", http=http, receive=True, buttons=True, owners=["ou_1"]) + assert ch.send(ALERT) + up = next(c for c in http.calls if "im/v1/images" in c["url"]) + assert up["body"] == {"image_type": "message"} and up["files"]["image"][1] == b"\x89PNG fake" + assert up["headers"]["Authorization"] == "Bearer t" + card = json.loads(http.bodies("messages")[0]["content"]) + assert "" in card["elements"][0]["text"]["content"] + assert card["elements"][1] == {"tag": "img", "img_key": "img_v3_1", "alt": {"tag": "plain_text", "content": "loss"}} + acts = card["elements"][2]["actions"] + assert acts[0]["value"] == {"cmd": "/ack c1"} and acts[1]["url"] == "https://d.example/" + bad = FakeHttp([("tenant_access_token", token_resp("t")), + ("im/v1/images", {"code": 99991672, "msg": "no im:resource"})]) + LarkAppChannel("a", "s", "oc", http=bad, receive=True, buttons=False).send(ALERT) + card = json.loads(bad.bodies("messages")[0]["content"]) + assert [e["tag"] for e in card["elements"]] == ["div", "action"] # URL button stays, command button dropped + assert "loss isl0: " in card["elements"][0]["text"]["content"] + + +def test_webhook_uses_sparklines_and_keeps_url_buttons(): + http = FakeHttp([("hook", {"code": 0})]) + LarkWebhookChannel("https://x/hook/y", http=http, owners=["ou_1"]).send(ALERT) + card = http.bodies("hook")[0]["card"] + body = card["elements"][0]["text"]["content"] + assert "loss isl0: " in body and "" in body + assert card["elements"][1]["actions"] == [ + {"tag": "button", "text": {"tag": "plain_text", "content": "Open dashboard"}, "type": "default", + "url": "https://d.example/"}] + + +def test_sparkline_and_text_lines(): + assert sparkline([1, 2, 3, 4, 5, 6, 7, 8]) == "\u2581\u2582\u2583\u2584\u2585\u2586\u2587\u2588" + assert sparkline([2, 2, 2]) == "\u2584" * 3 and sparkline([]) == "" + assert len(sparkline(list(range(100)), width=24)) == 24 + assert text_lines(Chart("gn", [("a", [1, 2], [1.0, float("nan")])])) == ["gn a: \u2584 1..1 (last 1)"] + assert charts_mod.render_png(CHART) is None # matplotlib "missing" (fixture) + + +def test_render_png_with_matplotlib(monkeypatch): + pytest.importorskip("matplotlib") + monkeypatch.undo() + png = charts_mod.render_png(Chart("loss", [("a", [1e9, 1e9 + 60, 1e9 + 120], [0.9, 0.8, 0.7])], time_axis=True, + hlines=[("ref", 0.75)])) + assert png.startswith(b"\x89PNG") + assert charts_mod.render_png(Chart("empty", [("a", [1], [1.0])])) is None + + +def test_multipart_body(): + body, ctype = multipart({"image_type": "message"}, {"image": ("c.png", b"PNGDATA", "image/png")}) + boundary = ctype.split("boundary=")[1] + assert ctype.startswith("multipart/form-data; boundary=yeto-") + assert b'name="image_type"\r\n\r\nmessage\r\n' in body + assert b'filename="c.png"\r\nContent-Type: image/png\r\n\r\nPNGDATA\r\n' in body + assert body.endswith(f"--{boundary}--\r\n".encode()) diff --git a/tests/test_bot_monitor.py b/tests/test_bot_monitor.py new file mode 100644 index 00000000..4f349e2b --- /dev/null +++ b/tests/test_bot_monitor.py @@ -0,0 +1,410 @@ +"""yeto.bot rules and alert policy: severity routing, transitions, cooldown, re-notify, +ack, mute, quiet hours, merging one poll into one card.""" + +import math +import time + +import pytest + +from yeto.bot.config import Policy, Thresholds, parse_clock, parse_hours +from yeto.bot.monitor import Monitor +from yeto.bot.rules import P0, P1, P2, EventTracker, conditions, missed_streaks +from yeto.dashboard.alerts import DEFAULT_THRESHOLDS + +from bot_helpers import MERGE_COLS, STEP_COLS, T0, FakeClock, clone, eval_row, relayout, round_row, row, view + + +def mon(clock=None, policy=None, **th): + return Monitor(Thresholds(**th), policy or Policy(), clock=clock or FakeClock()) + + +def keys(notices): + return sorted(n.key for n in notices) + + +def add_row(v, isl=0, **kw): + rows = v["sft"]["islands"][isl]["rows"] + rows.append(row(rows[-1][0] + 1, **kw)) + return v + + +def noisy(v, isl=0): + """Give an island's loss some spread (0.69 / 0.71) so the sigma rule has a baseline.""" + for i, r in enumerate(v["sft"]["islands"][isl]["rows"]): + r[2] = 0.69 if i % 2 else 0.71 + return v + + +# -- rules: severity routing ---------------------------------------------------------- + +def sev_of(v, key, **th): + c = conditions(v, Thresholds(**th)).get(key) + return None if c is None else c.sev + + +def test_island_condition_severities(): + v = view(1) + assert conditions(v, Thresholds()) == {} + st = clone(v) + st["sft"]["islands"][0]["status"] = "stale" + assert sev_of(st, ("0", "stale")) == P0 + assert sev_of(add_row(clone(v), loss=None, gn=None), ("0", "loss_nan")) == P0 # dashboard null = NaN + assert sev_of(add_row(clone(v), loss=math.nan), ("0", "loss_nan")) == P0 + assert sev_of(add_row(clone(v), loss=None, gn=None, lr=None), ("0", "loss_nan")) is None # not logged yet + assert sev_of(add_row(clone(v), hbm=53.9), ("0", "hbm")) == P0 # refs sft_hbm_limit_gib + assert sev_of(add_row(clone(v), hbm=53.0), ("0", "hbm_warn")) == P1 + assert sev_of(add_row(clone(v), hbm=53.0), ("0", "hbm"), hbm_limit_gib=52.5) == P0 # env override + assert sev_of(add_row(clone(v), hbm=52.1), ("0", "hbm_warn")) == P1 # > 0.965 x 53.9 = 52.01 + assert sev_of(add_row(clone(v), hbm=52.1), ("0", "hbm_warn"), hbm_warn_frac=0.98) is None + assert sev_of(add_row(clone(v), hbm=60.0), ("0", "hbm_warn"), hbm_limit_gib=80.0) is None # follows the limit + assert sev_of(add_row(clone(v), hbm=78.0), ("0", "hbm_warn"), hbm_limit_gib=80.0) == P1 + assert sev_of(add_row(clone(v), hbm=50.5), ("0", "hbm_warn"), hbm_warn_gib=50.0) == P1 # explicit line + no_limit = add_row(clone(v), hbm=200.0) + del no_limit["sft"]["refs"]["sft_hbm_limit_gib"] + assert conditions(no_limit, Thresholds()) == {} # no limit known: neither P0 nor the derived warning + assert sev_of(no_limit, ("0", "hbm_warn"), hbm_warn_gib=100.0) == P1 + assert sev_of(add_row(clone(v), gn=1.5), ("0", "grad_norm")) == P1 + assert sev_of(add_row(clone(v), gn=1.5), ("0", "grad_norm"), grad_norm_max=2.0) is None + done = clone(v) + done["sft"]["islands"][0]["status"] = "done" + assert sev_of(done, ("0", "done")) == P2 + + +def test_rules_resolve_columns_by_name(): + v = stale(add_row(add_row(view(2), gn=2.0, hbm=53.95), isl=1, loss=None, gn=None), isl=0) + v["sft"]["now"] = v["sft"]["syncer"]["outer"][-1][0] + 4000 # syncer stuck (merges carry the round times) + v["sft"]["syncer"]["outer"] = [] + want = conditions(v, Thresholds()) + assert {("0", "stale"), ("0", "grad_norm"), ("0", "hbm"), ("1", "loss_nan"), ("syncer", "stuck")} <= set(want) + shuffled = relayout(v, cols=list(reversed(STEP_COLS)) + ["new_metric"], merge_cols=list(reversed(MERGE_COLS))) + assert shuffled["sft"]["islands"][0]["rows"][0][-2:] == [1, None] # step moved to the end, then new_metric + got = conditions(shuffled, Thresholds()) + assert {k: (c.sev, c.lines) for k, c in got.items()} == {k: (c.sev, c.lines) for k, c in want.items()} + no_hbm = relayout(v, cols=[c for c in STEP_COLS if c != "hbm_gib"]) + got = conditions(no_hbm, Thresholds()) + assert ("0", "hbm") not in got and ("0", "grad_norm") in got + bare = relayout(v, cols=["step", "ts"], merge_cols=["ts"]) # almost nothing named: rules go quiet, no crash + assert set(conditions(bare, Thresholds())) == {("0", "stale")} + + +def test_loss_spike_mean_plus_three_sigma(): + v = noisy(view(1)) + assert sev_of(add_row(clone(v), loss=0.72), ("0", "loss_spike")) is None # inside 3 sd (0.01) + c = conditions(add_row(clone(v), loss=0.75), Thresholds())[("0", "loss_spike")] + assert c.sev == P1 and "> mean 0.7 + 3 sd (0.01) of the previous 10 steps" in c.lines[0] + assert sev_of(add_row(view(1), loss=0.9), ("0", "loss_spike")) is None # flat history: no sd + + +def test_syncer_stuck_is_median_round_plus_two_grace(): + v = view(outer_steps=8, round_s=600.0) + last = v["sft"]["syncer"]["outer"][-1][0] + v["sft"]["now"] = last + 600 + 2 * 600 - 1 + assert sev_of(v, ("syncer", "stuck")) is None + v["sft"]["now"] = last + 1801 + c = conditions(v, Thresholds())[("syncer", "stuck")] + assert c.sev == P0 and "limit 30min" in c.lines[0] + assert sev_of(v, ("syncer", "stuck"), grace_s=900) is None + assert sev_of(v, ("syncer", "stuck"), stuck_round_s=1000) == P0 + v["sft"]["syncer"]["outer_step"] = 960 # finished + assert sev_of(v, ("syncer", "stuck")) is None + + +def test_rss_host_disk_and_missed_rounds(): + assert sev_of(view(rss=1.5e12), ("syncer", "rss")) == P1 # refs sft_syncer_rss_limit_bytes 1.4e12 + assert sev_of(view(rss=1.5e12), ("syncer", "rss"), rss_warn_bytes=1.6e12) is None + v = view(1) + v["sft"]["islands"][0]["host"] = {"node-a": [[T0, 80e9, 1.5e12]]} + assert sev_of(v, ("0", "host", "node-a")) == P1 + v["sft"]["syncer"]["disk"] = [[T0, 4e9, 100e9]] # hook: fires once the view carries disk samples + assert sev_of(v, ("syncer", "disk")) == P0 + rounds = [round_row(1), round_row(2, missed=[1]), round_row(3, missed=[1]), round_row(4, missed=[0, 1])] + assert missed_streaks(rounds) == {"1": 3, "0": 1} + assert Thresholds().missed_rounds == DEFAULT_THRESHOLDS["missed_consecutive_warn"] + c = conditions(view(rounds=rounds), Thresholds()) + assert c[("1", "missed")].sev == P1 and ("0", "missed") not in c + assert "isl1 not merged for 3 rounds" == c[("1", "missed")].title + + +def test_forwarded_dashboard_alerts_and_oom(): + crit = {"sev": 0, "rule": "recovery_required", "title": "island 1 RECOVERY_REQUIRED", "detail": "x", "island": "1"} + oom = {"sev": 0, "rule": "recovery_required", "title": "island 2 RECOVERY_REQUIRED", + "detail": "CUDA out of memory", "island": "2"} + warn = {"sev": 1, "rule": "kl", "title": "island 1 KL", "detail": "y", "island": "1"} + c = conditions(view(alerts=[crit, oom, warn]), Thresholds()) + assert c[("dash", "recovery_required", "1")].sev == P0 + assert c[("oom", "2")].sev == P0 and c[("oom", "2")].title.startswith("OOM") + assert len(c) == 2 + assert conditions(view(alerts=[crit]), Thresholds(forward_dashboard_alerts=False)) == {} + + +def test_event_tracker_syncer_events_and_eval(): + tr = EventTracker() + base = [[T0, "rejected", 1, "log", None]] + assert tr.events(view(events=base, eval_rows=[eval_row(32)]), seed=True) == [] + new = base + [[T0 + 5, "rejected", 1, "log", None], [T0 + 6, "rejected", 2, "log", None], + [T0 + 7, "checkpoint", None, "log", 12.5], [T0 + 8, "joined", 0, "log", None], + [T0 + 9, "late", 0, "log", None], [T0 + 10, "checkpoint_failed", None, "log", None]] + evs = {e.key: e for e in tr.events(view(events=new, eval_rows=[eval_row(32)]))} + assert evs[("syncer", "rejected")].sev == P1 and evs[("syncer", "rejected")].count == 2 + assert evs[("syncer", "rejected")].lines == ["learner 1 x1, learner 2 x1"] + assert evs[("syncer", "checkpoint")].sev == P2 and "write took 12s" in evs[("syncer", "checkpoint")].lines + assert evs[("syncer", "joined")].sev == P2 and evs[("syncer", "checkpoint_failed")].sev == P0 + assert ("syncer", "late") not in evs + assert tr.events(view(events=new, eval_rows=[eval_row(32)])) == [] # nothing new + better = tr.events(view(events=new, eval_rows=[eval_row(32), eval_row(64, loss=0.69)])) + assert [(e.key, e.sev) for e in better] == [(("eval", "new"), P2)] + assert "vs previous point -0.0100" in better[0].lines[0] + worse = tr.events(view(events=new, eval_rows=[eval_row(32), eval_row(64, loss=0.69), eval_row(96, loss=0.70)])) + assert [(e.key, e.sev) for e in worse] == [(("eval", "worse"), P1)] + assert worse[0].title == "eval step 96 worse than step 64" + assert worse[0].lines[1:] == ["vs base -0.0500", "vs ckpt100 -0.0100"] + renamed = tr.events(view(events=new, eval_rows=[eval_row(96, loss=0.70), eval_row(128, loss=0.71, ref="w0")])) + assert renamed[0].key == ("eval", "worse") and renamed[0].lines[1:] == ["vs base -0.0500", "vs w0 -0.0100"] + + +# -- policy: transitions, re-notify, cooldown --------------------------------------------- + +def stale(v, isl=1): + v = clone(v) + v["sft"]["islands"][isl]["status"] = "stale" + return v + + +def test_first_poll_seeds_then_p0_renotifies_until_acked(): + clock = FakeClock() + m = mon(clock) + st = stale(view()) + assert m.evaluate(st) == [] # already stale at startup: seeded, not paged yet + clock.t += 1799 + assert m.evaluate(st) == [] + clock.t += 2 + n = m.evaluate(st) + assert [(x.key, x.kind) for x in n] == [(("1", "stale"), "renotify")] + clock.t += 60 + assert m.evaluate(st) == [] + assert m.ack(None, "alice") == [("1", "stale")] + clock.t += 1800 + assert m.evaluate(st) == [] # acked: muted for 2 h + clock.t += 7200 + assert [x.kind for x in m.evaluate(st)] == ["renotify"] # ack expired, still firing + clock.t += 60 + assert [(x.key, x.kind) for x in m.evaluate(view())] == [(("1", "stale"), "resolve")] + + +def test_p0_fire_then_resolve_and_flap_suppression(): + clock = FakeClock() + m = mon(clock) + ok = view() + m.evaluate(ok) + clock.t += 60 + assert [(x.sev, x.kind) for x in m.evaluate(stale(ok))] == [(P0, "fire")] + clock.t += 60 + assert [x.kind for x in m.evaluate(ok)] == ["resolve"] + clock.t += 60 + assert m.evaluate(stale(ok)) == [] # flapped back inside the re-notify window + clock.t += 1800 + assert [x.kind for x in m.evaluate(stale(ok))] == ["fire"] # still there after 30 min: paged + + +def test_p1_cooldown_per_subject(): + clock = FakeClock() + m = mon(clock) + v = view(1) + hot = add_row(clone(v), gn=2.0) + m.evaluate(v) + clock.t += 60 + n = m.evaluate(hot) + assert [(x.key, x.sev) for x in n] == [(("0", "grad_norm"), P1)] + clock.t += 60 + assert m.evaluate(hot) == [] # P1 never re-notifies while it holds + clock.t += 60 + assert [x.kind for x in m.evaluate(v)] == ["resolve"] + clock.t += 60 + assert m.evaluate(hot) == [] # within the 1 h cooldown + clock.t += 60 + assert m.evaluate(v) == [] # suppressed one is not resolved either + clock.t += 3600 + assert [x.key for x in m.evaluate(hot)] == [("0", "grad_norm")] + + +def test_p1_events_accumulate_through_cooldown(): + clock = FakeClock() + m = mon(clock) + ev = [[T0 + 1, "rejected", 1, "log", None]] + m.evaluate(view(events=ev)) + ev += [[T0 + 100, "rejected", 1, "log", None]] + clock.t += 60 + assert [(x.key, x.count) for x in m.evaluate(view(events=ev))] == [(("syncer", "rejected"), 1)] + ev += [[T0 + 200, "rejected", 2, "log", None]] + clock.t += 60 + assert m.evaluate(view(events=ev)) == [] + ev += [[T0 + 300, "rejected", 2, "log", None]] + clock.t += 60 + assert m.evaluate(view(events=ev)) == [] + clock.t += 3600 + n = m.evaluate(view(events=ev)) # flushed with no new event once cooled down + assert [(x.key, x.count) for x in n] == [(("syncer", "rejected"), 2)] + + +def test_p2_events_have_no_cooldown(): + clock = FakeClock() + m = mon(clock) + ev = [] + m.evaluate(view(events=ev)) + for i in range(3): + ev = ev + [[T0 + i, "checkpoint", None, "log", 10.0]] + clock.t += 60 + assert [(x.sev, x.kind) for x in m.evaluate(view(events=ev))] == [(P2, "event")] + + +def test_mute_and_quiet_hours(): + clock = FakeClock() + m = mon(clock) + v = view(1) + m.evaluate(v) + m.mute(3600) + clock.t += 60 + both = stale(add_row(clone(v), gn=2.0), isl=0) + assert [x.sev for x in m.evaluate(both)] == [P0] # P0 pages through a plain mute + m.unmute() + m2 = mon(clock) + m2.evaluate(v) + m2.mute(3600, include_p0=True) + clock.t += 60 + assert m2.evaluate(both) == [] + lt = time.localtime(clock.t) + now_min = lt.tm_hour * 60 + lt.tm_min + quiet = Policy(quiet_hours=(now_min, (now_min + 120) % 1440 or 1440)) + m3 = mon(clock, quiet) + m3.evaluate(v) + clock.t += 60 + assert [x.sev for x in m3.evaluate(both)] == [P0] # only P0 inside quiet hours + assert any(not h[5] and h[1] == P1 for h in m3.history) # the P1 is recorded as suppressed + assert m3.quiet(clock.t) and not m3.quiet(clock.t + 3 * 3600) + + +def test_p1_held_back_by_mute_goes_out_after_the_mute(): + clock = FakeClock() + m = mon(clock) + v = view(1) + hot = add_row(clone(v), gn=2.0) + m.evaluate(v) + m.mute(3600) + clock.t += 60 + assert m.evaluate(hot) == [] # appeared while muted + for _ in range(5): + clock.t += 300 + assert m.evaluate(hot) == [] + suppressed = [h for h in m.history if h[2] == ("0", "grad_norm")] + assert [(h[4], h[5]) for h in suppressed] == [("fire", False)] # recorded once, not every poll + clock.t += 3600 # mute over, still firing + [n] = m.evaluate(hot) + assert (n.key, n.sev, n.kind) == (("0", "grad_norm"), P1, "fire") + assert "held back until now" in n.lines[-1] and n.lines[-1].startswith("active for 85min") + clock.t += 60 + assert m.evaluate(hot) == [] # delivered once + clock.t += 60 + assert [x.kind for x in m.evaluate(v)] == ["resolve"] + + +def test_p1_held_back_by_quiet_hours_or_cooldown_and_dropped_if_resolved(): + clock = FakeClock(time.mktime((2026, 11, 7, 7, 0, 0, 0, 0, -1))) + m = mon(clock, Policy(quiet_hours=(23 * 60, 8 * 60))) + v = view(1) + hot = add_row(clone(v), gn=2.0) + m.evaluate(v) + clock.t += 60 + assert m.evaluate(hot) == [] # 07:01, quiet + clock.t += 3600 # 08:01 + assert [x.key for x in m.evaluate(hot)] == [("0", "grad_norm")] + # resolved, then back inside the cooldown and still active once the cooldown ends + clock.t += 60 + m.evaluate(v) + clock.t += 60 + assert m.evaluate(hot) == [] + clock.t += 3600 + assert [x.key for x in m.evaluate(hot)] == [("0", "grad_norm")] + # held back, then cleared before the hold ended: never announced, no resolve either + m2 = mon(FakeClock(), Policy()) + m2.evaluate(v) + m2.mute(3600) + m2.clock.t += 60 + assert m2.evaluate(hot) == [] + m2.clock.t += 60 + assert m2.evaluate(v) == [] + m2.clock.t += 7200 + assert m2.evaluate(v) == [] and m2.unsent == {} + + +def test_p1_active_at_startup_stays_unannounced(): + clock = FakeClock() + m = mon(clock) + hot = add_row(view(1), gn=2.0) + assert m.evaluate(hot) == [] + clock.t += 7200 + assert m.evaluate(hot) == [] + + +def test_parse_quiet_hours_and_clock(): + assert parse_hours("23-8") == (23 * 60, 8 * 60) + assert parse_hours("22:30-06:15") == (22 * 60 + 30, 6 * 60 + 15) + assert parse_hours("") is None + with pytest.raises(ValueError): + parse_hours("night") + assert parse_clock("09:00") == 540 and parse_clock("7") == 420 + m = mon(policy=Policy(quiet_hours=(23 * 60, 8 * 60))) + day = time.mktime((2026, 11, 7, 12, 0, 0, 0, 0, -1)) + night = time.mktime((2026, 11, 7, 2, 0, 0, 0, 0, -1)) + late = time.mktime((2026, 11, 7, 23, 30, 0, 0, 0, -1)) + assert not m.quiet(day) and m.quiet(night) and m.quiet(late) + + +def test_dashboard_down_after_n_failures_and_recovery(): + m = mon(dashboard_down_after=3) + m.evaluate(view()) + assert m.fetch_failed(OSError("refused")) == [] + assert m.fetch_failed(OSError("refused")) == [] + n = m.fetch_failed(OSError("refused")) + assert [(x.sev, x.title) for x in n] == [(P1, "dashboard unreachable")] + assert m.fetch_failed(OSError("refused")) == [] + assert [(x.kind, x.title) for x in m.evaluate(view())] == [("resolve", "dashboard reachable again")] + + +# -- one card per poll ---------------------------------------------------------------------- + +def test_several_alerts_merge_into_one_card(): + clock = FakeClock() + m = mon(clock) + v = view(2, events=[]) + m.evaluate(v) + clock.t += 60 + bad = stale(add_row(add_row(clone(v), isl=1, gn=2.0), isl=1, gn=2.5), isl=0) + bad["sft"]["syncer"]["events"] = [[T0 + 50, "checkpoint", None, "log", 9.0]] + n = m.evaluate(bad) + assert keys(n) == [("0", "stale"), ("1", "grad_norm"), ("syncer", "checkpoint")] + msg = m.card(n, link="https://dash.example/") + assert msg.level == "p0" and msg.mention + assert msg.title == "P0 isl0 stale (process dead or hung?) (+2 more)" + text = msg.plain() + assert text.index("P0 isl0 stale") < text.index("P1 isl1 grad_norm high") < text.index("P2 syncer checkpoint") + assert msg.buttons == [("Ack 2h", "/ack c1"), ("Details", "/details c1"), + ("Open dashboard", "https://dash.example/")] + assert [x.key for x in m.card_notices("c1")] == [x.key for x in n] + p2 = m.card([n[2]]) + assert p2.level == "p2" and not p2.mention and p2.buttons == [("Details", "/details c2")] + clock.t += 60 + res = m.card(m.evaluate(v)) + assert res.level == "ok" and res.title.startswith("resolved: ") and "(+1 more)" in res.title + assert m.card([]) is None + + +def test_alerts_message_lists_active_ack_and_history(): + clock = FakeClock() + m = mon(clock) + m.evaluate(view()) + clock.t += 60 + m.evaluate(stale(view())) + m.ack(None, "bob") + text = m.alerts_message().plain() + assert "Active (1)" in text and "isl1 stale" in text and "acked by bob" in text + assert "P0 1 | P1 0 | P2 0" in text diff --git a/yeto/bot/__init__.py b/yeto/bot/__init__.py new file mode 100644 index 00000000..1114ca4b --- /dev/null +++ b/yeto/bot/__init__.py @@ -0,0 +1,7 @@ +"""Training report / alert bot for Telegram and Lark (Feishu). + +A read-only client of the yeto dashboard (``yeto dashboard serve``): it polls +``/api/view`` over HTTP, answers chat commands and pushes alerts on state +transitions. It never connects to training nodes. Run ``python -m yeto.bot``; +configuration is environment variables only (``yeto.bot.config``). +""" diff --git a/yeto/bot/__main__.py b/yeto/bot/__main__.py new file mode 100644 index 00000000..46476297 --- /dev/null +++ b/yeto/bot/__main__.py @@ -0,0 +1,60 @@ +"""``python -m yeto.bot``: run the bot, or do one thing and exit.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time + +from .app import Bot, build_channels +from .channels import Command, log +from .config import Config +from .http import DashboardClient, request_json + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(prog="python -m yeto.bot", description=__doc__) + one = ap.add_mutually_exclusive_group() + one.add_argument("--once", action="store_true", + help="run one monitor pass (sends the progress card) and exit") + one.add_argument("--command", metavar="TEXT", + help="answer one command (e.g. '/status') into each channel's report chat and exit") + one.add_argument("--request-decision", metavar="JSON", + help='post a decision card and exit: {"title", "evidence": [...], "options": [...], ' + '"suggested"}; the running bot records the press in the audit log') + args = ap.parse_args(argv) + cfg = Config.from_env() + oneshot = bool(args.once or args.command or args.request_decision) + channels = build_channels(cfg, receive=not oneshot) + client = DashboardClient(cfg.dashboard_url, timeout=cfg.dashboard_timeout, http=request_json) + bot = Bot(cfg, client, channels) + if args.command: + for ch in channels: + if ch.report_chat: + # whoever can run this already holds the bot's credentials: trust it + bot.handle(ch, Command(ch.name, ch.report_chat, args.command, "cli", trusted=True)) + else: + log(f"{ch.name}: no report chat configured, skipping") + return 0 + if args.request_decision: + spec = json.loads(args.request_decision) + bot.request_decision(spec["title"], spec.get("evidence") or [], spec["options"], spec.get("suggested", "")) + return 0 + if args.once: + bot.tick() + return 0 + if not channels: + log("no channel configured (TELEGRAM_BOT_TOKEN / LARK_APP_ID+LARK_APP_SECRET / LARK_WEBHOOK_URL): idling") + while True: + time.sleep(3600) + log(f"started: dashboard {cfg.dashboard_url}, poll every {cfg.poll_interval:g}s, progress every " + f"{cfg.policy.progress_interval:g}s, channels " + + ", ".join(f"{c.name}({'rx' if c.can_receive else 'tx'}, chat {c.report_chat or 'unset'})" for c in channels) + + (" [dry run]" if cfg.dry_run else "")) + bot.run() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/yeto/bot/app.py b/yeto/bot/app.py new file mode 100644 index 00000000..2b17785b --- /dev/null +++ b/yeto/bot/app.py @@ -0,0 +1,419 @@ +"""Bot core: command routing with chat / user authorization, the alert monitor, +the scheduled progress card and daily digest, decision requests. + +Phase-2 hook: text that is not a slash command goes to ``Bot.nl`` (a callable +``(text, ReadOnlyTools) -> Message``) when one is installed, e.g. an LLM that +answers questions with the read-only dashboard tools; without it such text is +ignored. +""" + +from __future__ import annotations + +import json +import os +import re +import threading +import time +from collections import Counter +from dataclasses import replace +from typing import Callable + +from . import lark_ws +from .channels import Channel, Command, LarkAppChannel, LarkWebhookChannel, TelegramChannel, dry_run_http, log +from .config import Config +from .decisions import DecisionLog +from .http import DashboardClient, HttpFn, request_json +from .monitor import Monitor +from .render import (Message, b, ckpt_message, details_message, digest_message, eval_message, fdur, help_message, + island_message, merge_message, res_message, status_message) +from .rules import P0, P1, P2, SEV_NAME, hbm_limit, hbm_warn, rss_warn + +VIEW_COMMANDS = ("/status", "/island", "/eval", "/merge", "/res", "/ckpt", "/details") +LOCAL_COMMANDS = ("/help", "/start", "/alerts", "/ack", "/mute", "/unmute", "/decide") +# commands (and their buttons) that change state other people rely on: refused +# while the channel has no user allowlist, since then anyone in the report chat +# could run them. Value: why, for the refusal message. +ALERT_CONTROL_COMMANDS = {"/ack": "changes alerting for everyone", "/mute": "changes alerting for everyone", + "/unmute": "changes alerting for everyone", + "/decide": "records a decision the operator will act on"} +_DUR = re.compile(r"^(\d+(?:\.\d+)?)([smhd]?)$") + + +def parse_command(text: str) -> tuple[str, str]: + """``"/Island@yeto_bot 2"`` -> ``("/island", "2")``.""" + head, _, rest = text.strip().partition(" ") + return head.split("@", 1)[0].lower(), rest.strip() + + +def parse_duration(text: str) -> float | None: + m = _DUR.match(text.strip().lower()) + if not m: + return None + return float(m.group(1)) * {"": 60, "s": 1, "m": 60, "h": 3600, "d": 86400}[m.group(2)] + + +class ReadOnlyTools: + """What a phase-2 natural-language handler may call: dashboard reads and alert state, nothing else.""" + + def __init__(self, client: DashboardClient, monitor: Monitor): + self._client, self._monitor = client, monitor + + def view(self) -> dict: + return self._client.view() + + def bad_rounds(self) -> list[dict]: + return self._client.bad_rounds() + + def active_alerts(self) -> list[dict]: + return [{"key": list(k), "sev": SEV_NAME[c.sev], "title": c.title, "lines": c.lines} + for k, c in self._monitor.active.items()] + + +class Bot: + def __init__(self, cfg: Config, client: DashboardClient, channels: list[Channel], *, + clock: Callable[[], float] = time.time, + nl: Callable[[str, ReadOnlyTools], Message] | None = None): + self.cfg = cfg + self.client = client + self.channels = channels + self.clock = clock + self.nl = nl + self.monitor = Monitor(cfg.th, cfg.policy, clock=clock) + self.decisions = DecisionLog(cfg.audit_log, clock=clock) + self.last_progress = float("-inf") + self.digest_day: tuple | None = None + self.lock = threading.RLock() + self._load_state() + + # -- registered chats (optional state file) ------------------------------------- + def _load_state(self) -> None: + if not self.cfg.state_file or not os.path.exists(self.cfg.state_file): + return + try: + with open(self.cfg.state_file) as fh: + chats = json.load(fh).get("chats") or {} + except (OSError, ValueError) as exc: + log(f"state file unreadable, ignoring: {exc}") + return + for ch in self.channels: + if not ch.configured_chat and chats.get(ch.name): + ch.registered_chat = str(chats[ch.name]) + + def _save_state(self) -> None: + if not self.cfg.state_file: + return + chats = {ch.name: ch.registered_chat for ch in self.channels if ch.registered_chat} + tmp = self.cfg.state_file + ".tmp" + with open(tmp, "w") as fh: + json.dump({"chats": chats}, fh) + os.replace(tmp, self.cfg.state_file) + + # -- commands ------------------------------------------------------------------- + def alerts_line(self) -> str: + c = self.monitor.active_counts() + line = f"{b('Alerts')} active P0 {c[P0]} | P1 {c[P1]}" + now = self.clock() + if now < self.monitor.mute_until: + line += f" | muted for {fdur(self.monitor.mute_until - now)}" + return line + + def progress(self, view: dict) -> Message: + return status_message(view, alerts_line=self.alerts_line()) + + def reply_for(self, name: str, arg: str, cmd: Command) -> Message: + if name in ("/help", "/start"): + return help_message() + if name == "/alerts": + return self.monitor.alerts_message() + if name == "/ack": + return self._ack(arg, cmd) + if name == "/mute": + parts = arg.split() + secs = parse_duration(parts[0]) if parts else None + if secs is None: + return Message("usage", ["/mute <30m|2h|1d> [all] (all: P0 too)"], "warn") + everything = len(parts) > 1 and parts[1].lower() == "all" + self.monitor.mute(secs, include_p0=everything) + log(f"muted for {secs:.0f}s by {cmd.user or '?'} (P0 {'too' if everything else 'still pages'})") + return Message("muted", [f"{'all alerts' if everything else 'P1/P2'} muted for {fdur(secs)}" + f" by {cmd.user or '?'}" + ("" if everything else "; P0 still pages")], "ok") + if name == "/unmute": + self.monitor.unmute() + return Message("unmuted", [f"alerts resumed by {cmd.user or '?'}"], "ok") + if name == "/decide": + req_id, _, option = arg.partition(" ") + return self.decisions.decide(req_id, option.strip(), user=cmd.user, channel=cmd.channel, chat=cmd.chat, + user_id=cmd.user_id) + if name not in VIEW_COMMANDS: + return help_message(f"unknown command {name}") + view = self.client.view() + th = self.cfg.th + if name == "/status": + return self.progress(view) + if name == "/island": + return island_message(view, arg, hbm_limit=hbm_limit(view, th)) + if name == "/eval": + return eval_message(view) + if name == "/merge": + if not view.get("sft") and not view.get("rounds_syncer"): + try: + view = dict(view, rounds_syncer=self.client.bad_rounds()) + except Exception as exc: + log(f"/api/rounds failed: {exc}") + return merge_message(view) + if name == "/res": + return res_message(view, hbm_limit=hbm_limit(view, th), hbm_warn=hbm_warn(view, th), + rss_warn=rss_warn(view, th), host_min=th.host_avail_min_bytes) + if name == "/ckpt": + return ckpt_message(view) + notices = self.monitor.card_notices(arg) # /details cN + if notices is None: + return Message("unknown card", [f"no alert card {arg!r} (cards are kept in memory)"], "warn") + return details_message(view, notices, ref=arg) + + def _ack(self, arg: str, cmd: Command) -> Message: + ref = arg.strip().lower() + if ref in ("", "all"): + keys = self.monitor.ack(None, cmd.user) + titles = [self.monitor.active[k].title for k in keys] + else: + notices = self.monitor.card_notices(ref) + if notices is None: + return Message("unknown card", [f"no alert card {arg!r}"], "warn") + live = [n for n in notices if n.kind != "resolve"] + keys = self.monitor.ack([n.key for n in live], cmd.user) + titles = list(dict.fromkeys(n.title for n in live)) + if not keys: + return Message("nothing to ack", ["no active alert subjects"], "info") + log(f"acked {len(keys)} subject(s) by {cmd.user or '?'}: {titles}") + return Message(f"acked {len(keys)} alert(s) for {fdur(self.cfg.policy.ack_mute_s)}", + [f"by {cmd.user or '?'}", *titles], "ok") + + def send(self, ch: Channel, msg: Message, chat: str | None = None) -> bool: + if self.cfg.title_prefix: + msg = replace(msg, title=self.cfg.title_prefix + msg.title) + return ch.send(msg, chat) + + def _register(self, ch: Channel, cmd: Command) -> Message | None: + """``/start`` without a configured report chat: the chat becomes the report chat, + but only for an allowlisted user. Without an allowlist anyone who can reach the + bot could claim the run's alerts, so registration is refused.""" + who = f"chat {cmd.chat} user {cmd.user or '?'} ({cmd.user_id or 'no id'})" + if not ch.allowed: + log(f"{ch.name}: refusing /start from {who}: configure the report chat id or a user allowlist") + reply = Message("registration disabled", [ + "This bot has no report chat and no user allowlist, so it does not let a chat register itself.", + f"Set {ch.chat_env or 'the report chat id'} or {ch.allow_env or 'a user allowlist'} and restart."], + "warn") + ch.answer(cmd, "Not allowed") + self.send(ch, reply, cmd.chat) + return reply + if not ch.user_allowed(cmd): + log(f"{ch.name}: refusing /start from {who}: not in the allowlist") + ch.answer(cmd, "Not allowed") + return None + ch.registered_chat = cmd.chat + log(f"{ch.name}: report chat registered: {cmd.chat} (by {cmd.user or '?'} {cmd.user_id})") + self._save_state() + reply = help_message("Registered this chat for training alerts.") + ch.answer(cmd, "ok") + self.send(ch, reply, cmd.chat) + return reply + + def handle(self, ch: Channel, cmd: Command) -> Message | None: + """Authorize and answer one command; returns the reply that was sent (or None).""" + slash = cmd.text.strip().startswith("/") + name, arg = parse_command(cmd.text) if slash else ("", cmd.text.strip()) + target = ch.report_chat + if not target: + if not slash: + return None + if name == "/start": + return self._register(ch, cmd) + reply = Message("not registered", ["send /start in the chat that should receive alerts"], "info") + ch.answer(cmd, "ok") + self.send(ch, reply, cmd.chat) + return reply + # run data stays in the report chat; local CLI commands skip the user allowlist + if str(cmd.chat) != str(target) or not (cmd.trusted or ch.user_allowed(cmd)): + log(f"{ch.name}: ignoring {name or 'text'} from chat {cmd.chat} user {cmd.user or '?'}") + ch.answer(cmd, "Not allowed") + return None + if not slash: + if self.nl is None: + return None + try: + reply = self.nl(arg, ReadOnlyTools(self.client, self.monitor)) + except Exception as exc: + reply = Message("question failed", [str(exc)[:500]], "warn") + self.send(ch, reply, cmd.chat) + return reply + if name in ALERT_CONTROL_COMMANDS and not ch.allowed and not cmd.trusted: + log(f"{ch.name}: refusing {name} from user {cmd.user or '?'} ({cmd.user_id or 'no id'}):" + f" no user allowlist configured") + ch.answer(cmd, "Needs a user allowlist") + reply = Message("not allowed", [ + f"{name} {ALERT_CONTROL_COMMANDS[name]}, so it needs a user allowlist" + f" ({ch.allow_env or 'the channel allowlist'}); read-only commands still work."], "warn") + self.send(ch, reply, cmd.chat) + return reply + ch.answer(cmd, f"running {name}") + try: + with self.lock: # monitor state is shared with the monitor thread + reply = self.reply_for(name, arg, cmd) + except Exception as exc: + reply = Message("command failed", [f"{name}: {exc}"[:500]], "warn") + self.send(ch, reply, cmd.chat) + return reply + + # -- broadcasts ----------------------------------------------------------------- + def broadcast(self, msgs: list[Message]) -> None: + for msg in msgs: + for ch in self.channels: + if ch.report_chat: + self.send(ch, msg) + + def request_decision(self, title: str, evidence: list[str], options: list[str], suggested: str = "") -> Message: + """Post a decision card to every report chat; presses are only recorded.""" + msg = self.decisions.request(title, evidence, options, suggested).message() + self.broadcast([msg]) + return msg + + def _digest_due(self, now: float) -> bool: + at = self.cfg.policy.digest_at + if at is None: + return False + t = time.localtime(now) + day, minute = (t.tm_year, t.tm_yday), t.tm_hour * 60 + t.tm_min + if self.digest_day is None: # started after today's slot: first digest tomorrow + self.digest_day = day if minute >= at else (t.tm_year, t.tm_yday - 1) + if minute >= at and day != self.digest_day: + self.digest_day = day + return True + return False + + def digest(self, view: dict) -> Message: + hist = [h for h in self.monitor.history_since(self.clock() - 86400) if h[5] and h[4] != "resolve"] + counts = {SEV_NAME[s]: sum(1 for h in hist if h[1] == s) for s in (P0, P1, P2)} + top = [f"{n}x {t}" for t, n in Counter(h[3] for h in hist if h[1] <= P1).most_common(5)] + return digest_message(view, alert_counts=counts, top_alerts=top) + + def tick(self) -> list[Message]: + """One monitor pass (+ the progress card / digest when due); returns what was broadcast.""" + with self.lock: + msgs = self._tick() + self.broadcast(msgs) + return msgs + + def _tick(self) -> list[Message]: + try: + view = self.client.view() + except Exception as exc: + log(f"dashboard fetch failed: {exc}") + card = self.monitor.card(self.monitor.fetch_failed(exc), link=self.cfg.dashboard_link) + return [card] if card else [] + card = self.monitor.card(self.monitor.evaluate(view), link=self.cfg.dashboard_link) + msgs = [card] if card else [] + now = self.clock() + every = self.cfg.policy.progress_interval + if every > 0 and now - self.last_progress >= every: + self.last_progress = now + msgs.append(self.progress(view)) + if self._digest_due(now): + msgs.append(self.digest(view)) + return msgs + + # -- threads -------------------------------------------------------------------- + def _monitor_loop(self, stop: threading.Event) -> None: + while not stop.is_set(): + try: + self.tick() + except Exception as exc: + log(f"monitor pass failed: {exc}") + stop.wait(self.cfg.poll_interval) + + def _poll_loop(self, ch: Channel, stop: threading.Event) -> None: + while not stop.is_set(): + try: + cmds = ch.poll_commands(25) + except Exception as exc: + log(f"{ch.name}: poll failed: {exc}") + stop.wait(5) + continue + for cmd in cmds: + try: + self.handle(ch, cmd) + except Exception as exc: + log(f"{ch.name}: command {cmd.text!r} failed: {exc}") + + def allowlist_warnings(self) -> list[str]: + return [f"{ch.name}: no user allowlist ({ch.allow_env or 'allowlist'} unset): anyone in the report chat can" + f" run read-only commands; {', '.join(ALERT_CONTROL_COMMANDS)} are refused" + for ch in self.channels if ch.can_receive and not ch.allowed] + + def run(self, stop: threading.Event | None = None) -> None: + stop = stop or threading.Event() + for line in self.allowlist_warnings(): + log("warning: " + line) + threads = [threading.Thread(target=self._monitor_loop, args=(stop,), name="monitor", daemon=True)] + threads += [threading.Thread(target=self._poll_loop, args=(ch, stop), name=f"poll-{ch.name}", daemon=True) + for ch in self.channels if ch.can_receive] + for t in threads: + t.start() + try: + while not stop.is_set(): + stop.wait(3600) + except KeyboardInterrupt: + stop.set() + + +def build_channels(cfg: Config, http: HttpFn | None = None, *, receive: bool = True) -> list[Channel]: + """Every channel whose configuration is complete; the rest are skipped with one log line. + ``receive=False`` (one-shot CLI runs, dry run) never opens a long connection or polls.""" + http = http or (dry_run_http() if cfg.dry_run else request_json) + receive = receive and not cfg.dry_run + out: list[Channel] = [] + if cfg.telegram_token: + out.append(TelegramChannel(cfg.telegram_token, cfg.telegram_chat_id, api_base=cfg.telegram_api_base, + http=http, owners=cfg.telegram_owners, allowed=cfg.telegram_allowed)) + out[-1].can_receive = receive + else: + log("telegram: TELEGRAM_BOT_TOKEN not set, skipping") + lark = build_lark(cfg, http, receive) + if lark: + out.append(lark) + return out + + +def build_lark(cfg: Config, http: HttpFn, receive: bool = True) -> Channel | None: + has_app = bool(cfg.lark_app_id and cfg.lark_app_secret) + hook = (LarkWebhookChannel(cfg.lark_webhook_url, cfg.lark_webhook_secret, http=http, owners=cfg.lark_owners) + if cfg.lark_webhook_url else None) + mode = cfg.lark_mode + if mode == "webhook" or (mode == "auto" and not has_app): + if hook is None: + log("lark: neither LARK_APP_ID/LARK_APP_SECRET nor LARK_WEBHOOK_URL set, skipping") + else: + log("lark: group webhook (send-only)" + (" with signature" if cfg.lark_webhook_secret else "")) + return hook + if not has_app: + log(f"lark: LARK_MODE={mode} needs LARK_APP_ID and LARK_APP_SECRET, skipping") + return None + want_ws = mode in ("auto", "ws") and receive + ws_ok = want_ws and lark_ws.available() + if want_ws and not ws_ok: + log("lark: lark-oapi not installed (pip install lark-oapi): no command receiving") + if mode == "auto" and not cfg.lark_chat_id and hook is not None: + log("lark: no LARK_CHAT_ID for the send-only app bot, falling back to the group webhook") + return hook + ch = LarkAppChannel(cfg.lark_app_id, cfg.lark_app_secret, cfg.lark_chat_id, domain=cfg.lark_domain, http=http, + refresh_margin=cfg.lark_token_refresh_margin, receive=ws_ok, + buttons=cfg.lark_card_buttons, owners=cfg.lark_owners, allowed=cfg.lark_allowed) + if cfg.dry_run: # render the cards as the listening bot would + ch.buttons = cfg.lark_card_buttons + if ws_ok: + lark_ws.start(ch) + log("lark: app bot with long-connection events") + else: + log("lark: app bot, send-only" + ("" if cfg.lark_chat_id else " (set LARK_CHAT_ID to receive alerts)")) + return ch diff --git a/yeto/bot/channels.py b/yeto/bot/channels.py new file mode 100644 index 00000000..9642d43c --- /dev/null +++ b/yeto/bot/channels.py @@ -0,0 +1,482 @@ +"""Chat channels: Telegram (Bot API long polling), Lark app bot, Lark group webhook. + +Every channel implements the same small interface: + +``send(msg, chat=None) -> bool`` + post a :class:`~yeto.bot.render.Message` to ``chat`` or the report chat: + buttons (command buttons only where the press comes back to the bot, URL + buttons everywhere), charts (PNG when matplotlib and an upload path exist, + else a text sparkline) and an @ of the owners when ``msg.mention``; +``poll_commands(timeout) -> list[Command]`` + block up to ``timeout`` seconds for incoming commands and button presses + (empty for send-only channels); +``answer(cmd, text)`` + acknowledge a button press. + +All network calls go through an injected ``http`` callable (see ``yeto.bot.http``). +""" + +from __future__ import annotations + +import base64 +import hashlib +import hmac +import json +import queue +import re +import sys +import time +from collections import deque +from dataclasses import dataclass, replace +from typing import Any, Callable + +from . import charts as charts_mod +from .http import HttpError, HttpFn, redact_url +from .render import BOLD_CLOSE, BOLD_OPEN, LEVEL_ICON, Message, is_url, strip_bold + +TELEGRAM_MAX = 4000 # Bot API limit is 4096 characters +TELEGRAM_CAPTION_MAX = 1000 +LARK_MAX = 8000 # keep cards well under the 30 KB request limit +LARK_TOKEN_EXPIRED = (99991661, 99991663, 99991668) +LARK_TEMPLATE = {"p0": "red", "p1": "orange", "p2": "grey", "warn": "orange", "ok": "green", "info": "blue", + "report": "indigo", "decision": "purple"} + + +@dataclass +class Command: + channel: str + chat: str + text: str + user: str = "" # display handle for replies and logs only, never for authorization + button: bool = False + ref: Any = None # channel-specific handle (Telegram callback query id) + user_id: str = "" # stable id the allowlist is checked against (Telegram from.id, Lark open_id) + trusted: bool = False # issued locally (``python -m yeto.bot --command``), not by a chat user + + +def log(text: str) -> None: + print(f"[yeto-bot] {text}", file=sys.stderr, flush=True) + + +def truncate(text: str, limit: int) -> str: + if len(text) <= limit: + return text + return text[: limit - 20].rstrip() + "\n... (truncated)" + + +def with_chart_text(msg: Message, failed: list) -> Message: + """Inline sparklines for charts that could not be sent as images.""" + extra = [line for c in failed for line in charts_mod.text_lines(c)] + return replace(msg, lines=msg.lines + ([""] + extra if extra else [])) if extra else msg + + +class Channel: + name = "channel" + can_receive = False + chat_env = "" # environment variables that configure the report chat / the allowlist (for messages) + allow_env = "" + + def __init__(self, configured_chat: str = "", *, owners: list[str] | None = None, + allowed: list[str] | None = None): + self.configured_chat = str(configured_chat or "") + self.registered_chat = "" + self.owners = list(owners or []) # @ on P0 + self.allowed = list(allowed or []) # user ids who may command / press; empty: anyone in the chat + + @property + def report_chat(self) -> str: + return self.configured_chat or self.registered_chat + + def user_allowed(self, cmd: Command) -> bool: + """Only the platform's stable user id counts: display names are chosen by the + sender (Telegram first names are free text, usernames can be released and + taken by someone else).""" + return not self.allowed or (bool(cmd.user_id) and cmd.user_id in self.allowed) + + def send(self, msg: Message, chat: str | None = None) -> bool: + raise NotImplementedError + + def poll_commands(self, timeout: float = 25.0) -> list[Command]: + time.sleep(timeout) + return [] + + def answer(self, cmd: Command, text: str) -> None: + pass + + +# -- Telegram --------------------------------------------------------------------- + +_TG_SPECIAL = re.compile(r"([_*`\[])") + + +def tg_markdown(msg: Message, owners: list[str] | None = None) -> str: + """Legacy Telegram Markdown: bold spans become ``*x*`` (special characters inside + are replaced, since entities cannot contain escapes); everything else is escaped + with a backslash.""" + def conv(text: str) -> str: + out, pos = [], 0 + for m in re.finditer(f"{BOLD_OPEN}(.*?){BOLD_CLOSE}", text, re.S): + out.append(_TG_SPECIAL.sub(r"\\\1", text[pos:m.start()])) + inner = m.group(1).replace("*", "x").replace("_", "-").replace("`", "'").replace("[", "(") + out.append(f"*{inner}*") + pos = m.end() + out.append(_TG_SPECIAL.sub(r"\\\1", strip_bold(text[pos:]))) + return "".join(out) + + icon = LEVEL_ICON.get(msg.level, "") + head = conv(f"{icon} {BOLD_OPEN}{msg.title}{BOLD_CLOSE}".strip()) + lines = [head, *(conv(line) for line in msg.lines)] + if msg.mention and owners: + lines.append(conv("cc " + " ".join(f"@{o}" for o in owners))) + return truncate("\n".join(lines), TELEGRAM_MAX) + + +def tg_plain(msg: Message, owners: list[str] | None = None) -> str: + icon = LEVEL_ICON.get(msg.level, "") + text = f"{icon} {msg.plain()}".strip() + if msg.mention and owners: + text += "\ncc " + " ".join(f"@{o}" for o in owners) + return truncate(text, TELEGRAM_MAX) + + +def tg_keyboard(msg: Message) -> dict | None: + if not msg.buttons: + return None + row = [{"text": label, "url": target} if is_url(target) else {"text": label, "callback_data": target[:64]} + for label, target in msg.buttons] + return {"inline_keyboard": [row[i:i + 4] for i in range(0, len(row), 4)]} + + +def _tg_user(frm: dict | None) -> tuple[str, str]: + """(display name, numeric user id) of a Telegram ``from`` object; only the id is trusted.""" + frm = frm or {} + uid = str(frm.get("id") or "") + return str(frm.get("username") or frm.get("first_name") or uid), uid + + +class TelegramChannel(Channel): + name = "telegram" + can_receive = True + chat_env, allow_env = "TELEGRAM_CHAT_ID", "TELEGRAM_ALLOWED_USERS" + + def __init__(self, token: str, chat: str = "", *, api_base: str = "https://api.telegram.org", + http: HttpFn, owners: list[str] | None = None, allowed: list[str] | None = None): + super().__init__(chat, owners=owners, allowed=allowed) + self.token = token + self.api_base = api_base.rstrip("/") + self.http = http + self.offset = 0 + + def api(self, method: str, files: dict | None = None, **params: Any) -> Any: + url = f"{self.api_base}/bot{self.token}/{method}" + if files: + return self.http("POST", url, body={k: str(v) for k, v in params.items()}, files=files, timeout=60) + return self.http("POST", url, body=params, timeout=35) + + def send(self, msg: Message, chat: str | None = None) -> bool: + target = chat or self.report_chat + if not target: + return False + pngs = [(c, charts_mod.render_png(c)) for c in msg.charts] + msg = with_chart_text(msg, [c for c, png in pngs if png is None]) + extra = {} + kb = tg_keyboard(msg) + if kb: + extra["reply_markup"] = kb + ok = False + try: + self.api("sendMessage", chat_id=target, text=tg_markdown(msg, self.owners), parse_mode="Markdown", **extra) + ok = True + except Exception as exc: # bad markup: retry as plain text + log(f"telegram send failed, retrying as plain text: {exc}") + try: + self.api("sendMessage", chat_id=target, text=tg_plain(msg, self.owners), **extra) + ok = True + except Exception as exc2: + log(f"telegram plain send failed: {exc2}") + for chart, png in pngs: + if png is None: + continue + try: + self.api("sendPhoto", files={"photo": ("chart.png", png, "image/png")}, chat_id=target, + caption=chart.title[:TELEGRAM_CAPTION_MAX]) + except Exception as exc: + log(f"telegram sendPhoto failed: {exc}") + return ok + + def poll_commands(self, timeout: float = 25.0) -> list[Command]: + try: + res = self.api("getUpdates", offset=self.offset, timeout=int(timeout), + allowed_updates=["message", "callback_query"]) + except Exception as exc: + log(f"telegram getUpdates failed: {exc}") + time.sleep(5) + return [] + return self.parse_updates(res) + + def parse_updates(self, res: Any) -> list[Command]: + out = [] + for upd in (res or {}).get("result") or []: + self.offset = max(self.offset, int(upd.get("update_id", 0)) + 1) + cq = upd.get("callback_query") + if cq: + chat = ((cq.get("message") or {}).get("chat") or {}).get("id") + if chat is not None: + name, uid = _tg_user(cq.get("from")) + out.append(Command(self.name, str(chat), cq.get("data") or "", name, True, cq.get("id"), uid)) + continue + msg = upd.get("message") or {} + text = (msg.get("text") or "").strip() + chat = (msg.get("chat") or {}).get("id") + if chat is None or not text: + continue + name, uid = _tg_user(msg.get("from")) + out.append(Command(self.name, str(chat), text, name, user_id=uid)) + return out + + def answer(self, cmd: Command, text: str) -> None: + if cmd.button and cmd.ref: + try: + self.api("answerCallbackQuery", callback_query_id=cmd.ref, text=text[:190]) + except Exception as exc: + log(f"telegram answerCallbackQuery failed: {exc}") + + +# -- Lark --------------------------------------------------------------------------- + +_LARK_ESC = {"&": "&", "<": "<", ">": ">", "*": "*", "_": "_", "~": "~", + "`": "`", "[": "[", "]": "]"} + + +def lark_md(text: str) -> str: + esc = "".join(_LARK_ESC.get(ch, ch) for ch in text) + return esc.replace(BOLD_OPEN, "**").replace(BOLD_CLOSE, "**") + + +def lark_card(msg: Message, *, buttons: bool, owners: list[str] | None = None, + images: list[tuple[str, str]] | None = None) -> dict: + """Interactive card (message card JSON 1.0): coloured header by level, one + lark_md block, chart images, then buttons (command buttons only when + ``buttons``; URL buttons always).""" + icon = LEVEL_ICON.get(msg.level, "") + body = truncate("\n".join(lark_md(line) for line in msg.lines), LARK_MAX) or " " + if msg.mention and owners: + body += "\n" + " ".join(f"" for o in owners) + elements: list[dict] = [{"tag": "div", "text": {"tag": "lark_md", "content": body}}] + for title, key in images or []: + elements.append({"tag": "img", "img_key": key, "alt": {"tag": "plain_text", "content": title}}) + actions = [] + for label, target in msg.buttons: + if is_url(target): + actions.append({"tag": "button", "text": {"tag": "plain_text", "content": label}, "type": "default", + "url": target}) + elif buttons: + actions.append({"tag": "button", "text": {"tag": "plain_text", "content": label}, + "type": "primary" if not actions else "default", "value": {"cmd": target}}) + if actions: + elements.append({"tag": "action", "actions": actions}) + return {"config": {"wide_screen_mode": True}, + "header": {"template": LARK_TEMPLATE.get(msg.level, "blue"), + "title": {"tag": "plain_text", "content": f"{icon} {strip_bold(msg.title)}".strip()}}, + "elements": elements} + + +def lark_plain(msg: Message) -> str: + return truncate(f"{LEVEL_ICON.get(msg.level, '')} {msg.plain()}".strip(), LARK_MAX) + + +def _lark_code(resp: Any) -> int | None: + if not isinstance(resp, dict): + return None + code = resp.get("code", resp.get("StatusCode")) + return int(code) if isinstance(code, (int, float)) else None + + +def _lark_call(fn: Callable[[], Any]) -> tuple[int | None, Any]: + """Run one Lark request; HTTP errors carry the JSON body with the error code.""" + try: + resp = fn() + except HttpError as exc: + try: + resp = json.loads(exc.body) + except ValueError: + resp = {"code": -exc.status, "msg": exc.body[:200]} + return _lark_code(resp), resp + + +class LarkAppChannel(Channel): + """Self-built app bot: tenant_access_token + im/v1/messages (+ im/v1/images for + charts); commands arrive through the long-connection client (``yeto.bot.lark_ws``).""" + + name = "lark" + chat_env, allow_env = "LARK_CHAT_ID", "LARK_ALLOWED_OPEN_IDS" + + def __init__(self, app_id: str, app_secret: str, chat: str = "", *, domain: str = "https://open.larksuite.com", + http: HttpFn, refresh_margin: float = 600.0, clock: Callable[[], float] = time.time, + receive: bool = False, buttons: bool = False, owners: list[str] | None = None, + allowed: list[str] | None = None): + super().__init__(chat, owners=owners, allowed=allowed) + self.app_id, self.app_secret = app_id, app_secret + self.domain = domain.rstrip("/") + self.http = http + self.refresh_margin = refresh_margin + self.clock = clock + self.can_receive = receive + self.buttons = buttons and receive # a press is only answered over the long connection + self._token: str | None = None + self._expires = 0.0 + self.inbox: queue.Queue = queue.Queue() + self._seen: deque = deque(maxlen=2000) + + # token ------------------------------------------------------------------- + def token(self, force: bool = False) -> str: + now = self.clock() + if not force and self._token and now < self._expires - self.refresh_margin: + return self._token + code, resp = _lark_call(lambda: self.http( + "POST", f"{self.domain}/open-apis/auth/v3/tenant_access_token/internal", + body={"app_id": self.app_id, "app_secret": self.app_secret}, timeout=15)) + if code != 0 or not resp.get("tenant_access_token"): + raise RuntimeError(f"lark tenant_access_token failed: code {code} {resp.get('msg') if resp else ''}") + self._token = resp["tenant_access_token"] + self._expires = now + float(resp.get("expire") or 7200) + return self._token + + def _authed(self, fn: Callable[[str], Any]) -> tuple[int | None, Any]: + """Call with a token; on an expired-token code refresh once and retry.""" + code, resp = _lark_call(lambda: fn(self.token())) + if code in LARK_TOKEN_EXPIRED: + code, resp = _lark_call(lambda: fn(self.token(force=True))) + return code, resp + + def _post_message(self, chat: str, msg_type: str, content: dict) -> tuple[int | None, Any]: + body = {"receive_id": chat, "msg_type": msg_type, "content": json.dumps(content, ensure_ascii=False)} + url = f"{self.domain}/open-apis/im/v1/messages?receive_id_type=chat_id" + return self._authed(lambda tok: self.http("POST", url, body=body, + headers={"Authorization": f"Bearer {tok}"}, timeout=15)) + + def upload_image(self, png: bytes) -> str | None: + """im/v1/images (scope im:resource) -> image_key, or None.""" + url = f"{self.domain}/open-apis/im/v1/images" + code, resp = self._authed(lambda tok: self.http( + "POST", url, body={"image_type": "message"}, files={"image": ("chart.png", png, "image/png")}, + headers={"Authorization": f"Bearer {tok}"}, timeout=30)) + key = ((resp or {}).get("data") or {}).get("image_key") if code == 0 else None + if not key: + log(f"lark image upload failed (code {code}: {(resp or {}).get('msg')}); using a text sparkline") + return key + + def send(self, msg: Message, chat: str | None = None) -> bool: + target = chat or self.report_chat + if not target: + return False + try: + images, failed = [], [] + for c in msg.charts: + png = charts_mod.render_png(c) + key = self.upload_image(png) if png else None + if key: + images.append((c.title, key)) + else: + failed.append(c) + msg = with_chart_text(msg, failed) + card = lark_card(msg, buttons=self.buttons, owners=self.owners, images=images) + code, resp = self._post_message(target, "interactive", card) + if code == 0: + return True + log(f"lark card send failed (code {code}: {(resp or {}).get('msg')}), retrying as text") + code, resp = self._post_message(target, "text", {"text": lark_plain(msg)}) + if code == 0: + return True + log(f"lark text send failed (code {code}: {(resp or {}).get('msg')})") + except Exception as exc: + log(f"lark send failed: {exc}") + return False + + # commands (filled by the ws event client) ----------------------------------- + def push(self, cmd: Command, event_id: str = "") -> None: + if event_id: + if event_id in self._seen: # Lark redelivers events it thinks were not acked + return + self._seen.append(event_id) + self.inbox.put(cmd) + + def poll_commands(self, timeout: float = 25.0) -> list[Command]: + out = [] + try: + out.append(self.inbox.get(timeout=timeout)) + while True: + out.append(self.inbox.get_nowait()) + except queue.Empty: + pass + return out + + +def webhook_sign(secret: str, timestamp: int) -> str: + """Lark custom-bot signature: HMAC-SHA256 keyed by ``"{timestamp}\\n{secret}"`` + over an empty message, base64-encoded.""" + key = f"{timestamp}\n{secret}".encode("utf-8") + return base64.b64encode(hmac.new(key, b"", hashlib.sha256).digest()).decode("ascii") + + +class LarkWebhookChannel(Channel): + """Group custom bot: one-way push to a fixed group (no callbacks, no image upload).""" + + name = "lark-webhook" + + def __init__(self, url: str, secret: str = "", *, http: HttpFn, clock: Callable[[], float] = time.time, + owners: list[str] | None = None): + super().__init__("webhook", owners=owners) + self.url, self.secret = url, secret + self.http = http + self.clock = clock + + def _payload(self, extra: dict) -> dict: + if not self.secret: + return extra + ts = int(self.clock()) + return {"timestamp": str(ts), "sign": webhook_sign(self.secret, ts), **extra} + + def send(self, msg: Message, chat: str | None = None) -> bool: + msg = with_chart_text(msg, msg.charts) + try: + code, resp = _lark_call(lambda: self.http( + "POST", self.url, body=self._payload({"msg_type": "interactive", + "card": lark_card(msg, buttons=False, owners=self.owners)}), + timeout=15)) + if code == 0: + return True + log(f"lark webhook card failed (code {code}: {(resp or {}).get('msg')}), retrying as text") + code, resp = _lark_call(lambda: self.http( + "POST", self.url, body=self._payload({"msg_type": "text", "content": {"text": lark_plain(msg)}}), + timeout=15)) + if code == 0: + return True + log(f"lark webhook text failed (code {code}: {(resp or {}).get('msg')})") + except Exception as exc: + log(f"lark webhook send failed: {exc}") + return False + + +# -- dry run ------------------------------------------------------------------------ + +def dry_run_http(out: Callable[[str], None] = print) -> HttpFn: + """An ``http`` stand-in that prints each request and answers like the real API.""" + def fake(method: str, url: str, *, body: Any = None, headers: dict | None = None, timeout: float = 0, + files: dict | None = None) -> Any: + shown = dict(body or {}) + if "app_secret" in shown: + shown["app_secret"] = "" + for k, (fname, data, ctype) in (files or {}).items(): + shown[k] = f"<{ctype} {fname}, {len(data)} bytes>" + safe_url = redact_url(url) + if "getUpdates" in url: + time.sleep(min(timeout, 1.0)) + return {"ok": True, "result": []} + out(f"[dry-run] {method} {safe_url}\n{json.dumps(shown, ensure_ascii=False, indent=2)}") + if "tenant_access_token" in url: + return {"code": 0, "tenant_access_token": "dry-run", "expire": 7200} + if "im/v1/images" in url: + return {"code": 0, "data": {"image_key": "img_dry_run"}} + return {"ok": True, "code": 0, "result": {}} + return fake diff --git a/yeto/bot/charts.py b/yeto/bot/charts.py new file mode 100644 index 00000000..af62d797 --- /dev/null +++ b/yeto/bot/charts.py @@ -0,0 +1,99 @@ +"""Small line charts: PNG through matplotlib when it is installed (optional +dependency), otherwise a text sparkline the channels inline into the message.""" + +from __future__ import annotations + +import io +import math +import statistics +from dataclasses import dataclass, field + +SPARK = "▁▂▃▄▅▆▇█" + + +@dataclass +class Chart: + title: str + series: list[tuple[str, list[float], list[float]]] = field(default_factory=list) # (label, xs, ys) + xlabel: str = "" + time_axis: bool = False # xs are unix seconds + hlines: list[tuple[str, float]] = field(default_factory=list) # reference lines + + def empty(self) -> bool: + return not any(len(ys) >= 2 for _, _, ys in self.series) + + +def _clean(xs: list, ys: list) -> tuple[list[float], list[float]]: + pts = [(x, y) for x, y in zip(xs, ys, strict=False) if isinstance(y, (int, float)) and not isinstance(y, bool) + and math.isfinite(y) and isinstance(x, (int, float))] + return [p[0] for p in pts], [float(p[1]) for p in pts] + + +def sparkline(ys: list[float], width: int = 24) -> str: + vals = [y for y in ys if isinstance(y, (int, float)) and math.isfinite(y)] + if not vals: + return "" + if len(vals) > width: # bucket means + n = len(vals) + vals = [statistics.fmean(vals[i * n // width:(i + 1) * n // width]) for i in range(width)] + lo, hi = min(vals), max(vals) + if hi - lo < 1e-12: + return SPARK[3] * len(vals) + return "".join(SPARK[min(7, int((v - lo) / (hi - lo) * 7.999))] for v in vals) + + +def text_lines(chart: Chart) -> list[str]: + """Sparkline fallback: one line per series with its range.""" + out = [] + for label, xs, ys in chart.series: + _, ys = _clean(xs, ys) + if ys: + out.append(f"{chart.title} {label}: {sparkline(ys)} {min(ys):.4g}..{max(ys):.4g} (last {ys[-1]:.4g})") + return out + + +def available() -> bool: + try: + import matplotlib # noqa: F401 + except ImportError: + return False + return True + + +def render_png(chart: Chart) -> bytes | None: + """PNG bytes, or None without matplotlib or without data.""" + if chart.empty() or not available(): + return None + import matplotlib + + matplotlib.use("Agg") + from matplotlib.figure import Figure + + fig = Figure(figsize=(6.4, 2.6), dpi=110) + ax = fig.add_subplot(1, 1, 1) + for label, xs, ys in chart.series: + xs, ys = _clean(xs, ys) + if not ys: + continue + if chart.time_axis: + from datetime import datetime + + xs = [datetime.fromtimestamp(x) for x in xs] + ax.plot(xs, ys, label=label, linewidth=1.4, marker="o" if len(ys) <= 12 else None, markersize=3) + for label, y in chart.hlines: + ax.axhline(y, linestyle="--", linewidth=1.0, color="grey") + ax.annotate(label, xy=(0, y), xycoords=("axes fraction", "data"), fontsize=7, color="grey", + va="bottom") + ax.set_title(chart.title, fontsize=10) + if chart.xlabel: + ax.set_xlabel(chart.xlabel, fontsize=8) + ax.tick_params(labelsize=7) + ax.grid(alpha=0.3) + if len(chart.series) > 1: + ax.legend(fontsize=7, loc="best") + if chart.time_axis: + fig.autofmt_xdate() + fig.tight_layout() + buf = io.BytesIO() + fig.savefig(buf, format="png") + return buf.getvalue() diff --git a/yeto/bot/config.py b/yeto/bot/config.py new file mode 100644 index 00000000..5266da40 --- /dev/null +++ b/yeto/bot/config.py @@ -0,0 +1,200 @@ +"""Bot configuration, read from environment variables only. + +Thresholds left unset fall back to the dashboard: the HBM line and the syncer +RSS cap come from ``sft.refs`` in ``/api/view`` (the dashboard's +``--thresholds`` file changes both), the HBM warning sits at a fraction of +that HBM line, and the consecutive-missed-rounds count +reuses ``yeto.dashboard.alerts.DEFAULT_THRESHOLDS``. +""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass, field + +from yeto.dashboard.alerts import DEFAULT_THRESHOLDS + + +def _str(env: dict, key: str, default: str = "") -> str: + return (env.get(key) or default).strip() + + +def _float(env: dict, key: str, default: float | None) -> float | None: + raw = (env.get(key) or "").strip() + if not raw: + return default + try: + return float(raw) + except ValueError: + raise ValueError(f"{key} must be a number, got {raw!r}") from None + + +def _bool(env: dict, key: str, default: bool = False) -> bool: + raw = (env.get(key) or "").strip().lower() + if not raw: + return default + return raw in ("1", "true", "yes", "on") + + +def _list(env: dict, key: str) -> list[str]: + return [x.strip() for x in (env.get(key) or "").replace(";", ",").split(",") if x.strip()] + + +def _telegram_ids(items: list[str]) -> list[str]: + bad = [x for x in items if not re.fullmatch(r"\d+", x)] + if bad: + raise ValueError(f"TELEGRAM_ALLOWED_USERS takes numeric Telegram user ids (the from.id of a message)," + f" got {bad}; usernames and display names are not trusted") + return items + + +def parse_hours(spec: str) -> tuple[int, int] | None: + """``"23-8"`` / ``"23:00-08:00"`` -> (start minute, end minute) of the day; None when empty.""" + spec = (spec or "").strip() + if not spec: + return None + m = re.fullmatch(r"(\d{1,2})(?::(\d\d))?\s*-\s*(\d{1,2})(?::(\d\d))?", spec) + if not m: + raise ValueError(f"quiet hours must look like 23-8 or 23:00-08:00, got {spec!r}") + a, b = int(m.group(1)) * 60 + int(m.group(2) or 0), int(m.group(3)) * 60 + int(m.group(4) or 0) + if a >= 24 * 60 or b > 24 * 60: + raise ValueError(f"quiet hours out of range: {spec!r}") + return a, b + + +def parse_clock(spec: str) -> int | None: + """``"09:00"`` -> minute of the day; None when empty.""" + spec = (spec or "").strip() + if not spec: + return None + m = re.fullmatch(r"(\d{1,2})(?::(\d\d))?", spec) + if not m or int(m.group(1)) > 23: + raise ValueError(f"time of day must look like 09:00, got {spec!r}") + return int(m.group(1)) * 60 + int(m.group(2) or 0) + + +@dataclass +class Thresholds: + # None: resolved against the dashboard refs at evaluation time + hbm_limit_gib: float | None = None # P0 at or above (refs sft_hbm_limit_gib) + hbm_warn_gib: float | None = None # P1 above; None: hbm_warn_frac x the HBM limit (no limit: no warning) + # 0.965 x 53.9 GiB keeps the warning just above a healthy ~52 GiB peak on the reference card + hbm_warn_frac: float = 0.965 + rss_warn_bytes: float | None = None # P1 above (refs sft_syncer_rss_limit_bytes) + host_avail_min_bytes: float = 100e9 # P1 below (per node, when the dashboard has host samples) + disk_free_min_frac: float = 0.05 # P0 below (only if the view ever carries disk samples) + loss_window: int = 20 + loss_min_points: int = 5 + loss_sigma: float = 3.0 # P1: loss > mean + sigma * std of the window + grad_norm_max: float = 1.0 # P1 above + missed_rounds: int = int(DEFAULT_THRESHOLDS["missed_consecutive_warn"]) # P1: N rounds in a row + grace_s: float = 600.0 # syncer grace window + stuck_round_s: float | None = None # P0; None: median round interval + 2 x grace + dashboard_down_after: int = 3 + forward_dashboard_alerts: bool = True # dashboard sev 0 -> P0 + + +@dataclass +class Policy: + p0_renotify_s: float = 1800.0 + p1_cooldown_s: float = 3600.0 + ack_mute_s: float = 7200.0 + quiet_hours: tuple[int, int] | None = None # only P0 inside + progress_interval: float = 7200.0 + digest_at: int | None = 9 * 60 # minute of the local day; None disables + + +@dataclass +class Config: + dashboard_url: str = "http://127.0.0.1:8787" + dashboard_link: str = "" # URL for the "Open dashboard" button (e.g. through a tunnel) + dashboard_timeout: float = 10.0 + poll_interval: float = 60.0 + dry_run: bool = False + title_prefix: str = "" # e.g. "[staging] " to tell several bots / runs apart + state_file: str = "" + audit_log: str = "yeto-bot-audit.jsonl" + telegram_token: str = "" + telegram_chat_id: str = "" + telegram_api_base: str = "https://api.telegram.org" + telegram_owners: list[str] = field(default_factory=list) # usernames to @ on P0 + # numeric Telegram user ids (from.id); usernames are not accepted since they can be + # released and re-registered by someone else. Empty: anyone in the report chat. + telegram_allowed: list[str] = field(default_factory=list) + lark_app_id: str = "" + lark_app_secret: str = "" + lark_chat_id: str = "" + lark_domain: str = "https://open.larksuite.com" # Feishu (China): https://open.feishu.cn + lark_mode: str = "auto" # auto | ws | app | webhook + lark_card_buttons: bool = True # callback buttons need card.action.trigger on the long connection + lark_owners: list[str] = field(default_factory=list) # open_ids to @ on P0 + lark_allowed: list[str] = field(default_factory=list) # open_ids; empty: anyone in the chat + lark_webhook_url: str = "" + lark_webhook_secret: str = "" + lark_token_refresh_margin: float = 600.0 + th: Thresholds = field(default_factory=Thresholds) + policy: Policy = field(default_factory=Policy) + + @classmethod + def from_env(cls, env: dict | None = None) -> "Config": + e = dict(os.environ if env is None else env) + mode = _str(e, "LARK_MODE", "auto").lower() + if mode not in ("auto", "ws", "app", "webhook"): + raise ValueError(f"LARK_MODE must be auto, ws, app or webhook, got {mode!r}") + d = Thresholds() + th = Thresholds( + hbm_limit_gib=_float(e, "YETO_BOT_HBM_LIMIT_GIB", None), + hbm_warn_gib=_float(e, "YETO_BOT_HBM_WARN_GIB", None), + hbm_warn_frac=_float(e, "YETO_BOT_HBM_WARN_FRAC", d.hbm_warn_frac), + rss_warn_bytes=_float(e, "YETO_BOT_RSS_WARN_BYTES", None), + host_avail_min_bytes=_float(e, "YETO_BOT_HOST_AVAIL_MIN_BYTES", d.host_avail_min_bytes), + disk_free_min_frac=_float(e, "YETO_BOT_DISK_FREE_MIN_FRAC", d.disk_free_min_frac), + loss_window=int(_float(e, "YETO_BOT_LOSS_WINDOW", d.loss_window)), + loss_min_points=int(_float(e, "YETO_BOT_LOSS_MIN_POINTS", d.loss_min_points)), + loss_sigma=_float(e, "YETO_BOT_LOSS_SIGMA", d.loss_sigma), + grad_norm_max=_float(e, "YETO_BOT_GRAD_NORM_MAX", d.grad_norm_max), + missed_rounds=int(_float(e, "YETO_BOT_MISSED_ROUNDS", d.missed_rounds)), + grace_s=_float(e, "YETO_BOT_GRACE_S", d.grace_s), + stuck_round_s=_float(e, "YETO_BOT_STUCK_ROUND_S", None), + dashboard_down_after=int(_float(e, "YETO_BOT_DASHBOARD_DOWN_AFTER", d.dashboard_down_after)), + forward_dashboard_alerts=_bool(e, "YETO_BOT_FORWARD_DASHBOARD_ALERTS", True), + ) + p = Policy() + digest = _str(e, "YETO_BOT_DIGEST_AT", "09:00") + policy = Policy( + p0_renotify_s=_float(e, "YETO_BOT_P0_RENOTIFY_S", p.p0_renotify_s), + p1_cooldown_s=_float(e, "YETO_BOT_P1_COOLDOWN_S", p.p1_cooldown_s), + ack_mute_s=_float(e, "YETO_BOT_ACK_MUTE_S", p.ack_mute_s), + quiet_hours=parse_hours(_str(e, "YETO_BOT_QUIET_HOURS")), + progress_interval=_float(e, "YETO_BOT_PROGRESS_INTERVAL", p.progress_interval), + digest_at=None if digest.lower() in ("off", "none") else parse_clock(digest), + ) + return cls( + dashboard_url=_str(e, "YETO_DASHBOARD_URL", "http://127.0.0.1:8787").rstrip("/"), + dashboard_link=_str(e, "YETO_DASHBOARD_LINK"), + dashboard_timeout=_float(e, "YETO_DASHBOARD_TIMEOUT", 10.0), + poll_interval=_float(e, "YETO_BOT_POLL_INTERVAL", 60.0), + dry_run=_bool(e, "YETO_BOT_DRY_RUN"), + title_prefix=e.get("YETO_BOT_TITLE_PREFIX") or "", + state_file=_str(e, "YETO_BOT_STATE_FILE"), + audit_log=_str(e, "YETO_BOT_AUDIT_LOG", "yeto-bot-audit.jsonl"), + telegram_token=_str(e, "TELEGRAM_BOT_TOKEN"), + telegram_chat_id=_str(e, "TELEGRAM_CHAT_ID"), + telegram_api_base=_str(e, "TELEGRAM_API_BASE", "https://api.telegram.org").rstrip("/"), + telegram_owners=[x.lstrip("@") for x in _list(e, "TELEGRAM_OWNERS")], + telegram_allowed=_telegram_ids(_list(e, "TELEGRAM_ALLOWED_USERS")), + lark_app_id=_str(e, "LARK_APP_ID"), + lark_app_secret=_str(e, "LARK_APP_SECRET"), + lark_chat_id=_str(e, "LARK_CHAT_ID"), + lark_domain=_str(e, "LARK_DOMAIN", "https://open.larksuite.com").rstrip("/"), + lark_mode=mode, + lark_card_buttons=_bool(e, "LARK_CARD_BUTTONS", True), + lark_owners=_list(e, "LARK_OWNER_OPEN_IDS") or _list(e, "LARK_OWNER_OPEN_ID"), + lark_allowed=_list(e, "LARK_ALLOWED_OPEN_IDS"), + lark_webhook_url=_str(e, "LARK_WEBHOOK_URL"), + lark_webhook_secret=_str(e, "LARK_WEBHOOK_SECRET"), + lark_token_refresh_margin=_float(e, "LARK_TOKEN_REFRESH_MARGIN", 600.0), + th=th, + policy=policy, + ) diff --git a/yeto/bot/decisions.py b/yeto/bot/decisions.py new file mode 100644 index 00000000..6de24713 --- /dev/null +++ b/yeto/bot/decisions.py @@ -0,0 +1,103 @@ +"""Decision-request cards and their append-only JSONL audit log. + +A decision request is a card with a title, evidence lines, a suggested option +and one button per option. Pressing a button only appends a record to the audit +log and replies that the operator will act: the bot executes nothing. Requests +are written to the same log, so a request posted by one process (for example +``python -m yeto.bot --request-decision``) is answered by the running bot. + +Record types: ``{"type": "request", "id", "title", "evidence", "options", +"suggested", "ts"}`` and ``{"type": "decision", "id", "option", "user", +"user_id", "channel", "chat", "ts"}`` (``user`` is the display name, ``user_id`` +the platform id the allowlist checks). +""" + +from __future__ import annotations + +import json +import os +import secrets +import threading +import time +from dataclasses import dataclass, field + +from .render import Message, b + + +@dataclass +class DecisionRequest: + title: str + evidence: list[str] = field(default_factory=list) + options: list[str] = field(default_factory=list) + suggested: str = "" + id: str = "" + ts: float = 0.0 + + def message(self) -> Message: + lines = [b("Evidence"), *[f"- {x}" for x in self.evidence]] + if self.suggested: + lines += ["", f"{b('Suggested')}: {self.suggested}"] + lines += ["", "Pressing a button records the decision for the operator; the bot runs nothing."] + buttons = [(opt + (" (suggested)" if opt == self.suggested else ""), f"/decide {self.id} {opt}") + for opt in self.options] + return Message(f"decision needed: {self.title}", lines, "decision", buttons) + + +class DecisionLog: + def __init__(self, path: str, *, clock=time.time): + self.path = path + self.clock = clock + self.lock = threading.Lock() + + def _append(self, rec: dict) -> None: + with self.lock: + d = os.path.dirname(self.path) + if d: + os.makedirs(d, exist_ok=True) + with open(self.path, "a", encoding="utf-8") as fh: + fh.write(json.dumps(rec, ensure_ascii=False) + "\n") + fh.flush() + os.fsync(fh.fileno()) + + def records(self) -> list[dict]: + if not os.path.exists(self.path): + return [] + out = [] + with open(self.path, encoding="utf-8") as fh: + for line in fh: + try: + rec = json.loads(line) + except ValueError: + continue + if isinstance(rec, dict): + out.append(rec) + return out + + def request(self, title: str, evidence: list[str], options: list[str], suggested: str = "") -> DecisionRequest: + if not options or any(" " in o or not o for o in options): + raise ValueError("decision options must be non-empty single words") + if suggested and suggested not in options: + raise ValueError(f"suggested option {suggested!r} is not one of {options}") + req = DecisionRequest(title, list(evidence), list(options), suggested, "d" + secrets.token_hex(3), + self.clock()) + self._append({"type": "request", "id": req.id, "title": title, "evidence": req.evidence, + "options": req.options, "suggested": suggested, "ts": req.ts}) + return req + + def decide(self, req_id: str, option: str, *, user: str, channel: str, chat: str, + user_id: str = "") -> Message: + recs = self.records() + req = next((r for r in recs if r.get("type") == "request" and r.get("id") == req_id), None) + if req is None: + return Message("unknown decision", [f"no decision request {req_id!r}"], "warn") + if option not in (req.get("options") or []): + return Message("unknown option", [f"{option!r} is not one of {req.get('options')}"], "warn") + done = next((r for r in recs if r.get("type") == "decision" and r.get("id") == req_id), None) + if done is not None: + return Message("decision already recorded", + [f"{req.get('title')}: {done.get('option')} by {done.get('user') or '?'} at" + f" {time.strftime('%m-%d %H:%M', time.localtime(done.get('ts') or 0))}"], "info") + self._append({"type": "decision", "id": req_id, "option": option, "user": user, "user_id": user_id, + "channel": channel, "chat": chat, "ts": self.clock()}) + return Message("decision recorded", [f"{req.get('title')}: {b(option)} (by {user or '?'})", + "recorded; operator will act"], "ok") diff --git a/yeto/bot/http.py b/yeto/bot/http.py new file mode 100644 index 00000000..efce5e10 --- /dev/null +++ b/yeto/bot/http.py @@ -0,0 +1,88 @@ +"""Tiny JSON-over-HTTP helper (urllib) and the read-only dashboard client. + +Everything that talks to the network takes an ``http`` callable with the +signature of :func:`request_json`, so tests swap in a fake. +""" + +from __future__ import annotations + +import json +import re +import urllib.error +import urllib.request +import uuid +from typing import Any, Callable + +HttpFn = Callable[..., Any] + + +class HttpError(RuntimeError): + def __init__(self, status: int, body: str, url: str): + super().__init__(f"HTTP {status} from {url}: {body[:200]}") + self.status = status + self.body = body + + +def multipart(fields: dict, files: dict) -> tuple[bytes, str]: + """``files``: name -> (filename, bytes, content type). Returns (body, content type).""" + boundary = "yeto-" + uuid.uuid4().hex + out = [] + for k, v in (fields or {}).items(): + out.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"\r\n\r\n{v}\r\n'.encode()) + for k, (fname, data, ctype) in files.items(): + out.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"; filename="{fname}"\r\n' + f"Content-Type: {ctype}\r\n\r\n".encode() + data + b"\r\n") + out.append(f"--{boundary}--\r\n".encode()) + return b"".join(out), f"multipart/form-data; boundary={boundary}" + + +def request_json(method: str, url: str, *, body: Any = None, headers: dict | None = None, + timeout: float = 30.0, files: dict | None = None) -> Any: + """JSON request (or multipart form when ``files`` is given, ``body`` then holds the form fields).""" + data = None + hdrs = dict(headers or {}) + if files: + data, ctype = multipart(body or {}, files) + hdrs["Content-Type"] = ctype + elif body is not None: + data = json.dumps(body).encode("utf-8") + hdrs.setdefault("Content-Type", "application/json; charset=utf-8") + req = urllib.request.Request(url, data=data, headers=hdrs, method=method) + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + raw = resp.read() + except urllib.error.HTTPError as exc: + raise HttpError(exc.code, exc.read().decode("utf-8", errors="replace"), redact_url(url)) from None + return json.loads(raw) if raw else None + + +_URL_SECRETS = ( + (re.compile(r"/bot[^/?#]+"), "/bot"), # Telegram Bot API: /bot/ + (re.compile(r"/hook/[^/?#]+"), "/hook/"), # Lark custom bot: /open-apis/bot/v2/hook/ +) + + +def redact_url(url: str) -> str: + """``url`` with the credentials some APIs carry in the path replaced by ````. + Use it wherever a URL reaches a log line, an exception or the dry-run output.""" + for pat, repl in _URL_SECRETS: + url = pat.sub(repl, url) + return url + + +class DashboardClient: + """GET-only client of ``yeto dashboard serve`` (never touches nodes or the run).""" + + def __init__(self, base_url: str, *, timeout: float = 10.0, http: HttpFn = request_json): + self.base = base_url.rstrip("/") + self.timeout = timeout + self.http = http + + def _get(self, path: str) -> Any: + return self.http("GET", self.base + path, timeout=self.timeout) + + def view(self) -> dict: + return self._get("/api/view") + + def bad_rounds(self) -> list[dict]: + return (self._get("/api/rounds?only_bad=1") or {}).get("rounds") or [] diff --git a/yeto/bot/lark_ws.py b/yeto/bot/lark_ws.py new file mode 100644 index 00000000..313d9746 --- /dev/null +++ b/yeto/bot/lark_ws.py @@ -0,0 +1,137 @@ +"""Lark long-connection (WebSocket) event client, an optional dependency. + +The official SDK (``pip install lark-oapi``) opens an outbound WebSocket to the +Lark open platform, so the bot receives ``im.message.receive_v1`` (commands) and +``card.action.trigger`` (button presses) without a public HTTP endpoint. Without +the SDK the Lark app channel is send-only, or the bot falls back to the group +webhook (``LARK_WEBHOOK_URL``). + +Event parsing lives in plain functions over attribute objects so it is testable +without the SDK. +""" + +from __future__ import annotations + +import json +import re +import threading +import time +from typing import Any, Callable + +from .channels import Command, LarkAppChannel, log + +_MENTION = re.compile(r"@_user_\d+\s*") + + +def available() -> bool: + try: + import lark_oapi # noqa: F401 + except ImportError: + return False + return True + + +def _get(obj: Any, *path: str) -> Any: + for p in path: + if obj is None: + return None + obj = obj.get(p) if isinstance(obj, dict) else getattr(obj, p, None) + return obj + + +def parse_message_event(data: Any) -> tuple[Command, str] | None: + """``P2ImMessageReceiveV1`` -> (command, message id); None for non-text messages. + Text without a leading ``/`` is kept for the natural-language hook.""" + msg = _get(data, "event", "message") + if _get(msg, "message_type") != "text": + return None + try: + text = json.loads(_get(msg, "content") or "{}").get("text") or "" + except ValueError: + return None + text = _MENTION.sub("", text).strip() # group messages carry "@_user_1 /status" + chat = _get(msg, "chat_id") + if not chat or not text: + return None + user = str(_get(data, "event", "sender", "sender_id", "open_id") or "") + return Command("lark", str(chat), text, user, user_id=user), str(_get(msg, "message_id") or "") + + +def parse_card_event(data: Any) -> tuple[Command, str] | None: + """``P2CardActionTrigger`` -> (command, event id) for buttons built by ``lark_card``.""" + value = _get(data, "event", "action", "value") or {} + cmd = value.get("cmd") if isinstance(value, dict) else None + chat = _get(data, "event", "context", "open_chat_id") + if not cmd or not chat: + return None + user = str(_get(data, "event", "operator", "open_id") or "") + event_id = _get(data, "header", "event_id") or "" + return Command("lark", str(chat), str(cmd), user, button=True, user_id=user), str(event_id) + + +def card_toast(parsed: tuple[Command, str] | None) -> str: + """Toast shown on the press itself. Neutral on purpose: the allowlist is checked + later in ``Bot.handle``, whose reply says what ran or why it was refused.""" + return "received" if parsed else "unknown button" + + +def callbacks(channel: LarkAppChannel) -> tuple[Callable[[Any], None], Callable[[Any], Any]]: + """(message handler, card-action handler) that feed ``channel``'s inbox.""" + + def on_message(data: Any) -> None: + parsed = parse_message_event(data) + if parsed: + channel.push(*parsed) + + def on_card(data: Any) -> Any: + from lark_oapi.event.callback.model.p2_card_action_trigger import P2CardActionTriggerResponse + + parsed = parse_card_event(data) + if parsed: + channel.push(*parsed) + return P2CardActionTriggerResponse({"toast": {"type": "info", "content": card_toast(parsed)}}) + + return on_message, on_card + + +def build_handler(channel: LarkAppChannel) -> Any: + import lark_oapi as lark + + on_message, on_card = callbacks(channel) + builder = lark.EventDispatcherHandler.builder("", "").register_p2_im_message_receive_v1(on_message) + if hasattr(builder, "register_p2_card_action_trigger"): + builder = builder.register_p2_card_action_trigger(on_card) + else: + log("lark-oapi too old for card.action.trigger: buttons will not respond (upgrade lark-oapi)") + return builder.build() + + +def start(channel: LarkAppChannel, *, stop: threading.Event | None = None, + sleep: Callable[[float], None] = time.sleep) -> threading.Thread: + """Run the SDK's blocking ws client in a daemon thread, reconnecting with backoff.""" + + def run() -> None: + import asyncio + + import lark_oapi as lark + import lark_oapi.ws.client as ws_client + + # the SDK binds a module-level event loop at import; give this thread its own + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + ws_client.loop = loop + backoff = 5.0 + while stop is None or not stop.is_set(): + try: + cli = lark.ws.Client(channel.app_id, channel.app_secret, event_handler=build_handler(channel), + domain=channel.domain, log_level=lark.LogLevel.WARNING) + log("lark long connection starting") + cli.start() + except Exception as exc: + log(f"lark long connection failed: {exc}; retrying in {backoff:.0f}s") + sleep(backoff) + backoff = min(300.0, backoff * 2) + + t = threading.Thread(target=run, name="lark-ws", daemon=True) + t.start() + return t diff --git a/yeto/bot/monitor.py b/yeto/bot/monitor.py new file mode 100644 index 00000000..32d601b0 --- /dev/null +++ b/yeto/bot/monitor.py @@ -0,0 +1,291 @@ +"""Alert policy: transitions, severity routing, cooldown, re-notify, ack, mute, +quiet hours, and merging everything one poll produced into a single card. + +* P0 conditions notify on appearance and then every ``p0_renotify_s`` while + they hold, until acked; they also pass quiet hours and a plain ``/mute``. +* P1 conditions notify on appearance, at most once per ``p1_cooldown_s`` per + subject; P1 events held back by the cooldown accumulate and go out together. +* P2 notifies on appearance, never pages. +* A P1 / P2 appearance held back by a mute, quiet hours, an ack or the + cooldown is not lost: it goes out at the first poll after the hold ends, + if the condition is still active (its suppression is recorded once). +* ``/ack`` mutes the subjects of one card (or all active ones) for ``ack_mute_s``. +* The first snapshot only seeds: P0 conditions already true are re-notified + after ``p0_renotify_s`` unless acked; nothing else from before the start is + announced (the startup progress card shows it). +""" + +from __future__ import annotations + +import time +from collections import OrderedDict, deque +from dataclasses import dataclass, field +from typing import Callable + +from .config import Policy, Thresholds +from .render import Message, b, fdur +from .rules import P0, P1, P2, SEV_NAME, Condition, Event, EventTracker, conditions + +SEV_ICON = {P0: "\U0001f534", P1: "\U0001f7e0", P2: "⚪"} +SEV_LEVEL = {P0: "p0", P1: "p1", P2: "p2"} +HISTORY_MAX = 5000 +CARDS_MAX = 500 + + +@dataclass +class Notice: + sev: int + key: tuple + title: str + lines: list[str] = field(default_factory=list) + island: str | None = None + kind: str = "fire" # fire | renotify | resolve | event + count: int = 1 + since: float | None = None + held: bool = False # an earlier attempt was already recorded as suppressed + + +def key_str(key: tuple) -> str: + return "/".join(str(k) for k in key) + + +class Monitor: + def __init__(self, th: Thresholds, policy: Policy, *, clock: Callable[[], float] = time.time): + self.th = th + self.policy = policy + self.clock = clock + self.seeded = False + self.active: dict[tuple, Condition] = {} + self.since: dict[tuple, float] = {} + self.notified: dict[tuple, float] = {} # last delivered notice per subject (survives flaps) + self.alerted: set[tuple] = set() # active subjects whose appearance was delivered + # active P1 / P2 subjects whose appearance is not delivered yet -> suppression recorded + self.unsent: dict[tuple, bool] = {} + self.acked: dict[tuple, tuple[float, str]] = {} # subject -> (until, by) + self.pending: dict[tuple, Event] = {} # P1 events held back by the cooldown + self.tracker = EventTracker() + self.mute_until = 0.0 + self.mute_all = False + self.fail_streak = 0 + self.dashboard_down = False + self.history: deque = deque(maxlen=HISTORY_MAX) # (ts, sev, key, title, kind, delivered) + self.cards: OrderedDict[int, list[Notice]] = OrderedDict() + self._card_seq = 0 + + # -- suppression ------------------------------------------------------------ + def quiet(self, now: float) -> bool: + qh = self.policy.quiet_hours + if not qh: + return False + t = time.localtime(now) + m = t.tm_hour * 60 + t.tm_min + a, z = qh + return a <= m < z if a < z else (m >= a or m < z) + + def muted(self, sev: int, now: float) -> bool: + return now < self.mute_until and (sev > P0 or self.mute_all) + + def acked_now(self, key: tuple, now: float) -> bool: + ack = self.acked.get(key) + return ack is not None and now < ack[0] + + def deliverable(self, n: Notice, now: float) -> bool: + if n.kind in ("fire", "renotify", "event") and self.acked_now(n.key, now): + return False + if self.muted(n.sev, now): + return False + return not (n.sev > P0 and n.kind != "resolve" and self.quiet(now)) + + # -- one poll ----------------------------------------------------------------- + def evaluate(self, view: dict) -> list[Notice]: + now = self.clock() + out: list[Notice] = [] + if self.dashboard_down: + out.append(Notice(P1, ("dashboard", "down"), "dashboard reachable again", + [f"after {self.fail_streak} failed polls"], kind="resolve")) + self.fail_streak, self.dashboard_down = 0, False + cur = conditions(view, self.th) + if not self.seeded: + for key, c in cur.items(): + self.since[key] = now + if c.sev == P0: # re-notified after p0_renotify_s unless acked + self.notified[key] = now + self.alerted.add(key) + self.active = cur + self.tracker.events(view, seed=True) + self.seeded = True + return [] + candidates: list[Notice] = [] + for key, c in cur.items(): + if key not in self.active: + self.since[key] = now + if c.sev != P0: + self.unsent[key] = False + last = self.notified.get(key, float("-inf")) + if c.sev == P0: + due = now - last >= self.policy.p0_renotify_s + if due: + kind = "renotify" if key in self.alerted else "fire" + candidates.append(Notice(c.sev, key, c.title, c.lines, c.island, kind, since=self.since[key])) + elif key in self.unsent and (c.sev == P2 or now - last >= self.policy.p1_cooldown_s): + since = self.since[key] + lines = c.lines + ([f"active for {fdur(now - since)}, held back until now (mute, quiet hours," + f" ack or cooldown)"] if now > since else []) + candidates.append(Notice(c.sev, key, c.title, lines, c.island, "fire", since=since, + held=self.unsent[key])) + for key, c in self.active.items(): + if key not in cur: + self.since.pop(key, None) + self.acked.pop(key, None) + self.unsent.pop(key, None) + if key in self.alerted: + self.alerted.discard(key) + if c.resolve: + candidates.append(Notice(c.sev, key, c.title, c.lines, c.island, "resolve")) + self.active = cur + for ev in self.tracker.events(view): + candidates.extend(self._event(ev, now)) + for key in list(self.pending): # cooled-down P1 events still waiting + if now - self.notified.get(key, float("-inf")) >= self.policy.p1_cooldown_s: + ev = self.pending.pop(key) + candidates.append(Notice(ev.sev, key, ev.title, ev.lines, ev.island, "event", ev.count)) + delivered = [] + for n in out + candidates: + ok = self.deliverable(n, now) + if ok or not n.held: # a held-back appearance retried every poll is recorded once + self.history.append((now, n.sev, n.key, n.title, n.kind, ok)) + if not ok: + if n.key in self.unsent and n.kind == "fire": + self.unsent[n.key] = True + continue + if n.kind != "resolve": + self.notified[n.key] = now + if n.kind in ("fire", "renotify"): + self.alerted.add(n.key) + self.unsent.pop(n.key, None) + delivered.append(n) + return delivered + + def _event(self, ev: Event, now: float) -> list[Notice]: + if ev.sev == P1: + held = self.pending.get(ev.key) + if held is not None: + ev = Event(ev.sev, ev.key, ev.title, held.lines + ev.lines, ev.island, held.count + ev.count) + if now - self.notified.get(ev.key, float("-inf")) < self.policy.p1_cooldown_s: + self.pending[ev.key] = ev + return [] + self.pending.pop(ev.key, None) + return [Notice(ev.sev, ev.key, ev.title, ev.lines, ev.island, "event", ev.count)] + + def fetch_failed(self, exc: BaseException) -> list[Notice]: + self.fail_streak += 1 + if self.dashboard_down or self.fail_streak < self.th.dashboard_down_after: + return [] + self.dashboard_down = True + n = Notice(P1, ("dashboard", "down"), "dashboard unreachable", + [f"{self.fail_streak} failed polls in a row: {exc}"[:300]]) + now = self.clock() + ok = self.deliverable(n, now) + self.history.append((now, n.sev, n.key, n.title, n.kind, ok)) + return [n] if ok else [] + + # -- operator actions ----------------------------------------------------------- + def ack(self, keys: list[tuple] | None, by: str) -> list[tuple]: + """Mute subjects (active conditions or event subjects such as rejected pushes) for ``ack_mute_s``.""" + now = self.clock() + keys = [k for k, c in self.active.items() if c.resolve] if keys is None else list(dict.fromkeys(keys)) + for k in keys: + self.acked[k] = (now + self.policy.ack_mute_s, by) + self.pending.pop(k, None) + return keys + + def mute(self, seconds: float, *, include_p0: bool = False) -> None: + self.mute_until = self.clock() + seconds + self.mute_all = include_p0 + + def unmute(self) -> None: + self.mute_until, self.mute_all = 0.0, False + + # -- cards ------------------------------------------------------------------------ + def card(self, notices: list[Notice], *, link: str = "") -> Message | None: + """All notices of one poll as one message; registers the card for /ack and /details.""" + if not notices: + return None + self._card_seq += 1 + cid = self._card_seq + self.cards[cid] = list(notices) + while len(self.cards) > CARDS_MAX: + self.cards.popitem(last=False) + firing = sorted((n for n in notices if n.kind != "resolve"), key=lambda n: n.sev) + resolved = [n for n in notices if n.kind == "resolve"] + now = self.clock() + lines: list[str] = [] + for n in firing: + head = f"{SEV_ICON[n.sev]} {b(SEV_NAME[n.sev])} {n.title}" + if n.count > 1: + head += f" (+{n.count})" + if n.kind == "renotify": + head += f" - still firing for {fdur(now - (n.since or now))}, not acked" + lines.append(head) + lines += [f" {x}" for x in n.lines if x] + for n in resolved: + lines.append(f"✅ resolved: {n.title}") + if firing: + top = firing[0] + level = SEV_LEVEL[top.sev] + title = f"{SEV_NAME[top.sev]} {top.title}" + (f" (+{len(notices) - 1} more)" if len(notices) > 1 else "") + else: + level = "ok" + title = f"resolved: {resolved[0].title}" + (f" (+{len(resolved) - 1} more)" if len(resolved) > 1 else "") + buttons: list[tuple[str, str]] = [] + if any(n.sev <= P1 for n in firing): + buttons.append(("Ack 2h", f"/ack c{cid}")) + buttons.append(("Details", f"/details c{cid}")) + if link: + buttons.append(("Open dashboard", link)) + mention = any(n.sev == P0 and n.kind in ("fire", "renotify", "event") for n in firing) + return Message(title, lines, level, buttons, mention=mention) + + def card_notices(self, ref: str) -> list[Notice] | None: + ref = ref.strip().lower().lstrip("c") + return self.cards.get(int(ref)) if ref.isdigit() else None + + # -- summaries -------------------------------------------------------------------- + def active_counts(self) -> dict[int, int]: + out = {P0: 0, P1: 0, P2: 0} + for c in self.active.values(): + if c.sev in out and c.resolve: + out[c.sev] += 1 + return out + + def history_since(self, t0: float) -> list[tuple]: + return [h for h in self.history if h[0] >= t0] + + def alerts_message(self, *, window_s: float = 86400.0) -> Message: + now = self.clock() + lines = [] + act = sorted(((k, c) for k, c in self.active.items() if c.resolve), key=lambda kc: kc[1].sev) + lines.append(b(f"Active ({len(act)})")) + for k, c in act: + ack = self.acked.get(k) + tag = f" | acked by {ack[1]} until {time.strftime('%H:%M', time.localtime(ack[0]))}" \ + if ack and now < ack[0] else "" + age = fdur(now - self.since.get(k, now)) + lines.append(f"{SEV_ICON[c.sev]} {SEV_NAME[c.sev]} {c.title} | for {age}{tag}") + if not act: + lines.append("none") + hist = self.history_since(now - window_s) + counts = {s: sum(1 for h in hist if h[1] == s and h[4] != "resolve") for s in (P0, P1, P2)} + lines += ["", b(f"Last {fdur(window_s)}") + f": P0 {counts[P0]} | P1 {counts[P1]} | P2 {counts[P2]}" + f" | suppressed {sum(1 for h in hist if not h[5])}"] + for ts, sev, _key, title, kind, ok in list(hist)[-10:]: + lines.append(f"{time.strftime('%m-%d %H:%M', time.localtime(ts))} {SEV_NAME.get(sev, '?')} {kind}" + f" {title}" + ("" if ok else " (suppressed)")) + if now < self.mute_until: + lines.append(f"muted until {time.strftime('%m-%d %H:%M', time.localtime(self.mute_until))}" + + (" (P0 too)" if self.mute_all else " (P0 still pages)")) + if self.policy.quiet_hours: + a, z = self.policy.quiet_hours + lines.append(f"quiet hours {a // 60:02d}:{a % 60:02d}-{z // 60:02d}:{z % 60:02d} (P0 only)" + + (" - now quiet" if self.quiet(now) else "")) + return Message("yeto alerts", lines, "report", [("Ack all", "/ack all"), ("Status", "/status")]) diff --git a/yeto/bot/render.py b/yeto/bot/render.py new file mode 100644 index 00000000..c6e50e86 --- /dev/null +++ b/yeto/bot/render.py @@ -0,0 +1,607 @@ +"""Channel-neutral messages built from the dashboard's ``/api/view`` JSON. + +A :class:`Message` is a title, body lines, optional buttons (a command, or a +URL when it starts with ``http``), optional charts and a flag to @ the owners. +Bold spans are wrapped by :func:`b` in two control characters; each channel +turns them into its own markup (Telegram ``*x*``, Lark ``**x**``, plain text +drops them) and escapes everything else itself. +""" + +from __future__ import annotations + +import math +import statistics +import time +from dataclasses import dataclass, field +from typing import Any + +from yeto.dashboard.sft import MERGE_COLS, STEP_COLS + +from .charts import Chart + +BOLD_OPEN, BOLD_CLOSE = "\x02", "\x03" + +# sft.islands[].rows and sft.syncer.merges are positional; their column names come +# with the view (sft.cols, sft.syncer.merge_cols) and are resolved per view with +# step_cols() / merge_cols(). A column the view does not name resolves to None and +# reads as missing data. syncer.outer rows carry no names: +# [ts, step, fragment, step_seconds, timed_out, arrived]. +PROBLEMS = ("rejected", "timeout", "not_carried") +STATUS_DOT = {"ok": "\U0001f7e2", "warn": "\U0001f7e1", "stale": "\U0001f534", "done": "\U0001f3c1", + "unknown": "⚪"} +LEVEL_ICON = {"info": "ℹ️", "ok": "✅", "warn": "⚠️", "p0": "\U0001f6a8", "p1": "\U0001f7e0", + "p2": "⚪", "report": "\U0001f4ca", "decision": "\U0001f9ed"} + + +@dataclass +class Message: + title: str + lines: list[str] = field(default_factory=list) + level: str = "info" # info | ok | warn | p0 | p1 | p2 | report | decision + buttons: list[tuple[str, str]] = field(default_factory=list) # (label, command or URL) + charts: list[Chart] = field(default_factory=list) + mention: bool = False # @ the configured owners + + def plain(self) -> str: + return strip_bold("\n".join([self.title, *self.lines])) + + +def b(text: Any) -> str: + return f"{BOLD_OPEN}{text}{BOLD_CLOSE}" + + +def strip_bold(text: str) -> str: + return text.replace(BOLD_OPEN, "").replace(BOLD_CLOSE, "") + + +def is_url(target: str) -> bool: + return target.startswith(("http://", "https://")) + + +# -- number formatting ------------------------------------------------------------ + +def num(v: Any) -> float | None: + """A finite float, or None (bools, strings and NaN/Inf are None).""" + if isinstance(v, bool) or not isinstance(v, (int, float)): + return None + v = float(v) + return v if math.isfinite(v) else None + + +def nonfinite(v: Any) -> bool: + if isinstance(v, str): + return v.strip().lower() in ("nan", "inf", "-inf", "infinity", "-infinity") + return isinstance(v, float) and not math.isfinite(v) + + +def f3(v: Any) -> str: + v = num(v) + if v is None: + return "-" + return f"{v:.4f}" if abs(v) < 10 else f"{v:.3g}" + + +def fsig(v: Any, digits: int = 3) -> str: + v = num(v) + return "-" if v is None else f"{v:.{digits}g}" + + +def fsigned(v: Any) -> str: + v = num(v) + return "-" if v is None else f"{v:+.4f}" + + +def _g1(v: float) -> str: + """One decimal, without a trailing ``.0``.""" + out = f"{v:.1f}" + return out[:-2] if out.endswith(".0") else out + + +def fdur(s: Any) -> str: + s = num(s) + if s is None: + return "-" + if s < 120: + return f"{s:.0f}s" + if s < 7200: + return f"{_g1(s / 60)}min" + if s < 172800: + return f"{_g1(s / 3600)}h" + return f"{_g1(s / 86400)}d" + + +def fbytes(v: Any) -> str: + v = num(v) + if v is None: + return "-" + for unit, scale in (("TB", 1e12), ("GB", 1e9), ("MB", 1e6)): + if v >= scale: + return f"{v / scale:.2f} {unit}" + return f"{v:.0f} B" + + +def fclock(ts: Any) -> str: + ts = num(ts) + return "-" if ts is None else time.strftime("%m-%d %H:%M", time.localtime(ts)) + + +# -- view helpers --------------------------------------------------------------- + +def col_index(names: Any, fallback: tuple[str, ...]) -> dict[str, int]: + """Column name -> position; ``fallback`` when the view names no columns (an older + dashboard, whose layout was the fallback).""" + if not isinstance(names, (list, tuple)) or not names: + names = fallback + return {str(n): i for i, n in enumerate(names)} + + +def step_cols(sft: dict) -> dict[str, int]: + """Positions in ``sft.islands[].rows`` by name (``sft.cols``).""" + return col_index((sft or {}).get("cols"), STEP_COLS) + + +def merge_cols(sft: dict) -> dict[str, int]: + """Positions in ``sft.syncer.merges`` rows by name (``sft.syncer.merge_cols``).""" + return col_index(((sft or {}).get("syncer") or {}).get("merge_cols"), MERGE_COLS) + + +def cell(r: list, col: int | None) -> Any: + """``r[col]``, or None when the column is unknown or the row is short.""" + return r[col] if col is not None and 0 <= col < len(r) else None + + +def last_value(rows: list[list], col: int | None) -> Any: + """Newest non-None value of a column (the page does the same: the newest step + often has bench numbers before ms-swift logs its loss).""" + for r in reversed(rows): + v = cell(r, col) + if v is not None: + return v + return None + + +def column(rows: list[list], col: int | None) -> list[float]: + return [v for v in (num(cell(r, col)) for r in rows) if v is not None] + + +def step_p50(rows: list[list], ix: dict[str, int]) -> float | None: + vals = sorted(column(rows[-20:], ix.get("step_s"))) + return vals[max(0, math.ceil(0.5 * len(vals)) - 1)] if vals else None + + +def refs(view: dict) -> dict: + return ((view or {}).get("sft") or {}).get("refs") or {} + + +def problem_counts(sft: dict) -> dict[str, int]: + out = {c: 0 for c in PROBLEMS} + for e in (sft.get("syncer") or {}).get("events") or []: + if len(e) > 1 and e[1] in out: + out[e[1]] += 1 + return out + + +def last_round_ts(sft: dict) -> float | None: + syn = sft.get("syncer") or {} + m_ts = merge_cols(sft).get("ts") + ts = [num(r[0]) for r in (syn.get("outer") or []) if r] + [num(cell(r, m_ts)) for r in (syn.get("merges") or [])] + ts = [t for t in ts if t is not None] + return max(ts) if ts else None + + +def round_intervals(sft: dict, n: int = 16) -> list[float]: + """Seconds between consecutive outer steps (newest ``n``).""" + syn = sft.get("syncer") or {} + pts = sorted({(num(r[0]), r[1]) for r in (syn.get("outer") or []) if r and num(r[0]) is not None + and isinstance(r[1], int)}) + if len(pts) < 2: + mx = merge_cols(sft) + m_ts, m_step = mx.get("ts"), mx.get("step") + pts = sorted({(num(cell(r, m_ts)), cell(r, m_step)) for r in (syn.get("merges") or []) + if num(cell(r, m_ts)) is not None and isinstance(cell(r, m_step), int)}) + gaps = [q[0] - p[0] for p, q in zip(pts, pts[1:], strict=False) if q[1] > p[1] and q[0] > p[0]] + return gaps[-n:] + + +def find_island(sft: dict, key: str) -> dict | None: + k = key.strip().lower() + if k.startswith("island"): + k = k[len("island"):].strip(" -_") + for isl in sft.get("islands") or []: + if str(isl.get("id")).lower() == k or str(isl.get("name") or "").lower() in (key.strip().lower(), k): + return isl + return None + + +def latest_rss(sft: dict) -> float | None: + rss = (sft.get("syncer") or {}).get("rss") or [] + return num(rss[-1][1]) if rss and len(rss[-1]) > 1 else None + + +def island_label(isl: dict) -> str: + return isl.get("name") or f"island {isl.get('id')}" + + +def window(rows: list[list], now: float | None, seconds: float, ts_col: int | None) -> list[list]: + if now is None: + return rows + return [r for r in rows if num(cell(r, ts_col)) is not None and r[ts_col] >= now - seconds] + + +def loss_chart(sft: dict, seconds: float, title: str) -> Chart: + now = num(sft.get("now")) + ix = step_cols(sft) + ts, loss = ix.get("ts"), ix.get("loss") + series = [] + for isl in sft.get("islands") or []: + rows = window(isl.get("rows") or [], now, seconds, ts) + series.append((island_label(isl), [cell(r, ts) for r in rows], [cell(r, loss) for r in rows])) + return Chart(title, series, time_axis=True) + + +def merge_health(view: dict, n: int = 20) -> dict: + """Direct-in-round ratio (rounds merged with every expected push and no resend + or timeout), missed counts per island, round timing medians.""" + sft = (view or {}).get("sft") or {} + syn = sft.get("syncer") or {} + rounds = ((view or {}).get("rounds_syncer") or [])[-n:] + out: dict[str, Any] = {"rounds": len(rounds), "direct": None, "missed": {}, "weights": None} + if rounds: + out["direct"] = sum(1 for r in rounds if not r.get("missed") and not r.get("resend")) + for r in rounds: + for iid in r.get("missed") or []: + out["missed"][str(iid)] = out["missed"].get(str(iid), 0) + 1 + w = [r.get("weights") for r in rounds if isinstance(r.get("weights"), dict)] + out["weights"] = w[-1] if w else None # hook: per-island merge weight once exposed + else: + outer = (syn.get("outer") or [])[-n:] + out["rounds"] = len(outer) + out["direct"] = sum(1 for r in outer if not (len(r) > 4 and r[4])) if outer else None + merges = (syn.get("merges") or [])[-n:] + mx = merge_cols(sft) + for name in ("round_s", "quorum_s", "grace_s", "merge_s"): + vals = column(merges, mx.get(name)) + out[name] = statistics.median(vals) if vals else None + return out + + +def merge_line(view: dict) -> str: + mh = merge_health(view) + if not mh["rounds"]: + return f"{b('Merge')} no rounds yet" + parts = [f"{b('Merge')} direct-in-round {mh['direct']}/{mh['rounds']}" + f" ({100 * mh['direct'] / mh['rounds']:.0f}%)"] + if mh["round_s"] is not None: + parts.append(f"median round {fdur(mh['round_s'])} (quorum {fdur(mh['quorum_s'])}, grace {fdur(mh['grace_s'])}," + f" merge {fdur(mh['merge_s'])})") + if mh["missed"]: + parts.append("missed " + ", ".join(f"{k} x{v}" for k, v in sorted(mh["missed"].items()))) + return " | ".join(parts) + + +# -- commands ------------------------------------------------------------------- + +NAV_BUTTONS = [("Status", "/status"), ("Merge", "/merge"), ("Eval", "/eval"), ("Alerts", "/alerts")] + +HELP_LINES = [ + "/status - progress: outer step, ETA, islands, merge health, 6 h loss chart", + "/island - one island: last steps, loss / grad_norm / step-time charts", + "/eval - eval table, dloss vs references, chart", + "/merge - syncer rounds: direct-in-round ratio, missed islands, timing, problem events", + "/res - resources: HBM per island, host memory, syncer RSS", + "/ckpt - checkpoints: syncer writes and eval checkpoints", + "/alerts - active alerts and the last 24 h", + "/ack [cN|all] - mute the subjects of alert card N (or all) for 2 h", + "/mute <30m|2h|1d> [all] - mute P1/P2 (all: P0 too); /unmute", + "/help - this help; /start - register this chat for alerts", +] + + +def help_message(prefix: str | None = None) -> Message: + return Message("yeto training bot", ([prefix, ""] if prefix else []) + HELP_LINES, "info", NAV_BUTTONS) + + +def island_row(isl: dict, ix: dict[str, int]) -> str: + rows = isl.get("rows") or [] + st = isl.get("status") or "unknown" + step = last_value(rows, ix.get("step")) + return (f"{STATUS_DOT.get(st, '')} {b(island_label(isl))} {step if step is not None else '-'}/" + f"{isl.get('steps_total') or '?'} | loss {f3(last_value(rows, ix.get('loss')))}" + f" | {fdur(step_p50(rows, ix))}/step" + (f" | {st}" if st != "ok" else "")) + + +def syncer_line(sft: dict) -> str: + syn = sft.get("syncer") or {} + cur, total = syn.get("outer_step"), syn.get("total_steps") + pct = f" ({100 * cur / total:.1f}%)" if isinstance(cur, int) and isinstance(total, int) and total else "" + lrt, now = last_round_ts(sft), num(sft.get("now")) + line = f"{b('Outer')} {cur if cur is not None else '-'}/{total or '?'}{pct}" + if num(syn.get("eta_ts")): + line += f" | ETA {fclock(syn['eta_ts'])}" + if lrt is not None and now is not None: + line += f" | last round {fdur(now - lrt)} ago" + return line + + +def eval_refs(rows: list[dict]) -> list[str]: + """Reference names of eval rows: every ``vs_`` key is a comparison against + reference ```` (``vs_base``, a fixed checkpoint, ...), in first-seen order; + ``prev`` (the previous eval point) goes last.""" + names: list[str] = [] + for r in rows: + for k in r if isinstance(r, dict) else (): + if isinstance(k, str) and k.startswith("vs_") and len(k) > 3 and k[3:] not in names: + names.append(k[3:]) + return [n for n in names if n != "prev"] + (["prev"] if "prev" in names else []) + + +def ref_cmp(row: dict, name: str) -> dict: + c = row.get(f"vs_{name}") + return c if isinstance(c, dict) else {} + + +def eval_line(row: dict | None) -> str: + if not row: + return f"{b('Eval')} none yet" + parts = [f"{b('Eval')} step {row.get('step')} loss {f3(row.get('loss'))}"] + for name in eval_refs([row]): + d = num(ref_cmp(row, name).get("dloss")) + if d is not None: + parts.append(f"vs {name} {fsigned(d)}") + return " | ".join(parts) + + +def status_message(view: dict, *, alerts_line: str | None = None, chart_hours: float = 6.0) -> Message: + """The progress card (/status and the scheduled one).""" + ov = (view or {}).get("overview") or {} + sft = (view or {}).get("sft") + title = f"yeto progress: {ov.get('run') or 'run'}" + if not sft: + return rl_status_message(view, title) + lines = [syncer_line(sft)] + ix = step_cols(sft) + lines += [island_row(isl, ix) for isl in sft.get("islands") or []] or ["no SFT islands yet"] + lines += [merge_line(view), eval_line(((sft.get("eval") or {}).get("rows") or [None])[-1])] + if alerts_line: + lines.append(alerts_line) + chart = loss_chart(sft, chart_hours * 3600, f"loss, last {chart_hours:g} h") + return Message(title, lines, "report", NAV_BUTTONS, [chart] if not chart.empty() else []) + + +def rl_status_message(view: dict, title: str) -> Message: + ov = (view or {}).get("overview") or {} + gs = ov.get("global_status") or {} + lines = [f"{b('Fleet')} {gs.get('text') or '-'}"] + for c in ov.get("islands") or []: + lines.append(f"{STATUS_DOT.get(c.get('status'), '')} {b('island ' + str(c.get('id')))} {c.get('status')}" + f" | round {c.get('round', '-')} | last event {fdur(c.get('last_event_age_s'))} ago") + alerts = ov.get("alerts") or [] + if alerts: + lines.append(f"{len(alerts)} dashboard alert(s); first: {alerts[0].get('title')}") + return Message(title, lines, "report", NAV_BUTTONS) + + +def steps_table(rows: list[list], ix: dict[str, int], n: int = 10) -> list[str]: + out = [b("step | loss | grad_norm | step time | HBM GiB | tokens/s")] + for r in rows[-n:]: + c = {k: cell(r, ix.get(k)) for k in ("step", "loss", "grad_norm", "step_s", "hbm_gib", "tps")} + out.append(f"{c['step'] if c['step'] is not None else '-'} | {f3(c['loss'])} | {fsig(c['grad_norm'])}" + f" | {fdur(c['step_s'])} | {fsig(c['hbm_gib'])} | {fsig(c['tps'], 4)}") + return out + + +def unknown_island(sft: dict, key: str) -> Message: + names = ", ".join(f"{i.get('id')} ({i.get('name')})" for i in sft.get("islands") or []) or "none" + return Message("unknown island", [f"no island {key!r}; islands: {names}", "usage: /island "], "warn") + + +def island_message(view: dict, key: str, *, hbm_limit: float | None = None, chart_steps: int = 200) -> Message: + sft = (view or {}).get("sft") or {} + isl = find_island(sft, key) if key else None + if isl is None: + return unknown_island(sft, key) + rows = isl.get("rows") or [] + ix = step_cols(sft) + meta = [f"{k} {isl[k]}" for k in ("nodes", "layout", "cards", "dataset", "gbs") if isl.get(k) not in (None, "")] + lines = [island_row(isl, ix)] + if meta: + lines.append(" | ".join(meta)) + lines.append(f"HBM {fsig(last_value(rows, ix.get('hbm_gib')))}/{fsig(hbm_limit)} GiB" + f" | tokens/s {fsig(last_value(rows, ix.get('tps')), 4)}" + f" | MFU {fsig(last_value(rows, ix.get('mfu')))}%" + f" | lr {fsig(last_value(rows, ix.get('lr')))}" + f" | last step {fdur(isl.get('age_s'))} ago") + if isl.get("last_push"): + p = isl["last_push"] + lines.append(f"last push step {p.get('step')} at {fclock(p.get('ts'))} (export {fdur(p.get('export_s'))})") + if isl.get("last_apply"): + a = isl["last_apply"] + lines.append(f"last apply step {a.get('step')} at {fclock(a.get('ts'))} (apply {fdur(a.get('apply_s'))})") + lines += ["", *steps_table(rows, ix, 5)] + tail = rows[-chart_steps:] + xs = [cell(r, ix.get("step")) for r in tail] + charts = [Chart(f"{island_label(isl)} {name}", [(name, xs, [cell(r, ix.get(col)) for r in tail])], + xlabel="step") + for name, col in (("loss", "loss"), ("grad_norm", "grad_norm"), ("step time s", "step_s"))] + return Message(f"island {isl.get('id')}: {island_label(isl)}", lines, "report", + [("Status", "/status"), ("Resources", "/res")], [c for c in charts if not c.empty()]) + + +def merge_message(view: dict, limit: int = 8) -> Message: + sft = (view or {}).get("sft") or {} + syn = sft.get("syncer") or {} + mh = merge_health(view) + lines = [syncer_line(sft) if sft else "no syncer data", merge_line(view)] + if mh["weights"]: + lines.append("merge weight " + ", ".join(f"{k} {fsig(v)}" for k, v in sorted(mh["weights"].items()))) + timed_out = [r for r in syn.get("outer") or [] if len(r) > 4 and r[4]] + mx = merge_cols(sft) + + def mc(m: list, name: str) -> Any: + return cell(m, mx.get(name)) + + merges = [m for m in syn.get("merges") or [] if num(mc(m, "round_s")) is not None] + med = statistics.median([mc(m, "round_s") for m in merges]) if merges else None + slow = [m for m in merges if med and mc(m, "round_s") > 2 * med] + if timed_out: + lines.append(b(f"Timed-out rounds ({len(timed_out)})")) + lines += [f"outer {r[1]} frag {r[2]} at {fclock(r[0])}" for r in timed_out[-limit:]] + if slow: + lines.append(b(f"Slow rounds (> 2x median {fdur(med)}): {len(slow)}")) + lines += [f"outer {mc(m, 'step')} frag {mc(m, 'fragment')} round {fdur(mc(m, 'round_s'))}" + f" (quorum {fdur(mc(m, 'quorum_s'))}, grace {fdur(mc(m, 'grace_s'))}," + f" merge {fdur(mc(m, 'merge_s'))}) at {fclock(mc(m, 'ts'))}" + for m in slow[-limit:]] + events = [e for e in syn.get("events") or [] if len(e) > 1 and e[1] in PROBLEMS] + pc = problem_counts(sft) + lines.append(f"rejected {pc['rejected']} | timeout {pc['timeout']} | not carried {pc['not_carried']}") + if events: + lines += [f"{fclock(e[0])} {e[1]}" + (f" learner {e[2]}" if e[2] is not None else "") for e in events[-limit:]] + bad = [r for r in (view or {}).get("rounds_syncer") or [] if r.get("bad")] + for r in bad[-limit:]: + lines.append(f"round {r.get('round')}: missed {', '.join(map(str, r.get('missed') or [])) or '-'}" + f" | resend {r.get('resend', 0)}") + walls = [(mc(m, "ts"), mc(m, "round_s")) for m in merges[-200:]] + chart = Chart("syncer round time s", [("round", [w[0] for w in walls], [w[1] for w in walls])], time_axis=True) + return Message("yeto merge", lines, "report", [("Status", "/status")], [] if chart.empty() else [chart]) + + +def res_message(view: dict, *, hbm_limit: float | None, hbm_warn: float | None, rss_warn: float | None, + host_min: float) -> Message: + sft = (view or {}).get("sft") or {} + lines = [b(f"HBM peak (warn > {fsig(hbm_warn)}, limit {fsig(hbm_limit)} GiB)")] + c_hbm = step_cols(sft).get("hbm_gib") + for isl in sft.get("islands") or []: + rows = isl.get("rows") or [] + recent = column(rows[-20:], c_hbm) + lines.append(f"{island_label(isl)}: last {fsig(last_value(rows, c_hbm))} | max of 20 steps" + f" {fsig(max(recent) if recent else None)} GiB") + hosts = [(island_label(isl), node, ser[-1]) for isl in sft.get("islands") or [] + for node, ser in (isl.get("host") or {}).items() if ser] + lines.append(b(f"Host memory available (warn < {fbytes(host_min)})")) + lines += [f"{who} {node}: {fbytes(s[1])}" + (f" of {fbytes(s[2])}" if len(s) > 2 else "") + for who, node, s in hosts] or ["no host samples"] + lines.append(f"{b('Syncer RSS')} {fbytes(latest_rss(sft))} (warn > {fbytes(rss_warn)})") + rss = (sft.get("syncer") or {}).get("rss") or [] + chart = Chart("syncer RSS GB", [("rss", [r[0] for r in rss[-500:]], [r[1] / 1e9 for r in rss[-500:]])], + time_axis=True) + return Message("yeto resources", lines, "report", [("Status", "/status")], [] if chart.empty() else [chart]) + + +def ckpt_message(view: dict, limit: int = 8) -> Message: + sft = (view or {}).get("sft") or {} + syn = sft.get("syncer") or {} + writes = [e for e in syn.get("events") or [] if len(e) > 1 and e[1] in ("checkpoint", "checkpoint_failed")] + lines = [b(f"Syncer checkpoint writes ({len(writes)})")] + lines += [f"{fclock(e[0])} {'FAILED' if e[1] == 'checkpoint_failed' else 'written'}" + + (f" in {fdur(e[4])}" if len(e) > 4 and num(e[4]) is not None else "") for e in writes[-limit:]] \ + or ["none seen in syncer.log"] + ev = [r for r in (sft.get("eval") or {}).get("rows") or [] if r.get("ckpt_written")] + if ev: + lines.append(b("Evaluated checkpoints")) + lines += [f"step {r.get('step')} written {r.get('ckpt_written')} loss {f3(r.get('loss'))}" for r in ev[-limit:]] + isl_ck = [e for e in syn.get("events") or [] if len(e) > 1 and e[1] == "island_checkpoint"] + if isl_ck: # hook: island checkpoint events once the dashboard reports them + lines.append(b("Island checkpoints")) + lines += [f"{fclock(e[0])} learner {e[2]}" for e in isl_ck[-limit:]] + return Message("yeto checkpoints", lines, "report", [("Status", "/status")]) + + +def eval_message(view: dict, limit: int = 8) -> Message: + rows = (((view or {}).get("sft") or {}).get("eval") or {}).get("rows") or [] + if not rows: + return Message("yeto eval", ["no eval points yet"], "info") + last = rows[-1] + lines = eval_detail(last) + names = eval_refs(rows[-limit:]) + lines += ["", b(" | ".join(["step", "loss", *(f"vs {n}" for n in names)]))] + for r in rows[-limit:]: + d = [fsigned(ref_cmp(r, n).get("dloss")) for n in names] + lines.append(" | ".join([str(r.get("step")), f3(r.get("loss")), *d])) + return Message(f"yeto eval: step {last.get('step')}", lines, "report", [("Status", "/status")], + [c for c in [eval_chart(rows)] if not c.empty()]) + + +def eval_chart(rows: list[dict]) -> Chart: + xs = [r.get("step") for r in rows] + series = [("eval loss", xs, [r.get("loss") for r in rows])] + for name in eval_refs(rows): + if name != "prev": # the previous point's loss is the eval line itself, shifted + series.append((f"{name} ref", xs, [ref_cmp(r, name).get("ref_loss") for r in rows])) + return Chart("held-out eval loss", series, xlabel="syncer outer step") + + +def eval_detail(row: dict) -> list[str]: + lines = [f"loss {f3(row.get('loss'))} | ppl {fsig(row.get('ppl'), 4)}"] + for label in eval_refs([row]): + c = ref_cmp(row, label) + if num(c.get("dloss")) is None: + continue + better = num(c.get("better")) + se = num(c.get("se")) + lines.append(f"vs {label}: dloss {fsigned(c['dloss'])} (ref {f3(c.get('ref_loss'))})" + + (f" | paired se {fsig(se, 2)}" if se is not None else "") + + (f" | better on {100 * better:.0f}%" if better is not None else "")) + return lines + + +def details_message(view: dict, notices: list, *, ref: str) -> Message: + """Context for one alert card: its notices, the involved islands' last 10 steps, + the last syncer rounds when a syncer subject is involved.""" + sft = (view or {}).get("sft") or {} + ix, mx = step_cols(sft), merge_cols(sft) + lines = [] + for n in notices: + lines.append(f"{b(n.title)} ({n.kind})") + lines += [f" {x}" for x in n.lines if x] + islands = [] + for n in notices: + if n.island is not None and n.island not in islands: + islands.append(n.island) + for iid in islands: + isl = find_island(sft, str(iid)) + if isl: + head = f"{island_label(isl)} (status {isl.get('status')}, last step {fdur(isl.get('age_s'))} ago)" + lines += ["", b(head), *steps_table(isl.get("rows") or [], ix, 10)] + if any(n.key and n.key[0] in ("syncer", "dashboard") for n in notices): + merges = ((sft.get("syncer") or {}).get("merges") or [])[-10:] + lines += ["", syncer_line(sft) if sft else "", b("outer | frag | round | quorum | grace | merge | at")] + lines += [" | ".join([str(cell(m, mx.get("step"))), str(cell(m, mx.get("fragment")))] + + [fdur(cell(m, mx.get(k))) for k in ("round_s", "quorum_s", "grace_s", "merge_s")] + + [fclock(cell(m, mx.get("ts")))]) for m in merges] + return Message(f"details {ref}", lines or ["nothing recorded"], "report", [("Status", "/status")]) + + +def digest_message(view: dict, *, alert_counts: dict[str, int], top_alerts: list[str]) -> Message: + """Daily digest: 24 h progress, per-island throughput / MFU, eval, alerts.""" + sft = (view or {}).get("sft") or {} + ov = (view or {}).get("overview") or {} + now = num(sft.get("now")) or time.time() + lines = [syncer_line(sft) if sft else "no SFT data"] + syn = sft.get("syncer") or {} + outer_24 = [r for r in syn.get("outer") or [] if num(r[0]) is not None and r[0] >= now - 86400] + steps = [r[1] for r in outer_24 if isinstance(r[1], int)] + if steps: + lines.append(f"outer steps in 24 h: {max(steps) - min(steps) + 1}") + lines.append(b("Islands, last 24 h")) + ix = step_cols(sft) + for isl in sft.get("islands") or []: + rows = window(isl.get("rows") or [], now, 86400, ix.get("ts")) + loss = column(rows, ix.get("loss")) + lines.append(f"{STATUS_DOT.get(isl.get('status') or 'unknown', '')} {island_label(isl)}: {len(rows)} steps" + f" | loss {f3(loss[0] if loss else None)} -> {f3(loss[-1] if loss else None)}" + f" | {fdur(step_p50(rows, ix))}/step | tokens/s {fsig(_mean(column(rows, ix.get('tps'))), 4)}" + f" | MFU {fsig(_mean(column(rows, ix.get('mfu'))))}%") + rows = (sft.get("eval") or {}).get("rows") or [] + lines.append(eval_line(rows[-1] if rows else None)) + lines.append(merge_line(view)) + lines.append(f"{b('Alerts, last 24 h')} P0 {alert_counts.get('P0', 0)} | P1 {alert_counts.get('P1', 0)}" + f" | P2 {alert_counts.get('P2', 0)}") + lines += [f" {t}" for t in top_alerts[:5]] + charts = [loss_chart(sft, 86400, "loss, last 24 h"), eval_chart(rows)] + return Message(f"yeto daily digest: {ov.get('run') or 'run'} {time.strftime('%Y-%m-%d', time.localtime(now))}", + lines, "report", NAV_BUTTONS, [c for c in charts if not c.empty()]) + + +def _mean(vals: list[float]) -> float | None: + return statistics.fmean(vals) if vals else None diff --git a/yeto/bot/rules.py b/yeto/bot/rules.py new file mode 100644 index 00000000..e284f410 --- /dev/null +++ b/yeto/bot/rules.py @@ -0,0 +1,286 @@ +"""Alert rules over one ``/api/view`` snapshot. + +*Conditions* are states that hold while true (an island is stale, HBM is at the +limit); the monitor alerts on their transitions. *Events* are one-shot facts +discovered since the previous snapshot (new syncer log lines, a new eval point). + +Severities: 0 = P0 (page: red, @owner, re-notify until acked), 1 = P1 (orange, +cooldown per subject), 2 = P2 (grey, informational). + +Data the dashboard does not expose yet is wired as hooks that fire as soon as +the view carries it: ``disk`` samples (``[ts, free_bytes, total_bytes]`` under +``sft.syncer.disk`` or an island's ``disk``), syncer log categories ``oom`` / +``checkpoint_failed`` / ``island_checkpoint``, and per-island merge weights. +""" + +from __future__ import annotations + +import statistics +from dataclasses import dataclass, field + +from .config import Thresholds +from .render import (cell, column, eval_refs, fbytes, fdur, fsig, last_round_ts, latest_rss, nonfinite, num, + ref_cmp, refs, round_intervals, step_cols) + +P0, P1, P2 = 0, 1, 2 +SEV_NAME = {P0: "P0", P1: "P1", P2: "P2"} + +# syncer log / tape event categories -> (severity, label) +EVENT_RULES: dict[str, tuple[int, str]] = { + "oom": (P0, "out of memory"), + "checkpoint_failed": (P0, "checkpoint write failed"), + "rejected": (P1, "pushes rejected"), + "timeout": (P1, "quorum timeouts"), + "not_carried": (P1, "pushes not carried"), + "error": (P1, "syncer ERROR lines"), + "pool_leave": (P1, "island left the pool"), + "checkpoint": (P2, "syncer checkpoint written"), + "island_checkpoint": (P2, "island checkpoint saved"), + "joined": (P2, "learner (re)joined"), + "pool_join": (P2, "island joined the pool"), +} + + +@dataclass +class Condition: + sev: int + title: str + lines: list[str] = field(default_factory=list) + island: str | None = None + resolve: bool = True # announce when it clears + + +@dataclass +class Event: + sev: int + key: tuple + title: str + lines: list[str] = field(default_factory=list) + island: str | None = None + count: int = 1 + + +def island_name(isl: dict) -> str: + return isl.get("name") or f"island {isl.get('id')}" + + +def hbm_limit(view: dict, th: Thresholds) -> float | None: + return th.hbm_limit_gib if th.hbm_limit_gib is not None else num(refs(view).get("sft_hbm_limit_gib")) + + +def hbm_warn(view: dict, th: Thresholds) -> float | None: + """Explicit ``hbm_warn_gib``, else ``hbm_warn_frac`` of the HBM limit; None without a limit.""" + if th.hbm_warn_gib is not None: + return th.hbm_warn_gib + limit = hbm_limit(view, th) + return None if limit is None else th.hbm_warn_frac * limit + + +def rss_warn(view: dict, th: Thresholds) -> float | None: + return th.rss_warn_bytes if th.rss_warn_bytes is not None else num(refs(view).get("sft_syncer_rss_limit_bytes")) + + +def stuck_limit(sft: dict, th: Thresholds) -> float | None: + """Explicit ``stuck_round_s``, else the median recent round interval + 2 x grace.""" + if th.stuck_round_s is not None: + return th.stuck_round_s + gaps = round_intervals(sft) + if len(gaps) >= 3: + return statistics.median(gaps) + 2 * th.grace_s + return None + + +def missed_streaks(rounds: list[dict]) -> dict[str, int]: + """Island -> rounds missed in a row, counted back from the newest round.""" + streak: dict[str, int] = {} + for r in rounds: + missed = {str(x) for x in r.get("missed") or []} + for iid in list(streak): + if iid not in missed: + streak.pop(iid) + for iid in missed: + streak[iid] = streak.get(iid, 0) + 1 + return streak + + +def conditions(view: dict, th: Thresholds) -> dict[tuple, Condition]: + out: dict[tuple, Condition] = {} + sft = (view or {}).get("sft") or {} + hbm, warn = hbm_limit(view, th), hbm_warn(view, th) + ix = step_cols(sft) + names = {} + for isl in sft.get("islands") or []: + names[str(isl.get("id"))] = island_name(isl) + _island(isl, ix, hbm, warn, th, out) + if sft: + _syncer(sft, rss_warn(view, th), th, out) + for iid, n in missed_streaks((view or {}).get("rounds_syncer") or []).items(): + if n >= th.missed_rounds: + out[(iid, "missed")] = Condition(P1, f"{names.get(iid, 'island ' + iid)} not merged for {n} rounds", + [f"missed the last {n} syncer rounds in a row"], iid) + if th.forward_dashboard_alerts: + for a in ((view or {}).get("overview") or {}).get("alerts") or []: + text = f"{a.get('title')} {a.get('detail')}".lower() + if "out of memory" in text or " oom" in f" {text}": + out[("oom", a.get("island"))] = Condition(P0, f"OOM: {a.get('title')}", [str(a.get("detail") or "")], + a.get("island")) + elif a.get("sev") == 0: + out[("dash", a.get("rule"), a.get("island"))] = Condition( + P0, f"dashboard: {a.get('title')}", [str(a.get("detail") or "")], a.get("island")) + return out + + +def _island(isl: dict, ix: dict[str, int], hbm: float | None, hbm_warn_at: float | None, th: Thresholds, + out: dict) -> None: + """Island conditions; a rule whose column the view does not carry stays silent.""" + iid, name, rows = str(isl.get("id")), island_name(isl), isl.get("rows") or [] + c_step, c_loss, c_gn, c_lr, c_hbm = (ix.get(k) for k in ("step", "loss", "grad_norm", "lr", "hbm_gib")) + st = isl.get("status") + last_step = cell(rows[-1], c_step) if rows else None + last_step = "-" if last_step is None else last_step + if st == "done": + out[(iid, "done")] = Condition(P2, f"{name} finished", [f"step {last_step}/{isl.get('steps_total')}"], iid, + resolve=False) + return + if st == "stale": + out[(iid, "stale")] = Condition(P0, f"{name} stale (process dead or hung?)", + [f"no new step for {fdur(isl.get('age_s'))} (last step {last_step})"], iid) + # ms-swift logged the step (lr is set); a dashboard null loss on such a step is NaN + logged = [r for r in rows if cell(r, c_lr) is not None] if c_loss is not None else [] + if logged: + last = logged[-1] + step, loss = cell(last, c_step), cell(last, c_loss) + if loss is None or nonfinite(loss): + out[(iid, "loss_nan")] = Condition(P0, f"{name} loss NaN/Inf", + [f"step {step} logged a non-finite loss" + f" (grad_norm {fsig(cell(last, c_gn))})"], iid) + else: + prev = column(logged[:-1], c_loss)[-th.loss_window:] + if len(prev) >= th.loss_min_points: + mean, sd = statistics.fmean(prev), statistics.pstdev(prev) + limit = mean + th.loss_sigma * sd + if num(loss) is not None and loss > limit and sd > 0: + out[(iid, "loss_spike")] = Condition( + P1, f"{name} loss spike", [f"step {step} loss {fsig(loss, 4)} > mean" + f" {fsig(mean, 4)} + {th.loss_sigma:g} sd ({fsig(sd, 2)}) of the" + f" previous {len(prev)} steps"], iid) + gn = num(cell(last, c_gn)) + if gn is not None and gn > th.grad_norm_max: + out[(iid, "grad_norm")] = Condition(P1, f"{name} grad_norm high", + [f"step {step} grad_norm {fsig(gn)} > {th.grad_norm_max:g}"], + iid) + hbm_rows = [r for r in rows if num(cell(r, c_hbm)) is not None] + if hbm_rows: + step, peak = cell(hbm_rows[-1], c_step), hbm_rows[-1][c_hbm] + if hbm is not None and peak >= hbm: + out[(iid, "hbm")] = Condition(P0, f"{name} HBM peak at the limit", + [f"step {step} peak {fsig(peak)} GiB >= {fsig(hbm)} GiB"], iid) + elif hbm_warn_at is not None and peak > hbm_warn_at: + out[(iid, "hbm_warn")] = Condition(P1, f"{name} HBM peak high", + [f"step {step} peak {fsig(peak)} GiB > {fsig(hbm_warn_at)} GiB" + + (f" (limit {fsig(hbm)})" if hbm is not None else "")], iid) + for node, ser in (isl.get("host") or {}).items(): + if ser and num(ser[-1][1]) is not None and ser[-1][1] < th.host_avail_min_bytes: + out[(iid, "host", node)] = Condition(P1, f"{name} host {node} memory low", + [f"available {fbytes(ser[-1][1])} < {fbytes(th.host_avail_min_bytes)}" + + (f" of {fbytes(ser[-1][2])}" if len(ser[-1]) > 2 else "")], iid) + _disk(isl.get("disk"), (iid, "disk"), name, th, out, iid) + + +def _disk(samples, key: tuple, who: str, th: Thresholds, out: dict, island: str | None = None) -> None: + """Hook: fires once the dashboard exposes ``[ts, free_bytes, total_bytes]`` samples.""" + if not samples: + return + last = samples[-1] + free, total = (num(last[1]), num(last[2])) if len(last) > 2 else (None, None) + if free is not None and total and free / total < th.disk_free_min_frac: + out[key] = Condition(P0, f"{who} disk almost full", + [f"free {fbytes(free)} of {fbytes(total)} ({free / total:.1%})"], island) + + +def _syncer(sft: dict, rss_limit: float | None, th: Thresholds, out: dict) -> None: + syn = sft.get("syncer") or {} + cur, total = syn.get("outer_step"), syn.get("total_steps") + finished = isinstance(cur, int) and isinstance(total, int) and cur >= total + lrt, now = last_round_ts(sft), num(sft.get("now")) + limit = stuck_limit(sft, th) + if not finished and lrt is not None and now is not None and limit is not None and now - lrt > limit: + out[("syncer", "stuck")] = Condition(P0, "syncer outer step not advancing", + [f"outer step {cur}/{total or '?'} unchanged for {fdur(now - lrt)}" + f" (limit {fdur(limit)} = median round + 2 x grace)"]) + rss = latest_rss(sft) + if rss is not None and rss_limit and rss > rss_limit: + out[("syncer", "rss")] = Condition(P1, "syncer RSS high", [f"RSS {fbytes(rss)} > {fbytes(rss_limit)}"]) + _disk(syn.get("disk"), ("syncer", "disk"), "syncer host", th, out) + + +class EventTracker: + """High-water marks for one-shot events (syncer events by timestamp, eval by step).""" + + def __init__(self) -> None: + self.event_hw: float | None = None + self.eval_step: int | None = None + + def events(self, view: dict, *, seed: bool = False) -> list[Event]: + sft = (view or {}).get("sft") or {} + out = self._syncer_events(sft, seed) + self._eval_events(sft, seed) + return [] if seed else out + + def _syncer_events(self, sft: dict, seed: bool) -> list[Event]: + evs = [e for e in (sft.get("syncer") or {}).get("events") or [] + if len(e) > 1 and e[1] in EVENT_RULES and num(e[0]) is not None] + newest = max((e[0] for e in evs), default=None) + hw = self.event_hw + if newest is not None and (hw is None or newest > hw): + self.event_hw = newest + if seed or hw is None and not evs: + return [] + fresh = [e for e in evs if hw is None or e[0] > hw] + by_cat: dict[str, list] = {} + for e in fresh: + by_cat.setdefault(e[1], []).append(e) + out = [] + for cat, items in by_cat.items(): + sev, label = EVENT_RULES[cat] + who: dict = {} + for e in items: + lid = e[2] if len(e) > 2 and e[2] is not None else None + if lid is not None: + who[lid] = who.get(lid, 0) + 1 + lines = [", ".join(f"learner {k} x{v}" for k, v in sorted(who.items(), key=lambda kv: str(kv[0])))] \ + if who else [] + if cat == "checkpoint": + secs = [num(e[4]) for e in items if len(e) > 4 and num(e[4]) is not None] + if secs: + lines.append(f"write took {fdur(secs[-1])}") + out.append(Event(sev, ("syncer", cat), label, lines, count=len(items))) + return out + + def _eval_events(self, sft: dict, seed: bool) -> list[Event]: + rows = (sft.get("eval") or {}).get("rows") or [] + steps = [r["step"] for r in rows if isinstance(r.get("step"), int)] + if not steps: + return [] + newest, prev_seen = max(steps), self.eval_step + if prev_seen is not None and newest <= prev_seen: + return [] + self.eval_step = newest + if seed: + return [] + row = next(r for r in rows if r.get("step") == newest) + older = [r for r in rows if isinstance(r.get("step"), int) and r["step"] < newest] + prev = max(older, key=lambda r: r["step"]) if older else None + # "worse" compares the loss with the previous eval point, not with any named reference + d_prev = num(ref_cmp(row, "prev").get("dloss")) + if d_prev is None and prev is not None and None not in (num(row.get("loss")), num(prev.get("loss"))): + d_prev = row["loss"] - prev["loss"] + lines = [f"loss {fsig(row.get('loss'), 4)}" + (f" | vs previous point {d_prev:+.4f}" if d_prev is not None + else "")] + for name in eval_refs([row]): + d = num(ref_cmp(row, name).get("dloss")) + if d is not None and name != "prev": + lines.append(f"vs {name} {d:+.4f}") + if d_prev is not None and d_prev > 0: + return [Event(P1, ("eval", "worse"), f"eval step {newest} worse than step {prev['step'] if prev else '?'}", + lines)] + return [Event(P2, ("eval", "new"), f"new eval point: step {newest}", lines)]