stack/dev/scripts/marimo_latest.py

"""Keep the notebooks image on marimo's latest release.

Run nightly by the ``marimo-latest`` systemd user timer, from the main
checkout:

1. Find the newest stable marimo release on GitHub that is at least
   ``--cooldown-days`` old (default 7: the stack's release cooldown, as in
   renovate.json's minimumReleaseAge and pyproject's exclude-newer).
2. If it is newer than ``ARG MARIMO_VERSION`` in
   infra/images/notebooks.Dockerfile: bump the pin, build the image, gate it
   with nb_fe_smoke.py --image (the overlay's own guards already fail the
   build if marimo's chrome controls come back), commit the pin on main and
   push, then roll the notebooks container (HEAVY_TAG = the new commit) and
   probe the live service.
3. On failure the pin is restored (a live failure also rolls HEAVY_TAG
   back), the version is recorded so later runs skip it until a newer
   release, and a Gitea issue is filed.

A no-op when the pin is current; ``--dry-run`` only reports. Rolling ends
open marimo sessions, hence the early-morning timer.

Units live in ~/.config/systemd/user/ (like bib-refresh-iom):

  marimo-latest.service  Type=oneshot, WorkingDirectory=/home/kert/stack,
                         EnvironmentFile=/home/kert/stack/.env,
                         Environment=DOCKER_HOST=unix:///run/user/1000/docker.sock,
                         ExecStart=/usr/bin/python3 dev/scripts/marimo_latest.py
  marimo-latest.timer    OnCalendar=*-*-* 05:15:00, Persistent=true

Stdlib only.
"""

from __future__ import annotations

import argparse
import datetime as dt
import fcntl
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
DOCKERFILE = "infra/images/notebooks.Dockerfile"
# Everything the notebooks image COPYs; the build reads the working tree, so
# these must match HEAD or the image would carry edits the commit doesn't.
IMAGE_INPUTS = [DOCKERFILE, "pyproject.toml", "infra/marimo/theme"]
ZOTERO_INPUTS = ["infra/images/zotero.Dockerfile", "infra/images/zotero"]
STATE = ROOT / ".state" / "marimo-latest"
RELEASES_URL = "https://api.github.com/repos/marimo-team/marimo/releases?per_page=30"
ISSUES_URL = "https://git.fhirworx.io/api/v1/repos/homelab/stack/issues"
PIN_RE = re.compile(r"^(ARG MARIMO_VERSION=)(\S+)$", re.M)
STABLE_RE = re.compile(r"^v?(\d+)\.(\d+)\.(\d+)$")


class Failed(Exception):
    def __init__(self, stage: str, detail: str):
        super().__init__(f"{stage}: {detail}")
        self.stage, self.detail = stage, detail


def log(msg: str) -> None:
    print(f"{dt.datetime.now():%Y-%m-%d %H:%M:%S} {msg}", flush=True)


def run(
    cmd: list[str], *, stage: str = "", check: bool = True, **kw
) -> subprocess.CompletedProcess:
    r = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True, **kw)
    if check and r.returncode:
        raise Failed(stage or " ".join(cmd[:2]), (r.stdout + r.stderr)[-3000:])
    return r


def git(*args: str, check: bool = True) -> str:
    return run(["git", *args], check=check).stdout.strip()


def ver(tag: str) -> tuple[int, ...]:
    return tuple(int(x) for x in STABLE_RE.match(tag).groups())


def latest_release(cooldown_days: int) -> str:
    req = urllib.request.Request(
        RELEASES_URL, headers={"Accept": "application/vnd.github+json"}
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        releases = json.load(resp)
    cutoff = dt.datetime.now(dt.UTC) - dt.timedelta(days=cooldown_days)
    tags = [
        r["tag_name"].lstrip("v")
        for r in releases
        if not (r["draft"] or r["prerelease"])
        and STABLE_RE.match(r["tag_name"])
        and dt.datetime.fromisoformat(r["published_at"]) <= cutoff
    ]
    if not tags:
        raise Failed("github", "no stable release past the cooldown")
    return max(tags, key=ver)


def read_pin() -> str:
    return PIN_RE.search((ROOT / DOCKERFILE).read_text()).group(2)


def write_pin(version: str) -> None:
    p = ROOT / DOCKERFILE
    p.write_text(PIN_RE.sub(rf"\g<1>{version}", p.read_text(), count=1))


def commit_pin(message: str) -> str:
    # --only commits just the Dockerfile, whatever else sits in the shared
    # index (concurrent sessions stage work in this checkout).
    run(["git", "commit", "--only", "-m", message, "--", DOCKERFILE], stage="commit")
    sha = git("rev-parse", "--short=7", "HEAD")
    if run(["git", "push", "gitea", "main"], check=False).returncode:
        log(f"WARN push failed; {sha} stays local until the next push")
    return sha


def set_heavy_tag(tag: str) -> None:
    env = ROOT / ".env"
    env.write_text(
        re.sub(r"^HEAVY_TAG=.*$", f"HEAVY_TAG={tag}", env.read_text(), flags=re.M)
    )


def roll_notebooks(tag: str) -> None:
    # Pass HEAVY_TAG explicitly: compose prefers the process environment over
    # .env, and the systemd unit loads .env as EnvironmentFile at start, so
    # without this compose keeps resolving the old tag and recreates nothing.
    run(
        [
            "docker",
            "compose",
            "up",
            "-d",
            "--no-build",
            "--pull",
            "never",
            "--no-deps",
            "notebooks",
        ],
        stage="roll",
        env={**os.environ, "HEAVY_TAG": tag},
    )
    image = run(["docker", "inspect", "-f", "{{.Config.Image}}", "notebooks"]).stdout
    if not image.strip().endswith(f"/notebooks:{tag}"):
        raise Failed("roll", f"notebooks runs {image.strip()}, expected tag {tag}")
    for _ in range(60):
        status = run(
            ["docker", "inspect", "-f", "{{.State.Health.Status}}", "notebooks"],
            check=False,
        )
        if status.stdout.strip() == "healthy":
            return
        time.sleep(3)
    raise Failed("roll", "notebooks never reported healthy")


def retag_zotero(old: str, new: str) -> None:
    """HEAVY_TAG drives zotero too; carry its image over when its inputs match."""
    same = (
        run(
            ["git", "diff", "--quiet", old, "HEAD", "--", *ZOTERO_INPUTS], check=False
        ).returncode
        == 0
    )
    have = (
        run(
            ["docker", "image", "inspect", f"fhirworx/zotero:{old}"], check=False
        ).returncode
        == 0
    )
    if same and have:
        run(["docker", "tag", f"fhirworx/zotero:{old}", f"fhirworx/zotero:{new}"])
    else:
        log(
            f"WARN zotero:{new} not tagged (inputs changed since {old} or image missing); build it before rolling zotero"
        )


def file_issue(version: str, err: Failed) -> None:
    token = os.environ.get("GITEA_TOKEN")
    if not token:
        log("WARN no GITEA_TOKEN; issue not filed")
        return
    body = json.dumps(
        {
            "title": f"marimo-latest: {version} failed at {err.stage}",
            "body": (
                f"`dev/scripts/marimo_latest.py` could not move the notebooks image to marimo "
                f"{version}; the pin and the running image stay on the previous version, and "
                f"later runs skip {version} until a newer release.\n\n```\n{err.detail[-2500:]}\n```"
            ),
        }
    )
    r = run(
        [
            "curl",
            "-sf",
            "-X",
            "POST",
            "-H",
            f"Authorization: token {token}",
            "-H",
            "Content-Type: application/json",
            "--data-binary",
            "@-",
            ISSUES_URL,
        ],
        check=False,
        input=body,
    )
    log(
        "filed Gitea issue"
        if r.returncode == 0
        else f"WARN issue filing failed: {r.stderr[:200]}"
    )


def preflight() -> None:
    if git("rev-parse", "--abbrev-ref", "HEAD") != "main":
        raise Failed("preflight", "checkout is not on main")
    if git("status", "--porcelain", "--", *IMAGE_INPUTS):
        raise Failed("preflight", f"uncommitted changes in {', '.join(IMAGE_INPUTS)}")
    git("fetch", "-q", "gitea", "main")
    if git("rev-list", "--count", "HEAD..gitea/main") != "0":
        # Renovate automerges onto the server's main. Fast-forward only: git
        # refuses if that would touch files someone has modified here.
        r = run(["git", "merge", "--ff-only", "gitea/main"], check=False)
        if r.returncode:
            raise Failed(
                "preflight", f"can't fast-forward to gitea/main: {r.stderr[-300:]}"
            )
        log(f"fast-forwarded main to {git('rev-parse', '--short=7', 'HEAD')}")


def upgrade(old: str, new: str) -> None:
    image = f"fhirworx/notebooks:marimo-{new}"
    write_pin(new)
    try:
        log(f"building {image}")
        run(
            ["docker", "build", "-f", DOCKERFILE, "-t", image, "."],
            stage="build",
            timeout=3600,
        )
        log("gating the image (nb_fe_smoke --image)")
        run(
            [
                "python3",
                "dev/scripts/nb_fe_smoke.py",
                "--image",
                image,
                "--network",
                "gateway",
            ],
            stage="image gate",
        )
        sha = commit_pin(
            f"chore(notebooks): marimo {old} → {new}\n\n"
            "Bumped, built, gated and rolled by dev/scripts/marimo_latest.py."
        )
    except Exception:
        write_pin(old)
        raise

    prev = re.search(r"^HEAVY_TAG=(.*)$", (ROOT / ".env").read_text(), re.M).group(1)
    shutil.copy(ROOT / ".env", STATE / "env.bak")
    run(["docker", "tag", image, f"fhirworx/notebooks:{sha}"])
    retag_zotero(prev, sha)
    log(f"rolling notebooks to {sha} (was {prev})")
    set_heavy_tag(sha)
    try:
        roll_notebooks(sha)
        run(
            [
                "python3",
                "dev/scripts/nb_fe_smoke.py",
                "--url",
                "http://notebooks:2718",
                "--network",
                "gateway",
            ],
            stage="live probe",
        )
    except Exception:
        log(f"live probe failed; rolling back to {prev}")
        set_heavy_tag(prev)
        roll_notebooks(prev)
        write_pin(old)
        commit_pin(
            f"revert(notebooks): marimo {new} → {old}\n\n"
            f"The live probe failed after rolling {sha}; see the marimo-latest issue."
        )
        raise
    log(f"notebooks on marimo {new} ({sha})")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--cooldown-days", type=int, default=7)
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    STATE.mkdir(parents=True, exist_ok=True)
    lock = open(STATE / "lock", "w")
    fcntl.flock(lock, fcntl.LOCK_EX)
    failed_path = STATE / "failed.json"
    failed = json.loads(failed_path.read_text()) if failed_path.exists() else {}

    new, old = latest_release(args.cooldown_days), read_pin()
    if ver(new) <= ver(old):
        log(f"marimo {old} is current (latest past cooldown: {new})")
        return 0
    if new in failed:
        log(f"skipping marimo {new}: failed before at {failed[new]}")
        return 0
    log(f"marimo {old} → {new}" + (" (dry run)" if args.dry_run else ""))
    if args.dry_run:
        return 0
    try:
        preflight()
        upgrade(old, new)
    except Failed as err:
        log(f"FAILED {err.stage}: {err.detail[-800:]}")
        if (
            err.stage != "preflight"
        ):  # a dirty checkout is transient, not the release's fault
            failed[new] = err.stage
            failed_path.write_text(json.dumps(failed, indent=2))
            file_issue(new, err)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())