gpushift 0.1.0__py3-none-any.whl

gpushift/__init__.py

@@ -0,0 +1,7 @@
+"""gpushift: a shared GPU job queue.
+
+Keep this module free of imports: ``gpushift.checkpoint`` is imported by user
+training scripts and must stay dependency-free.
+"""
+
+__version__ = "0.1.0"

gpushift/__main__.py

@@ -0,0 +1,4 @@
+from gpushift._entry import main
+
+if __name__ == "__main__":
+    main()

gpushift/_entry.py

@@ -0,0 +1,14 @@
+"""Console entry point. `gpushift hook` runs before every agent shell command, so
+it skips importing the full CLI (psutil, sqlite, ...)."""
+
+import sys
+
+
+def main() -> None:
+    if len(sys.argv) > 1 and sys.argv[1] == "hook":
+        from gpushift.hook import main as hook_main
+
+        sys.exit(hook_main(sys.argv[2] if len(sys.argv) > 2 else "pretooluse"))
+    from gpushift.cli import main as cli_main
+
+    cli_main()

gpushift/_proc.py

@@ -0,0 +1,152 @@
+"""Small cross-platform process helpers (spawning detached, suspending trees)."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import psutil
+
+IS_WINDOWS = os.name == "nt"
+NO_WINDOW = 0x08000000 if IS_WINDOWS else 0  # CREATE_NO_WINDOW
+NEW_GROUP = 0x00000200 if IS_WINDOWS else 0  # CREATE_NEW_PROCESS_GROUP
+BREAKAWAY = 0x01000000 if IS_WINDOWS else 0  # CREATE_BREAKAWAY_FROM_JOB
+
+
+def console_python() -> str:
+    """A python executable that has a (hidden) console: under pythonw, child
+    console programs would otherwise pop up their own windows."""
+    exe = Path(sys.executable)
+    if IS_WINDOWS and exe.name.lower() == "pythonw.exe":
+        cand = exe.with_name("python.exe")
+        if cand.exists():
+            return str(cand)
+    return str(exe)
+
+
+def windowless_python() -> str:
+    exe = Path(sys.executable)
+    if IS_WINDOWS and exe.name.lower() == "python.exe":
+        cand = exe.with_name("pythonw.exe")
+        if cand.exists():
+            return str(cand)
+    return str(exe)
+
+
+def spawn_detached(argv: list[str], **kw) -> subprocess.Popen:
+    """Start a process that outlives its parent (and the terminal/agent that
+    started it). On Windows, try to break away from the caller's job object so
+    an agent's shell cleanup cannot kill it."""
+    kw.setdefault("stdin", subprocess.DEVNULL)
+    kw.setdefault("stdout", subprocess.DEVNULL)
+    kw.setdefault("stderr", subprocess.DEVNULL)
+    if IS_WINDOWS:
+        base = NO_WINDOW | NEW_GROUP
+        try:
+            return subprocess.Popen(argv, creationflags=base | BREAKAWAY, **kw)
+        except OSError:
+            return subprocess.Popen(argv, creationflags=base, **kw)
+    return subprocess.Popen(argv, start_new_session=True, **kw)
+
+
+def get_proc(pid: int | None, ctime: float | None = None) -> psutil.Process | None:
+    """The process with this pid, if it is alive and (when ctime is given) is
+    the same process we started (guards against pid reuse)."""
+    if not pid:
+        return None
+    try:
+        p = psutil.Process(pid)
+        if ctime and abs(p.create_time() - ctime) > 1.0:
+            return None
+        if p.status() == psutil.STATUS_ZOMBIE:
+            return None
+        return p
+    except psutil.Error:
+        return None
+
+
+def alive(pid: int | None, ctime: float | None = None) -> bool:
+    return get_proc(pid, ctime) is not None
+
+
+def create_time(pid: int) -> float | None:
+    try:
+        return psutil.Process(pid).create_time()
+    except psutil.Error:
+        return None
+
+
+def tree(p: psutil.Process) -> list[psutil.Process]:
+    try:
+        return [p, *p.children(recursive=True)]
+    except psutil.Error:
+        return [p]
+
+
+def suspend_tree(pid: int | None, ctime: float | None = None) -> bool:
+    p = get_proc(pid, ctime)
+    if p is None:
+        return False
+    # Parent first so it cannot spawn new children while we work down the tree.
+    for q in tree(p):
+        try:
+            q.suspend()
+        except psutil.Error:
+            pass
+    return True
+
+
+def resume_tree(pid: int | None, ctime: float | None = None) -> bool:
+    p = get_proc(pid, ctime)
+    if p is None:
+        return False
+    for q in reversed(tree(p)):
+        try:
+            q.resume()
+        except psutil.Error:
+            pass
+    return True
+
+
+def kill(pid: int | None, ctime: float | None = None, timeout: float = 5.0) -> bool:
+    """Kill this one process, leaving its children running."""
+    p = get_proc(pid, ctime)
+    if p is None:
+        return False
+    try:
+        p.kill()
+        p.wait(timeout)
+    except (psutil.Error, psutil.TimeoutExpired):
+        pass
+    return True
+
+
+def kill_tree(pid: int | None, ctime: float | None = None, timeout: float = 5.0) -> bool:
+    p = get_proc(pid, ctime)
+    if p is None:
+        return False
+    procs = tree(p)
+    for q in reversed(procs):
+        try:
+            q.kill()
+        except psutil.Error:
+            pass
+    psutil.wait_procs(procs, timeout=timeout)
+    return True
+
+
+def terminate(pid: int | None, ctime: float | None = None) -> bool:
+    """Polite stop request: SIGTERM on POSIX. No-op on Windows (gpushift uses a
+    stop file there, which the checkpoint helper polls)."""
+    if IS_WINDOWS:
+        return False
+    p = get_proc(pid, ctime)
+    if p is None:
+        return False
+    try:
+        p.terminate()
+        return True
+    except psutil.Error:
+        return False

gpushift/autostart.py

@@ -0,0 +1,97 @@
+"""Start the daemon at login: HKCU Run key on Windows, a systemd user service
+on Linux, a LaunchAgent on macOS. None of these need admin rights."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+from . import _proc
+
+NAME = "gpushift"
+UNIT = Path.home() / ".config" / "systemd" / "user" / "gpushift.service"
+PLIST = Path.home() / "Library" / "LaunchAgents" / "com.gpushift.daemon.plist"
+
+
+def _cmd() -> list[str]:
+    return [_proc.windowless_python(), "-m", "gpushift", "daemon"]
+
+
+def install() -> str:
+    if sys.platform == "win32":
+        import winreg
+
+        cmd = subprocess.list2cmdline(_cmd())
+        with winreg.OpenKey(
+            winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Windows\CurrentVersion\Run", 0, winreg.KEY_SET_VALUE
+        ) as k:
+            winreg.SetValueEx(k, NAME, 0, winreg.REG_SZ, cmd)
+        return f"registered at login (HKCU\\...\\Run\\{NAME}): {cmd}"
+    if sys.platform == "darwin":
+        args = "".join(f"<string>{a}</string>" for a in _cmd())
+        PLIST.parent.mkdir(parents=True, exist_ok=True)
+        PLIST.write_text(
+            '<?xml version="1.0" encoding="UTF-8"?>\n'
+            '<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">\n'
+            f'<plist version="1.0"><dict><key>Label</key><string>com.gpushift.daemon</string>'
+            f"<key>ProgramArguments</key><array>{args}</array>"
+            "<key>RunAtLoad</key><true/><key>KeepAlive</key><true/></dict></plist>\n"
+        )
+        subprocess.run(["launchctl", "load", "-w", str(PLIST)], check=False)
+        return f"installed LaunchAgent {PLIST}"
+    exec_start = " ".join(f'"{a}"' if " " in a else a for a in _cmd())
+    UNIT.parent.mkdir(parents=True, exist_ok=True)
+    UNIT.write_text(
+        "[Unit]\nDescription=gpushift GPU job queue\n\n"
+        f"[Service]\nExecStart={exec_start}\nRestart=on-failure\nRestartSec=5\n\n"
+        "[Install]\nWantedBy=default.target\n"
+    )
+    subprocess.run(["systemctl", "--user", "daemon-reload"], check=False)
+    r = subprocess.run(["systemctl", "--user", "enable", "--now", "gpushift.service"], capture_output=True, text=True)
+    if r.returncode != 0:
+        return f"wrote {UNIT}, but enabling failed: {r.stderr.strip()}"
+    return (
+        f"installed systemd user service {UNIT}. To keep the queue running after you log out, "
+        "run: loginctl enable-linger $USER"
+    )
+
+
+def uninstall() -> str:
+    if sys.platform == "win32":
+        import winreg
+
+        try:
+            with winreg.OpenKey(
+                winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Windows\CurrentVersion\Run", 0, winreg.KEY_SET_VALUE
+            ) as k:
+                winreg.DeleteValue(k, NAME)
+        except FileNotFoundError:
+            return "not registered"
+        return "removed from login startup"
+    if sys.platform == "darwin":
+        if PLIST.exists():
+            subprocess.run(["launchctl", "unload", "-w", str(PLIST)], check=False)
+            PLIST.unlink()
+            return "removed LaunchAgent"
+        return "not registered"
+    if UNIT.exists():
+        subprocess.run(["systemctl", "--user", "disable", "--now", "gpushift.service"], check=False)
+        UNIT.unlink()
+        return "removed systemd user service"
+    return "not registered"
+
+
+def status() -> str | None:
+    """Description of the registered autostart, or None."""
+    if sys.platform == "win32":
+        import winreg
+
+        try:
+            with winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Windows\CurrentVersion\Run") as k:
+                return winreg.QueryValueEx(k, NAME)[0]
+        except FileNotFoundError:
+            return None
+    if sys.platform == "darwin":
+        return str(PLIST) if PLIST.exists() else None
+    return str(UNIT) if UNIT.exists() else None

gpushift/checkpoint.py

@@ -0,0 +1,140 @@
+"""gpushift.checkpoint: make a training script evictable. Dependency-free.
+
+Jobs started by gpushift can always ``import gpushift.checkpoint``, even in an
+environment where gpushift is not installed (gpushift puts this file on PYTHONPATH).
+
+Usage::
+
+    from gpushift import checkpoint as ckpt
+
+    state = load(ckpt.latest_checkpoint("ckpts"))   # None on a fresh start
+    for step in range(start, total):
+        train_step()
+        if step % 500 == 0 or ckpt.should_stop():
+            save(f"ckpts/step{step}.pt")
+        if ckpt.should_stop():
+            break                                   # exit cleanly; gpushift requeues the job
+
+When ``gpushift free`` needs the GPUs back, it asks the job to stop (SIGTERM on
+Linux/macOS, a stop file on every platform). ``should_stop()`` turns True;
+the script saves and exits; gpushift puts the job back at the front of the queue
+with ``GPUSHIFT_RESUME=1`` so it continues from its latest checkpoint.
+
+Calling ``should_stop()`` once marks the job as resumable, so gpushift evicts it
+instead of freezing it. Scripts that never call it are frozen (their VRAM stays
+allocated) during ``gpushift free``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import signal
+import threading
+import time
+from pathlib import Path
+
+__all__ = [
+    "atomic_path",
+    "install",
+    "job_id",
+    "latest_checkpoint",
+    "managed",
+    "resuming",
+    "should_stop",
+]
+
+_stop = False
+_installed = False
+_last_poll = 0.0
+
+
+def managed() -> bool:
+    """True when running as a gpushift job."""
+    return "GPUSHIFT_JOB_ID" in os.environ
+
+
+def job_id() -> int | None:
+    v = os.environ.get("GPUSHIFT_JOB_ID")
+    return int(v) if v else None
+
+
+def _on_signal(signum, frame):
+    global _stop
+    _stop = True
+
+
+def install() -> None:
+    """Install the stop handlers and mark this job resumable. Idempotent;
+    ``should_stop()`` calls it for you."""
+    global _installed
+    if _installed:
+        return
+    _installed = True
+    if threading.current_thread() is threading.main_thread():
+        for name in ("SIGTERM", "SIGBREAK"):
+            sig = getattr(signal, name, None)
+            if sig is not None:
+                with contextlib.suppress(ValueError, OSError):
+                    signal.signal(sig, _on_signal)
+    d = os.environ.get("GPUSHIFT_JOB_DIR")
+    if d:
+        with contextlib.suppress(OSError):
+            Path(d).mkdir(parents=True, exist_ok=True)
+            Path(d, "resumable").touch()
+
+
+def should_stop(poll_interval: float = 1.0) -> bool:
+    """True once gpushift (or anyone sending SIGTERM) wants the job to stop.
+    Cheap enough to call every step: the stop file is checked at most once per
+    ``poll_interval`` seconds."""
+    global _stop, _last_poll
+    install()
+    if _stop:
+        return True
+    f = os.environ.get("GPUSHIFT_STOP_FILE")
+    if f:
+        now = time.monotonic()
+        if now - _last_poll >= poll_interval:
+            _last_poll = now
+            if os.path.exists(f):
+                _stop = True
+    return _stop
+
+
+def resuming() -> bool:
+    """True when gpushift restarted this job after an eviction."""
+    return os.environ.get("GPUSHIFT_RESUME") == "1"
+
+
+def latest_checkpoint(directory, pattern: str = "*") -> Path | None:
+    """Newest file or directory in ``directory`` matching ``pattern`` (by
+    modification time), or None. Temp files from ``atomic_path`` are ignored."""
+    d = Path(directory)
+    if not d.is_dir():
+        return None
+    cands = [p for p in d.glob(pattern) if not p.name.startswith(".") and not p.name.endswith(".gpushift-tmp")]
+    if not cands:
+        return None
+    return max(cands, key=lambda p: p.stat().st_mtime)
+
+
+@contextlib.contextmanager
+def atomic_path(path):
+    """Write a checkpoint without ever leaving a half-written file behind::
+
+    with ckpt.atomic_path("ckpts/last.pt") as tmp:
+        torch.save(state, tmp)
+    """
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_name(path.name + ".gpushift-tmp")
+    yield tmp
+    os.replace(tmp, path)
+
+
+def _reset_for_tests() -> None:
+    global _stop, _installed, _last_poll
+    _stop = False
+    _installed = False
+    _last_poll = 0.0