Bind a worker's name and labels into its token; list and revoke tokens by jti

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013GAeg4Jo6xjZzjwy9QDWMQ
This commit is contained in:
2026-10-07 08:42:14 +02:00
co-authored by Claude Opus 5.5
parent 71b0f986e9
commit 21537e2eb6
16 changed files with 470 additions and 38 deletions
@@ -0,0 +1,37 @@
"""worker_token
Every worker credential a superuser mints, by the `jti` it carries, so one can
be listed and revoked without rotating the key that signs them all.
Revision ID: b7d2e9c4a1f3
Revises: f8c3e5a91d27
Create Date: 2026-10-06
"""
import sqlalchemy as sa
import sqlmodel.sql.sqltypes
from alembic import op
# revision identifiers, used by Alembic.
revision = "b7d2e9c4a1f3"
down_revision = "f8c3e5a91d27"
branch_labels = None
depends_on = None
def upgrade():
op.create_table(
"worker_token",
sa.Column("jti", sqlmodel.sql.sqltypes.AutoString(length=32), nullable=False),
sa.Column("name", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=False),
sa.Column("labels", sa.JSON(), nullable=True),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
sa.Column("revoked_at", sa.DateTime(timezone=True), nullable=True),
sa.PrimaryKeyConstraint("jti"),
)
def downgrade():
op.drop_table("worker_token")
+20 -1
View File
@@ -20,7 +20,7 @@ from fluksio.flow.artifacts import is_reference
from fluksio.flow.controller import FlowController
from fluksio.flow.dashboards import DashboardDef, DashboardStore
from fluksio.flow.workers import PythonWorkerPool
from fluksio.models import OAuthClient, TokenPayload, User
from fluksio.models import OAuthClient, TokenPayload, User, WorkerToken
reusable_oauth2 = OAuth2PasswordBearer(
tokenUrl=f"{settings.API_V1_STR}/login/access-token"
@@ -300,6 +300,25 @@ def oauth_client_lives(client_id: str) -> bool:
return session.get(OAuthClient, key) is not None
def worker_claims(token: str) -> dict[str, Any]:
"""A worker credential's claims, if it is still good. Blocking.
Raises ``InvalidTokenError`` for one that does not verify, one minted
before tokens carried a `jti` and their labels, and one a superuser has
revoked. A token whose `jti` has no row is one the engine minted for
itself, a provisioned machine's, and is good until it expires.
"""
claims = security.decode_worker_token(token)
jti = claims.get("jti")
if not jti:
raise InvalidTokenError("this worker token predates revocation; mint a new one")
with Session(engine) as session:
row = session.get(WorkerToken, str(jti))
if row is not None and row.revoked_at is not None:
raise InvalidTokenError("this worker token has been revoked")
return claims
def _gate_agent(payload: dict[str, Any]) -> dict[str, Any]:
"""Refuse a token whose agent has been revoked."""
if not oauth_client_lives(str(payload.get("client_id") or "")):
+2 -3
View File
@@ -17,8 +17,7 @@ from pydantic import BaseModel
from sqlmodel import Session
from starlette.concurrency import run_in_threadpool
from fluksio.api.deps import user_from_token
from fluksio.core import security
from fluksio.api.deps import user_from_token, worker_claims
from fluksio.core.config import settings
from fluksio.core.db import engine
from fluksio.flow.artifacts import ArtifactStore
@@ -42,7 +41,7 @@ def artifact_caller(request: Request) -> str:
if not token:
raise HTTPException(status_code=401, detail="Not authenticated")
try:
claims = security.decode_worker_token(token)
claims = worker_claims(token)
except InvalidTokenError:
pass
else:
+111 -15
View File
@@ -2,24 +2,39 @@
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, says what it can do, and then answers calls on the socket
it opened.
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
from datetime import timedelta
import uuid
from datetime import UTC, datetime, timedelta
from typing import Any
from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
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 get_current_active_superuser, get_current_user
from fluksio.api.deps import (
SessionDep,
get_current_active_superuser,
get_current_user,
worker_claims,
)
from fluksio.core import security
from fluksio.flow.remote import (
PROTOCOL,
@@ -27,6 +42,7 @@ from fluksio.flow.remote import (
RemoteWorkerHub,
WorkerInventory,
)
from fluksio.models import WorkerToken
logger = logging.getLogger(__name__)
@@ -35,6 +51,9 @@ 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]] = {}
class WorkerInfo(BaseModel):
name: str
@@ -51,12 +70,16 @@ class WorkerInfo(BaseModel):
class TokenRequest(BaseModel):
name: str
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
@@ -145,14 +168,58 @@ def read_resources(request: Request) -> Any:
response_model=TokenIssued,
dependencies=[Depends(get_current_active_superuser)],
)
def issue_token(body: TokenRequest) -> Any:
def issue_token(session: SessionDep, body: TokenRequest) -> Any:
"""Mint the credential a worker presents when it dials in.
Shown once. It is signed with the same keypair the agent tokens use, so
rotating that key revokes every worker along with them.
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.
"""
token = security.create_worker_token(body.name, timedelta(days=TOKEN_DAYS))
return TokenIssued(name=body.name, token=token)
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(
@@ -177,14 +244,29 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
does: a websocket handshake carries no headers of its own.
"""
try:
claims = security.decode_worker_token(token)
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
@@ -196,13 +278,27 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
await websocket.close(code=1002)
return
# The token names the worker; what it calls itself is a suggestion, so two
# hosts cannot fight over one identity by claiming the same name.
# 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=[str(label) for label in (hello.get("labels") or [])],
labels=labels,
send=websocket.send_json,
loop=asyncio.get_running_loop(),
max_parallel=max(1, int(hello.get("max_parallel") or 1)),
+13 -5
View File
@@ -132,13 +132,19 @@ def create_oauth_access_token(
)
def create_worker_token(name: str, expires_delta: timedelta) -> str:
def create_worker_token(
name: str,
expires_delta: timedelta,
labels: list[str] | None = None,
jti: str | None = None,
) -> str:
"""A credential a remote worker presents when it dials in.
Signed with the same keypair the agent tokens use, so the whole set can be
revoked by rotating one key, and told apart from them by its audience: a
worker's token grants no API access, and an agent's opens no worker
connection.
Signed with the same keypair the agent tokens use, and told apart from them
by its audience: a worker's token grants no API access, and an agent's
opens no worker connection. It names the worker and the labels it may
carry, so a worker cannot claim work meant for another machine, and its
`jti` is what revoking one token rather than the whole key goes by.
"""
now = datetime.now(UTC)
payload = {
@@ -148,6 +154,8 @@ def create_worker_token(name: str, expires_delta: timedelta) -> str:
"iat": now,
"exp": now + expires_delta,
"scope": WORKER_SCOPE,
"labels": sorted(set(labels or [])),
"jti": jti or uuid.uuid4().hex,
}
return jwt.encode(
payload,
+3 -1
View File
@@ -199,7 +199,9 @@ class SlurmProvisioner:
def _submit(self, profile: SlurmProfile) -> None:
worker = f"{self.name}-{profile.name}-{uuid.uuid4().hex[:8]}"
token = security.create_worker_token(worker, timedelta(hours=TOKEN_HOURS))
token = security.create_worker_token(
worker, timedelta(hours=TOKEN_HOURS), labels=profile.labels
)
try:
result = self._ssh(["sbatch", "--parsable"], self._payload(profile, token))
except (OSError, subprocess.SubprocessError) as exc:
+22
View File
@@ -566,3 +566,25 @@ class ConnectorState(SQLModel, table=True):
key: str = Field(primary_key=True, max_length=200)
value: str
updated_at: datetime = Field(default_factory=get_datetime_utc, sa_type=UTCDateTime)
# -----------------------------------------------------------------------------
# Worker tokens
#
# What a superuser minted for a remote worker, so it can be listed and revoked.
# The token itself is not stored: it carries its own name and labels, and this
# row is how one is refused before it expires.
# -----------------------------------------------------------------------------
class WorkerToken(SQLModel, table=True):
"""One worker credential, by the id it carries as `jti`."""
__tablename__ = "worker_token"
jti: str = Field(primary_key=True, max_length=32)
name: str = Field(max_length=255)
labels: list[str] = Field(sa_column=Column(JSON), default_factory=list)
created_at: datetime = Field(default_factory=get_datetime_utc, sa_type=UTCDateTime)
expires_at: datetime = Field(sa_type=UTCDateTime)
revoked_at: datetime | None = Field(default=None, sa_type=UTCDateTime)
+99
View File
@@ -0,0 +1,99 @@
"""A worker is what its token says, and a superuser can take the token back."""
from datetime import timedelta
import jwt
import pytest
from fastapi.testclient import TestClient
from starlette.websockets import WebSocketDisconnect
from fluksio.core import security
from fluksio.core.config import settings
from fluksio.flow.remote import PROTOCOL
PREFIX = f"{settings.API_V1_STR}/workers"
MISSING = "sha256:" + "0" * 64
def mint(client: TestClient, headers: dict[str, str]) -> dict[str, str]:
answer = client.post(
f"{PREFIX}/tokens", headers=headers, json={"name": "gpu-box", "labels": ["gpu"]}
)
assert answer.status_code == 200
return answer.json()
def hello(*labels: str) -> dict[str, object]:
return {"op": "hello", "protocol": PROTOCOL, "name": "elsewhere", "labels": labels}
def test_the_token_names_the_worker_and_its_labels(
client: TestClient, superuser_token_headers: dict[str, str]
) -> None:
issued = mint(client, superuser_token_headers)
with client.websocket_connect(f"{PREFIX}/attach?token={issued['token']}") as ws:
ws.send_json(hello())
assert ws.receive_json()["name"] == "gpu-box"
attached = client.get(PREFIX, headers=superuser_token_headers).json()
assert [(w["name"], w["labels"]) for w in attached] == [("gpu-box", ["gpu"])]
def test_a_worker_cannot_claim_a_label_its_token_lacks(
client: TestClient, superuser_token_headers: dict[str, str]
) -> None:
issued = mint(client, superuser_token_headers)
with client.websocket_connect(f"{PREFIX}/attach?token={issued['token']}") as ws:
ws.send_json(hello("gpu", "vault"))
assert ws.receive_json()["op"] == "refused"
def test_a_revoked_token_is_dropped_and_refused(
client: TestClient,
superuser_token_headers: dict[str, str],
normal_user_token_headers: dict[str, str],
) -> None:
issued = mint(client, superuser_token_headers)
url = f"{PREFIX}/attach?token={issued['token']}"
bearer = {"Authorization": f"Bearer {issued['token']}"}
assert (
client.get(
f"{settings.API_V1_STR}/artifacts/{MISSING}", headers=bearer
).status_code
== 404
)
revoke = f"{PREFIX}/tokens/{issued['jti']}/revoke"
assert client.post(revoke, headers=normal_user_token_headers).status_code == 403
with client.websocket_connect(url) as ws:
ws.send_json(hello("gpu"))
ws.receive_json()
assert client.post(revoke, headers=superuser_token_headers).json()["revoked_at"]
with pytest.raises(WebSocketDisconnect):
ws.receive_json()
with pytest.raises(WebSocketDisconnect), client.websocket_connect(url):
pass
assert (
client.get(
f"{settings.API_V1_STR}/artifacts/{MISSING}", headers=bearer
).status_code
== 401
)
listed = client.get(f"{PREFIX}/tokens", headers=superuser_token_headers).json()
assert [row for row in listed if row["jti"] == issued["jti"]][0]["revoked_at"]
assert all("token" not in row for row in listed)
def test_a_token_minted_before_revocation_is_refused(client: TestClient) -> None:
claims = security.decode_worker_token(
security.create_worker_token("old", timedelta(days=1))
)
del claims["jti"], claims["labels"]
legacy = jwt.encode(
claims, security.oauth_key(), algorithm=security.OAUTH_ALGORITHM
)
with (
pytest.raises(WebSocketDisconnect),
client.websocket_connect(f"{PREFIX}/attach?token={legacy}"),
):
pass
+3 -1
View File
@@ -229,7 +229,9 @@ the process is dead.
| Method | Path | What |
|---|---|---|
| `GET` | `/workers` | what is attached, its labels and how busy it is |
| `POST` | `/workers/tokens` | mint a worker credential (superuser; shown once) |
| `POST` | `/workers/tokens` | mint a worker credential with its name and labels (superuser; shown once) |
| `GET` | `/workers/tokens` | the worker credentials minted, without the tokens (superuser) |
| `POST` | `/workers/tokens/{jti}/revoke` | refuse one from now on and drop a worker attached with it (superuser) |
| `GET` | `/workers/runtime` | the node runner's source, for a host without pip |
| `WS` | `/workers/attach` | where a worker dials in |
+23 -5
View File
@@ -30,8 +30,8 @@ driver in order to run a training step. An engine host already has it, and
|---|---|---|
| `--url` | **required** | `wss://…/api/v1/workers/attach` |
| `--token` | `$FLUKSIO_WORKER_TOKEN` | the credential, minted on the engine |
| `--name` | this host's name | how it shows up in the worker list |
| `--labels` | none | comma-separated; what a node's `device` matches |
| `--name` | this host's name | sent along; the worker list shows the token's name |
| `--labels` | the token's | comma-separated; must match the labels the token was minted with |
| `--python` | this interpreter | the interpreter node code runs on |
| `--parallel` | `1` | how many node calls it will take at once |
| `--artifact-url` | derived from `--url` | where the artifact store is, if not beside the socket |
@@ -84,12 +84,26 @@ On the engine, as a superuser:
```sh
curl -X POST $FLUKSIO/workers/tokens -H "Authorization: Bearer $TOKEN" \
-H 'Content-Type: application/json' -d '{"name": "gpu-dev"}'
-H 'Content-Type: application/json' \
-d '{"name": "gpu-dev", "labels": ["gpu", "cuda12"]}'
```
Shown once, valid for a year, since a worker is a machine somebody sets up and
leaves running. It is signed with the same keypair agent tokens use, so
rotating that key revokes every worker along with them.
leaves running. The name and labels are signed into it: labels decide which
nodes, with their source and secrets, are sent to a worker, so a worker cannot
choose its own. It shows up under the token's name, and a worker whose
`--labels` differ from the token's is refused.
```sh
curl -s $FLUKSIO/workers/tokens -H "Authorization: Bearer $TOKEN" | jq
curl -X POST $FLUKSIO/workers/tokens/$JTI/revoke -H "Authorization: Bearer $TOKEN"
```
The first lists what was minted, by `jti` (the tokens themselves are not
kept); the second refuses one from then on, here and at the artifact
endpoint, and drops a worker attached with it. Rotating the keypair agent
tokens are signed with still revokes every worker at once. A token minted
before tokens carried their labels is refused: mint a new one.
??? note "A host where pip is not an option"
@@ -181,6 +195,10 @@ will never accept it; `pip install -U fluksio-worker` on that host is the whole
upgrade. Nothing changed in the runner served at `GET /api/v1/workers/runtime`,
so a host that copies its two files copies the same one as before.
A token minted before tokens carried their labels and a `jti` is refused. Mint
each existing worker host a new one with `POST /workers/tokens`, naming the
labels it should carry, and swap it in for `FLUKSIO_WORKER_TOKEN`.
## Machines from a batch scheduler
A cluster is not a machine that attaches and stays; it is a queue somebody else
+1 -1
View File
@@ -264,7 +264,7 @@ nothing should expose Redis across hosts. Install it on the box, mint it a
token, and start it:
```bash
curl -X POST $FLUKSIO/workers/tokens -d '{"name": "gpu-dev"}' # once, as an admin
curl -X POST $FLUKSIO/workers/tokens -d '{"name": "gpu-dev", "labels": ["gpu", "cuda12"]}' # once, as an admin
pip install fluksio-worker
+4 -1
View File
@@ -252,9 +252,12 @@ On the engine, mint a token:
```sh
curl -X POST https://api.${DOMAIN}/api/v1/workers/tokens \
-H "Authorization: Bearer $TOKEN" -d '{"name": "shed-pi"}'
-H "Authorization: Bearer $TOKEN" -d '{"name": "shed-pi", "labels": ["shed", "gpio"]}'
```
The labels are signed into it, so the worker can repeat them with `--labels`
but not change them.
On the other machine:
```sh
+69 -1
View File
@@ -3860,6 +3860,17 @@ export const TokenIssuedSchema = {
type: 'string',
title: 'Token'
},
jti: {
type: 'string',
title: 'Jti'
},
labels: {
items: {
type: 'string'
},
type: 'array',
title: 'Labels'
},
expires_days: {
type: 'integer',
title: 'Expires Days',
@@ -3867,7 +3878,7 @@ export const TokenIssuedSchema = {
}
},
type: 'object',
required: ['name', 'token'],
required: ['name', 'token', 'jti'],
title: 'TokenIssued'
} as const;
@@ -3875,7 +3886,16 @@ export const TokenRequestSchema = {
properties: {
name: {
type: 'string',
maxLength: 255,
minLength: 1,
title: 'Name'
},
labels: {
items: {
type: 'string'
},
type: 'array',
title: 'Labels'
}
},
type: 'object',
@@ -4520,6 +4540,54 @@ export const WorkerInfoSchema = {
title: 'WorkerInfo'
} as const;
export const WorkerTokenSchema = {
properties: {
jti: {
type: 'string',
maxLength: 32,
title: 'Jti'
},
name: {
type: 'string',
maxLength: 255,
title: 'Name'
},
labels: {
items: {
type: 'string'
},
type: 'array',
title: 'Labels'
},
created_at: {
type: 'string',
format: 'date-time',
title: 'Created At'
},
expires_at: {
type: 'string',
format: 'date-time',
title: 'Expires At'
},
revoked_at: {
anyOf: [
{
type: 'string',
format: 'date-time'
},
{
type: 'null'
}
],
title: 'Revoked At'
}
},
type: 'object',
required: ['jti', 'name', 'expires_at'],
title: 'WorkerToken',
description: 'One worker credential, by the id it carries as `jti`.'
} as const;
export const fluksio__api__routes__dashboards__PublishRequestSchema = {
properties: {
version: {
File diff suppressed because one or more lines are too long
+23
View File
@@ -1291,11 +1291,14 @@ export type Token = {
export type TokenIssued = {
name: string;
token: string;
jti: string;
labels?: Array<(string)>;
expires_days?: number;
};
export type TokenRequest = {
name: string;
labels?: Array<(string)>;
};
export type TriggerRequest = {
@@ -1474,6 +1477,18 @@ export type WorkerInfo = {
ram_mb?: (number | null);
};
/**
* One worker credential, by the id it carries as `jti`.
*/
export type WorkerToken = {
jti: string;
name: string;
labels?: Array<(string)>;
created_at?: string;
expires_at: string;
revoked_at?: (string | null);
};
export type AiReadAiConfigResponse = (AIConfigPublic);
export type AiUpdateAiConfigData = {
@@ -2207,10 +2222,18 @@ export type WorkersReadWorkersResponse = (Array<WorkerInfo>);
export type WorkersReadResourcesResponse = (ResourcesSnapshot);
export type WorkersReadTokensResponse = (Array<WorkerToken>);
export type WorkersIssueTokenData = {
requestBody: TokenRequest;
};
export type WorkersIssueTokenResponse = (TokenIssued);
export type WorkersRevokeTokenData = {
jti: string;
};
export type WorkersRevokeTokenResponse = (WorkerToken);
export type WorkersReadRuntimeResponse = (string);
+2 -1
View File
@@ -16,7 +16,8 @@ fluksio-worker \
`--python` is the interpreter node code runs on, which is how this machine keeps
its own wheels without the engine ever installing them. Mint a token from the
engine with `POST /api/v1/workers/tokens`.
engine with `POST /api/v1/workers/tokens`; it carries this worker's name and
labels, and `--labels` may repeat those but not change them.
A plain `ws://` address is refused unless the engine is on this machine, since
the token and node secrets would cross the network in clear; `--insecure`