stroblme
ae9a9792c7
It expects worker_main.py beside it, which is true on a worker host and not in a clone — so trying a remote worker started with copying files around. It now looks one directory over as well.
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())
|