stroblmeandClaude Opus 5
60d7ec81c0
A wheel whose top-level module is `app` collides with anything else in a user's venv, so the package that is about to be published takes the name it is published under. Only the Python package moves; the repo, the Docker WORKDIR and the compose project keep theirs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
309 lines
11 KiB
Python
309 lines
11 KiB
Python
"""The agent that runs Fluksio nodes on a machine the engine cannot reach.
|
|
|
|
Copy this file and ``worker_main.py`` onto the box with the GPU, point it at
|
|
the engine, and it dials in::
|
|
|
|
pip install websockets
|
|
python fluksio_worker.py --url wss://api.example.com/api/v1/workers/attach \\
|
|
--token "$FLUKSIO_WORKER_TOKEN" --labels gpu --python /opt/venv/bin/python
|
|
|
|
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.
|
|
|
|
Deliberately one file with one dependency. Nothing of the app is imported
|
|
here; a worker host installs Python, ``websockets``, and whatever the nodes
|
|
themselves need.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
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
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
log = logging.getLogger("fluksio-worker")
|
|
|
|
PROTOCOL = 1
|
|
#: Sent while a call is running, so the engine can tell working from wedged.
|
|
HEARTBEAT_S = 10.0
|
|
#: Reconnection backs off to this and no further.
|
|
MAX_BACKOFF_S = 30.0
|
|
|
|
|
|
def _worker_main() -> Path:
|
|
"""The user-code runner: beside this file once deployed, or in a checkout.
|
|
|
|
On a worker host the two files sit together, which is what the install
|
|
instructions say. Run straight from a clone and the runner is one directory
|
|
over, in the engine's own package — worth finding, so trying this out does
|
|
not start with copying files around.
|
|
"""
|
|
here = Path(__file__).resolve().parent
|
|
for candidate in (here / "worker_main.py", here.parent / "flow" / "worker_main.py"):
|
|
if candidate.exists():
|
|
return candidate
|
|
return here / "worker_main.py"
|
|
|
|
|
|
WORKER_MAIN = _worker_main()
|
|
|
|
|
|
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,
|
|
)
|
|
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:
|
|
with contextlib.suppress(OSError):
|
|
self.proc.send_signal(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 = dict(os.environ)
|
|
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] = {}
|
|
|
|
async def serve_forever(self) -> None:
|
|
backoff = 1.0
|
|
while True:
|
|
try:
|
|
await self._session()
|
|
backoff = 1.0
|
|
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:
|
|
url = f"{self.args.url}?token={self.args.token}"
|
|
async with websockets.connect(url, max_size=None, ping_interval=20) as socket:
|
|
await socket.send(
|
|
json.dumps(
|
|
{
|
|
"op": "hello",
|
|
"protocol": PROTOCOL,
|
|
"name": self.args.name,
|
|
"labels": self.args.labels,
|
|
"python": self.args.python,
|
|
"max_parallel": self.args.parallel,
|
|
"venv_digest": _venv_digest(self.args.python),
|
|
}
|
|
)
|
|
)
|
|
welcome = json.loads(await socket.recv())
|
|
if welcome.get("op") != "welcome":
|
|
raise RuntimeError(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,
|
|
)
|
|
|
|
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"))
|
|
task = asyncio.create_task(self._run(socket, message))
|
|
# 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)
|
|
elif op == "cancel":
|
|
self._cancel(str(message.get("call_id") or ""))
|
|
|
|
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()
|
|
worker = Subprocess(self.args.python, self.env)
|
|
self.running[call_id] = worker
|
|
|
|
async def beat() -> None:
|
|
while True:
|
|
await asyncio.sleep(HEARTBEAT_S)
|
|
with contextlib.suppress(Exception):
|
|
await socket.send(
|
|
json.dumps({"call_id": call_id, "event": "heartbeat"})
|
|
)
|
|
|
|
heartbeat = asyncio.create_task(beat())
|
|
try:
|
|
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)
|
|
worker.kill()
|
|
|
|
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()
|
|
|
|
|
|
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 _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 _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() -> int:
|
|
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 device matches",
|
|
)
|
|
parser.add_argument(
|
|
"--python",
|
|
default=sys.executable,
|
|
help="the interpreter node code runs on; point it at the venv with torch",
|
|
)
|
|
parser.add_argument("--parallel", type=int, default=1)
|
|
parser.add_argument("--artifact-url", default="")
|
|
args = parser.parse_args()
|
|
|
|
if not args.token:
|
|
parser.error("a token is required (--token or FLUKSIO_WORKER_TOKEN)")
|
|
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:
|
|
log.info("stopping")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|