stroblmeandClaude Opus 5
40f8ad378d
Slurm is not a machine that attaches and stays; it is a queue somebody else owns. So nothing here submits a node to it. It submits a job whose payload is an ordinary worker dialling back in, and everything downstream — the protocol, the artifacts, cancellation, the books — already worked and did not have to learn what Slurm is. The alternative, which Covalent takes, is to stage a serialized call and a runner onto the login node, poll squeue and copy the result back: a second way of running a node beside the one that exists. The cost of not doing that is one assumption, that a compute node can open a connection outward. Where that is false, _payload is the single method a staged variant would replace. Clusters are configured in provisioners.json beside the alerts, since this is infrastructure an operator writes rather than anything a flow says. The script is generated with the system ssh and no new dependency, and prerun owns the environment — deliberately no pip install, because what is on a cluster is somebody's decision. One outstanding request per profile, cancelled if it never attaches and on the way out. Nothing autoscales. The run gate needed the same hook: a run held before it starts never reaches the placer's own wait, so it would have queued forever on a machine nothing had asked for. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW
356 lines
14 KiB
Python
356 lines
14 KiB
Python
import asyncio
|
|
import contextlib
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
|
|
|
import sentry_sdk
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.concurrency import run_in_threadpool
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi.routing import APIRoute
|
|
from fluksio_worker.worker_main import ARTIFACT_DIR_ENV
|
|
from starlette.middleware.cors import CORSMiddleware
|
|
|
|
from fluksio.api.main import api_router
|
|
from fluksio.api.routes.alerts import read_config as read_alerts_config
|
|
from fluksio.cloud import config as cloud_config
|
|
from fluksio.core import security
|
|
from fluksio.core.config import settings
|
|
from fluksio.core.db import engine as db_engine
|
|
from fluksio.core.db import prepare
|
|
from fluksio.flow import logs, modules
|
|
from fluksio.flow.alerts import AlertManager
|
|
from fluksio.flow.artifacts import ArtifactStore
|
|
from fluksio.flow.controller import FlowController, RebuildBusy
|
|
from fluksio.flow.dashboards import DashboardStore
|
|
from fluksio.flow.events import event_bus
|
|
from fluksio.flow.executor import ExecutionService
|
|
from fluksio.flow.metrics import MetricsCollector
|
|
from fluksio.flow.nodes.http import close_shared_client
|
|
from fluksio.flow.pipeline import ValueSource
|
|
from fluksio.flow.placement import Placer
|
|
from fluksio.flow.plugins import load_plugins
|
|
from fluksio.flow.provision import load_provisioners
|
|
from fluksio.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue
|
|
from fluksio.flow.remote import RemoteWorkerHub
|
|
from fluksio.flow.resources import ResourceAccountant, fair_share_env
|
|
from fluksio.flow.runs import RUN_STATE_TTL, RunService, sweep_artifacts
|
|
from fluksio.flow.secrets import init_secrets
|
|
from fluksio.flow.state import MemoryState, RedisState, StateBackend
|
|
from fluksio.flow.store import FlowStore
|
|
from fluksio.flow.watchdog import LoopWatchdog
|
|
from fluksio.flow.workers import PythonWorkerPool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def custom_generate_unique_id(route: APIRoute) -> str:
|
|
return f"{route.tags[0]}-{route.name}"
|
|
|
|
|
|
if settings.SENTRY_DSN and settings.ENVIRONMENT != "local":
|
|
# `enable_tracing` was removed in sentry-sdk 2.x; this is what it meant.
|
|
sentry_sdk.init(dsn=str(settings.SENTRY_DSN), traces_sample_rate=1.0)
|
|
|
|
|
|
def _state_backend() -> StateBackend:
|
|
if settings.REDIS_HOST:
|
|
return RedisState(host=settings.REDIS_HOST, port=settings.REDIS_PORT)
|
|
return MemoryState()
|
|
|
|
|
|
def _work_queue(namespace: str = "queue") -> WorkQueue:
|
|
"""Redis makes queued work survive the process; memory does not pretend to."""
|
|
if settings.REDIS_HOST:
|
|
return RedisWorkQueue(
|
|
host=settings.REDIS_HOST, port=settings.REDIS_PORT, namespace=namespace
|
|
)
|
|
return MemoryWorkQueue()
|
|
|
|
|
|
def _run_state(namespace: str) -> StateBackend:
|
|
"""A state backend of a run's own, which is what isolates it.
|
|
|
|
It expires: a finished run's messages are read out into its result, and
|
|
what is left is only worth keeping while someone might look at it.
|
|
"""
|
|
if settings.REDIS_HOST:
|
|
return RedisState(
|
|
host=settings.REDIS_HOST,
|
|
port=settings.REDIS_PORT,
|
|
namespace=namespace,
|
|
ttl=RUN_STATE_TTL,
|
|
)
|
|
return MemoryState()
|
|
|
|
|
|
async def _sweep_artifacts(store: ArtifactStore, controller: FlowController) -> None:
|
|
"""Take unreferenced artifact bytes off the disk, on a slow loop.
|
|
|
|
A flow streaming media writes one artifact per frame, so a store nothing
|
|
prunes only grows. Runs in a thread: it walks a directory and reads state.
|
|
"""
|
|
interval = settings.ARTIFACT_GC_INTERVAL_S
|
|
if interval <= 0:
|
|
return
|
|
while True:
|
|
await asyncio.sleep(interval)
|
|
try:
|
|
await run_in_threadpool(
|
|
sweep_artifacts,
|
|
store,
|
|
controller.state,
|
|
settings.ARTIFACT_GC_GRACE_S,
|
|
)
|
|
except Exception:
|
|
logger.exception("Artifact sweep failed")
|
|
|
|
|
|
def _mcp_sessions() -> AbstractAsyncContextManager[None]:
|
|
"""The MCP session manager's run scope, or nothing when MCP is off."""
|
|
if not settings.MCP_ENABLED:
|
|
return contextlib.nullcontext()
|
|
from fluksio.mcp.server import mcp as mcp_server
|
|
|
|
return mcp_server.session_manager.run()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
"""Start the flow engine alongside the API."""
|
|
# The schema and the first superuser, before anything reads either. It is
|
|
# idempotent, so a deployment that ran this from its own prestart step
|
|
# pays a version check for it and nothing else.
|
|
await run_in_threadpool(prepare, db_engine)
|
|
event_bus.bind(asyncio.get_running_loop())
|
|
# Node code is user code, and `print` is how it says things.
|
|
logs.install()
|
|
init_secrets(settings.SECRETS_FILE, settings.SECRET_KEY)
|
|
# Connectors register their node types before any flow is built with them.
|
|
load_plugins()
|
|
|
|
alerts = AlertManager(event_bus, config=read_alerts_config())
|
|
execution = ExecutionService(
|
|
queue=_work_queue(),
|
|
max_workers=settings.FLOW_MAX_WORKERS,
|
|
events=event_bus,
|
|
max_cascades=settings.FLOW_MAX_CASCADES,
|
|
)
|
|
store = FlowStore(settings.FLOWS_DIR)
|
|
# The packages node code imports, before anything tries to import them.
|
|
await run_in_threadpool(modules.reconcile, store)
|
|
# Beside the flows rather than in them: an artifact is what a run produced,
|
|
# not something anyone wrote, so it has no business in the git repository.
|
|
artifacts = ArtifactStore(settings.FLOWS_DIR.parent / "artifacts")
|
|
app.state.artifact_store = artifacts
|
|
accountant = ResourceAccountant(cpus=settings.FLOW_CPUS, gpus=settings.FLOW_GPUS)
|
|
app.state.resources = accountant
|
|
# Every machine a node could run on: this one, and whatever attaches.
|
|
placer = Placer(local=accountant, events=event_bus)
|
|
# Where more machines can be asked for when nothing attached will do.
|
|
placer.provisioners = load_provisioners(settings.PROVISIONERS_FILE, event_bus)
|
|
app.state.placer = placer
|
|
pool = PythonWorkerPool(
|
|
python=modules.venv_python(),
|
|
size=settings.FLOW_MAX_WORKERS,
|
|
events=event_bus,
|
|
# A worker in this container writes to the store directly; a remote one
|
|
# is given a URL instead. Node code calls the same two functions.
|
|
env={
|
|
ARTIFACT_DIR_ENV: str(artifacts.root),
|
|
# Every slot can be busy at once, so a worker left to size its own
|
|
# thread pool to the machine means as many processes as there are
|
|
# slots, each believing it has the whole of it. A node that says
|
|
# what it needs overrides this; one that says nothing gets a share.
|
|
**fair_share_env(accountant.cpus, settings.FLOW_MAX_WORKERS),
|
|
},
|
|
)
|
|
pool.start()
|
|
app.state.worker_pool = pool
|
|
# Assigned rather than passed both ways: the hub tells the placer when a
|
|
# machine comes or goes, and the placer needs the hub to know what is there.
|
|
worker_hub = RemoteWorkerHub(on_change=placer.wake)
|
|
placer.hub = worker_hub
|
|
app.state.worker_hub = worker_hub
|
|
controller = FlowController(
|
|
store=store,
|
|
state=_state_backend(),
|
|
events=event_bus,
|
|
max_workers=settings.FLOW_MAX_WORKERS,
|
|
fastapi_app=app,
|
|
execution=execution,
|
|
alerts=alerts,
|
|
workers=pool,
|
|
remote=worker_hub,
|
|
resources=accountant,
|
|
placer=placer,
|
|
)
|
|
app.state.flow_controller = controller
|
|
# A "dashboard" alert channel puts its alert into the graph. Bound here
|
|
# rather than passed in: the manager is built before the controller is.
|
|
alerts.publish = lambda name, value: controller.publish_message(
|
|
name, value, ValueSource(kind="api", id="alerts", label="Alerts")
|
|
)
|
|
dashboards = DashboardStore(controller.store)
|
|
app.state.dashboard_store = dashboards
|
|
controller.dashboards = dashboards
|
|
# Charts need a deeper series than the default; tell the engine
|
|
# before it starts recording.
|
|
controller.set_history_limits(dashboards.history_requirements())
|
|
# Runs read from a stream of their own: a burst of sweep runs must not
|
|
# stand between the automations and their work, and a run that takes an
|
|
# hour must not be judged by the cascade reaper's timings.
|
|
run_service = RunService(
|
|
controller=controller,
|
|
queue=_work_queue("run"),
|
|
state_factory=_run_state,
|
|
artifacts=artifacts,
|
|
)
|
|
app.state.run_service = run_service
|
|
watchdog = LoopWatchdog(event_bus)
|
|
app.state.watchdog = watchdog
|
|
watchdog_task = asyncio.create_task(watchdog.run(), name="loop-watchdog")
|
|
alerts_task = asyncio.create_task(alerts.run(), name="alert-manager")
|
|
metrics_task = asyncio.create_task(
|
|
MetricsCollector(event_bus).run(), name="metrics-collector"
|
|
)
|
|
gc_task = asyncio.create_task(
|
|
_sweep_artifacts(artifacts, controller), name="artifact-gc"
|
|
)
|
|
await controller.start()
|
|
run_service.start()
|
|
# Optional, and off unless someone enrolled this installation: the
|
|
# connector dials the portal, nothing dials in.
|
|
cloud_task: asyncio.Task[None] | None = None
|
|
app.state.cloud_connector = None
|
|
app.state.cloud_task = None
|
|
from fluksio.cloud import connector as cloud_connector
|
|
|
|
if cloud_config.exists():
|
|
cloud_connector.start(app)
|
|
cloud_task = app.state.cloud_task
|
|
# Watched whether or not one exists now: enrolling from the CLI writes the
|
|
# config from another process entirely, and an engine already serving
|
|
# should pick it up rather than need restarting.
|
|
enrol_task = asyncio.create_task(
|
|
cloud_connector.watch_enrolment(app), name="cloud-enrolment-watch"
|
|
)
|
|
try:
|
|
# A mounted sub-app gets no lifespan of its own, so the MCP session
|
|
# manager is entered here; without it every /mcp request fails.
|
|
async with _mcp_sessions():
|
|
yield
|
|
finally:
|
|
watchdog_task.cancel()
|
|
alerts_task.cancel()
|
|
metrics_task.cancel()
|
|
gc_task.cancel()
|
|
enrol_task.cancel()
|
|
# Re-read from app.state: enrolling at runtime replaces this.
|
|
running_cloud = getattr(app.state, "cloud_task", None) or cloud_task
|
|
if running_cloud is not None:
|
|
running_cloud.cancel()
|
|
await run_in_threadpool(run_service.stop)
|
|
await controller.stop()
|
|
pool.stop()
|
|
# A machine asked for and not yet arrived would hold an allocation
|
|
# nobody is going to use.
|
|
for provisioner in placer.provisioners:
|
|
await run_in_threadpool(provisioner.shutdown)
|
|
close_shared_client()
|
|
if settings.MCP_ENABLED:
|
|
from fluksio.mcp.http import aclose
|
|
|
|
await aclose()
|
|
|
|
|
|
# The schema enumerates every endpoint this installation serves, including the
|
|
# paths trigger nodes mount at runtime. That is exactly what a developer wants
|
|
# and exactly what an internet-facing deployment should not hand out, so it
|
|
# follows the environment — the same rule the portal's backend uses. The
|
|
# generated client is built from a local run, not from the deployed host.
|
|
_docs_enabled = settings.ENVIRONMENT != "production"
|
|
|
|
app = FastAPI(
|
|
title=settings.PROJECT_NAME,
|
|
openapi_url=f"{settings.API_V1_STR}/openapi.json" if _docs_enabled else None,
|
|
docs_url="/docs" if _docs_enabled else None,
|
|
redoc_url="/redoc" if _docs_enabled else None,
|
|
generate_unique_id_function=custom_generate_unique_id,
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
# Set all CORS enabled origins
|
|
if settings.all_cors_origins:
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=settings.all_cors_origins,
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
app.include_router(api_router, prefix=settings.API_V1_STR)
|
|
|
|
|
|
@app.exception_handler(RebuildBusy)
|
|
async def rebuild_busy(request: Request, exc: RebuildBusy) -> JSONResponse: # noqa: ARG001
|
|
"""Every route that deploys something answers a wedged rebuild the same way.
|
|
|
|
The request was fine and retrying it may well work, so this is the engine
|
|
saying it is busy rather than the request having gone wrong.
|
|
"""
|
|
return JSONResponse(status_code=503, content={"detail": str(exc)})
|
|
|
|
|
|
# Tagged because the operation-id builder reads the first tag; the route
|
|
# itself stays out of the schema.
|
|
@app.get(
|
|
"/.well-known/oauth-authorization-server",
|
|
include_in_schema=False,
|
|
tags=["oauth"],
|
|
)
|
|
def oauth_authorization_server() -> JSONResponse:
|
|
"""RFC 8414 metadata, so an agent can find its way in unaided.
|
|
|
|
The authorization endpoint is the dashboard rather than the API: approving
|
|
a client needs a signed-in human, and the browser session lives there.
|
|
"""
|
|
issuer = settings.oauth_issuer
|
|
return JSONResponse(
|
|
content={
|
|
"issuer": issuer,
|
|
"authorization_endpoint": (
|
|
f"{settings.FRONTEND_HOST.rstrip('/')}/oauth/authorize"
|
|
),
|
|
"token_endpoint": f"{issuer}{settings.API_V1_STR}/oauth/token",
|
|
"registration_endpoint": f"{issuer}{settings.API_V1_STR}/oauth/register",
|
|
"jwks_uri": f"{issuer}/.well-known/jwks.json",
|
|
"response_types_supported": ["code"],
|
|
"grant_types_supported": ["authorization_code", "refresh_token"],
|
|
"code_challenge_methods_supported": ["S256"],
|
|
"token_endpoint_auth_methods_supported": ["none"],
|
|
"scopes_supported": [security.MCP_SCOPE],
|
|
},
|
|
headers={"Cache-Control": "public, max-age=3600"},
|
|
)
|
|
|
|
|
|
@app.get("/.well-known/jwks.json", include_in_schema=False, tags=["oauth"])
|
|
def jwks() -> JSONResponse:
|
|
"""The public half of the MCP signing key."""
|
|
return JSONResponse(
|
|
content=security.public_jwks(),
|
|
headers={"Cache-Control": "public, max-age=3600"},
|
|
)
|
|
|
|
|
|
# Mounted last, and at the root: the SDK serves both /mcp and the protected
|
|
# resource metadata that has to sit beside it, so mounting under /mcp would put
|
|
# that metadata somewhere no client looks for it.
|
|
if settings.MCP_ENABLED:
|
|
from fluksio.mcp.http import build_http_app
|
|
|
|
app.mount("/", build_http_app(app))
|