From 06d518567e123292ade7cdfa26c7e35ba880cec3 Mon Sep 17 00:00:00 2001 From: Giovanni Ferri Date: Sat, 12 Sep 2026 02:09:37 +0100 Subject: [PATCH] feat: gate Imp runtime OnDelete rollouts --- scripts/rollout_imp_runtime.py | 155 ++++++++++++++++++++++++++++ scripts/test_rollout_imp_runtime.py | 56 ++++++++++ 2 files changed, 211 insertions(+) create mode 100644 scripts/rollout_imp_runtime.py create mode 100644 scripts/test_rollout_imp_runtime.py diff --git a/scripts/rollout_imp_runtime.py b/scripts/rollout_imp_runtime.py new file mode 100644 index 0000000..182fdb4 --- /dev/null +++ b/scripts/rollout_imp_runtime.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Safely replace stale imp-runtime OnDelete DaemonSet pods. + +Use only after Argo CD has synced a new immutable runtime image. The default is +read-only. --apply deletes at most one stale pod at a time and waits for its +replacement to be Ready on the same node with the desired image digest. + +The gate fails closed when any ImpVM is nonterminal. This prevents a runtime +restart while a one-use runner may be registering, claimed, or executing work. +Scale the runner pool to zero through GitOps and wait for all ImpVMs to finish +before rerunning with --apply. +""" + +import argparse +import json +import re +import subprocess +import sys +import time +from typing import Any + +DIGEST_RE = re.compile(r"sha256:[0-9a-f]{64}") + + +def kubectl(args: list[str], context: str | None, check: bool = True) -> subprocess.CompletedProcess[str]: + command = ["kubectl"] + if context: + command += ["--context", context] + command += args + result = subprocess.run(command, capture_output=True, text=True) + if check and result.returncode: + raise RuntimeError(result.stderr.strip() or "kubectl command failed") + return result + + +def get_json(args: list[str], context: str | None) -> dict[str, Any]: + return json.loads(kubectl([*args, "-o", "json"], context).stdout) + + +def digest(image: str) -> str: + match = DIGEST_RE.search(image) + if not match: + raise RuntimeError(f"expected an immutable sha256 image reference, got {image!r}") + return match.group(0) + + +def ready(pod: dict[str, Any]) -> bool: + return any( + condition.get("type") == "Ready" and condition.get("status") == "True" + for condition in pod.get("status", {}).get("conditions", []) + ) + + +def pod_digest(pod: dict[str, Any]) -> str | None: + statuses = pod.get("status", {}).get("containerStatuses", []) + if not statuses: + return None + image_id = statuses[0].get("imageID", "") + match = DIGEST_RE.search(image_id) + return match.group(0) if match else None + + +def selector_args(match_labels: dict[str, str]) -> list[str]: + if not match_labels: + raise RuntimeError("DaemonSet selector has no matchLabels; refusing to select pods") + return ["-l", ",".join(f"{key}={value}" for key, value in sorted(match_labels.items()))] + + +def nonterminal_vms(namespace: str, context: str | None) -> list[str]: + items = get_json(["-n", namespace, "get", "impvms.imp.dev"], context).get("items", []) + blocked = [] + for vm in items: + status = vm.get("status", {}) + phase = status.get("phase", "Unknown") + exit_code = status.get("runnerExitCode") + if phase != "Terminated" and exit_code is None: + blocked.append(f"{vm['metadata']['name']} (phase={phase})") + return blocked + + +def replacement_ready( + namespace: str, + selector: list[str], + node: str, + expected_digest: str, + context: str | None, + timeout_seconds: int, +) -> bool: + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + pods = get_json(["-n", namespace, "get", "pods", *selector], context).get("items", []) + candidates = [pod for pod in pods if pod.get("spec", {}).get("nodeName") == node] + if any(ready(pod) and pod_digest(pod) == expected_digest for pod in candidates): + return True + time.sleep(2) + return False + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--context", help="kubectl context; defaults to the current context") + parser.add_argument("--namespace", default="imp-system") + parser.add_argument("--daemonset", default="imp-runtime") + parser.add_argument("--timeout", type=int, default=300, help="per-pod replacement timeout in seconds") + parser.add_argument("--apply", action="store_true", help="delete stale pods after all safety gates pass") + args = parser.parse_args() + + try: + daemonset = get_json(["-n", args.namespace, "get", "daemonset", args.daemonset], args.context) + if daemonset.get("spec", {}).get("updateStrategy", {}).get("type") != "OnDelete": + raise RuntimeError("runtime DaemonSet is not OnDelete; refusing this manual rollout path") + + containers = daemonset.get("spec", {}).get("template", {}).get("spec", {}).get("containers", []) + if len(containers) != 1: + raise RuntimeError("expected exactly one runtime container; refusing ambiguous image verification") + expected_digest = digest(containers[0].get("image", "")) + selector = selector_args(daemonset["spec"]["selector"].get("matchLabels", {})) + pods = get_json(["-n", args.namespace, "get", "pods", *selector], args.context).get("items", []) + stale = [pod for pod in pods if pod_digest(pod) != expected_digest] + + blocked = nonterminal_vms(args.namespace, args.context) + if blocked: + raise RuntimeError( + "refusing runtime rollout while ImpVMs are nonterminal: " + + ", ".join(blocked) + + ". Scale the runner pool to zero through GitOps and wait for cleanup first." + ) + + if not stale: + print(f"OK: every {args.daemonset} pod already uses {expected_digest}") + return 0 + + names = ", ".join(pod["metadata"]["name"] for pod in stale) + if not args.apply: + print(f"DRY RUN: would replace stale pods one at a time: {names}") + return 0 + + for pod in stale: + name = pod["metadata"]["name"] + node = pod.get("spec", {}).get("nodeName") + if not node: + raise RuntimeError(f"pod {name} has no assigned node") + print(f"Replacing {name} on {node} ...") + kubectl(["-n", args.namespace, "delete", "pod", name, "--wait=false"], args.context) + if not replacement_ready(args.namespace, selector, node, expected_digest, args.context, args.timeout): + raise RuntimeError(f"replacement on {node} did not become Ready with {expected_digest}") + print(f"OK: replacement on {node} is Ready with {expected_digest}") + except (KeyError, RuntimeError, json.JSONDecodeError) as exc: + print(f"FAIL: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/test_rollout_imp_runtime.py b/scripts/test_rollout_imp_runtime.py new file mode 100644 index 0000000..c327892 --- /dev/null +++ b/scripts/test_rollout_imp_runtime.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +import contextlib +import importlib.util +import io +import pathlib +import sys +import unittest +from unittest.mock import patch + +SCRIPT = pathlib.Path(__file__).with_name("rollout_imp_runtime.py") +spec = importlib.util.spec_from_file_location("rollout_imp_runtime", SCRIPT) +if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {SCRIPT}") +rollout = importlib.util.module_from_spec(spec) +spec.loader.exec_module(rollout) + + +class RolloutGateTests(unittest.TestCase): + def test_refuses_to_delete_when_a_runner_vm_is_nonterminal(self): + daemonset = { + "spec": { + "updateStrategy": {"type": "OnDelete"}, + "selector": {"matchLabels": {"app": "imp-runtime"}}, + "template": { + "spec": { + "containers": [ + {"image": "ghcr.io/example/imp@sha256:" + "a" * 64} + ] + } + }, + } + } + stale_pod = { + "metadata": {"name": "imp-runtime-old"}, + "spec": {"nodeName": "runner-node"}, + "status": {"containerStatuses": [{"imageID": "sha256:" + "b" * 64}]}, + } + running_vm = { + "metadata": {"name": "github-runner-active"}, + "status": {"phase": "Running", "runnerHandoffAccepted": True}, + } + responses = iter([daemonset, {"items": [stale_pod]}, {"items": [running_vm]}]) + + with ( + patch.object(rollout, "get_json", side_effect=lambda *_: next(responses)), + patch.object(sys, "argv", ["rollout_imp_runtime.py", "--apply"]), + patch.object(rollout, "kubectl") as kubectl, + contextlib.redirect_stderr(io.StringIO()), + ): + self.assertEqual(rollout.main(), 1) + + kubectl.assert_not_called() + + +if __name__ == "__main__": + unittest.main()