From 31eb65dedeed2ee89b1440d5dd186d47c3c47051 Mon Sep 17 00:00:00 2001 From: Konrad Rokicki Date: Thu, 15 Oct 2026 09:44:27 -0400 Subject: [PATCH] mia-data: poller in its own module, job builders take JobArgs, checks split - poll.py: run_poll walks issues and PRs; deliver_label, wake_dataset_workflow and pass_status_without_snippet each handle one item and fail alone. Same behaviour, same log lines; the poll tests in test_mia_workflows cover it. - jobspec builders take the submit activity's JobArgs instead of unpacking it into up to twelve positional parameters; binding.yaml / record.yaml are module constants since every workspace uses those names. - cmd_check's csub policy check becomes _csub_problems with _mount_problems and _queue_problems; config._finish calls five named checks. Co-Authored-By: Claude Fable 5.1 --- mia_data/activities.py | 143 ++++-------------------------------- mia_data/cli.py | 123 ++++++++++++++++++------------- mia_data/config.py | 28 +++++-- mia_data/jobspec.py | 63 +++++++--------- mia_data/poll.py | 150 ++++++++++++++++++++++++++++++++++++++ tests/test_mia_jobspec.py | 31 ++++---- 6 files changed, 303 insertions(+), 235 deletions(-) create mode 100644 mia_data/poll.py diff --git a/mia_data/activities.py b/mia_data/activities.py index 0a17da3..06251c9 100644 --- a/mia_data/activities.py +++ b/mia_data/activities.py @@ -18,7 +18,6 @@ import re import shutil from collections.abc import Awaitable, Callable -from datetime import UTC, datetime, timedelta from pathlib import Path from typing import TYPE_CHECKING, Any, cast @@ -28,7 +27,6 @@ from autodev.activities import PERMANENT_CSUB, neutralize, permanent_errors from autodev.jobs import JobSpec, read_untrusted -from autodev.review import WRITE_PERMISSIONS from . import catalog from . import publish as crop_publish @@ -54,7 +52,6 @@ GalleryResult, JobArgs, JobDone, - LabelEvent, MatchResult, Pins, PrUpdate, @@ -62,6 +59,7 @@ SmokeResult, VerifyResult, ) +from .poll import run_poll from .publish import CROP_ID_RE, Conflict from .seed import ID_RE, MARKER_RE from .status import render_binding, render_crops_table, render_links, render_stage @@ -78,7 +76,6 @@ TEST_PR_RE = re.compile(r"i0-e(\d+)-[0-9a-f]{7}") # `mia-data test-pr`'s round id: PR at a head, no data issue ROUND_RE = re.compile(rf"i\d+-e\d+|{TEST_PR_RE.pattern}") -SNIPPET_ISSUE_RE = re.compile(r"") # open_snippet_issue's marker KIND_RE = re.compile(r"[a-z][a-z0-9-]{0,60}") # a job kind, optionally with a suffix (crop-) HEX64 = re.compile(r"[0-9a-f]{64}") NAME_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") # a snippet directory @@ -467,37 +464,24 @@ async def probe_hosts(self) -> list[str]: def _spec(self, kind: str, args: JobArgs, snapshot_dir: Path | None) -> JobSpec: """Build the JobSpec for one job kind from its arguments.""" - workspace, pins, cfg = Path(args.workspace), args.pins, self.cfg + cfg = self.cfg if kind == "match": - return match_job(cfg, workspace, pins, "record.yaml") - if kind in ("smoke", "crop") and snapshot_dir is None: - raise ValueError(f"{kind} job needs a snippets snapshot") - existing = self._dataset_dir(args.round_id, pins) / "manifest.json" if args.existing else None + return match_job(cfg, args) + existing = self._dataset_dir(args.round_id, args.pins) / "manifest.json" if args.existing else None existing = existing if existing and existing.is_file() else None - if kind == "smoke" and snapshot_dir: - return smoke_job(cfg, workspace, pins, "binding.yaml", snapshot_dir, "record.yaml") - if kind == "crop" and snapshot_dir: - return crop_job( - cfg, - workspace, - pins, - "binding.yaml", - snapshot_dir, - "record.yaml", - args.n_crops, - args.seed, - args.round_id, - args.batch, - existing, - args.trusted, - ) + if kind in ("smoke", "crop"): + if snapshot_dir is None: + raise ValueError(f"{kind} job needs a snippets snapshot") + if kind == "smoke": + return smoke_job(cfg, args, snapshot_dir) + return crop_job(cfg, args, snapshot_dir, existing) if kind == "verify": if args.crop_workspace is None: raise ValueError("verify job needs crop_workspace") crop_workspace = Path(args.crop_workspace) - if crop_workspace.parent != workspace.parent: + if crop_workspace.parent != Path(args.workspace).parent: raise ValueError(f"crop workspace {args.crop_workspace} is not of this round") - return verify_job(cfg, workspace, pins, self._crops_dir(crop_workspace), "binding.yaml", existing) + return verify_job(cfg, args, self._crops_dir(crop_workspace), existing) raise ValueError(f"unknown job kind {kind!r}") @activity.defn(name="submit") @@ -846,103 +830,6 @@ async def adopt_pr_binding(self, round_id: str, pr: int, sha: str) -> dict: @activity.defn(name="poll") @mia_permanent_errors async def poll(self) -> list[str]: - """One poll: deliver `data:go`/`data:approve` labels as signals, wake dataset workflows whose snippet - PR changed, and pass the `mia-data/crops` status on PRs that add no snippet. - - One bad issue or PR never skips the rest; the next poll retries it. Returns a log of what was done.""" - from temporalio.service import RPCError, RPCStatusCode - - cfg, gh_data, gh_snippets, done = self.cfg, self.gh_data, self.gh_snippets, [] - for kind in HUMAN_LABELS: - label = cfg.label(kind) - for issue in await gh_data.list_issues(labels=label, state="open"): - number = issue["number"] - try: - actor, event_id, labeled_at = await gh_data.label_event(number, label) - if not actor or await gh_data.permission(actor) not in WRITE_PERMISSIONS: - await gh_data.remove_label(number, label) - marker = cfg.marker(f"ignored {label} {event_id}") - await self._comment_once( - gh_data, - number, - f"Ignoring `{label}`: it was added by `{actor or 'unknown'}`, who does " - f"not have write access to this repository.\n\n{marker}", - marker, - ) - elif not (event_id and labeled_at): - log.warning("poll: #%s %s: label event has no id or time", number, label) - elif kind == "go" and "licence-excluded" in {lb["name"] for lb in issue.get("labels", [])}: - await self.refuse_label( - number, - label, - "the dataset's licence is not in `crops.licences` (`licence-excluded`)", - event_id, - ) - else: - # start_signal: starts the workflow if needed and delivers the label as a signal either way - await self.temporal.start_workflow( - "MiaDataDataset", - DatasetState(number), - id=cfg.workflow_id(number), - task_queue=cfg.task_queue, - start_signal=kind, - start_signal_args=[LabelEvent(kind, event_id, actor, labeled_at)], - ) - await gh_data.remove_label(number, label) # only once the signal is accepted - done.append(f"{kind} #{number}") - except Exception: - log.exception("poll: issue #%s %s failed", number, label) - recent = datetime.now(UTC) - timedelta(days=1) - # ponytail: the first page (50 most recently updated PRs); paginate if more than 50 are active - for pr_data in await gh_snippets.list_prs(state="all"): - try: - head = pr_data["head"] - if pr_data["state"] == "open" and not any( - status.get("context") == "mia-data/crops" for status in await gh_snippets.statuses(head["sha"]) - ): - files = await gh_snippets.pr_files(pr_data["number"]) - if len(files) < 100 and not any(file.startswith("snippets/") for file in files): # 100: a full page - await gh_snippets.set_status( - head["sha"], "mia-data/crops", "success", "no snippet in this PR", None - ) - done.append(f"status PR #{pr_data['number']}") - match = re.fullmatch(r"snippets/issue-(\d+)", head["ref"]) - if ( - not match - or (head.get("repo") or {}).get("full_name") != gh_snippets.repo - or (pr_data["state"] != "open" and datetime.fromisoformat(pr_data["updated_at"]) < recent) - ): - continue - snippet_issue = await gh_snippets.get_issue(int(match.group(1))) - issue_marker = SNIPPET_ISSUE_RE.match(snippet_issue.get("body") or "") - if ( - not issue_marker or (snippet_issue.get("user") or {}).get("login") != self.bot - ): # an outsider can copy the marker - continue - merged = bool(pr_data.get("merged_at")) - update = PrUpdate( - pr_data["number"], - head["sha"], - "merged" if merged else pr_data["state"], - pr_data.get("merge_commit_sha") if merged else None, - snippet_issue["number"], - "snippets:needs-human" in {lb["name"] for lb in snippet_issue.get("labels", [])}, - pr_data.get("html_url"), - ) - fingerprint = (update.head_sha, update.state, update.merge_sha, update.needs_human) - if self._sent.get(update.pr) == fingerprint: - continue # signalled already: a wake-up per change, not per poll - try: - await self.temporal.get_workflow_handle(cfg.workflow_id(int(issue_marker.group(1)))).signal( - "pr_update", update - ) - self._sent[update.pr] = fingerprint - done.append(f"pr_update PR #{pr_data['number']}") - except RPCError as error: - if error.status != RPCStatusCode.NOT_FOUND: # no dataset workflow (any more): nothing to wake - raise - except Exception: - log.exception("poll: PR #%s failed", pr_data.get("number")) - if done: - log.info("poll: %s", ", ".join(done)) - return done + """One scheduled pass over GitHub: labels become signals, changed snippet PRs wake their workflows + (see poll.py). Returns a log of what was done.""" + return await run_poll(self) diff --git a/mia_data/cli.py b/mia_data/cli.py index 6746634..ffda616 100644 --- a/mia_data/cli.py +++ b/mia_data/cli.py @@ -63,12 +63,81 @@ def _gh(cfg: Config, repo: str, tokens: Callable[[], str | None] | None = None) return GitHub(repo, tokens or token_fn(cfg), api_url=cfg.github.api_url) +def _wanted_hosts(cfg: Config) -> list[str]: + """Every host a mia-data job may need: package indexes, redirect sources and targets, the data hosts.""" + network = cfg.network + return [ + *network.package_hosts, + *network.redirect_hosts, + *(target for targets in network.redirect_hosts.values() for target in targets), + *DATA_HOSTS, + ] + + +def _csub_problems(cfg: Config, probe: Any, lsf_queues: set[str] | None) -> list[str]: + """Everything in the broker policy (`csub probe`) that would stop mia-data's jobs: hosts, the agent, env + vars, mounts, queues. Each message says what to change in broker.toml.""" + from .binding import missing_hosts + from .jobspec import HARNESS_ENV + + problems = [f"host not allowed: {host}" for host in missing_hosts(_wanted_hosts(cfg), list(probe.allowed_hosts))] + if "claude" not in probe.agents: + problems.append( + 'agent "claude" not permitted (broker.toml: add [agents] claude = ["api.anthropic.com", ' + '"claude.ai", "platform.claude.com"])' + ) + if missing := [var for var in HARNESS_ENV if not any(fnmatchcase(var, pattern) for pattern in probe.env_allow)]: + problems.append( + f'env not allowed: {", ".join(missing)} (broker.toml [broker]: add "MIA_HARNESS_*" to env_allow)' + ) + problems += _mount_problems(cfg, probe) + problems += _queue_problems(probe, lsf_queues) + return problems + + +def _mount_problems(cfg: Config, probe: Any) -> list[str]: + """The directories every job mounts must be under the policy's allowed_roots and away from $HOME, which + csub never mounts.""" + mounted = [cfg.harness.dir, cfg.paths.work_root, cfg.paths.storage_root] + problems = [] + if outside := [ + str(directory) + for directory in mounted + if not any(directory.is_relative_to(root) for root in probe.allowed_roots) + ]: + problems.append( + f"outside allowed_roots: {', '.join(outside)} (broker.toml [broker]: add " + f"{', '.join(json.dumps(directory) for directory in outside)} to allowed_roots)" + ) + home = Path.home().resolve() + if inside := [ + str(directory) for directory in mounted if directory.is_relative_to(home) or home.is_relative_to(directory) + ]: + problems.append( + f"under or containing the home directory {home}, which csub never mounts: {', '.join(inside)} " + "(move it out of $HOME)" + ) + return problems + + +def _queue_problems(probe: Any, lsf_queues: set[str] | None) -> list[str]: + """A policy queue table named after a queue LSF does not have fails every submit with "No such queue". + Skipped (None) where bqueues is not installed.""" + if lsf_queues is None: + return [] + if missing := [name for name, queue in probe.queues.items() if not queue.get("gpu") and name not in lsf_queues]: + return [ + f"not an LSF queue: {', '.join(missing)} (broker.toml: name the CPU [queues.] tables and " + f"[limits] cpu_short_queue/cpu_default_queue after real queues from `bqueues`, e.g. short, local)" + ] + return [] + + async def cmd_check(cfg: Config, args: argparse.Namespace) -> int: """Try every external dependency and print one ok/BAD line per check.""" from autodev.jobs import csub_client - from .binding import missing_hosts - from .jobspec import HARNESS_ENV, harness_env, harness_pin + from .jobspec import harness_env, harness_pin rows: list[tuple[str, bool, str]] = [] @@ -97,55 +166,11 @@ async def temporal(): return f"{cfg.temporal.address} ns={cfg.temporal.namespace}" async def csub(): - """Everything in the broker policy that would stop mia-data's jobs: hosts, the agent, env vars, mounts, queues.""" probe = await asyncio.to_thread(csub_client(cfg.deployment).probe) - network = cfg.network - want = [ - *network.package_hosts, - *network.redirect_hosts, - *(target for targets in network.redirect_hosts.values() for target in targets), - *DATA_HOSTS, - ] - bad = [f"host not allowed: {host}" for host in missing_hosts(want, list(probe.allowed_hosts))] - if "claude" not in probe.agents: - bad.append( - 'agent "claude" not permitted (broker.toml: add [agents] claude = ["api.anthropic.com", ' - '"claude.ai", "platform.claude.com"])' - ) - if missing := [var for var in HARNESS_ENV if not any(fnmatchcase(var, pattern) for pattern in probe.env_allow)]: - bad.append( - f'env not allowed: {", ".join(missing)} (broker.toml [broker]: add "MIA_HARNESS_*" to env_allow)' - ) - mounted = [cfg.harness.dir, cfg.paths.work_root, cfg.paths.storage_root] # every job mounts under these - if missing := [ - str(directory) - for directory in mounted - if not any(directory.is_relative_to(root) for root in probe.allowed_roots) - ]: - bad.append( - f"outside allowed_roots: {', '.join(missing)} (broker.toml [broker]: add " - f"{', '.join(json.dumps(directory) for directory in missing)} to allowed_roots)" - ) - home = Path.home().resolve() - if inside := [ - str(directory) for directory in mounted if directory.is_relative_to(home) or home.is_relative_to(directory) - ]: - bad.append( - f"under or containing the home directory {home}, which csub never mounts: {', '.join(inside)} " - "(move it out of $HOME)" - ) - # a policy queue table named after a queue LSF does not have fails every submit with "No such queue" - if (lsf_queues := await asyncio.to_thread(_lsf_queues)) is not None: - if missing := [ - name for name, queue in probe.queues.items() if not queue.get("gpu") and name not in lsf_queues - ]: - bad.append( - f"not an LSF queue: {', '.join(missing)} (broker.toml: name the CPU [queues.] tables and " - f"[limits] cpu_short_queue/cpu_default_queue after real queues from `bqueues`, e.g. short, local)" - ) - if bad: + lsf_queues = await asyncio.to_thread(_lsf_queues) + if bad := _csub_problems(cfg, probe, lsf_queues): raise RuntimeError("; ".join(bad)) - return f"broker user {probe.user}, {len(want)} hosts ok" + return f"broker user {probe.user}, {len(_wanted_hosts(cfg))} hosts ok" async def dirs(): for directory in (cfg.paths.state_dir, cfg.paths.work_root, cfg.paths.storage_root): diff --git a/mia_data/config.py b/mia_data/config.py index 355e906..5f7fdb0 100644 --- a/mia_data/config.py +++ b/mia_data/config.py @@ -143,33 +143,52 @@ def _finish(self) -> Config: """Cross-field checks and path normalization that pydantic's per-field validation cannot do.""" if not re.fullmatch(r"[a-z][a-z0-9-]{1,30}", self.deployment): raise ValueError(f"deployment {self.deployment!r}: use 2-31 chars of a-z, 0-9, '-'") + self._check_github() + self._resolve_paths() + self._check_limits() + self._check_agent() + if not self.crops.licences: + raise ValueError("crops.licences must not be empty") + self._check_network() + return self + + def _check_github(self) -> None: github = self.github if github.auth == "app" and not (github.app_id and github.installation_id and github.private_key_path): raise ValueError("github.auth=app needs github.app_id, installation_id and private_key_path") for attr in ("private_key_path", "token_file"): if getattr(github, attr): setattr(github, attr, getattr(github, attr).expanduser()) + + def _resolve_paths(self) -> None: + """Make every path absolute, and keep state_dir (secrets, never mounted) apart from the mounted roots.""" paths = self.paths for section in (paths, self.harness): for name, value in vars(section).items(): if isinstance(value, Path): setattr(section, name, value.expanduser().resolve()) - for other in (paths.work_root, paths.storage_root): # state_dir holds secrets and is never mounted + for other in (paths.work_root, paths.storage_root): if paths.state_dir == other or paths.state_dir in other.parents or other in paths.state_dir.parents: raise ValueError(f"paths.state_dir {paths.state_dir} overlaps {other}") + + def _check_limits(self) -> None: for name, value in vars(self.limits).items(): if not (math.isfinite(value) and value > 0): raise ValueError(f"limits.{name} must be finite and positive") if self.limits.max_stage_attempts != 1: raise ValueError("limits.max_stage_attempts: only 1 is supported (a failed job ends the round)") - for name in ("allow_hosts", "extra_mounts", "api_key_env"): # autodev's; mia-data's jobs never read them + + def _check_agent(self) -> None: + """The agent section is autodev's; mia-data's jobs never read these fields, so setting them is a mistake.""" + for name in ("allow_hosts", "extra_mounts", "api_key_env"): if getattr(self.agent, name): raise ValueError( f"agent.{name} is not used by mia-data (job hosts come from the binding and network.*, " "mounts from the job builders); remove it" ) - if not self.crops.licences: - raise ValueError("crops.licences must not be empty") + + def _check_network(self) -> None: + """No GitHub host may be allowed to a job, and redirect targets must be exact hosts.""" network = self.network redirect_targets = [target for targets in network.redirect_hosts.values() for target in targets] for host in network.package_hosts + list(network.redirect_hosts) + redirect_targets: @@ -177,7 +196,6 @@ def _finish(self) -> Config: raise ValueError(f"network: {host!r} is a GitHub host") if bad := [target for target in redirect_targets if "*" in target]: raise ValueError(f"network.redirect_hosts: {bad} must be exact hosts, no '*'") - return self @property def task_queue(self) -> str: diff --git a/mia_data/jobspec.py b/mia_data/jobspec.py index a951246..e377f22 100644 --- a/mia_data/jobspec.py +++ b/mia_data/jobspec.py @@ -38,7 +38,7 @@ if TYPE_CHECKING: from .config import Config - from .models import Pins + from .models import JobArgs, Pins SHA_RE = re.compile(r"[0-9a-f]{40}") MATCH_TOOLS = "Bash(harness smoke:*),Read,Write,Glob,Grep" @@ -51,6 +51,8 @@ "MIA_HARNESS_CACHE_WRITABLE", ) # every job's env; csub's env_allow must cover it quote = shlex.quote +BINDING_FILE = "binding.yaml" # the round's checked binding, copied into the workspace by prepare_job +RECORD_FILE = "record.yaml" # the dataset record, copied into every workspace class PinMoved(RuntimeError): @@ -291,13 +293,14 @@ def _spec( ) -def match_job(cfg: Config, workspace: Path, pins: Pins, record_path: str = "record.yaml") -> JobSpec: +def match_job(cfg: Config, args: JobArgs) -> JobSpec: """The match agent: Claude Code with the harness's `smoke` command as its only tool besides reading. The caller has written `prompt.md` (match_prompt), `record.yaml`, `downloads.json`, any `conversion.json` / `manifest.json`, `catalog/` and `SPEC.md` into the workspace, and made the snapshot of `pins.snippets_commit`. Anthropic hosts come from the csub policy's [agents] table, never from here.""" - record = yaml.safe_load(Path(workspace, record_path).read_text()) or {} + workspace, pins = Path(args.workspace), args.pins + record = yaml.safe_load((workspace / RECORD_FILE).read_text()) or {} snapshot_dir = _snapshot_dir(cfg, pins.snippets_commit) body = ( f"ln -sfn {quote(str(snapshot_dir))} snapshot && \\\n" @@ -320,62 +323,46 @@ def match_job(cfg: Config, workspace: Path, pins: Pins, record_path: str = "reco ) -def smoke_job(cfg: Config, workspace: Path, pins: Pins, binding_path: str, snapshot: Path, record_path: str) -> JobSpec: +def smoke_job(cfg: Config, args: JobArgs, snapshot: Path) -> JobSpec: """The worker's own smoke of a binding, so the agent's claim is checked. Reviewed code: cache read-write.""" - record, urls = _binding(workspace, binding_path) - body = ( - f"harness smoke {quote(binding_path)} --snippets {quote(str(snapshot))} " - f"--record {quote(record_path)} --out smoke.json" - ) - return _spec(cfg, "smoke", workspace, pins, body, hosts(urls, cfg), [snapshot], record=record, writable=True) + workspace = Path(args.workspace) + record, urls = _binding(workspace, BINDING_FILE) + body = f"harness smoke {BINDING_FILE} --snippets {quote(str(snapshot))} --record {RECORD_FILE} --out smoke.json" + return _spec(cfg, "smoke", workspace, args.pins, body, hosts(urls, cfg), [snapshot], record=record, writable=True) -def crop_job( - cfg: Config, - workspace: Path, - pins: Pins, - binding_path: str, - snapshot: Path, - record_path: str, - n_crops: int, - seed: str, - round_id: str, - batch: int, - existing: Path | None, - trusted: bool, -) -> JobSpec: - """`harness crop` into `/crops`. +def crop_job(cfg: Config, args: JobArgs, snapshot: Path, existing: Path | None) -> JobSpec: + """`harness crop` into `/crops`, `args.n_crops` of them from `args.seed`. - `trusted` (a matched or merged SHA, or production) mounts the dataset cache read-write; a PR head gets it - read-only and keeps a private cache. `existing` is a published manifest.json to avoid duplicating.""" - record, urls = _binding(workspace, binding_path) + `args.trusted` (a matched or merged SHA, or production) mounts the dataset cache read-write; a PR head gets + it read-only and keeps a private cache. `existing` is a published manifest.json to avoid duplicating.""" + workspace = Path(args.workspace) + record, urls = _binding(workspace, BINDING_FILE) body = ( - f"harness crop {quote(binding_path)} --snippets {quote(str(snapshot))} --record {quote(record_path)} " - f"--n {int(n_crops)} --seed {quote(seed)} --round {quote(round_id)} --batch {int(batch)} --out crops" - + (f" --existing {quote(str(existing))}" if existing else "") + f"harness crop {BINDING_FILE} --snippets {quote(str(snapshot))} --record {RECORD_FILE} " + f"--n {int(args.n_crops)} --seed {quote(args.seed)} --round {quote(args.round_id)} " + f"--batch {int(args.batch)} --out crops" + (f" --existing {quote(str(existing))}" if existing else "") ) return _spec( cfg, "crop", workspace, - pins, + args.pins, body, hosts(urls, cfg), [snapshot, _parent(existing)], record=record, - writable=trusted, + writable=args.trusted, ) -def verify_job( - cfg: Config, workspace: Path, pins: Pins, crops_dir: Path, binding_path: str, existing: Path | None -) -> JobSpec: +def verify_job(cfg: Config, args: JobArgs, crops_dir: Path, existing: Path | None) -> JobSpec: """`harness verify` then `harness preview` on a crop job's output, mounted read-only. Loads no snippet code: no network, no dataset cache.""" crops = quote(str(crops_dir)) body = ( - f"harness verify {crops} --binding {quote(binding_path)}" + f"harness verify {crops} --binding {BINDING_FILE}" + (f" --existing {quote(str(existing))}" if existing else "") + f" --out verify.json; v=$?\nharness preview {crops} --out preview && (exit $v)" ) # overlays even if a crop fails - return _spec(cfg, "verify", workspace, pins, body, [], [crops_dir, _parent(existing)]) + return _spec(cfg, "verify", Path(args.workspace), args.pins, body, [], [crops_dir, _parent(existing)]) diff --git a/mia_data/poll.py b/mia_data/poll.py new file mode 100644 index 0000000..c017bb8 --- /dev/null +++ b/mia_data/poll.py @@ -0,0 +1,150 @@ +"""The poller: one scheduled pass that turns GitHub state into workflow signals. + +Three jobs, each done per item so one bad issue or PR never skips the rest (the next poll retries it): + +- deliver a `data:go` / `data:approve` label on a dataset issue as a signal, starting the dataset workflow if needed; +- wake the dataset workflow whose snippet PR changed (head, state, merge, needs-human); +- pass the `mia-data/crops` status on snippet PRs that add no snippet, so it never blocks them. + +Runs as the `poll` activity (Activities.poll) and uses the activity object's clients and helpers. +""" + +from __future__ import annotations + +import logging +import re +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING + +from autodev.review import WRITE_PERMISSIONS + +from .config import HUMAN_LABELS +from .models import DatasetState, LabelEvent, PrUpdate + +if TYPE_CHECKING: + from .activities import Activities + +log = logging.getLogger(__name__) + +SNIPPET_ISSUE_RE = re.compile(r"") # open_snippet_issue's marker + + +async def run_poll(activities: Activities) -> list[str]: + """One pass over the data repo's labelled issues and the snippets repo's recent PRs. Returns a log of + what was done.""" + done: list[str] = [] + for kind in HUMAN_LABELS: + label = activities.cfg.label(kind) + for issue in await activities.gh_data.list_issues(labels=label, state="open"): + try: + if await deliver_label(activities, kind, label, issue): + done.append(f"{kind} #{issue['number']}") + except Exception: + log.exception("poll: issue #%s %s failed", issue["number"], label) + recent = datetime.now(UTC) - timedelta(days=1) + # ponytail: the first page (50 most recently updated PRs); paginate if more than 50 are active + for pr_data in await activities.gh_snippets.list_prs(state="all"): + try: + if pr_data["state"] == "open" and await pass_status_without_snippet(activities, pr_data): + done.append(f"status PR #{pr_data['number']}") + if await wake_dataset_workflow(activities, pr_data, recent): + done.append(f"pr_update PR #{pr_data['number']}") + except Exception: + log.exception("poll: PR #%s failed", pr_data.get("number")) + if done: + log.info("poll: %s", ", ".join(done)) + return done + + +async def deliver_label(activities: Activities, kind: str, label: str, issue: dict) -> bool: + """Turn one human label on a dataset issue into a signal, or remove it with a comment saying why not. + + Only a person with write access counts; `go` on a licence-excluded dataset is refused. The label is removed + only once Temporal has accepted the signal, so a crash in between leaves it for the next poll. Returns True + when a signal was delivered.""" + cfg, gh_data, number = activities.cfg, activities.gh_data, issue["number"] + actor, event_id, labeled_at = await gh_data.label_event(number, label) + if not actor or await gh_data.permission(actor) not in WRITE_PERMISSIONS: + await gh_data.remove_label(number, label) + marker = cfg.marker(f"ignored {label} {event_id}") + await activities._comment_once( + gh_data, + number, + f"Ignoring `{label}`: it was added by `{actor or 'unknown'}`, who does " + f"not have write access to this repository.\n\n{marker}", + marker, + ) + return False + if not (event_id and labeled_at): + log.warning("poll: #%s %s: label event has no id or time", number, label) + return False + if kind == "go" and "licence-excluded" in {entry["name"] for entry in issue.get("labels", [])}: + await activities.refuse_label( + number, label, "the dataset's licence is not in `crops.licences` (`licence-excluded`)", event_id + ) + return False + # start_signal: starts the workflow if needed and delivers the label as a signal either way + await activities.temporal.start_workflow( + "MiaDataDataset", + DatasetState(number), + id=cfg.workflow_id(number), + task_queue=cfg.task_queue, + start_signal=kind, + start_signal_args=[LabelEvent(kind, event_id, actor, labeled_at)], + ) + await gh_data.remove_label(number, label) # only once the signal is accepted + return True + + +async def pass_status_without_snippet(activities: Activities, pr_data: dict) -> bool: + """Set `mia-data/crops` to success on an open PR that touches nothing under snippets/, if it has no such + status yet. Returns True when the status was set.""" + gh_snippets, head_sha = activities.gh_snippets, pr_data["head"]["sha"] + if any(status.get("context") == "mia-data/crops" for status in await gh_snippets.statuses(head_sha)): + return False + files = await gh_snippets.pr_files(pr_data["number"]) + if len(files) >= 100 or any(file.startswith("snippets/") for file in files): # 100: a full page, unknown rest + return False + await gh_snippets.set_status(head_sha, "mia-data/crops", "success", "no snippet in this PR", None) + return True + + +async def wake_dataset_workflow(activities: Activities, pr_data: dict, recent: datetime) -> bool: + """Signal `pr_update` to the dataset workflow behind a `snippets/issue-N` PR when something about the PR + changed since the last poll. Closed PRs older than `recent` are skipped. Returns True when a signal was sent.""" + from temporalio.service import RPCError, RPCStatusCode + + gh_snippets, head = activities.gh_snippets, pr_data["head"] + match = re.fullmatch(r"snippets/issue-(\d+)", head["ref"]) + if ( + not match + or (head.get("repo") or {}).get("full_name") != gh_snippets.repo + or (pr_data["state"] != "open" and datetime.fromisoformat(pr_data["updated_at"]) < recent) + ): + return False + snippet_issue = await gh_snippets.get_issue(int(match.group(1))) + issue_marker = SNIPPET_ISSUE_RE.match(snippet_issue.get("body") or "") + if not issue_marker or (snippet_issue.get("user") or {}).get("login") != activities.bot: + return False # not one of ours: an outsider can copy the marker, but not post as the bot + merged = bool(pr_data.get("merged_at")) + update = PrUpdate( + pr_data["number"], + head["sha"], + "merged" if merged else pr_data["state"], + pr_data.get("merge_commit_sha") if merged else None, + snippet_issue["number"], + "snippets:needs-human" in {entry["name"] for entry in snippet_issue.get("labels", [])}, + pr_data.get("html_url"), + ) + fingerprint = (update.head_sha, update.state, update.merge_sha, update.needs_human) + if activities._sent.get(update.pr) == fingerprint: + return False # signalled already: a wake-up per change, not per poll + try: + workflow_id = activities.cfg.workflow_id(int(issue_marker.group(1))) + await activities.temporal.get_workflow_handle(workflow_id).signal("pr_update", update) + except RPCError as error: + if error.status != RPCStatusCode.NOT_FOUND: # no dataset workflow (any more): nothing to wake + raise + return False + activities._sent[update.pr] = fingerprint + return True diff --git a/tests/test_mia_jobspec.py b/tests/test_mia_jobspec.py index 9fe5579..423b31e 100644 --- a/tests/test_mia_jobspec.py +++ b/tests/test_mia_jobspec.py @@ -25,7 +25,7 @@ snapshot, verify_job, ) -from mia_data.models import Pins +from mia_data.models import JobArgs, Pins RECORD = {"data": {"download_url": "https://data.example.org/a.zip"}} BINDING = { @@ -54,6 +54,11 @@ def _pins(cfg, record="rec-1"): return Pins(record, "c" * 40, "a" * 40, *harness_pin(cfg)) +def _args(ws, pins, **job_args): + """A submit activity's JobArgs for workspace `ws` in round i1-e2.""" + return JobArgs(str(ws), "i1-e2", pins, **job_args) + + def _ws(tmp_path): ws = tmp_path / "ws" ws.mkdir() @@ -146,7 +151,7 @@ def test_a_new_harness_commit_does_not_touch_an_admitted_round(tmp_path): run_git("add", "-A", cwd=cfg.harness.dir) run_git(*GIT_ID, "commit", "-qm", "moved", cwd=cfg.harness.dir) assert_pin(cfg, pins) # HEAD moved, env didn't - s = smoke_job(cfg, ws, pins, "binding.yaml", tmp_path / "snap", "record.yaml") + s = smoke_job(cfg, _args(ws, pins), tmp_path / "snap") assert f"PYTHONPATH={snap}/src" in s.command and str(snap) in s.ro_mounts assert (snap / "src" / "mia_harness" / "cli.py").read_text() == "def main(): return 0\n" # still the admitted one assert harness_snapshot(cfg, harness_pin(cfg)[0]) != snap # the new commit gets its own @@ -171,7 +176,7 @@ def test_assert_pin_raises_when_lock_changes(tmp_path): def test_match_job(tmp_path): cfg, ws = _cfg(tmp_path), _ws(tmp_path) pins = _pins(cfg) - s = match_job(cfg, ws, pins, "record.yaml") + s = match_job(cfg, _args(ws, pins)) assert s.allow_hosts == sorted(["data.example.org", "cdn.example.net", *cfg.network.package_hosts]) assert not any("github" in h or "anthropic" in h or "claude" in h for h in s.allow_hosts) assert s.agent == "claude" and f"--max-budget-usd {cfg.limits.agent_max_budget_usd}" in s.command @@ -187,9 +192,7 @@ def test_match_job(tmp_path): @pytest.mark.parametrize("trusted", [False, True]) def test_crop_job_cache_mode(tmp_path, trusted): cfg, ws = _cfg(tmp_path), _ws(tmp_path) - s = crop_job( - cfg, ws, _pins(cfg), "binding.yaml", tmp_path / "snap", "record.yaml", 4, "s-t", "i1-e2", 0, None, trusted - ) + s = crop_job(cfg, _args(ws, _pins(cfg), n_crops=4, seed="s-t", batch=0, trusted=trusted), tmp_path / "snap", None) cache = str(cfg.paths.work_root / "_cache" / "rec-1") assert (cache in s.rw_mounts, cache in s.ro_mounts) == (trusted, not trusted) assert s.env["MIA_HARNESS_CACHE_WRITABLE"] == ("1" if trusted else "0") and Path(cache).is_dir() @@ -199,7 +202,7 @@ def test_crop_job_cache_mode(tmp_path, trusted): def test_crop_job_mounts_the_existing_manifests_dir(tmp_path): cfg, ws, m = _cfg(tmp_path), _ws(tmp_path), tmp_path / "ds" / "manifest.json" - s = crop_job(cfg, ws, _pins(cfg), "binding.yaml", tmp_path / "snap", "record.yaml", 4, "s", "i1-e2", 1, m, True) + s = crop_job(cfg, _args(ws, _pins(cfg), n_crops=4, seed="s", batch=1, trusted=True), tmp_path / "snap", m) body = next(line for line in s.command.split("\n") if line.startswith("harness crop")) assert body.endswith(f"--existing {m}") assert str(m.parent) in s.ro_mounts and str(m) not in s.ro_mounts # csub mounts directories only @@ -209,7 +212,7 @@ def test_verify_job(tmp_path): cfg, ws = _cfg(tmp_path), _ws(tmp_path) crops = tmp_path / "crop-ws" / "crops" m = tmp_path / "ds" / "manifest.json" - s = verify_job(cfg, ws, _pins(cfg), crops, "binding.yaml", m) + s = verify_job(cfg, _args(ws, _pins(cfg)), crops, m) assert s.allow_hosts == [] and str(crops) in s.ro_mounts and f"--existing {m} " in s.command assert str(m.parent) in s.ro_mounts and str(m) not in s.ro_mounts # csub mounts directories only assert s.rw_mounts == [] and s.env["MIA_HARNESS_CACHE_DIR"] == "" and "harness preview" in s.command @@ -218,7 +221,7 @@ def test_verify_job(tmp_path): def test_check_hosts_allowed(tmp_path): cfg, ws = _cfg(tmp_path), _ws(tmp_path) - s = smoke_job(cfg, ws, _pins(cfg), "binding.yaml", tmp_path / "snap", "record.yaml") + s = smoke_job(cfg, _args(ws, _pins(cfg)), tmp_path / "snap") check_hosts_allowed(s, ["*.example.org", "cdn.example.net", *cfg.network.package_hosts]) with pytest.raises(HostsNotAllowed, match="data.example.org") as e: check_hosts_allowed(s, ["cdn.example.net", "*.pypi.org", *cfg.network.package_hosts]) @@ -228,12 +231,10 @@ def test_check_hosts_allowed(tmp_path): def _builders(cfg, ws, pins, tmp_path): snap = tmp_path / "snap" return [ - lambda: match_job(cfg, ws, pins, "record.yaml"), - lambda: smoke_job(cfg, ws, pins, "binding.yaml", snap, "record.yaml"), - lambda: crop_job( - cfg, ws, pins, "binding.yaml", snap, "record.yaml", 4, "s", "i1-e2", 1, tmp_path / "m.json", True - ), - lambda: verify_job(cfg, ws, pins, tmp_path / "crops", "binding.yaml", None), + lambda: match_job(cfg, _args(ws, pins)), + lambda: smoke_job(cfg, _args(ws, pins), snap), + lambda: crop_job(cfg, _args(ws, pins, n_crops=4, seed="s", batch=1, trusted=True), snap, tmp_path / "m.json"), + lambda: verify_job(cfg, _args(ws, pins), tmp_path / "crops", None), ]