Melvin Strobl
85dad39973
- `service` node type: command, port, worker label, health path, env and
weights (an artifact digest the worker fetches as {weights}). Its flow
starts it; it restarts with backoff, reads degraded while starting or
without a worker and down after three starts that never got healthy, and
publishes http://<worker>:<port> once /health answers 200, "" when it ends.
- Worker protocol 3: a `serve` call, every call in its own process group,
killed when the engine ends the call, the socket drops or the agent gets
SIGTERM. `--service-host` overrides the address the engine sees.
- Services and compiles take no `--parallel` slot, and a flow's detail is
built off the event loop, so a save no longer stalls the engine.
692 lines
26 KiB
Python
692 lines
26 KiB
Python
"""The agent that runs Fluksio nodes on a machine the engine cannot reach.
|
|
|
|
Point it at the engine on the box with the GPU and it dials in::
|
|
|
|
pip install fluksio-worker
|
|
fluksio-worker --url wss://api.example.com/api/v1/workers/attach --labels gpu
|
|
|
|
Install it into the environment the training code already runs in and node
|
|
code runs on that: ``--python`` defaults to the interpreter this was started
|
|
with. Point it elsewhere only when the two are meant to differ.
|
|
|
|
It connects *out*, so the engine needs no route back and nothing has to expose
|
|
Redis. What it then does is what the engine's own worker pool does: hold a few
|
|
subprocesses running ``worker_main.py``, hand each call to one, and pass back
|
|
everything that comes out — including the metrics a training loop reports
|
|
while it is still running. A service node's call runs a command instead, a
|
|
server kept up until the engine ends the call.
|
|
|
|
Every call runs in a process group of its own, and ends with it: when the
|
|
engine ends the call, when the socket drops, or when this agent stops. A
|
|
process left behind would keep its port and its card.
|
|
|
|
Deliberately two files with one dependency. Nothing of the engine is imported
|
|
here; a worker host installs Python, ``websockets``, and whatever the nodes
|
|
themselves need. Where pip is not an option, copying this file and
|
|
``worker_main.py`` into one directory and running ``python agent.py`` is the
|
|
same thing — the engine serves the runner at ``GET /api/v1/workers/runtime``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import contextlib
|
|
import glob
|
|
import importlib.metadata
|
|
import ipaddress
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import urlsplit
|
|
|
|
try:
|
|
import websockets
|
|
except ImportError: # pragma: no cover - the one dependency, named plainly
|
|
print(
|
|
"This needs the 'websockets' package: pip install websockets", file=sys.stderr
|
|
)
|
|
raise SystemExit(1) from None
|
|
|
|
log = logging.getLogger("fluksio-worker")
|
|
|
|
PROTOCOL = 3
|
|
#: Sent while a call is running, so the engine can tell working from wedged.
|
|
HEARTBEAT_S = 10.0
|
|
#: How long a stopped service gets to shut down before its group is killed.
|
|
STOP_GRACE_S = 3.0
|
|
#: How much of what a service printed last is kept, to say why it ended.
|
|
TAIL_BYTES = 4096
|
|
#: Reconnection backs off to this and no further.
|
|
MAX_BACKOFF_S = 30.0
|
|
#: Environment node code is not given; the engine's lists, see `_node_env`.
|
|
ENV_DENY_PREFIXES = ("DATABASE_URL", "FIRST_SUPERUSER", "SENTRY_DSN")
|
|
ENV_DENY_WORDS = ("PASSWORD", "SECRET", "TOKEN")
|
|
|
|
|
|
class Refused(RuntimeError):
|
|
"""The engine will not have this worker, and retrying will not change it."""
|
|
|
|
|
|
class Idle(Exception):
|
|
"""Nothing has run here for long enough that this worker is done."""
|
|
|
|
|
|
#: The user-code runner, beside this file — installed together, or copied
|
|
#: together onto a host where pip is not an option.
|
|
WORKER_MAIN = Path(__file__).resolve().with_name("worker_main.py")
|
|
|
|
|
|
class Subprocess:
|
|
"""One user-code process and the framing of one call over its pipes."""
|
|
|
|
def __init__(self, python: str, env: dict[str, str]) -> None:
|
|
self.proc = subprocess.Popen(
|
|
[python, str(WORKER_MAIN)],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
close_fds=True,
|
|
env=env,
|
|
start_new_session=True,
|
|
)
|
|
self._buffer = bytearray()
|
|
|
|
def alive(self) -> bool:
|
|
return self.proc.poll() is None
|
|
|
|
def send(self, request: dict[str, Any]) -> None:
|
|
assert self.proc.stdin is not None
|
|
self.proc.stdin.write((json.dumps(request) + "\n").encode())
|
|
self.proc.stdin.flush()
|
|
|
|
def read_line(self) -> str:
|
|
"""One line, blocking. Empty when the process is gone.
|
|
|
|
Read at the file-descriptor level so a report arriving mid-call is
|
|
passed on the moment it is written rather than when a buffer fills.
|
|
"""
|
|
assert self.proc.stdout is not None
|
|
fd = self.proc.stdout.fileno()
|
|
while True:
|
|
end = self._buffer.find(b"\n")
|
|
if end >= 0:
|
|
line = bytes(self._buffer[: end + 1])
|
|
del self._buffer[: end + 1]
|
|
return line.decode(errors="replace")
|
|
try:
|
|
chunk = os.read(fd, 65536)
|
|
except OSError:
|
|
chunk = b""
|
|
if not chunk:
|
|
self._buffer.clear()
|
|
return ""
|
|
self._buffer += chunk
|
|
|
|
def kill(self) -> None:
|
|
"""SIGKILL to the whole group, so whatever the node started goes too."""
|
|
_signal_group(self.proc.pid, signal.SIGKILL)
|
|
with contextlib.suppress(subprocess.TimeoutExpired):
|
|
self.proc.wait(timeout=5)
|
|
|
|
|
|
class Agent:
|
|
"""Holds the connection, and one subprocess per call in flight."""
|
|
|
|
def __init__(self, args: argparse.Namespace) -> None:
|
|
self.args = args
|
|
self.env = _node_env()
|
|
self.env["FLUKSIO_ARTIFACT_URL"] = args.artifact_url or _artifacts_from(
|
|
args.url
|
|
)
|
|
self.env["FLUKSIO_ARTIFACT_TOKEN"] = args.token
|
|
self.running: dict[str, Subprocess] = {}
|
|
#: Service calls, by call id. Cancelling one stops its server.
|
|
self.serving: dict[str, asyncio.Task[None]] = {}
|
|
#: Every call of this session, so a dropped socket can end them.
|
|
self._calls: set[asyncio.Task[None]] = set()
|
|
self.last_done = time.monotonic()
|
|
self._idled = False
|
|
|
|
async def serve_forever(self) -> None:
|
|
# A stop is a Ctrl-C: the calls end, and their process groups with them.
|
|
main = asyncio.current_task()
|
|
assert main is not None
|
|
asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, main.cancel)
|
|
backoff = 1.0
|
|
while True:
|
|
try:
|
|
await self._session()
|
|
backoff = 1.0
|
|
except Idle:
|
|
log.info("nothing to do for %.0fs — stopping", self.args.max_idle)
|
|
return
|
|
except Refused:
|
|
raise
|
|
except Exception as exc:
|
|
log.warning("disconnected: %s — retrying in %.0fs", exc, backoff)
|
|
await asyncio.sleep(backoff)
|
|
backoff = min(MAX_BACKOFF_S, backoff * 2)
|
|
|
|
async def _session(self) -> None:
|
|
# A header rather than the query string, which access logs keep.
|
|
async with websockets.connect(
|
|
self.args.url,
|
|
additional_headers={"Authorization": f"Bearer {self.args.token}"},
|
|
max_size=None,
|
|
ping_interval=20,
|
|
) as socket:
|
|
await socket.send(
|
|
json.dumps(
|
|
{
|
|
"op": "hello",
|
|
"protocol": PROTOCOL,
|
|
"version": _version(),
|
|
"name": self.args.name,
|
|
"labels": self.args.labels,
|
|
"python": self.args.python,
|
|
"max_parallel": self.args.parallel,
|
|
"venv_digest": _venv_digest(self.args.python),
|
|
"inventory": _inventory(self.args),
|
|
"service_host": self.args.service_host,
|
|
}
|
|
)
|
|
)
|
|
welcome = json.loads(await socket.recv())
|
|
if welcome.get("op") != "welcome":
|
|
# Nothing about this machine will change the answer, so
|
|
# reconnecting would only be a quiet loop in the job's log.
|
|
raise Refused(str(welcome.get("reason") or "refused"))
|
|
log.info(
|
|
"attached to %s as '%s' with labels %s",
|
|
self.args.url,
|
|
welcome.get("name"),
|
|
self.args.labels,
|
|
)
|
|
|
|
self.last_done = time.monotonic()
|
|
watchdog = (
|
|
asyncio.create_task(self._idle_watch(socket))
|
|
if self.args.max_idle > 0
|
|
else None
|
|
)
|
|
try:
|
|
async for raw in socket:
|
|
message = json.loads(raw)
|
|
op = message.get("op")
|
|
if op in ("run", "compile"):
|
|
# Compiling is loading the source, which is the same
|
|
# trip through a subprocess a call is — and has to
|
|
# happen here rather than on the engine, because "does
|
|
# this import" is a question about *this* machine's
|
|
# packages.
|
|
log.info("%s %s", op, message.get("call_id"))
|
|
self._spawn(self._run(socket, message))
|
|
elif op == "serve":
|
|
log.info("serve %s", message.get("call_id"))
|
|
self._spawn(self._serve(socket, message))
|
|
elif op == "cancel":
|
|
self._cancel(str(message.get("call_id") or ""))
|
|
except Exception:
|
|
# The watchdog closing the socket is a stop, not a disconnect.
|
|
if not self._idled:
|
|
raise
|
|
finally:
|
|
if watchdog is not None:
|
|
watchdog.cancel()
|
|
# The engine has failed these calls already, and a process
|
|
# left running would keep its port for the next one.
|
|
await self._end_calls()
|
|
if self._idled:
|
|
raise Idle
|
|
|
|
def _spawn(self, call: Any) -> None:
|
|
task = asyncio.create_task(call)
|
|
self._calls.add(task)
|
|
task.add_done_callback(self._calls.discard)
|
|
# Without this a failure in the task is only noticed when it is
|
|
# garbage collected, which reads as a call that vanished.
|
|
task.add_done_callback(_report_failure)
|
|
|
|
async def _end_calls(self) -> None:
|
|
calls = list(self._calls)
|
|
for task in calls:
|
|
task.cancel()
|
|
await asyncio.gather(*calls, return_exceptions=True)
|
|
|
|
async def _idle_watch(self, socket: Any) -> None:
|
|
"""Close the socket once nothing has run here for ``--max-idle``.
|
|
|
|
A worker a batch scheduler started for one node has no other way to
|
|
know it is finished, and a job that idles until its walltime is a job
|
|
somebody else was queued behind.
|
|
"""
|
|
tick = min(30.0, max(1.0, self.args.max_idle / 4))
|
|
while True:
|
|
await asyncio.sleep(tick)
|
|
if self.running or self.serving:
|
|
continue
|
|
if time.monotonic() - self.last_done >= self.args.max_idle:
|
|
self._idled = True
|
|
with contextlib.suppress(Exception):
|
|
await socket.close()
|
|
return
|
|
|
|
async def _run(self, socket: Any, request: dict[str, Any]) -> None:
|
|
"""Execute one call in a subprocess, streaming what it says back."""
|
|
call_id = str(request.get("call_id") or "")
|
|
loop = asyncio.get_running_loop()
|
|
# What the engine allocated this call: thread caps and the GPUs it may
|
|
# see. Taken off the request because a library reads those once, when
|
|
# it is imported — which for this agent is a process that does not
|
|
# exist yet, so there is nothing to reconfigure and no pool to key.
|
|
extra = {str(k): str(v) for k, v in (request.pop("env", None) or {}).items()}
|
|
mask = self.env.get("CUDA_VISIBLE_DEVICES")
|
|
if mask and extra.get("CUDA_VISIBLE_DEVICES"):
|
|
extra["CUDA_VISIBLE_DEVICES"] = _remap_cuda(
|
|
mask, extra["CUDA_VISIBLE_DEVICES"]
|
|
)
|
|
worker = Subprocess(self.args.python, {**self.env, **extra})
|
|
self.running[call_id] = worker
|
|
|
|
heartbeat = asyncio.create_task(_beat(socket, call_id))
|
|
try:
|
|
# A subprocess that died before it could be written to is reported
|
|
# by the read below, which says so and ends the call — rather than
|
|
# raising here and leaving the engine waiting out its silence.
|
|
with contextlib.suppress(OSError):
|
|
await loop.run_in_executor(None, worker.send, request)
|
|
while True:
|
|
line = await loop.run_in_executor(None, worker.read_line)
|
|
if not line:
|
|
await socket.send(
|
|
json.dumps(
|
|
{
|
|
"call_id": call_id,
|
|
"ok": False,
|
|
"error": {
|
|
"type": "NodeCancelled"
|
|
if worker.proc.returncode
|
|
else "RemoteError",
|
|
"message": "the node process stopped",
|
|
"short": "the node process stopped",
|
|
"traceback": "",
|
|
},
|
|
}
|
|
)
|
|
)
|
|
return
|
|
await socket.send(line.strip())
|
|
# Anything without an `event` is the answer; the call is over.
|
|
if not json.loads(line).get("event"):
|
|
return
|
|
finally:
|
|
heartbeat.cancel()
|
|
self.running.pop(call_id, None)
|
|
self.last_done = time.monotonic()
|
|
worker.kill()
|
|
|
|
async def _serve(self, socket: Any, request: dict[str, Any]) -> None:
|
|
"""Run one server until it exits or the call is ended, then say why.
|
|
|
|
The engine checks its health and restarts it; this only runs it. What
|
|
it prints goes to this agent's log, and the last of it into the reply.
|
|
"""
|
|
call_id = str(request.get("call_id") or "")
|
|
task = asyncio.current_task()
|
|
assert task is not None
|
|
self.serving[call_id] = task
|
|
env = {
|
|
**self.env,
|
|
**{str(k): str(v) for k, v in (request.get("env") or {}).items()},
|
|
}
|
|
argv = [str(part) for part in request.get("argv") or []]
|
|
heartbeat = asyncio.create_task(_beat(socket, call_id))
|
|
proc: asyncio.subprocess.Process | None = None
|
|
tail = b""
|
|
try:
|
|
weights = str(request.get("weights") or "")
|
|
if weights:
|
|
path = await self._fetch(weights, env)
|
|
argv = [part.replace("{weights}", path) for part in argv]
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*argv,
|
|
stdin=subprocess.DEVNULL,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
start_new_session=True,
|
|
)
|
|
assert proc.stdout is not None
|
|
while chunk := await proc.stdout.read(65536):
|
|
sys.stderr.buffer.write(chunk)
|
|
sys.stderr.flush()
|
|
tail = (tail + chunk)[-TAIL_BYTES:]
|
|
reason = f"the service exited with code {await proc.wait()}"
|
|
last = tail.decode(errors="replace").strip().splitlines()
|
|
if last:
|
|
reason += f": {last[-1][:300]}"
|
|
except asyncio.CancelledError:
|
|
# Ended by the engine or by the socket going: stopping it is the
|
|
# job, and the reply below still goes if the socket is there.
|
|
reason = "the service was stopped"
|
|
except Exception as exc:
|
|
reason = f"{type(exc).__name__}: {exc}"
|
|
finally:
|
|
heartbeat.cancel()
|
|
self.serving.pop(call_id, None)
|
|
self.last_done = time.monotonic()
|
|
if proc is not None:
|
|
await _stop_group(proc)
|
|
log.info("%s ended: %s", call_id, reason)
|
|
with contextlib.suppress(Exception):
|
|
await socket.send(
|
|
json.dumps(
|
|
{
|
|
"call_id": call_id,
|
|
"ok": False,
|
|
"error": {
|
|
"type": "ServiceStopped",
|
|
"message": reason,
|
|
"short": reason,
|
|
"traceback": "",
|
|
},
|
|
}
|
|
)
|
|
)
|
|
|
|
async def _fetch(self, digest: str, env: dict[str, str]) -> str:
|
|
"""A local path for an artifact, fetched as node code fetches one."""
|
|
proc = await asyncio.create_subprocess_exec(
|
|
self.args.python,
|
|
str(WORKER_MAIN),
|
|
"--fetch",
|
|
digest,
|
|
env=env,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
out, err = await proc.communicate()
|
|
if proc.returncode:
|
|
problem = err.decode(errors="replace").strip().splitlines()
|
|
raise RuntimeError(
|
|
f"could not fetch {digest}: {problem[-1] if problem else proc.returncode}"
|
|
)
|
|
return out.decode().strip()
|
|
|
|
def _cancel(self, call_id: str) -> None:
|
|
worker = self.running.get(call_id)
|
|
if worker is not None:
|
|
log.info("cancelling %s", call_id)
|
|
worker.kill()
|
|
service = self.serving.get(call_id)
|
|
if service is not None:
|
|
log.info("stopping %s", call_id)
|
|
service.cancel()
|
|
|
|
|
|
async def _beat(socket: Any, call_id: str) -> None:
|
|
"""Say this call is alive, for as long as it runs."""
|
|
while True:
|
|
await asyncio.sleep(HEARTBEAT_S)
|
|
with contextlib.suppress(Exception):
|
|
await socket.send(json.dumps({"call_id": call_id, "event": "heartbeat"}))
|
|
|
|
|
|
def _signal_group(pid: int, sig: int) -> None:
|
|
"""Signal a call's process group, which it leads. Gone already is fine."""
|
|
with contextlib.suppress(OSError):
|
|
os.killpg(pid, sig)
|
|
|
|
|
|
async def _stop_group(proc: asyncio.subprocess.Process) -> None:
|
|
"""SIGTERM to the group, a moment to shut down, then SIGKILL to the rest."""
|
|
_signal_group(proc.pid, signal.SIGTERM)
|
|
try:
|
|
await asyncio.wait_for(proc.wait(), STOP_GRACE_S)
|
|
except TimeoutError:
|
|
pass
|
|
finally:
|
|
_signal_group(proc.pid, signal.SIGKILL)
|
|
|
|
|
|
def _report_failure(task: asyncio.Task[Any]) -> None:
|
|
if not task.cancelled() and task.exception() is not None:
|
|
log.exception("call failed", exc_info=task.exception())
|
|
|
|
|
|
def _node_env() -> dict[str, str]:
|
|
"""This host's environment with its credentials taken out, for node code.
|
|
|
|
The engine's own denylist (``fluksio.flow.workers.worker_env``), copied
|
|
because nothing of the engine is imported here: ``PATH``, ``HOME`` and
|
|
whatever a venv needs still arrive, a ``*_TOKEN`` or ``*_PASSWORD`` does
|
|
not — this worker's own token included. Not a sandbox: node code runs as
|
|
the same user on the same machine.
|
|
"""
|
|
return {
|
|
key: value
|
|
for key, value in os.environ.items()
|
|
if not key.startswith(ENV_DENY_PREFIXES)
|
|
and not any(word in key.upper() for word in ENV_DENY_WORDS)
|
|
}
|
|
|
|
|
|
def _encrypted_or_local(url: str) -> bool:
|
|
"""Whether the token and node secrets sent to *url* stay off the wire.
|
|
|
|
TLS, or this machine. An engine across a network on plain ``ws://`` hands
|
|
both to anything on the path.
|
|
"""
|
|
parts = urlsplit(url)
|
|
if parts.scheme in ("wss", "https"):
|
|
return True
|
|
host = parts.hostname or ""
|
|
if host == "localhost":
|
|
return True
|
|
try:
|
|
return ipaddress.ip_address(host).is_loopback
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def _artifacts_from(url: str) -> str:
|
|
"""The artifact endpoint beside the socket, so one URL configures both."""
|
|
base = url.replace("wss://", "https://").replace("ws://", "http://")
|
|
return base.rsplit("/workers/attach", 1)[0] + "/artifacts"
|
|
|
|
|
|
def _detect_cpus() -> int:
|
|
"""Cores this worker may use — what the batch job was given, or the box."""
|
|
given = os.environ.get("SLURM_CPUS_ON_NODE", "")
|
|
return int(given) if given.isdigit() else (os.cpu_count() or 1)
|
|
|
|
|
|
def _detect_ram_mb() -> int | None:
|
|
"""Memory in MB, or None where it cannot be asked — which is allowed."""
|
|
given = os.environ.get("SLURM_MEM_PER_NODE", "")
|
|
if given.isdigit():
|
|
return int(given)
|
|
try:
|
|
return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") // 2**20
|
|
except (AttributeError, ValueError, OSError):
|
|
return None
|
|
|
|
|
|
def _detect_gpus() -> int:
|
|
"""What this worker was *given*, else what the box appears to have.
|
|
|
|
No vendor tool is asked — that would make the one dependency two. A batch
|
|
scheduler says what the job was given in the environment and that always
|
|
wins, since a node with eight cards may have granted this job one. With
|
|
nothing said, NVIDIA's device nodes are counted, which is what the engine
|
|
does for its own machine.
|
|
"""
|
|
for name in ("SLURM_GPUS_ON_NODE", "FLUKSIO_WORKER_GPUS"):
|
|
given = os.environ.get(name, "")
|
|
if given.isdigit():
|
|
return int(given)
|
|
listed = os.environ.get("SLURM_JOB_GPUS", "")
|
|
if listed.strip():
|
|
return len([part for part in listed.split(",") if part.strip()])
|
|
return len(glob.glob("/dev/nvidia[0-9]*"))
|
|
|
|
|
|
def _inventory(args: argparse.Namespace) -> dict[str, Any]:
|
|
"""What this machine holds, for the engine to schedule against."""
|
|
return {
|
|
"cpus": args.cpus or _detect_cpus(),
|
|
"gpus": _detect_gpus() if args.gpus is None else args.gpus,
|
|
"ram_mb": args.ram_mb or _detect_ram_mb(),
|
|
}
|
|
|
|
|
|
def _remap_cuda(mask: str, requested: str) -> str:
|
|
"""Turn engine-assigned device numbers into the ones this job may see.
|
|
|
|
The engine counts GPUs from the inventory reported here, so it asks for
|
|
device 0. Under a scheduler's gres this agent already runs with something
|
|
like ``CUDA_VISIBLE_DEVICES=2,3``, where 0 is a device the job was never
|
|
given — so each index is read as a position in that mask.
|
|
"""
|
|
devices = [part.strip() for part in mask.split(",") if part.strip()]
|
|
out = []
|
|
for part in requested.split(","):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
try:
|
|
out.append(devices[int(part)])
|
|
except (ValueError, IndexError):
|
|
out.append(part)
|
|
return ",".join(out)
|
|
|
|
|
|
def _version() -> str:
|
|
"""The installed fluksio-worker, or empty when this file was copied over."""
|
|
try:
|
|
return importlib.metadata.version("fluksio-worker")
|
|
except importlib.metadata.PackageNotFoundError:
|
|
return ""
|
|
|
|
|
|
def _venv_digest(python: str) -> str:
|
|
"""What is installed here, so the engine can say when it has drifted."""
|
|
try:
|
|
listing = subprocess.run(
|
|
[python, "-m", "pip", "freeze"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=60,
|
|
check=False,
|
|
).stdout
|
|
except (OSError, subprocess.SubprocessError):
|
|
return ""
|
|
import hashlib
|
|
|
|
return hashlib.sha256(listing.encode()).hexdigest()[:16]
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
|
|
)
|
|
parser = argparse.ArgumentParser(description="Run Fluksio nodes on this machine.")
|
|
parser.add_argument("--url", required=True, help="wss://…/api/v1/workers/attach")
|
|
parser.add_argument(
|
|
"--token",
|
|
default=os.environ.get("FLUKSIO_WORKER_TOKEN", ""),
|
|
help="issued by POST /api/v1/workers/tokens",
|
|
)
|
|
parser.add_argument("--name", default=os.uname().nodename)
|
|
parser.add_argument(
|
|
"--labels",
|
|
default="",
|
|
help="comma-separated, e.g. gpu,cuda12 — what a node's worker matches",
|
|
)
|
|
parser.add_argument(
|
|
"--python",
|
|
default=sys.executable,
|
|
help=(
|
|
"the interpreter node code runs on (default: the one running this, "
|
|
"so installing into the venv with torch in it is enough)"
|
|
),
|
|
)
|
|
parser.add_argument("--parallel", type=int, default=1)
|
|
parser.add_argument("--artifact-url", default="")
|
|
parser.add_argument(
|
|
"--service-host",
|
|
default="",
|
|
help=(
|
|
"the address the engine reaches this machine's services on "
|
|
"(default: the address the engine sees this worker connect from)"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--insecure",
|
|
action="store_true",
|
|
help=(
|
|
"allow ws:// and http:// to an engine off this machine — the token "
|
|
"and every secret a node is handed then cross the network in clear"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--cpus",
|
|
type=int,
|
|
default=0,
|
|
help="cores to advertise (default: what the job or the machine has)",
|
|
)
|
|
parser.add_argument(
|
|
"--gpus",
|
|
type=int,
|
|
default=None,
|
|
help="GPUs to advertise — nothing is probed, so say so here or in the job",
|
|
)
|
|
parser.add_argument(
|
|
"--ram-mb",
|
|
type=int,
|
|
default=0,
|
|
help="memory in MB to advertise (default: what the job or the machine has)",
|
|
)
|
|
parser.add_argument(
|
|
"--max-idle",
|
|
type=float,
|
|
default=0.0,
|
|
help=(
|
|
"stop after this many seconds with nothing running — for a worker "
|
|
"a batch scheduler started for one job (default: never)"
|
|
),
|
|
)
|
|
args = parser.parse_args(argv)
|
|
|
|
if not args.token:
|
|
parser.error("a token is required (--token or FLUKSIO_WORKER_TOKEN)")
|
|
for url in (args.url, args.artifact_url):
|
|
if url and not args.insecure and not _encrypted_or_local(url):
|
|
parser.error(
|
|
f"{url} is not encrypted: use wss:// (https:// for --artifact-url), "
|
|
"or --insecure on a network you trust"
|
|
)
|
|
args.labels = [part.strip() for part in args.labels.split(",") if part.strip()]
|
|
if not WORKER_MAIN.exists():
|
|
parser.error(f"{WORKER_MAIN} is missing — copy it beside this file")
|
|
|
|
agent = Agent(args)
|
|
try:
|
|
asyncio.run(agent.serve_forever())
|
|
except (KeyboardInterrupt, asyncio.CancelledError):
|
|
log.info("stopping")
|
|
except Refused as exc:
|
|
log.error("the engine refused this worker: %s", exc)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|