diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index cb397463..e824eb8b 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -2,6 +2,8 @@ import asyncio import hmac import io import os, time +import random +from contextlib import asynccontextmanager from typing import AsyncGenerator import json import cv2 @@ -45,46 +47,19 @@ from aare.daq.automation_runner import PersistentWorkflowRunner from aare.daq.automation_api_router import router as workflow_router, set_workflow_dependencies logger = setup_logger("aareDAQ") -app = FastAPI() -register_exception_handlers(app) # OAuth2 setup oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") -bl = mx_beamline() -cfg = BeamlineConfig(bl) -daq = AareDAQ(cfg, bl) - -# ───────────────────────────────────────────── -# Initialize workflow system -# ───────────────────────────────────────────── +# ── Per-worker state: populated inside the lifespan, after fork ── +bl = None +cfg = None +daq = None +workflow_redis_manager = None +workflow_runner = None USE_SIMULATED_WORKFLOW = os.getenv("WORKFLOW_SIMULATION", "0") == "1" -workflow_redis_manager = WorkflowRedisManager( - client=cfg._BeamlineConfig__client, # Reuse existing Redis connection - beamline=bl.value, -) - -workflow_runner = PersistentWorkflowRunner( - redis_manager=workflow_redis_manager, - registry=STATE_REGISTRY, - handlers=SIMULATED_HANDLER_REGISTRY if USE_SIMULATED_WORKFLOW else HANDLER_REGISTRY, -) - -if USE_SIMULATED_WORKFLOW: - logger.warning("⚠️ Workflow system running in SIMULATION mode - no actual DAQ operations") - -set_workflow_dependencies(workflow_redis_manager, workflow_runner, cfg) - -# Include the workflow router -app.include_router(workflow_router) - -try: - daq.sync_current_sample_from_tell(force=True) -except Exception as e: - logger.warning(f"Initial sample sync from TELL failed: {e}") - _all_pgroups_cache: dict[str, tuple[list[str], float]] = {} _ALL_PGROUPS_TTL_S = 60.0 # adjust TTL as needed @@ -97,6 +72,60 @@ _face_detection_state: dict = { } _face_detection_state_lock = asyncio.Lock() +@asynccontextmanager +async def lifespan(application: FastAPI): + """ + Runs once per worker process, AFTER fork() and BEFORE serving requests. + All stateful / connection-opening initialisation belongs here so that + each worker gets its own fresh Redis, BEC, EPICS, and TELL connections. + """ + global bl, cfg, daq, workflow_redis_manager, workflow_runner + + await asyncio.sleep(random.uniform(0.5, 3.0)) + + logger.info(f"Worker {os.getpid()} starting initialisation...") + + # ── Core objects (Redis, EPICS PVs, BEC, TELL, JFJoch, etc.) ── + bl = mx_beamline() + cfg = BeamlineConfig(bl) + daq = AareDAQ(cfg, bl) + + # ── Workflow system ── + # workflow_redis_manager = WorkflowRedisManager( + # client=cfg._BeamlineConfig__client, # Reuse this worker's Redis connection + # beamline=bl.value, + # ) + # workflow_runner = PersistentWorkflowRunner( + # redis_manager=workflow_redis_manager, + # registry=STATE_REGISTRY, + # handlers=SIMULATED_HANDLER_REGISTRY if USE_SIMULATED_WORKFLOW else HANDLER_REGISTRY, + # ) + # if USE_SIMULATED_WORKFLOW: + # logger.warning("⚠️ Workflow system running in SIMULATION mode - no actual DAQ operations") + # + # set_workflow_dependencies(workflow_redis_manager, workflow_runner, cfg) + + # ── Initial TELL sync ── + try: + daq.sync_current_sample_from_tell(force=True) + except Exception as e: + logger.warning(f"Initial sample sync from TELL failed: {e}") + + # ── Wire callbacks ── + daq.set_face_detection_progress_callback(_push_face_detection_progress) + + logger.info(f"Worker {os.getpid()} initialised successfully.") + + yield # ── application serves requests here ── + + # Shutdown: add cleanup here if needed + logger.info(f"Worker {os.getpid()} shutting down.") + +app = FastAPI(lifespan=lifespan) +register_exception_handlers(app) + +app.include_router(workflow_router) + def _required_recovery_code() -> str: code = os.getenv("AARE_RECOVERY_CODE", "").strip() if not code: @@ -148,8 +177,6 @@ async def face_detection_event_stream() -> AsyncGenerator[str, None]: except asyncio.CancelledError: return -daq.set_face_detection_progress_callback(_push_face_detection_progress) - @app.post("/token") async def login(form_data: OAuth2PasswordRequestForm = Depends()): data = await run_in_threadpool(auth.authenticate_user, cfg, form_data) @@ -1045,15 +1072,12 @@ LOGGING_CONFIG = { }, } -urllib3.disable_warnings() - - def main(): # Remove in production! urllib3.disable_warnings() # Run the application using uvicorn - uvicorn.run("aare.daq.server:app", host="0.0.0.0", port=5210, workers=4, log_config=LOGGING_CONFIG) + uvicorn.run("aare.daq.server:app", host="0.0.0.0", port=5210, workers=2, log_config=LOGGING_CONFIG) if __name__ == "__main__":