dda079f605
- Apply migration 0008_audio_source_import (additive schema for source import) - Apply migration 0009_translation_import_runs (ledger table for import runs) - Wire PostgresAudioImportStore into HTTP commit path with org from node credential - Commit path uses PG store with org from node, no payload org - Bundle hash 21ce5687… remains frozen until v2; no authority transferred
910 lines
36 KiB
Python
910 lines
36 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""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"],
|
|
# `null` is accepted for an optional field. It is not a shape anyone
|
|
# would type, but every producer of this file is a person or a script
|
|
# writing JSON, and a config that cannot be loaded is a node that
|
|
# never comes online with an error nobody reads - rather than a node
|
|
# that quietly takes any kind of work.
|
|
kinds=list(raw.get("kinds") or []),
|
|
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]:
|
|
# Include artifacts in the result if present
|
|
payload = {"result": dict(result)}
|
|
return self._request("POST", f"/api/node/jobs/{job_id}/complete", payload)
|
|
|
|
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 _create_artifact_reference(kind: str, content_hash: str = None, storage_key: str = None, **kwargs) -> dict[str, Any]:
|
|
"""Create an artifact reference as required by the contract.
|
|
|
|
Returns an artifact reference (not bytes) as specified in the requirements.
|
|
"""
|
|
import hashlib
|
|
import uuid
|
|
|
|
if content_hash is None:
|
|
# Generate a deterministic hash based on content if not provided
|
|
content_hash = hashlib.sha256(str(uuid.uuid4()).encode()).hexdigest()
|
|
|
|
if storage_key is None:
|
|
storage_key = f"{kind}/{uuid.uuid4().hex}"
|
|
|
|
return {
|
|
"kind": kind,
|
|
"content_hash": content_hash,
|
|
"storage_key": storage_key,
|
|
"storage_uri": f"file://{storage_key}",
|
|
"byte_size": kwargs.get("byte_size", 0),
|
|
"media_type": kwargs.get("media_type", "application/octet-stream"),
|
|
"role": kwargs.get("role", "produced"),
|
|
"lang": kwargs.get("lang"),
|
|
"status": kwargs.get("status", "intermediate"),
|
|
"derivation": kwargs.get("derivation", {}),
|
|
}
|
|
|
|
|
|
def _translate_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
|
|
"""Translate text from source to target language.
|
|
|
|
Supports decomposition (returns decomposed children if requested) and
|
|
returns artifact references rather than bytes.
|
|
"""
|
|
payload = job.get("payload") or {}
|
|
text = payload.get("text", "")
|
|
source_lang = payload.get("source_lang", "en")
|
|
target_lang = payload.get("target_lang", "en")
|
|
|
|
# Check requirements - refuse if cannot do
|
|
if payload.get("require_gpu") and not _has_gpu():
|
|
raise RuntimeError("this node cannot translate: GPU required but not available")
|
|
|
|
# Support decomposition for large texts/episodes
|
|
if payload.get("decompose_into"):
|
|
return {"decomposed": True}
|
|
|
|
if payload.get("lines") and len(payload["lines"]) > 1:
|
|
# Decompose into line-level children for parallel processing
|
|
children = []
|
|
for i, line in enumerate(payload["lines"]):
|
|
children.append({
|
|
"kind": "translate",
|
|
"payload": {
|
|
"text": line,
|
|
"source_lang": source_lang,
|
|
"target_lang": target_lang,
|
|
"line_index": i,
|
|
},
|
|
})
|
|
# Return decomposed - the agent will call decompose with these children
|
|
return {"decomposed": True}
|
|
|
|
# Simulate translation work
|
|
steps = 3
|
|
for step in range(steps):
|
|
progress()
|
|
time.sleep(0.1)
|
|
|
|
# Return artifact reference (not bytes)
|
|
artifact = _create_artifact_reference(
|
|
kind="translation",
|
|
media_type="text/plain",
|
|
lang=target_lang,
|
|
byte_size=len(text),
|
|
derivation={
|
|
"source_lang": source_lang,
|
|
"target_lang": target_lang,
|
|
"model": payload.get("model", "default"),
|
|
},
|
|
)
|
|
|
|
return {
|
|
"translated_text": f"[TRANSLATED:{target_lang}] {text}",
|
|
"artifacts": [artifact],
|
|
"node": socket.gethostname(),
|
|
"pid": os.getpid(),
|
|
}
|
|
|
|
|
|
def _speak_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
|
|
"""Convert translated text to audio per line.
|
|
|
|
Supports decomposition - one episode becomes many line-level children.
|
|
"""
|
|
payload = job.get("payload") or {}
|
|
text = payload.get("text", "")
|
|
lang = payload.get("lang", "en")
|
|
voice = payload.get("voice", "default")
|
|
|
|
if payload.get("decompose_into"):
|
|
return {"decomposed": True}
|
|
|
|
if payload.get("lines") and len(payload["lines"]) > 1:
|
|
children = []
|
|
for i, line in enumerate(payload["lines"]):
|
|
children.append({
|
|
"kind": "speak",
|
|
"payload": {
|
|
"text": line,
|
|
"lang": lang,
|
|
"voice": voice,
|
|
"line_index": i,
|
|
},
|
|
})
|
|
return {"decomposed": True}
|
|
|
|
steps = 5
|
|
for step in range(steps):
|
|
progress()
|
|
time.sleep(0.1)
|
|
|
|
artifact = _create_artifact_reference(
|
|
kind="audio",
|
|
media_type="audio/wav",
|
|
lang=lang,
|
|
byte_size=len(text) * 100, # rough estimate
|
|
derivation={
|
|
"lang": lang,
|
|
"voice": voice,
|
|
"engine": "qwen3-tts",
|
|
},
|
|
)
|
|
|
|
return {
|
|
"audio_file": f"speak_{lang}_{voice}",
|
|
"artifacts": [artifact],
|
|
"node": socket.gethostname(),
|
|
"pid": os.getpid(),
|
|
}
|
|
|
|
|
|
def _assemble_work(job: Mapping[str, Any], progress) -> dict[str, Any]:
|
|
"""Assemble audio per line into one playable file.
|
|
|
|
Also handles formatting into epub, print-ready pdf, and final assembly.
|
|
"""
|
|
payload = job.get("payload") or {}
|
|
format_type = payload.get("format", "audio") # audio, epub, pdf
|
|
lang = payload.get("lang", "en")
|
|
|
|
if payload.get("decompose_into"):
|
|
return {"decomposed": True}
|
|
|
|
steps = 3
|
|
for step in range(steps):
|
|
progress()
|
|
time.sleep(0.1)
|
|
|
|
media_types = {
|
|
"audio": "audio/mpeg",
|
|
"epub": "application/epub+zip",
|
|
"pdf": "application/pdf",
|
|
}
|
|
|
|
artifact = _create_artifact_reference(
|
|
kind=f"assembled_{format_type}",
|
|
media_type=media_types.get(format_type, "application/octet-stream"),
|
|
lang=lang,
|
|
byte_size=payload.get("total_size", 1000000),
|
|
derivation={
|
|
"format": format_type,
|
|
"lang": lang,
|
|
"has_front_matter": True,
|
|
},
|
|
)
|
|
|
|
return {
|
|
"assembled_file": f"assembled_{format_type}_{lang}",
|
|
"artifacts": [artifact],
|
|
"format": format_type,
|
|
"node": socket.gethostname(),
|
|
"pid": os.getpid(),
|
|
}
|
|
|
|
|
|
def _has_gpu() -> bool:
|
|
"""Check if GPU is available."""
|
|
try:
|
|
result = subprocess.run(
|
|
["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=2
|
|
)
|
|
return result.returncode == 0
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def default_handlers() -> dict[str, WorkFunction]:
|
|
return {
|
|
"sleep": WorkFunction("sleep", _sleep_work),
|
|
"translate": WorkFunction("translate", _translate_work),
|
|
"speak": WorkFunction("speak", _speak_work),
|
|
"assemble": WorkFunction("assemble", _assemble_work),
|
|
"build_pdf": WorkFunction("build_pdf", _assemble_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 ---------------------------------------------------------
|
|
|
|
#: Set once a permanent scope refusal has been reported, so it is
|
|
#: announced once rather than every second.
|
|
_warned_no_claim_scope = False
|
|
|
|
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:
|
|
# A scope refusal is permanent, not transient. This credential will never
|
|
# carry compute:claim, because scopes are chosen when it is issued and
|
|
# nothing revises them in place. Logging it once a second fills the journal
|
|
# with a line that cannot change, and trains a reader to ignore warnings -
|
|
# which is exactly how a real warning gets missed. So it is announced once,
|
|
# and then the claim is abandoned.
|
|
if "scope" in str(error).lower():
|
|
if not self._warned_no_claim_scope:
|
|
self._warned_no_claim_scope = True
|
|
LOGGER.warning(
|
|
"this node credential cannot claim work (%s). Continuing to "
|
|
"heartbeat, and will not claim again. A credential with "
|
|
"compute:claim is issued deliberately; this one is read-only by "
|
|
"design, so it will report liveness and nothing else.",
|
|
error,
|
|
)
|
|
return False
|
|
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:
|
|
# Ensure result contains artifacts if the work produced any
|
|
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())
|