diff --git a/scripts/ci/e2e_runner_pool.py b/scripts/ci/e2e_runner_pool.py index 9c5a19dbef69..cd3f265b60af 100644 --- a/scripts/ci/e2e_runner_pool.py +++ b/scripts/ci/e2e_runner_pool.py @@ -120,6 +120,7 @@ import argparse import dataclasses +import json import datetime as dt import os from collections.abc import Callable, Mapping, Sequence @@ -240,8 +241,9 @@ def ui_owned_runner(label: str | None, *, test_filter: str | None, owned: str | """A UI run's pool, moved off Blacksmith onto an owned pool with machines (see the module docstring). Any other run, or an owned label, comes back unchanged, as does every - label when no owned pool of the lane's Xcode pin has a slot count: a - fleet drained by zeroing CI_OWNED_POOL_SLOTS keeps UI runs off it. + label when no owned pool of the lane's Xcode pin has machines in + `owned_slots` (the online runners when main() read them, else + CI_OWNED_POOL_SLOTS): a drained fleet, its runners offline, keeps UI runs off it. """ if (not ui_run(test_filter) or not label or pr_runner_pool.persistent(label) or (owned or "").strip() != "1" or not owned_target(test_filter, owned_ui)): @@ -496,28 +498,39 @@ def resolve( owned_slots=owned_slots, pr_xcode_app=pr_xcode_app, log=log) or label -def read_live_owned(repo: str, env: Mapping[str, str], owned: str | None, - pr_xcode_app: str | None) -> tuple[dict[str, int], dict[str, int]] | None: - """Idle and online runners per owned label (and its root label) from the runners API, or None. +def read_live_runners(repo: str, env: Mapping[str, str], owned: str | None, + pr_xcode_app: str | None) -> list[Mapping[str, Any]] | None: + """The repository's runners from the runners API, or None. - Needs the org App's token (ROUTE_TOKEN) and owned pools on; any error - leaves the snapshot to decide. + Needs the org App's token (ROUTE_TOKEN), owned pools on and a pin that + names an owned pool; any error leaves the snapshot and CI_OWNED_POOL_SLOTS to decide. """ token = (env.get("ROUTE_TOKEN") or "").strip() - if not token or not repo or (owned or "").strip() != "1": - return None - labels = pr_runner_pool.owned_pools(pr_xcode_app) - labels += tuple(pr_runner_pool.root_label(label) for label in labels) - if not labels: + if not token or not repo or (owned or "").strip() != "1" or not pr_runner_pool.owned_pools(pr_xcode_app): return None try: - runners = pr_runner_pool.GitHub(token, repo).runners() + return pr_runner_pool.GitHub(token, repo).runners() except Exception as error: # noqa: BLE001 - the snapshot path still decides print(f"could not list runners ({error}); using the snapshot", file=sys.stderr) return None + + +def live_owned(runners: Sequence[Mapping[str, Any]] | None, + pr_xcode_app: str | None) -> tuple[dict[str, int], dict[str, int]] | None: + """Idle and online runners per owned label (and its root label), or None without a listing.""" + if runners is None: + return None + labels = pr_runner_pool.owned_pools(pr_xcode_app) + labels += tuple(pr_runner_pool.root_label(label) for label in labels) return pr_runner_pool.live_owned_free(runners, labels), pr_runner_pool.live_online(runners, labels) +def read_live_owned(repo: str, env: Mapping[str, str], owned: str | None, + pr_xcode_app: str | None) -> tuple[dict[str, int], dict[str, int]] | None: + """Idle and online runners per owned label (and its root label) from the runners API, or None.""" + return live_owned(read_live_runners(repo, env, owned, pr_xcode_app), pr_xcode_app) + + def main(argv: Sequence[str] | None = None, env: Mapping[str, str] | None = None) -> int: env = os.environ if env is None else env parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) @@ -544,12 +557,18 @@ def main(argv: Sequence[str] | None = None, env: Mapping[str, str] | None = None run_id = (env.get("GITHUB_RUN_ID") or "").strip() now = dt.datetime.now(dt.timezone.utc) + # A run no owned pool may take (a UI filter without owned_ui) reads no runners. + runners = read_live_runners(repo, env, args.owned if owned_target(args.test_filter, args.owned_ui) else "", + args.pr_xcode_app) + # Which owned labels route (a pool, its root label): the online runners when they were read, + # CI_OWNED_POOL_SLOTS only when they could not be (pr_runner_pool.routing_slots()). + owned_slots = (args.owned_slots if runners is None else + json.dumps(pr_runner_pool.routing_slots(args.owned_slots, args.pr_xcode_app, runners))) + def measure() -> PoolLoad | None: if not token or not repo: raise RuntimeError("GH_TOKEN and GH_REPO are required") - # A run no owned pool may take (a UI filter without owned_ui) reads no runners. - owned = args.owned if owned_target(args.test_filter, args.owned_ui) else "" - idle, online = read_live_owned(repo, env, owned, args.pr_xcode_app) or (None, None) + idle, online = live_owned(runners, args.pr_xcode_app) or (None, None) return measure_load(pr_runner_pool.GitHub(token, repo), now=now, exclude_run_id=int(run_id) if run_id.isdigit() else None, live_owned=idle, live_online=online) @@ -557,7 +576,7 @@ def measure() -> PoolLoad | None: print(resolve( args.requested, args.variable, overflow=args.overflow, order=args.order, max_queued=args.max_queued, - owned=args.owned, owned_slots=args.owned_slots, pr_xcode_app=args.pr_xcode_app, + owned=args.owned, owned_slots=owned_slots, pr_xcode_app=args.pr_xcode_app, test_filter=args.test_filter, owned_ui=args.owned_ui, queue_rounds=args.queue_rounds, measure=measure, now=now, log=lambda message: print(message, file=sys.stderr), diff --git a/tests/test_run_e2e.py b/tests/test_run_e2e.py index 6197aec45887..d0f38aa1c388 100644 --- a/tests/test_run_e2e.py +++ b/tests/test_run_e2e.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 """Exercise the focused-run launcher against a fake GitHub CLI.""" import importlib.util +import io import json import re import os @@ -1714,6 +1715,23 @@ def test_live_owned_needs_the_token_and_owned_pools(self): mock.patch("sys.stderr"): self.assertIsNone(read("manaflow-ai/cmux", {"ROUTE_TOKEN": "t"}, "1", "/Applications/Xcode_26.6.app")) + def test_main_routes_by_the_online_runners_not_the_slot_variable(self): + root = self.pool.pr_runner_pool.root_label(MINI) + runners = [{"status": "online", "busy": False, "labels": [{"name": MINI}, {"name": root}]}] + argv = ["--requested", MINI, "--owned", "1", "--owned-slots", '{"std": 40}', + "--pr-xcode-app", "/Applications/Xcode_26.6.app"] + env = {"ROUTE_TOKEN": "t", "GITHUB_REPOSITORY": "manaflow-ai/cmux"} + + def run(listing): + with mock.patch.object(self.pool.pr_runner_pool.GitHub, "runners", **listing), \ + mock.patch("sys.stderr"), mock.patch("sys.stdout", new_callable=io.StringIO) as out: + self.assertEqual(self.pool.main(argv, env), 0) + return out.getvalue().strip() + # An online root runner turns root routing on without a root count in CI_OWNED_POOL_SLOTS. + self.assertEqual(run({"return_value": runners}), root) + # No listing: the variable decides, and it has no root count. + self.assertEqual(run({"side_effect": RuntimeError("403")}), MINI) + def test_the_workflow_mints_the_routing_token_for_auto_only(self): steps = self.jobs[next(name for name, job in self.jobs.items() if any(step.get("id") == "pool" for step in job.get("steps", [])))]["steps"]