368 lines
15 KiB
Python
368 lines
15 KiB
Python
|
"""Driving the training host over SSH, including getting the GPU back from the LLMs.
|
||
|
|
|
||
|
|
Training runs on a separate machine -- a Ryzen AI MAX+ 395 with a Radeon 8060S, whose
|
||
|
|
124 GB of memory is *unified*: the iGPU addresses it through GTT rather than owning
|
||
|
|
dedicated VRAM. That is why the box can train on frames this large at all, and also why
|
||
|
|
it can hold nothing else at the time: llama.cpp keeps 115.6 GB of the 124 GB resident, so
|
||
|
|
a training run does not get a smaller allocation, it gets ``Memory in use`` and dies.
|
||
|
|
|
||
|
|
So something has to give up the GPU for the duration. These services belong to the user,
|
||
|
|
not to us, which sets the bar for this module: stop as little as will do, put back
|
||
|
|
whatever it stopped -- including when the training run crashes -- and be able to say
|
||
|
|
afterwards whether that succeeded.
|
||
|
|
|
||
|
|
Stopping as little as will do is worth the extra code. Measured on this host, almost
|
||
|
|
all of that 115.6 GB is a 120B model that ``llama-swap`` had loaded on demand; with the
|
||
|
|
unit stopped, GTT falls to 24.8 GB and roughly 99 GB is free -- ample for training,
|
||
|
|
without touching anything else. So :func:`stop_llms` works in stages and escalates only
|
||
|
|
if the first stage leaves too little.
|
||
|
|
|
||
|
|
The two kinds of server are not alike:
|
||
|
|
|
||
|
|
* ``llama-swap.service``, a systemd --user unit. Stopping and starting it is exact, and
|
||
|
|
it is where the large on-demand models live. Almost always sufficient on its own.
|
||
|
|
* a **standalone** ``llama-server``, launched by hand and reparented to init. There is
|
||
|
|
no unit to restart, so the only way back is to record its argv and working directory
|
||
|
|
before killing it and re-exec them. This is the fragile one, only touched when the
|
||
|
|
first stage did not free enough, and the reason :func:`llms_paused` verifies the
|
||
|
|
restore rather than assuming it.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import shlex
|
||
|
|
import subprocess
|
||
|
|
import time
|
||
|
|
from contextlib import contextmanager
|
||
|
|
|
||
|
|
HOST = os.environ.get("SUVI_GPU_HOST", "htpc@192.168.1.66")
|
||
|
|
#: Read from a file rather than embedded, so the credential is not in the source tree.
|
||
|
|
PASSWORD_FILE = os.environ.get(
|
||
|
|
"SUVI_GPU_PASSWORD_FILE",
|
||
|
|
os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||
|
|
"gpu_system_ssh_info"),
|
||
|
|
)
|
||
|
|
#: systemd --user units that hold the GPU.
|
||
|
|
UNITS = ("llama-swap.service",)
|
||
|
|
#: Seconds to wait for GTT to drain / refill before giving up.
|
||
|
|
SETTLE_TIMEOUT = 120
|
||
|
|
#: Unified memory the iGPU can address, from mem_info_gtt_total on this host.
|
||
|
|
GTT_TOTAL = 115 * 1024**3
|
||
|
|
#: Free unified memory a training run needs before it will start.
|
||
|
|
GTT_FREE_BYTES = 48 * 1024**3
|
||
|
|
GTT_USED = "/sys/class/drm/card1/device/mem_info_gtt_used"
|
||
|
|
|
||
|
|
|
||
|
|
def _password():
|
||
|
|
with open(PASSWORD_FILE) as handle:
|
||
|
|
lines = [line.strip() for line in handle if line.strip()]
|
||
|
|
if len(lines) < 2:
|
||
|
|
raise RuntimeError(f"{PASSWORD_FILE}: expected an ssh line and a password line")
|
||
|
|
return lines[1]
|
||
|
|
|
||
|
|
|
||
|
|
def run(command, check=True, timeout=600, capture=True):
|
||
|
|
"""Run a shell command on the training host."""
|
||
|
|
argv = [
|
||
|
|
"sshpass", "-p", _password(), "ssh",
|
||
|
|
"-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
|
||
|
|
HOST, command,
|
||
|
|
]
|
||
|
|
result = subprocess.run(
|
||
|
|
argv, timeout=timeout,
|
||
|
|
stdout=subprocess.PIPE if capture else None,
|
||
|
|
stderr=subprocess.STDOUT if capture else None,
|
||
|
|
)
|
||
|
|
output = result.stdout.decode(errors="replace") if capture else ""
|
||
|
|
if check and result.returncode != 0:
|
||
|
|
raise RuntimeError(f"remote command failed ({result.returncode}): {command}\n{output}")
|
||
|
|
return output
|
||
|
|
|
||
|
|
|
||
|
|
def gtt_used():
|
||
|
|
"""Bytes of unified memory currently mapped to the GPU."""
|
||
|
|
try:
|
||
|
|
return int(run(f"cat {GTT_USED}", check=False).strip() or 0)
|
||
|
|
except ValueError:
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
# --------------------------------------------------------------- pausing the LLMs
|
||
|
|
|
||
|
|
|
||
|
|
def _standalone_servers():
|
||
|
|
"""Every llama-server not owned by one of our units, with enough to restart it.
|
||
|
|
|
||
|
|
A server spawned *by* llama-swap comes back when llama-swap does, so recording it
|
||
|
|
would restart it twice. Only the ones reparented to init are ours to restore.
|
||
|
|
"""
|
||
|
|
script = r"""
|
||
|
|
for p in $(pgrep -x llama-server); do
|
||
|
|
ppid=$(ps -o ppid= -p $p 2>/dev/null | tr -d ' ')
|
||
|
|
[ "$ppid" = "1" ] || continue
|
||
|
|
cwd=$(readlink /proc/$p/cwd)
|
||
|
|
argv=$(tr '\0' '\n' < /proc/$p/cmdline | sed 's/"/\\"/g' | awk '{printf "\"%s\",", $0}')
|
||
|
|
echo "{\"pid\": $p, \"cwd\": \"$cwd\", \"argv\": [${argv%,}]}"
|
||
|
|
done
|
||
|
|
"""
|
||
|
|
found = []
|
||
|
|
for line in run(script, check=False).splitlines():
|
||
|
|
line = line.strip()
|
||
|
|
if line.startswith("{"):
|
||
|
|
try:
|
||
|
|
found.append(json.loads(line))
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
continue
|
||
|
|
return found
|
||
|
|
|
||
|
|
|
||
|
|
def _wait_for_gtt(below, timeout=SETTLE_TIMEOUT):
|
||
|
|
deadline = time.time() + timeout
|
||
|
|
while time.time() < deadline:
|
||
|
|
if gtt_used() < below:
|
||
|
|
return True
|
||
|
|
time.sleep(2)
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def stop_llms(needed_bytes=None):
|
||
|
|
"""Free the GPU, stopping as little as will do.
|
||
|
|
|
||
|
|
Stage one stops the units, which is where the big on-demand models live and which
|
||
|
|
restarts exactly. Only if that leaves less than `needed_bytes` free does stage two
|
||
|
|
kill hand-launched servers, whose restoration is a re-exec rather than a restart.
|
||
|
|
|
||
|
|
Returns the state :func:`start_llms` needs to undo this.
|
||
|
|
"""
|
||
|
|
needed = GTT_TOTAL - (needed_bytes or GTT_FREE_BYTES)
|
||
|
|
state = {"units": [], "standalone": []}
|
||
|
|
|
||
|
|
for unit in UNITS:
|
||
|
|
active = run(f"systemctl --user is-active {unit}", check=False).strip()
|
||
|
|
state["units"].append({"unit": unit, "was_active": active == "active"})
|
||
|
|
if active == "active":
|
||
|
|
run(f"systemctl --user stop {unit}", check=False)
|
||
|
|
if _wait_for_gtt(needed, timeout=30):
|
||
|
|
return state
|
||
|
|
|
||
|
|
# Stage two: the units alone were not enough.
|
||
|
|
state["standalone"] = _standalone_servers()
|
||
|
|
for server in state["standalone"]:
|
||
|
|
run(f"kill {int(server['pid'])}", check=False)
|
||
|
|
if not _wait_for_gtt(needed):
|
||
|
|
for server in state["standalone"]:
|
||
|
|
run(f"kill -9 {int(server['pid'])}", check=False)
|
||
|
|
_wait_for_gtt(needed)
|
||
|
|
return state
|
||
|
|
|
||
|
|
|
||
|
|
def start_llms(state):
|
||
|
|
"""Put back exactly what :func:`stop_llms` took away.
|
||
|
|
|
||
|
|
Reports what it could not restore rather than raising: this runs in a `finally`,
|
||
|
|
and masking the training error with a restore error would lose the more useful of
|
||
|
|
the two.
|
||
|
|
"""
|
||
|
|
failures = []
|
||
|
|
for index, server in enumerate(state.get("standalone", [])):
|
||
|
|
argv = " ".join(shlex.quote(arg) for arg in server["argv"])
|
||
|
|
# systemd-run rather than nohup: a backgrounded process still holds the SSH
|
||
|
|
# channel open, so ssh blocks until it exits -- which for a server is never.
|
||
|
|
# A transient unit forks away cleanly and returns at once, and --collect means
|
||
|
|
# it leaves nothing behind when it stops.
|
||
|
|
started = run(
|
||
|
|
f"systemd-run --user --collect --unit=suvi-restored-llama-{index} "
|
||
|
|
f"--property=WorkingDirectory={shlex.quote(server['cwd'])} {argv}",
|
||
|
|
check=False, timeout=60,
|
||
|
|
)
|
||
|
|
if "Running as unit" not in started and "Failed" in started:
|
||
|
|
failures.append(f"standalone {os.path.basename(server['argv'][0])}")
|
||
|
|
|
||
|
|
for entry in state.get("units", []):
|
||
|
|
if not entry["was_active"]:
|
||
|
|
continue
|
||
|
|
run(f"systemctl --user start {entry['unit']}", check=False)
|
||
|
|
if run(f"systemctl --user is-active {entry['unit']}", check=False).strip() != "active":
|
||
|
|
failures.append(entry["unit"])
|
||
|
|
return failures
|
||
|
|
|
||
|
|
|
||
|
|
#: States stopped but not yet restored. Module level so the exit hooks can reach them.
|
||
|
|
_paused = []
|
||
|
|
_hooks_installed = False
|
||
|
|
|
||
|
|
|
||
|
|
def _restore_all(verbose=True):
|
||
|
|
"""Put back every outstanding pause. Safe to call more than once."""
|
||
|
|
while _paused:
|
||
|
|
state = _paused.pop()
|
||
|
|
# Our own containers first: restarting the LLMs while a training container still
|
||
|
|
# holds 60 GB of GTT just moves the memory exhaustion onto the user's services.
|
||
|
|
try:
|
||
|
|
stop_containers(verbose=verbose)
|
||
|
|
except (RuntimeError, OSError):
|
||
|
|
pass
|
||
|
|
failures = start_llms(state)
|
||
|
|
if failures:
|
||
|
|
print(f" WARNING: could not restart: {', '.join(failures)}")
|
||
|
|
print(" The GPU host is missing services it had before this run.")
|
||
|
|
elif verbose:
|
||
|
|
print(f" restored LLM services; GTT {gtt_used() / 1024**3:.1f} GiB used")
|
||
|
|
|
||
|
|
|
||
|
|
def _install_hooks():
|
||
|
|
"""Restore on the ways out that `finally` does not cover.
|
||
|
|
|
||
|
|
A `finally` block handles a return or an exception. It does not handle SIGTERM,
|
||
|
|
which is what a timeout, a `kill`, or a parent shell giving up actually sends -- the
|
||
|
|
default disposition terminates the interpreter without unwinding. That is not
|
||
|
|
hypothetical: it happened here, and left the user's LLM proxy stopped with nothing
|
||
|
|
scheduled to start it again. So the same belt-and-braces the archive traversals
|
||
|
|
use (see :class:`suvi.vfs.Reliever`): atexit for orderly exits, explicit handlers
|
||
|
|
for the signals.
|
||
|
|
"""
|
||
|
|
global _hooks_installed
|
||
|
|
if _hooks_installed:
|
||
|
|
return
|
||
|
|
import atexit
|
||
|
|
import signal
|
||
|
|
|
||
|
|
atexit.register(_restore_all, verbose=False)
|
||
|
|
for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
|
||
|
|
previous = signal.getsignal(signum)
|
||
|
|
|
||
|
|
def handler(number, frame, previous=previous):
|
||
|
|
_restore_all()
|
||
|
|
if callable(previous) and previous not in (signal.SIG_IGN, signal.SIG_DFL):
|
||
|
|
return previous(number, frame)
|
||
|
|
raise SystemExit(128 + number)
|
||
|
|
|
||
|
|
try:
|
||
|
|
signal.signal(signum, handler)
|
||
|
|
except (ValueError, OSError): # not the main thread, or no such signal
|
||
|
|
pass
|
||
|
|
_hooks_installed = True
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def llms_paused(needed_bytes=None, verbose=True):
|
||
|
|
"""Hold the GPU for the duration of the block, then give it back.
|
||
|
|
|
||
|
|
Restoration is registered before the block runs and happens on every exit path --
|
||
|
|
return, exception, or signal. What it cannot restore it names, because a
|
||
|
|
hand-launched server that quietly failed to come back would otherwise be discovered
|
||
|
|
by the user, later, as a service that is simply gone.
|
||
|
|
"""
|
||
|
|
_install_hooks()
|
||
|
|
state = stop_llms(needed_bytes)
|
||
|
|
_paused.append(state)
|
||
|
|
if verbose:
|
||
|
|
units = [e["unit"] for e in state["units"] if e["was_active"]]
|
||
|
|
print(f" paused: units {units or 'none'}"
|
||
|
|
+ (f", {len(state['standalone'])} standalone server(s)"
|
||
|
|
if state["standalone"] else " (no standalone servers touched)"))
|
||
|
|
print(f" GTT now {gtt_used() / 1024**3:.1f} GiB used")
|
||
|
|
try:
|
||
|
|
yield state
|
||
|
|
finally:
|
||
|
|
if state in _paused:
|
||
|
|
_paused.remove(state)
|
||
|
|
stop_containers(verbose=verbose)
|
||
|
|
failures = start_llms(state)
|
||
|
|
if failures:
|
||
|
|
print(f" WARNING: could not restart: {', '.join(failures)}")
|
||
|
|
print(" The GPU host is missing services it had before this run.")
|
||
|
|
elif verbose:
|
||
|
|
print(f" restored LLM services; GTT {gtt_used() / 1024**3:.1f} GiB used")
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ the container
|
||
|
|
|
||
|
|
|
||
|
|
#: ROCm+PyTorch plus this project's dependencies. Bazzite's root is read-only and its
|
||
|
|
#: Python is 3.14, which has no torch wheels, so the toolchain is containerised
|
||
|
|
#: regardless; the derived tag adds zstandard, opencv, scikit-image and astropy.
|
||
|
|
IMAGE = os.environ.get("SUVI_GPU_IMAGE", "localhost/suvi-train:latest")
|
||
|
|
#: Prefix for containers this module starts, so they can be found and stopped again.
|
||
|
|
CONTAINER_PREFIX = "suvi-run-"
|
||
|
|
|
||
|
|
|
||
|
|
def stop_containers(verbose=True):
|
||
|
|
"""Stop every container this module started.
|
||
|
|
|
||
|
|
Necessary because a container does *not* die with the SSH session that launched it.
|
||
|
|
A run that times out client-side leaves the container running, holding tens of GB of
|
||
|
|
unified memory; two such orphans were what put GTT at 89.6 GB with the GPU idle.
|
||
|
|
"""
|
||
|
|
names = [
|
||
|
|
line.strip()
|
||
|
|
for line in run(
|
||
|
|
f"podman ps --filter name={CONTAINER_PREFIX} --format '{{{{.Names}}}}'",
|
||
|
|
check=False,
|
||
|
|
).splitlines()
|
||
|
|
if line.strip().startswith(CONTAINER_PREFIX)
|
||
|
|
]
|
||
|
|
for name in names:
|
||
|
|
if verbose:
|
||
|
|
print(f" stopping orphaned container {name}")
|
||
|
|
run(f"podman stop -t 10 {shlex.quote(name)}", check=False, timeout=60)
|
||
|
|
return names
|
||
|
|
|
||
|
|
|
||
|
|
def launch(command, mounts=(), name=None, log="/var/home/htpc/suvi/run.log"):
|
||
|
|
"""Start a container detached on the host and return (name, log path).
|
||
|
|
|
||
|
|
Long jobs must not be tied to an SSH channel. Streaming a training run through
|
||
|
|
``ssh`` means the run dies with the connection, or -- worse, and observed here --
|
||
|
|
the container finishes while the local client stays blocked reading a pipe that
|
||
|
|
will never close. A transient systemd unit is owned by the host, survives the
|
||
|
|
client entirely, and can be polled with :func:`tail`.
|
||
|
|
"""
|
||
|
|
name = name or f"{CONTAINER_PREFIX}{int(time.time())}"
|
||
|
|
binds = " ".join(f"-v {shlex.quote(src)}:{shlex.quote(dst)}" for src, dst in mounts)
|
||
|
|
inner = (
|
||
|
|
f"podman run --rm --name {shlex.quote(name)} "
|
||
|
|
f"--device /dev/kfd --device /dev/dri "
|
||
|
|
f"--security-opt seccomp=unconfined --ipc=host {binds} {IMAGE} {command}"
|
||
|
|
)
|
||
|
|
run(f"rm -f {shlex.quote(log)}", check=False)
|
||
|
|
run(
|
||
|
|
f"systemd-run --user --collect --unit={shlex.quote(name)} "
|
||
|
|
f"--property=StandardOutput=append:{shlex.quote(log)} "
|
||
|
|
f"--property=StandardError=append:{shlex.quote(log)} "
|
||
|
|
f"/bin/sh -c {shlex.quote(inner)}",
|
||
|
|
timeout=60,
|
||
|
|
)
|
||
|
|
return name, log
|
||
|
|
|
||
|
|
|
||
|
|
def tail(log, lines=40):
|
||
|
|
return run(f"tail -n {int(lines)} {shlex.quote(log)} 2>/dev/null", check=False)
|
||
|
|
|
||
|
|
|
||
|
|
def running(name):
|
||
|
|
return run(f"podman ps --filter name={shlex.quote(name)} --format '{{{{.Names}}}}'",
|
||
|
|
check=False).strip() != ""
|
||
|
|
|
||
|
|
|
||
|
|
def podman(command, mounts=(), timeout=None, capture=True, name=None):
|
||
|
|
"""Run `command` inside the training container on the GPU host.
|
||
|
|
|
||
|
|
``/dev/kfd`` and ``/dev/dri`` are both world-accessible on this host, so no group
|
||
|
|
mapping is needed; ``seccomp=unconfined`` is what ROCm needs to issue its ioctls.
|
||
|
|
|
||
|
|
The container is named and stopped in a `finally`, so a client-side timeout cannot
|
||
|
|
leave it running -- ``--rm`` only covers containers that actually exit.
|
||
|
|
"""
|
||
|
|
name = name or f"{CONTAINER_PREFIX}{os.getpid()}"
|
||
|
|
binds = " ".join(f"-v {shlex.quote(src)}:{shlex.quote(dst)}" for src, dst in mounts)
|
||
|
|
try:
|
||
|
|
return run(
|
||
|
|
f"podman run --rm --name {shlex.quote(name)} "
|
||
|
|
f"--device /dev/kfd --device /dev/dri "
|
||
|
|
f"--security-opt seccomp=unconfined --ipc=host {binds} {IMAGE} {command}",
|
||
|
|
timeout=timeout, capture=capture,
|
||
|
|
)
|
||
|
|
finally:
|
||
|
|
run(f"podman stop -t 10 {shlex.quote(name)}", check=False, timeout=60)
|