stroblmeandClaude Fable 5
eb2d098d7c
The engine runs where the automations are and the GPU is somewhere else, usually behind a different network — so the worker connects out and the engine answers over the socket it was given. Nothing has to expose Redis, and the same connection works through the tunnel the hosted access will use. What travels is the protocol the local pool already speaks, so a node cannot tell which kind of worker it is on. A node declares device: gpu and device_policy, the label is resolved per call (a worker attaching later needs no rebuild), and a run whose labels nothing carries waits in the queue saying what it waits for rather than failing — submit from the couch, the GPU box picks it up when it is switched on. Two things had to move with it. Compiling now happens on the machine that will run the node: a node importing torch is correct on the GPU box and a missing module on the engine, so checking it here failed nodes that were fine. And the artifact endpoint accepts a worker's own credential, because storing a checkpoint is exactly what that credential is for — and only that. Verified against the real split: the training ran on this host (its checkpoint names the machine and a numpy the engine does not have), streamed 40 metric points back mid-run, and the evaluate node read the checkpoint on the engine. Cancel kills the remote training; pulling the worker fails the run in six seconds instead of waiting out its ten-minute timeout. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01AD8SfVhzXBG2nAfFcVh3iD
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
"""Artifacts over HTTP: the one way bytes get in and out of the store.
|
|
|
|
A node on this host could reach the directory itself, but a node on a remote
|
|
worker cannot — and having one path rather than two is what keeps a flow's
|
|
code the same wherever it runs.
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
|
from fastapi.responses import StreamingResponse
|
|
from jwt.exceptions import InvalidTokenError
|
|
from pydantic import BaseModel
|
|
from sqlmodel import Session
|
|
|
|
from app.api.deps import user_from_token
|
|
from app.core import security
|
|
from app.core.db import engine
|
|
from app.flow.artifacts import ArtifactStore
|
|
|
|
|
|
def artifact_caller(request: Request) -> str:
|
|
"""Who may move artifacts: a signed-in person, or an attached worker.
|
|
|
|
A worker's node stores its checkpoints through this endpoint, so its own
|
|
credential has to open it — and only it. The token is no use anywhere else
|
|
in the API, which is why this check is here rather than in the shared
|
|
dependency every other route uses.
|
|
"""
|
|
header = request.headers.get("Authorization", "")
|
|
token = header[7:] if header.lower().startswith("bearer ") else ""
|
|
if not token:
|
|
raise HTTPException(status_code=401, detail="Not authenticated")
|
|
try:
|
|
claims = security.decode_worker_token(token)
|
|
except InvalidTokenError:
|
|
pass
|
|
else:
|
|
return f"worker:{claims.get('sub')}"
|
|
with Session(engine) as session:
|
|
user = user_from_token(session, token)
|
|
if user is None:
|
|
raise HTTPException(status_code=401, detail="Not authenticated")
|
|
return user.email
|
|
|
|
|
|
router = APIRouter(
|
|
prefix="/artifacts", tags=["artifacts"], dependencies=[Depends(artifact_caller)]
|
|
)
|
|
|
|
|
|
class ArtifactRef(BaseModel):
|
|
digest: str
|
|
size: int
|
|
media_type: str
|
|
name: str = ""
|
|
|
|
|
|
def _store(request: Request) -> ArtifactStore:
|
|
store: ArtifactStore | None = getattr(request.app.state, "artifact_store", None)
|
|
if store is None:
|
|
raise HTTPException(status_code=503, detail="The artifact store is not ready")
|
|
return store
|
|
|
|
|
|
@router.put("", response_model=ArtifactRef)
|
|
async def put_artifact(
|
|
request: Request,
|
|
name: str = Query(default=""),
|
|
media_type: str = Query(default=""),
|
|
) -> Any:
|
|
"""Store the request body and answer with the reference to it."""
|
|
store = _store(request)
|
|
body = await request.body()
|
|
return store.put([body], name=name, media_type=media_type)
|
|
|
|
|
|
@router.get("/{digest}")
|
|
def get_artifact(digest: str, request: Request) -> Any:
|
|
"""Stream one artifact back."""
|
|
store = _store(request)
|
|
path = store.path(digest)
|
|
if path is None:
|
|
raise HTTPException(status_code=404, detail="No such artifact")
|
|
return StreamingResponse(
|
|
store.read(digest),
|
|
media_type="application/octet-stream",
|
|
headers={"Content-Length": str(path.stat().st_size)},
|
|
)
|