app/backend/fluksio/api/routes/workers.py

"""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 `device` 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
    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")),
    )
    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)