stroblme
34243e6902
save_artifact/load_artifact outside a node use the worker's own store code in $FLUKSIO_ARTIFACT_DIR or <tmp>/fluksio-artifacts, so a node can be called by hand. --follow prints status changes and each node's resource wait; the placer's waiting snapshot now carries the run. fluksio cancel <run>...
361 lines
12 KiB
Python
361 lines
12 KiB
Python
"""Remote workers: how one attaches, and what is attached right now.
|
|
|
|
A worker dials in rather than being dialled: the GPU box and the engine are
|
|
usually on different networks, and only one of them can be reached. It presents
|
|
a token minted here, which says who it is and which labels it carries, says
|
|
what else it can do, and then answers calls on the socket it opened.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
|
|
from fastapi import (
|
|
APIRouter,
|
|
Depends,
|
|
HTTPException,
|
|
Request,
|
|
WebSocket,
|
|
WebSocketDisconnect,
|
|
)
|
|
from fastapi.responses import PlainTextResponse
|
|
from fluksio_worker import worker_main
|
|
from jwt.exceptions import InvalidTokenError
|
|
from pydantic import BaseModel, Field
|
|
from sqlmodel import col, select
|
|
from starlette.concurrency import run_in_threadpool
|
|
|
|
from fluksio.api.deps import (
|
|
SessionDep,
|
|
get_current_active_superuser,
|
|
get_current_user,
|
|
worker_claims,
|
|
)
|
|
from fluksio.api.socket_auth import socket_token
|
|
from fluksio.core import security
|
|
from fluksio.flow.remote import (
|
|
PROTOCOL,
|
|
RemoteWorker,
|
|
RemoteWorkerHub,
|
|
WorkerInventory,
|
|
)
|
|
from fluksio.models import WorkerToken
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/workers", tags=["workers"])
|
|
|
|
#: Long, because a worker is a machine somebody set up once and left running.
|
|
TOKEN_DAYS = 365
|
|
|
|
#: The sockets attached with each token, by `jti`, so revoking one drops them.
|
|
_attached: dict[str, set[WebSocket]] = {}
|
|
|
|
#: Workers refused for an older protocol, by name, until one attaches again.
|
|
#: Such an agent exits on the refusal, so this is the one place left that says
|
|
#: which host still needs `pip install -U fluksio-worker`.
|
|
_behind: dict[str, WorkerInfo] = {}
|
|
|
|
|
|
class WorkerInfo(BaseModel):
|
|
name: str
|
|
labels: list[str] = Field(default_factory=list)
|
|
max_parallel: int = 1
|
|
in_flight: int = 0
|
|
attached_at: float = 0.0
|
|
last_seen: float = 0.0
|
|
python: str = ""
|
|
venv_digest: str = ""
|
|
cpus: int = 1
|
|
gpus: int = 0
|
|
ram_mb: int | None = None
|
|
protocol: int = 0
|
|
#: The fluksio-worker it runs; empty for an agent too old to say.
|
|
version: str = ""
|
|
#: Refused for speaking an older protocol than this engine, so not attached.
|
|
behind: bool = False
|
|
|
|
|
|
class TokenRequest(BaseModel):
|
|
name: str = Field(min_length=1, max_length=255)
|
|
#: What a node's `worker` may match on this worker; it cannot add to them.
|
|
labels: list[str] = Field(default_factory=list)
|
|
|
|
|
|
class TokenIssued(BaseModel):
|
|
name: str
|
|
token: str
|
|
jti: str
|
|
labels: list[str] = Field(default_factory=list)
|
|
expires_days: int = TOKEN_DAYS
|
|
|
|
|
|
def _hub(app: Any) -> RemoteWorkerHub:
|
|
hub: RemoteWorkerHub | None = getattr(app.state, "worker_hub", None)
|
|
if hub is None:
|
|
raise HTTPException(status_code=503, detail="Remote workers are not available")
|
|
return hub
|
|
|
|
|
|
@router.get(
|
|
"", response_model=list[WorkerInfo], dependencies=[Depends(get_current_user)]
|
|
)
|
|
def read_workers(request: Request) -> Any:
|
|
"""What is attached, and how busy it is."""
|
|
return [
|
|
WorkerInfo(
|
|
name=worker.name,
|
|
labels=sorted(worker.labels),
|
|
max_parallel=worker.max_parallel,
|
|
in_flight=worker.in_flight,
|
|
attached_at=worker.attached_at,
|
|
last_seen=worker.last_seen,
|
|
python=str(worker.info.get("python") or ""),
|
|
venv_digest=str(worker.info.get("venv_digest") or ""),
|
|
cpus=worker.inventory.cpus,
|
|
gpus=worker.inventory.gpus,
|
|
ram_mb=worker.inventory.ram_mb,
|
|
protocol=int(worker.info.get("protocol") or 0),
|
|
version=str(worker.info.get("version") or ""),
|
|
)
|
|
for worker in _hub(request.app).workers()
|
|
] + list(_behind.values())
|
|
|
|
|
|
class ResourceLevel(BaseModel):
|
|
total: int
|
|
free: int
|
|
|
|
|
|
class WaitingNode(BaseModel):
|
|
node: str
|
|
#: The run it belongs to, so a client can tell its own wait from another's.
|
|
run: str = ""
|
|
reason: str
|
|
seconds: float
|
|
|
|
|
|
class TargetResources(BaseModel):
|
|
"""One machine: this engine, or a worker attached to it."""
|
|
|
|
target: str
|
|
cpus: ResourceLevel
|
|
gpus: ResourceLevel
|
|
#: Absent where the machine did not say how much memory it has.
|
|
ram_mb: ResourceLevel | None = None
|
|
labels: list[str] = Field(default_factory=list)
|
|
in_flight: int = 0
|
|
|
|
|
|
class ResourcesSnapshot(BaseModel):
|
|
#: This engine's own figures, kept where they have always been.
|
|
cpus: ResourceLevel
|
|
gpus: ResourceLevel
|
|
waiting: list[WaitingNode] = Field(default_factory=list)
|
|
targets: list[TargetResources] = Field(default_factory=list)
|
|
provisioners: list[dict[str, Any]] = Field(default_factory=list)
|
|
|
|
|
|
@router.get(
|
|
"/resources",
|
|
response_model=ResourcesSnapshot,
|
|
dependencies=[Depends(get_current_user)],
|
|
)
|
|
def read_resources(request: Request) -> Any:
|
|
"""Every machine, what is free of it, and which nodes are queued.
|
|
|
|
A node waiting its turn looks exactly like a node that has hung — the run
|
|
sits at `running` and says nothing — so what is waiting, and for what, has
|
|
to be readable somewhere.
|
|
"""
|
|
placer = getattr(request.app.state, "placer", None)
|
|
if placer is None:
|
|
raise HTTPException(status_code=503, detail="Resources are not accounted here")
|
|
return placer.snapshot()
|
|
|
|
|
|
@router.post(
|
|
"/tokens",
|
|
response_model=TokenIssued,
|
|
dependencies=[Depends(get_current_active_superuser)],
|
|
)
|
|
def issue_token(session: SessionDep, body: TokenRequest) -> Any:
|
|
"""Mint the credential a worker presents when it dials in.
|
|
|
|
Shown once. The worker's name and labels are signed into it, so a worker is
|
|
what it was minted as whatever it says about itself, and only its `jti` is
|
|
kept here, to revoke it by.
|
|
"""
|
|
labels = sorted({label.strip() for label in body.labels if label.strip()})
|
|
row = WorkerToken(
|
|
jti=uuid.uuid4().hex,
|
|
name=body.name,
|
|
labels=labels,
|
|
expires_at=datetime.now(UTC) + timedelta(days=TOKEN_DAYS),
|
|
)
|
|
token = security.create_worker_token(
|
|
body.name, timedelta(days=TOKEN_DAYS), labels=labels, jti=row.jti
|
|
)
|
|
session.add(row)
|
|
session.commit()
|
|
return TokenIssued(name=body.name, token=token, jti=row.jti, labels=labels)
|
|
|
|
|
|
@router.get(
|
|
"/tokens",
|
|
response_model=list[WorkerToken],
|
|
dependencies=[Depends(get_current_active_superuser)],
|
|
)
|
|
def read_tokens(session: SessionDep) -> Any:
|
|
"""Every worker token minted here, newest first, without the tokens."""
|
|
return session.exec(
|
|
select(WorkerToken).order_by(col(WorkerToken.created_at).desc())
|
|
).all()
|
|
|
|
|
|
@router.post(
|
|
"/tokens/{jti}/revoke",
|
|
response_model=WorkerToken,
|
|
dependencies=[Depends(get_current_active_superuser)],
|
|
)
|
|
async def revoke_token(session: SessionDep, jti: str) -> Any:
|
|
"""Refuse this token from now on, and drop a worker attached with it."""
|
|
row = session.get(WorkerToken, jti)
|
|
if row is None:
|
|
raise HTTPException(status_code=404, detail="No such worker token")
|
|
if row.revoked_at is None:
|
|
row.revoked_at = datetime.now(UTC)
|
|
session.add(row)
|
|
session.commit()
|
|
session.refresh(row)
|
|
for socket in list(_attached.get(jti, ())):
|
|
await socket.close(code=1008)
|
|
return row
|
|
|
|
|
|
@router.get(
|
|
"/runtime",
|
|
response_class=PlainTextResponse,
|
|
dependencies=[Depends(get_current_user)],
|
|
)
|
|
def read_runtime() -> str:
|
|
"""The worker's own code, so a fresh host installs by fetching one file.
|
|
|
|
It is the same module the engine's local workers run — deliberately
|
|
standard library only, and with nothing of the engine importable in it.
|
|
"""
|
|
return worker_main.__file__ and open(worker_main.__file__).read()
|
|
|
|
|
|
@router.websocket("/attach")
|
|
async def attach(websocket: WebSocket, token: str = "") -> None:
|
|
"""A worker's connection, for as long as it holds.
|
|
|
|
The token comes as `Authorization: Bearer` on the handshake. `?token=` is
|
|
still read for one release, so the engine can be upgraded before its
|
|
workers.
|
|
"""
|
|
token = socket_token(websocket.headers, token=token)
|
|
try:
|
|
claims = await run_in_threadpool(worker_claims, token)
|
|
except InvalidTokenError:
|
|
await websocket.close(code=1008)
|
|
return
|
|
|
|
jti = str(claims["jti"])
|
|
await websocket.accept()
|
|
_attached.setdefault(jti, set()).add(websocket)
|
|
try:
|
|
await _attached_session(websocket, claims)
|
|
finally:
|
|
sockets = _attached.get(jti, set())
|
|
sockets.discard(websocket)
|
|
if not sockets:
|
|
_attached.pop(jti, None)
|
|
|
|
|
|
async def _attached_session(websocket: WebSocket, claims: dict[str, Any]) -> None:
|
|
"""The hello, then calls, until the socket goes or the token is revoked."""
|
|
try:
|
|
hello = await asyncio.wait_for(websocket.receive_json(), timeout=30)
|
|
except WebSocketDisconnect:
|
|
return
|
|
except (TimeoutError, ValueError):
|
|
await websocket.close(code=1002)
|
|
return
|
|
|
|
protocol = int(hello.get("protocol", 0))
|
|
if hello.get("op") != "hello" or protocol != PROTOCOL:
|
|
if hello.get("op") == "hello" and protocol < PROTOCOL:
|
|
name = str(claims.get("sub") or hello.get("name") or "worker")
|
|
_behind[name] = WorkerInfo(
|
|
name=name,
|
|
labels=sorted(str(label) for label in claims.get("labels") or []),
|
|
last_seen=time.time(),
|
|
python=str(hello.get("python") or ""),
|
|
protocol=protocol,
|
|
version=str(hello.get("version") or ""),
|
|
behind=True,
|
|
)
|
|
await websocket.send_json(
|
|
{"op": "refused", "reason": f"this engine speaks protocol {PROTOCOL}"}
|
|
)
|
|
await websocket.close(code=1002)
|
|
return
|
|
|
|
# The token names the worker and its labels. What it calls itself is a
|
|
# suggestion, so two hosts cannot fight over one identity by claiming the
|
|
# same name; labels it may repeat but not change, since they decide which
|
|
# node code and secrets come its way.
|
|
labels = sorted({str(label) for label in claims.get("labels") or []})
|
|
claimed = sorted({str(label) for label in hello.get("labels") or []})
|
|
if claimed and claimed != labels:
|
|
await websocket.send_json(
|
|
{
|
|
"op": "refused",
|
|
"reason": f"this token carries the labels {labels}, not {claimed}",
|
|
}
|
|
)
|
|
await websocket.close(code=1008)
|
|
return
|
|
|
|
name = str(claims.get("sub") or hello.get("name") or "worker")
|
|
hub = _hub(websocket.app)
|
|
worker = RemoteWorker(
|
|
name=name,
|
|
labels=labels,
|
|
send=websocket.send_json,
|
|
loop=asyncio.get_running_loop(),
|
|
max_parallel=max(1, int(hello.get("max_parallel") or 1)),
|
|
info={
|
|
"python": hello.get("python"),
|
|
"venv_digest": hello.get("venv_digest"),
|
|
"protocol": protocol,
|
|
"version": hello.get("version"),
|
|
},
|
|
inventory=WorkerInventory.from_hello(hello.get("inventory")),
|
|
host=str(
|
|
hello.get("service_host")
|
|
or (websocket.client.host if websocket.client else "")
|
|
),
|
|
)
|
|
hub.attach(worker)
|
|
_behind.pop(name, None)
|
|
await websocket.send_json({"op": "welcome", "protocol": PROTOCOL, "name": name})
|
|
|
|
try:
|
|
while True:
|
|
message = await websocket.receive_json()
|
|
worker.deliver(message)
|
|
except Exception:
|
|
# Any way this ends is the same thing: the socket is gone, and whatever
|
|
# was waiting on it has to be told rather than left hanging.
|
|
logger.info("Worker '%s' disconnected", name)
|
|
finally:
|
|
hub.detach(name)
|