Follows the portal: the noun is "instance" everywhere the app says it — UI strings, CLI output, error details, docs and comments. The wire keys (`instance_id`, `instance_token`) and the hub route this calls move with it. An existing cloud.json is adopted rather than refused: without the key alias the dataclass fails to parse, which the caller swallows and reads as "never enrolled" instead of "reconnect". `instance_key` on a node type becomes `target_key`. It means the outside thing a node points at, which is a different sense of the word, and keeping both would put two meanings of "instance" in one codebase. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015YrQnKV3bnQd4K342y8tKj
438 lines
18 KiB
Python
438 lines
18 KiB
Python
import asyncio
|
|
import contextlib
|
|
import inspect
|
|
import logging
|
|
from collections.abc import AsyncIterator, Callable, Coroutine
|
|
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
|
from typing import Any
|
|
|
|
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 import __version__
|
|
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":
|
|
# Imported here rather than at the top: it is a `fluksio[server]` extra, so
|
|
# a pip install without one has no sentry to import — and no DSN either.
|
|
import sentry_sdk
|
|
|
|
# `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())
|
|
|
|
# Everything acquired below registers how to release it, and the `finally`
|
|
# at the end walks the list backwards. It used to be a fixed block of
|
|
# shutdown steps *after* every acquisition — so a failure part-way through
|
|
# startup reached none of them and left the worker pool's subprocesses and
|
|
# every background task behind. Under `--reload` that is once per bad edit.
|
|
started: list[Callable[[], Any]] = []
|
|
|
|
async def _release() -> None:
|
|
for close in reversed(started):
|
|
try:
|
|
result = close()
|
|
if inspect.isawaitable(result):
|
|
await result
|
|
except Exception:
|
|
logger.exception("Error while shutting down")
|
|
|
|
def _background(
|
|
coro: Coroutine[Any, Any, None], name: str, once: bool = False
|
|
) -> asyncio.Task[None]:
|
|
"""Start a background task that says something if it ever stops.
|
|
|
|
The loops catch their own exceptions *inside* the loop, so one raised
|
|
anywhere else simply ended the task — an engine that went on serving
|
|
with no metrics, no alerts or no artifact sweep and nothing anywhere
|
|
saying so. ``once`` is for a task that is meant to finish.
|
|
"""
|
|
task = asyncio.create_task(coro, name=name)
|
|
|
|
def _finished(done: asyncio.Task[None]) -> None:
|
|
if done.cancelled():
|
|
return
|
|
exc = done.exception()
|
|
if exc is None:
|
|
if not once:
|
|
logger.warning("Background task '%s' stopped on its own", name)
|
|
return
|
|
logger.error("Background task '%s' died: %s", name, exc, exc_info=exc)
|
|
event_bus.publish(
|
|
{"type": "engine_degraded", "detail": f"{name} stopped: {exc}"}
|
|
)
|
|
|
|
task.add_done_callback(_finished)
|
|
started.append(task.cancel)
|
|
return task
|
|
|
|
# The startup itself is inside the try, so a step that raises part-way
|
|
# through still reaches the release below.
|
|
try:
|
|
# 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. In the background, because a first
|
|
# install of anything substantial takes minutes and `uv` is given five
|
|
# of them twice over — held here, the container never answers its
|
|
# healthcheck, autoheal restarts it, and it never finishes installing at
|
|
# all. A node whose import is not there yet fails and is retried, which
|
|
# is recoverable; a restart loop is not.
|
|
_background(
|
|
run_in_threadpool(modules.reconcile, store), "module-reconcile", once=True
|
|
)
|
|
# 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()
|
|
# On a thread: stopping a worker waits up to five seconds on each child,
|
|
# and with a pool per declared environment that is a shutdown the event
|
|
# loop should not be holding.
|
|
started.append(lambda: run_in_threadpool(pool.stop))
|
|
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,
|
|
parallel=settings.FLOW_MAX_RUNS,
|
|
)
|
|
app.state.run_service = run_service
|
|
# Said out loud because it is the only way to tell that a settings file was
|
|
# read at all — the numbers are what someone raising them is looking for.
|
|
logger.info(
|
|
"Engine limits: workers=%s cascades=%s runs=%s",
|
|
settings.FLOW_MAX_WORKERS,
|
|
execution.max_cascades,
|
|
run_service.parallel,
|
|
)
|
|
watchdog = LoopWatchdog(event_bus)
|
|
app.state.watchdog = watchdog
|
|
|
|
_background(watchdog.run(), "loop-watchdog")
|
|
_background(alerts.run(), "alert-manager")
|
|
_background(MetricsCollector(event_bus).run(), "metrics-collector")
|
|
_background(_sweep_artifacts(artifacts, controller), "artifact-gc")
|
|
await controller.start()
|
|
started.append(controller.stop)
|
|
run_service.start()
|
|
started.append(lambda: run_in_threadpool(run_service.stop))
|
|
# A machine asked for and not yet arrived would hold an allocation nobody
|
|
# is going to use.
|
|
started.append(
|
|
lambda: asyncio.gather(
|
|
*(
|
|
run_in_threadpool(provisioner.shutdown)
|
|
for provisioner in placer.provisioners
|
|
)
|
|
)
|
|
)
|
|
started.append(lambda: run_in_threadpool(close_shared_client))
|
|
# Optional, and off unless someone enrolled this instance: 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
|
|
|
|
def _stop_cloud() -> None:
|
|
# Re-read from app.state: enrolling at runtime replaces this.
|
|
running = getattr(app.state, "cloud_task", None) or cloud_task
|
|
if running is not None:
|
|
running.cancel()
|
|
|
|
started.append(_stop_cloud)
|
|
# 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.
|
|
_background(cloud_connector.watch_enrolment(app), "cloud-enrolment-watch")
|
|
# 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
|
|
except Exception:
|
|
# A startup that fails part-way lands here too, because everything
|
|
# since `event_bus.bind` registered its own release.
|
|
logger.exception("The engine could not start")
|
|
raise
|
|
finally:
|
|
await _release()
|
|
if settings.MCP_ENABLED:
|
|
from fluksio.mcp.http import aclose
|
|
|
|
await aclose()
|
|
|
|
|
|
# The schema enumerates every endpoint this instance 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,
|
|
version=__version__,
|
|
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))
|