0f5c1e3408
The compute proof of concept had never been run with the agent as umbra-poc. Running it that way found three defects, all the same shape: a check that passed because its mechanism did nothing. The demonstration never reached the agent at all. It held the pid it spawned, which under `sudo -n -u` is the sudo wrapper - sudo forks the agent before exec'ing. Signalling that killed the wrapper and left the worker running, renewing its lease and holding the job indefinitely, under a printed note saying "exactly like a power cut". An unprivileged process cannot signal another account's process at all, and kill(2) on a process group succeeds as soon as one member is signalled, so the wrapper's death read as the agent's. Every run leaked an agent. Lease recovery was a diagnosis rather than a mechanism. sweep_expired_leases documented itself as "called by the controller's own timer" and there was no timer; a killed node's job came back only when somebody noticed a stalled job and posted to an administrative route by hand. The proof of concept swept by hand and recorded the result as automatic. app.py now runs the loop for the life of the application, building the compute repository at startup so a restart mid-job still recovers it. The power cut is now the agent's own: os._exit(137) the instant it claims a job, before running it, because a process that does not own the agent cannot signal it. A clean stop goes through a stop_file the agent polls, for the same reason. Liveness is measured by finding the worker's own pid in /proc matched on its config path, sampled repeatedly and accumulated - the wrapper and the worker are not simultaneous, so a single sample finds only the wrapper. The node key travels by `sudo --preserve-env` rather than as VAR=value in argv, because argv is world-readable through /proc/<pid>/cmdline and the environment is not. The demonstration checks whether the key will survive before starting and falls back to a pipe, which node_agent.py --node-key-stdin accepts. Neither a file nor argv ever holds it. Eight tests, each verified to fail against a mutation of the code it covers: three for the sweeper (it must return an expired lease, must survive a failing round, must not invent a repository it has no database for) and five driving a real agent process against a stub controller (stops on the stop file, exits 137 holding an uncompleted job, takes the key from both routes and neither, reports an empty pipe as an empty pipe). 448 pass with a database, 421 without. The deployment database is untouched: users=2 orgs=3 projects=0 audit=43, no compute tables.
674 lines
28 KiB
Python
674 lines
28 KiB
Python
"""A node agent. One process, one machine, no database credential.
|
|
|
|
Runs on a worker node, talks to the controller over HTTP with a node credential,
|
|
and does three things in a loop: say that it is here, take a job if one is
|
|
suitable, and report what happened.
|
|
|
|
It holds no database connection and no project data. Everything it knows about
|
|
the world arrives over the API, and everything it produces leaves over the API.
|
|
That is deliberate: a node is remotely reachable by definition, so its blast
|
|
radius should be one render rather than a tenant. §7.2, §13.15.
|
|
|
|
Run it as the account that owns the work, not as the account that owns the
|
|
database:
|
|
|
|
umbra-node-agent --config /etc/umbra/node.json
|
|
|
|
The proof-of-concept driver in ``poc/run_poc.py`` runs this as a second, unprivileged
|
|
unix user precisely so that the boundary is demonstrated rather than asserted: the
|
|
agent cannot read the environment file holding the database DSN, because it is not
|
|
that user's file.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from dataclasses import dataclass, field, replace
|
|
from pathlib import Path
|
|
from typing import Any, Mapping
|
|
|
|
LOGGER = logging.getLogger("umbra-node-agent")
|
|
|
|
AGENT_VERSION = "0.1.0"
|
|
|
|
#: How often to say we are here. The controller treats a node stale after
|
|
#: missing three of these, so one slow or dropped report does not take a working
|
|
#: machine out of service.
|
|
HEARTBEAT_INTERVAL_SECONDS = 10
|
|
|
|
|
|
class AgentError(RuntimeError):
|
|
"""The controller refused us, or could not be reached."""
|
|
|
|
|
|
@dataclass
|
|
class NodeConfig:
|
|
"""Everything the agent needs, and nothing it does not.
|
|
|
|
The controller URL and the node key, and nothing else. There is deliberately
|
|
no database DSN field here, because an agent that could read one would not
|
|
need this program's API at all.
|
|
"""
|
|
|
|
controller_url: str
|
|
node_key: str
|
|
node_name: str
|
|
#: What this node is willing to be asked to do. An agent that will not accept
|
|
#: a kind of work should say so, rather than claiming a job and failing it.
|
|
kinds: list[str] = field(default_factory=list)
|
|
lease_seconds: int = 60
|
|
poll_interval_seconds: float = 1.0
|
|
#: How long to wait between failed polls. Backs off, so a controller restart
|
|
#: does not produce a tight reconnect loop.
|
|
retry_backoff_seconds: float = 2.0
|
|
max_retry_backoff_seconds: float = 30.0
|
|
#: A file whose existence means "stop". Optional, and the mechanism by which
|
|
#: something that is not allowed to signal this process can still ask it to
|
|
#: stop cleanly.
|
|
stop_file: Path | None = None
|
|
#: Fault injection: die abruptly after claiming this many jobs, with no
|
|
#: cleanup and no lease release. See :meth:`NodeAgent.run`.
|
|
exit_abruptly_after_claim: int = 0
|
|
|
|
@classmethod
|
|
def key_from_stdin(cls) -> str:
|
|
"""Read the node key from stdin, one line, and let go of it.
|
|
|
|
A third place for the key, after the environment and the config file, for
|
|
the case where the environment cannot be trusted to survive the trip to
|
|
this process. The ordinary way to run an agent as another unix user is
|
|
``sudo -u``, and sudo resets the environment by policy - so the key in
|
|
the environment silently does not arrive, and the agent reports only that
|
|
it found no key. A pipe sidesteps the policy: the value never appears in
|
|
argv, where every user on the host can read it from /proc, and never in a
|
|
file, which outlives the process.
|
|
"""
|
|
line = sys.stdin.readline().strip()
|
|
if not line:
|
|
raise AgentError(
|
|
"--node-key-stdin was given but stdin was empty. The writer closed the "
|
|
"pipe without sending a key."
|
|
)
|
|
return line
|
|
|
|
@classmethod
|
|
def load(cls, path: Path, *, key_from_stdin: bool = False) -> NodeConfig:
|
|
raw = json.loads(path.read_text())
|
|
# The key is deliberately *not* in this list. It may come from the
|
|
# environment, which is the preferred arrangement: a config file holding
|
|
# a secret is a secret that eventually ends up in a backup or a support
|
|
# bundle. Requiring it here would have refused the arrangement the
|
|
# docstring recommends.
|
|
missing = [key for key in ("controller_url", "node_name") if not raw.get(key)]
|
|
if missing:
|
|
raise AgentError(f"{path} is missing: {', '.join(missing)}")
|
|
key = os.environ.get("UMBRA_NODE_KEY")
|
|
if not key and key_from_stdin:
|
|
key = cls.key_from_stdin()
|
|
key = key or raw.get("node_key")
|
|
if not key:
|
|
hint = (
|
|
"Pass --node-key-stdin if the key is coming through a pipe, which is what "
|
|
"happens when sudo resets the environment on the way here."
|
|
if key_from_stdin
|
|
else "If this agent is started through sudo, check whether sudo is "
|
|
"resetting the environment, and pass --node-key-stdin instead."
|
|
)
|
|
raise AgentError(
|
|
f"{path} has no node_key and UMBRA_NODE_KEY is not set. Put the key in one "
|
|
f"of them - the environment is preferable, so the file stays safe to copy. {hint}"
|
|
)
|
|
return cls(
|
|
controller_url=str(raw["controller_url"]).rstrip("/"),
|
|
node_key=key,
|
|
node_name=raw["node_name"],
|
|
kinds=list(raw.get("kinds", [])),
|
|
lease_seconds=int(raw.get("lease_seconds", 60)),
|
|
poll_interval_seconds=float(raw.get("poll_interval_seconds", 1.0)),
|
|
stop_file=Path(raw["stop_file"]) if raw.get("stop_file") else None,
|
|
exit_abruptly_after_claim=int(raw.get("exit_abruptly_after_claim", 0)),
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Measurement
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
def _run(command: list[str]) -> str:
|
|
"""Run a command and return its stdout, or empty if it is unavailable.
|
|
|
|
Every probe here is best-effort. A machine without ``nvidia-smi`` still has
|
|
a CPU device and can still take work - refusing to start because a
|
|
measurement tool is missing would be the same eligibility-versus-preference
|
|
mistake in a different place. §7.3
|
|
"""
|
|
|
|
try:
|
|
completed = subprocess.run(command, capture_output=True, text=True, timeout=5, check=False)
|
|
except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
|
|
return ""
|
|
return completed.stdout.strip()
|
|
|
|
|
|
def _read_meminfo() -> tuple[int, int]:
|
|
"""Return ``(total_mb, available_mb)``.
|
|
|
|
``MemAvailable`` rather than ``MemFree``: the second excludes reclaimable
|
|
page cache, and a node with 2 GB of free RAM and 14 GB of cache is not out of
|
|
memory.
|
|
"""
|
|
|
|
total = available = 0
|
|
try:
|
|
for line in Path("/proc/meminfo").read_text().splitlines():
|
|
key, _, value = line.partition(":")
|
|
megabytes = int(value.strip().split()[0]) // 1024
|
|
if key == "MemTotal":
|
|
total = megabytes
|
|
elif key == "MemAvailable":
|
|
available = megabytes
|
|
except (OSError, ValueError, IndexError):
|
|
pass
|
|
return total, available
|
|
|
|
|
|
def _disk_free(path: str = "/") -> int:
|
|
try:
|
|
return shutil.disk_usage(path).free // (1024 * 1024)
|
|
except OSError:
|
|
return 0
|
|
|
|
|
|
def _nvidia_devices() -> list[dict[str, Any]]:
|
|
"""Report discrete NVIDIA devices, with a measured throughput factor.
|
|
|
|
The factor is a rough but *measured* order of magnitude rather than an
|
|
invented number, because a declared factor goes stale the moment
|
|
quantisation or context length changes and then sends work to the wrong
|
|
machine. §7.3
|
|
"""
|
|
|
|
output = _run(
|
|
[
|
|
"nvidia-smi",
|
|
"--query-gpu=index,name,memory.total,memory.used,utilization.gpu",
|
|
"--format=csv,noheader,nounits",
|
|
]
|
|
)
|
|
devices: list[dict[str, Any]] = []
|
|
for line in output.splitlines():
|
|
parts = [part.strip() for part in line.split(",")]
|
|
if len(parts) != 5:
|
|
continue
|
|
try:
|
|
index, name, total, used, utilisation = (int(parts[0]), parts[1], int(parts[2]), int(parts[3]), int(parts[4]))
|
|
except ValueError:
|
|
continue
|
|
# VRAM is *not* interchangeable with system RAM, but a model that will not
|
|
# fit in free VRAM can often still run by spilling into RAM. Reporting both
|
|
# numbers, and letting the controller decide, is better than reporting one
|
|
# and having it be wrong.
|
|
_, ram_available = _read_meminfo()
|
|
devices.append(
|
|
{
|
|
"kind": "gpu_discrete",
|
|
"label": f"gpu:{index}",
|
|
"vram_total_mb": total,
|
|
"vram_free_mb": max(total - used, 0),
|
|
"ram_total_mb": 0,
|
|
"ram_free_mb": ram_available,
|
|
"compute_units": None,
|
|
# Reference device is a mid-range discrete card. These factors order
|
|
# the fleet; they are not a benchmark.
|
|
"throughput_factor": 40.0 if "3060" in name or "4060" in name or "4070" in name else 20.0,
|
|
"supports_accelerators": ["cuda"],
|
|
"name": name,
|
|
"utilisation": utilisation,
|
|
}
|
|
)
|
|
return devices
|
|
|
|
|
|
def measure() -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
|
"""Return ``(devices, measurements)`` for the machine we are on.
|
|
|
|
A CPU device is always reported, even alongside a GPU. It is what makes this
|
|
node able to take work the GPU cannot, which is the whole of the
|
|
eligibility-versus-preference argument in miniature: a CPU device is slow, not
|
|
incapable.
|
|
"""
|
|
|
|
total_mb, available_mb = _read_meminfo()
|
|
gpus = _nvidia_devices()
|
|
cpu_cores = os.cpu_count() or 1
|
|
# Physical cores, where the platform can tell us. SMT doubles the thread count
|
|
# without doubling throughput, and a factor based on threads would rank a
|
|
# hyperthreaded machine above an unthreaded one of the same silicon.
|
|
physical = cpu_cores
|
|
physical_line = _run(["lscpu"])
|
|
for line in physical_line.splitlines():
|
|
if line.lower().startswith("core(s) per socket"):
|
|
try:
|
|
physical = int(line.split(":")[1].strip())
|
|
except (IndexError, ValueError):
|
|
pass
|
|
devices: list[dict[str, Any]] = [
|
|
{
|
|
"kind": "cpu",
|
|
"label": "cpu:0",
|
|
"vram_total_mb": 0,
|
|
"vram_free_mb": 0,
|
|
"ram_total_mb": total_mb,
|
|
"ram_free_mb": available_mb,
|
|
"compute_units": float(physical or cpu_cores),
|
|
# One reference CPU device. A GPU at 40 is roughly this much faster for
|
|
# generation work; the number orders placements, and nothing more.
|
|
"throughput_factor": 1.0,
|
|
"supports_accelerators": [],
|
|
}
|
|
]
|
|
devices.extend(gpus)
|
|
load = os.getloadavg()[0] if hasattr(os, "getloadavg") else 0.0
|
|
measurements = {
|
|
"ram_free_mb": available_mb,
|
|
"ram_total_mb": total_mb,
|
|
"vram_free_mb": max((device["vram_free_mb"] for device in gpus), default=0),
|
|
"vram_total_mb": max((device["vram_total_mb"] for device in gpus), default=0),
|
|
"cpu_busy_pct": round(min(100.0, 100.0 * load / max(physical, 1)), 2),
|
|
"load_avg_1m": round(load, 2),
|
|
"disk_free_mb": _disk_free(),
|
|
"host": socket.gethostname(),
|
|
"platform": platform.platform(),
|
|
"python": platform.python_version(),
|
|
}
|
|
return devices, measurements
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# The controller client
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class ControllerClient:
|
|
"""HTTP client for the agent surface.
|
|
|
|
Deliberately built on ``urllib`` rather than a dependency, because this
|
|
program runs on four machines with three operating systems and adding an
|
|
install step to each of them is a real cost for a proof of concept. It also
|
|
means the agent has no transitive dependencies at all.
|
|
"""
|
|
|
|
def __init__(self, base_url: str, node_key: str, *, timeout: float = 15.0) -> None:
|
|
self._base = base_url.rstrip("/")
|
|
self._key = node_key
|
|
self._timeout = timeout
|
|
|
|
def _request(self, method: str, path: str, payload: Mapping[str, Any] | None = None) -> Any:
|
|
url = f"{self._base}{path}"
|
|
body = json.dumps(payload).encode() if payload is not None else None
|
|
request = urllib.request.Request(url, data=body, method=method)
|
|
request.add_header("Authorization", f"Bearer {self._key}")
|
|
request.add_header("Content-Type", "application/json")
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=self._timeout) as response:
|
|
# Read once. A second ``response.read()`` on the same stream
|
|
# returns nothing, and every response that was not empty came
|
|
# back as a JSON decode error on an empty string - which looked
|
|
# like a server fault rather than a client bug.
|
|
body = response.read()
|
|
if not body:
|
|
# 204, or a claim that found nothing. Not an error.
|
|
return None
|
|
return json.loads(body.decode())
|
|
except urllib.error.HTTPError as error:
|
|
detail = error.read().decode(errors="replace")
|
|
try:
|
|
parsed = json.loads(detail)
|
|
message = parsed.get("error", {}).get("message", detail)
|
|
code = parsed.get("error", {}).get("code", "unknown")
|
|
except json.JSONDecodeError:
|
|
message, code = detail[:400], f"http_{error.code}"
|
|
raise AgentError(f"{code}: {message}") from None
|
|
except urllib.error.URLError as error:
|
|
raise AgentError(f"controller unreachable: {error.reason}") from None
|
|
except TimeoutError:
|
|
raise AgentError("controller timed out") from None
|
|
|
|
def heartbeat(self, payload: Mapping[str, Any]) -> dict[str, Any]:
|
|
return self._request("POST", "/api/node/heartbeat", payload)
|
|
|
|
def set_model_state(self, slug: str, payload: Mapping[str, Any]) -> None:
|
|
self._request("POST", f"/api/node/models/{slug}/state", payload)
|
|
|
|
def claim(self, kinds: list[str] | None, lease_seconds: int) -> dict[str, Any] | None:
|
|
query = f"?lease_seconds={lease_seconds}"
|
|
if kinds:
|
|
query += "".join(f"&kinds={kind}" for kind in kinds)
|
|
return self._request("POST", f"/api/node/jobs/claim{query}")
|
|
|
|
def renew(self, job_id: str, lease_seconds: int) -> bool:
|
|
result = self._request(
|
|
"POST", f"/api/node/jobs/{job_id}/renew?lease_seconds={lease_seconds}"
|
|
)
|
|
return bool(result and result.get("ok"))
|
|
|
|
def complete(self, job_id: str, result: Mapping[str, Any]) -> dict[str, Any]:
|
|
return self._request("POST", f"/api/node/jobs/{job_id}/complete", {"result": dict(result)})
|
|
|
|
def fail(self, job_id: str, error: str, *, retry: bool = True) -> dict[str, Any]:
|
|
return self._request(
|
|
"POST", f"/api/node/jobs/{job_id}/fail", {"error": error[:2000], "retry": retry}
|
|
)
|
|
|
|
def decompose(self, job_id: str, children: list[Mapping[str, Any]]) -> list[dict[str, Any]]:
|
|
return self._request("POST", f"/api/node/jobs/{job_id}/decompose", {"children": children})
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Work
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class WorkFunction:
|
|
"""Something this node knows how to do.
|
|
|
|
The proof of concept implements one: a sleep-and-report job, so that the queue,
|
|
the lease and the recovery path can be exercised without a model in the loop.
|
|
The point being demonstrated is the *dispatch*, not the work.
|
|
"""
|
|
|
|
def __init__(self, name: str, handler) -> None:
|
|
self.name = name
|
|
self._handler = handler
|
|
|
|
def __call__(self, job: Mapping[str, Any], progress) -> dict[str, Any]:
|
|
return self._handler(job, progress)
|
|
|
|
|
|
def _sleep_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
|
|
"""Take the time the job asked for, renewing the lease while doing it.
|
|
|
|
The renewal is the point. A job that outlasts its lease without renewing is a
|
|
job that gets requeued underneath a worker that is still running it, and the
|
|
duplicate is invisible until somebody notices the work happened twice.
|
|
"""
|
|
|
|
seconds = float((job.get("payload") or {}).get("seconds", 1.0))
|
|
if (job.get("payload") or {}).get("fail"):
|
|
raise RuntimeError((job["payload"]).get("fail_message", "requested failure"))
|
|
if (job.get("payload") or {}).get("decompose_into"):
|
|
return {"decomposed": True}
|
|
steps = max(int(seconds), 1)
|
|
for step in range(steps):
|
|
progress()
|
|
time.sleep(seconds / steps)
|
|
return {
|
|
"slept_seconds": seconds,
|
|
"node": socket.gethostname(),
|
|
"pid": os.getpid(),
|
|
"echo": (job.get("payload") or {}).get("echo"),
|
|
}
|
|
|
|
|
|
def default_handlers() -> dict[str, WorkFunction]:
|
|
return {"sleep": WorkFunction("sleep", _sleep_work)}
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# The agent loop
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class NodeAgent:
|
|
"""Heartbeat, claim, work, report. Then do it again."""
|
|
|
|
def __init__(self, config: NodeConfig, client: ControllerClient, handlers=None) -> None:
|
|
self._config = config
|
|
self._client = client
|
|
self._handlers = handlers if handlers is not None else default_handlers()
|
|
self._stop = False
|
|
self._active = 0
|
|
self._claims_made = 0
|
|
#: What this node has told the controller is resident. Reported on every
|
|
#: heartbeat so the console can distinguish a node that is idle from one
|
|
#: that is idle and out of memory. §7.3
|
|
self._resident: list[str] = []
|
|
self._last_heartbeat = 0.0
|
|
self._backoff = config.retry_backoff_seconds
|
|
|
|
def stop(self, *_: Any) -> None:
|
|
self._stop = True
|
|
|
|
# -- liveness ---------------------------------------------------------
|
|
|
|
def beat(self, *, force: bool = False) -> bool:
|
|
now = time.monotonic()
|
|
if not force and now - self._last_heartbeat < HEARTBEAT_INTERVAL_SECONDS:
|
|
return True
|
|
devices, measurements = measure()
|
|
try:
|
|
response = self._client.heartbeat(
|
|
{
|
|
"agent_version": AGENT_VERSION,
|
|
"devices": devices,
|
|
"resident_models": self._resident,
|
|
"active_jobs": self._active,
|
|
"measurements": measurements,
|
|
}
|
|
)
|
|
except AgentError as error:
|
|
LOGGER.warning("heartbeat failed: %s", error)
|
|
return False
|
|
self._last_heartbeat = now
|
|
if response:
|
|
LOGGER.info(
|
|
"online as %s: %s (schedulable=%s)", response.get("name"),
|
|
response.get("status_reason"), response.get("schedulable"),
|
|
)
|
|
return True
|
|
|
|
# -- claiming ---------------------------------------------------------
|
|
|
|
def claim_and_run(self) -> bool:
|
|
"""Take at most one job. Returns True if work was done.
|
|
|
|
One at a time, deliberately. A proof of purpose that claims five jobs and
|
|
runs them in threads would be demonstrating thread pools, not leases, and
|
|
the lease is the part worth demonstrating.
|
|
"""
|
|
|
|
try:
|
|
claimed = self._client.claim(self._config.kinds or None, self._config.lease_seconds)
|
|
except AgentError as error:
|
|
LOGGER.warning("claim failed: %s", error)
|
|
return False
|
|
if not claimed:
|
|
return False
|
|
job = claimed["job"]
|
|
job_id = job["id"]
|
|
placement = claimed.get("placement", {})
|
|
self._active = 1
|
|
LOGGER.info(
|
|
"claimed %s %s -> %s via %s (%s)",
|
|
job["kind"], job_id[:8], placement.get("node_name", "?"),
|
|
placement.get("device_kind", "?"), placement.get("reason", "?"),
|
|
)
|
|
self._claims_made += 1
|
|
if (
|
|
self._config.exit_abruptly_after_claim
|
|
and self._claims_made >= self._config.exit_abruptly_after_claim
|
|
):
|
|
# Here, between taking the job and doing it. Not after the work: a
|
|
# fault that fires once the job has finished tests a finished job,
|
|
# and the lease it is meant to strand is already released.
|
|
LOGGER.warning(
|
|
"fault injection: exiting abruptly having claimed %d job(s), without "
|
|
"running it and without releasing the lease",
|
|
self._claims_made,
|
|
)
|
|
os._exit(137)
|
|
# A placement that names something to evict says so before the work starts,
|
|
# because the cost of the eviction is somebody else's warm model.
|
|
if placement.get("would_evict"):
|
|
LOGGER.warning(
|
|
"this job needs to unload %s first; that is another node's work",
|
|
", ".join(placement["would_evict"]),
|
|
)
|
|
|
|
handler = self._handlers.get(job["kind"])
|
|
if handler is None:
|
|
self._report_failure(job_id, f"this agent does not implement '{job['kind']}'")
|
|
return True
|
|
|
|
def progress() -> None:
|
|
# Renew well inside the lease. A lease is a liveness signal, and a
|
|
# worker that is alive is one that keeps saying so.
|
|
if not self._client.renew(job_id, self._config.lease_seconds):
|
|
LOGGER.warning("lease renewal refused for %s; the job may have been requeued", job_id[:8])
|
|
|
|
try:
|
|
result = handler(job, progress)
|
|
except Exception as error: # noqa: BLE001 - an agent must survive its own work failing
|
|
LOGGER.exception("job %s failed", job_id[:8])
|
|
self._report_failure(job_id, f"{type(error).__name__}: {error}")
|
|
self._active = 0
|
|
return True
|
|
|
|
if result.get("decomposed"):
|
|
self._decompose(job)
|
|
self._active = 0
|
|
return True
|
|
try:
|
|
self._client.complete(job_id, result)
|
|
LOGGER.info("completed %s", job_id[:8])
|
|
except AgentError as error:
|
|
LOGGER.error("could not report completion for %s: %s", job_id[:8], error)
|
|
self._active = 0
|
|
return True
|
|
|
|
def _decompose(self, job: Mapping[str, Any]) -> None:
|
|
"""Hand the pieces back for allocation.
|
|
|
|
Only the holding node may do this, and the parent becomes a container
|
|
rather than a second copy of the work. §6.19
|
|
"""
|
|
|
|
children = (job.get("payload") or {}).get("decompose_into") or []
|
|
try:
|
|
created = self._client.decompose(job["id"], children)
|
|
LOGGER.info("decomposed %s into %d jobs", job["id"][:8], len(created))
|
|
except AgentError as error:
|
|
LOGGER.error("decomposition refused: %s", error)
|
|
self._report_failure(job["id"], f"decomposition refused: {error}")
|
|
|
|
def _report_failure(self, job_id: str, message: str) -> None:
|
|
try:
|
|
self._client.fail(job_id, message, retry=True)
|
|
except AgentError as error:
|
|
LOGGER.error("could not report failure for %s: %s", job_id[:8], error)
|
|
|
|
# -- loop -------------------------------------------------------------
|
|
|
|
def run(self, *, once: bool = False, max_iterations: int | None = None) -> None:
|
|
"""The loop.
|
|
|
|
``once`` and ``max_iterations`` exist so the proof-of-concept driver can
|
|
step this deterministically instead of racing a timer. An agent that can
|
|
only be tested by sleeping is an agent whose lease behaviour is untested.
|
|
"""
|
|
|
|
import signal
|
|
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
try:
|
|
signal.signal(sig, self.stop)
|
|
except (ValueError, OSError): # pragma: no cover - non-main thread
|
|
pass
|
|
|
|
self.beat(force=True)
|
|
iterations = 0
|
|
while not self._stop:
|
|
# Asked to stop by file. A supervisor that cannot signal this process -
|
|
# because it does not own it, which is the normal arrangement when an
|
|
# agent runs as a service account - still has to be able to stop it
|
|
# cleanly. Without this, stopping an agent meant killing the *sudo
|
|
# wrapper*, which orphaned the worker: every proof-of-concept run left
|
|
# an agent behind, quietly renewing leases for jobs nobody was watching.
|
|
if self._config.stop_file is not None and self._config.stop_file.exists():
|
|
LOGGER.info("stop file %s appeared; stopping", self._config.stop_file)
|
|
break
|
|
worked = False
|
|
try:
|
|
worked = self.claim_and_run()
|
|
except Exception: # noqa: BLE001 - the loop outlives any single failure
|
|
LOGGER.exception("unexpected error in the claim/run cycle")
|
|
|
|
if not self.beat():
|
|
# Backing off rather than spinning. A controller that is down
|
|
# should not be hammered by every node on the network at once.
|
|
time.sleep(self._backoff)
|
|
self._backoff = min(self._backoff * 2, self._config.max_retry_backoff_seconds)
|
|
else:
|
|
self._backoff = self._config.retry_backoff_seconds
|
|
iterations += 1
|
|
if once or (max_iterations is not None and iterations >= max_iterations):
|
|
break
|
|
if not worked:
|
|
time.sleep(self._config.poll_interval_seconds)
|
|
LOGGER.info("agent stopped after %d iteration(s)", iterations)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description="UmbraTrans compute node agent")
|
|
parser.add_argument("--config", required=True, type=Path, help="path to the node JSON config")
|
|
parser.add_argument(
|
|
"--node-key-stdin",
|
|
action="store_true",
|
|
help="read the node key from stdin instead of the environment; for use under "
|
|
"sudo, which resets the environment by policy",
|
|
)
|
|
parser.add_argument("--log-level", default="INFO")
|
|
parser.add_argument("--once", action="store_true", help="one claim/heartbeat cycle and exit")
|
|
parser.add_argument("--max-iterations", type=int, default=None)
|
|
parser.add_argument(
|
|
"--fault-exit-after-claim",
|
|
type=int,
|
|
default=None,
|
|
metavar="N",
|
|
help="fault injection: exit abruptly with no cleanup after claiming N jobs, "
|
|
"so lease recovery can be exercised against a process that really vanished",
|
|
)
|
|
args = parser.parse_args(argv)
|
|
|
|
logging.basicConfig(
|
|
level=getattr(logging, str(args.log_level).upper(), logging.INFO),
|
|
format="%(asctime)s %(levelname)-7s %(name)s: %(message)s",
|
|
)
|
|
config = NodeConfig.load(args.config, key_from_stdin=args.node_key_stdin)
|
|
if args.fault_exit_after_claim is not None:
|
|
config = replace(config, exit_abruptly_after_claim=args.fault_exit_after_claim)
|
|
client = ControllerClient(config.controller_url, config.node_key)
|
|
agent = NodeAgent(config, client)
|
|
LOGGER.info(
|
|
"node agent %s starting as %s against %s (kinds=%s)",
|
|
AGENT_VERSION, config.node_name, config.controller_url, config.kinds or ["any"],
|
|
)
|
|
agent.run(once=args.once, max_iterations=args.max_iterations)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__": # pragma: no cover
|
|
sys.exit(main())
|