scripts/campaign_daemon.py

#!/usr/bin/env python
"""Start/stop/status for the long-running datagen campaign, with a PID file.

**Why this exists rather than a shell one-liner.** The campaign runs for many
hours, so it has to survive the session that started it, and it has to be
stoppable *precisely*. Both were done with ``ps | grep | xargs kill`` before,
and that is how the campaign was killed by accident: the grep pattern matched
the very shell whose command line contained the script name, so the launcher
killed itself. A PID file removes the ambiguity -- ``stop`` signals exactly the
process ``start`` recorded, and nothing else.

``start`` detaches with ``os.setsid`` so the campaign is not in the launching
session's process group and does not die with it.

Usage::

    python scripts/campaign_daemon.py start
    python scripts/campaign_daemon.py status
    python scripts/campaign_daemon.py stop
"""

from __future__ import annotations

import argparse
import os
import signal
import subprocess
import sys
import time
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
PID_FILE = REPO_ROOT / "logs" / "campaign.pid"
LOG_FILE = REPO_ROOT / "logs" / "datagen_campaign.log"


def _read_pid() -> int | None:
    if not PID_FILE.exists():
        return None
    try:
        return int(PID_FILE.read_text().strip())
    except ValueError:
        return None


def _alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
    except (ProcessLookupError, PermissionError):
        return False
    return True


def start(extra: list[str]) -> int:
    pid = _read_pid()
    if pid and _alive(pid):
        print(f"already running: pid {pid}")
        return 1

    env = dict(os.environ)
    nvshim = REPO_ROOT / ".nvshim"
    if nvshim.is_dir():
        env["LD_LIBRARY_PATH"] = os.pathsep.join(
            [str(nvshim), env.get("LD_LIBRARY_PATH", "")]
        ).rstrip(os.pathsep)
    env["PYTHONPATH"] = os.pathsep.join(
        [str(REPO_ROOT / "src"), env.get("PYTHONPATH", "")]
    ).rstrip(os.pathsep)

    LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        sys.executable, str(REPO_ROOT / "scripts" / "run_datagen_until_converged.py"),
        *extra,
    ]
    with open(LOG_FILE, "a") as log:
        proc = subprocess.Popen(
            cmd, cwd=str(REPO_ROOT), env=env,
            stdout=log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL,
            start_new_session=True,  # os.setsid: survives the launching session
        )
    PID_FILE.write_text(str(proc.pid))
    time.sleep(3)
    if not _alive(proc.pid):
        print(f"started but exited immediately; see {LOG_FILE}")
        return 1
    print(f"started: pid {proc.pid}, log {LOG_FILE}")
    return 0


def stop() -> int:
    pid = _read_pid()
    if not pid:
        print("no pid file; nothing to stop")
        return 1
    if not _alive(pid):
        print(f"pid {pid} not running; clearing stale pid file")
        PID_FILE.unlink(missing_ok=True)
        return 0
    # Signal the whole process group: the campaign spawns worker processes,
    # and leaving those alive would keep GPU memory pinned.
    try:
        os.killpg(os.getpgid(pid), signal.SIGTERM)
    except ProcessLookupError:
        pass
    for _ in range(20):
        if not _alive(pid):
            break
        time.sleep(1)
    if _alive(pid):
        os.killpg(os.getpgid(pid), signal.SIGKILL)
    PID_FILE.unlink(missing_ok=True)
    print(f"stopped pid {pid}")
    return 0


def status() -> int:
    pid = _read_pid()
    if pid and _alive(pid):
        print(f"running: pid {pid}")
    else:
        print("not running")
    return 0


def main() -> int:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("action", choices=["start", "stop", "status"])
    args, extra = p.parse_known_args()
    return {"start": lambda: start(extra), "stop": stop, "status": status}[args.action]()


if __name__ == "__main__":
    raise SystemExit(main())