[Tinker] Keep SQLite off the forwarded-future hot path

skyrl/tinker/api.py

@@ -6,9 +6,9 @@
 import shutil
 import signal
 import threading
-from contextlib import asynccontextmanager, suppress
+from contextlib import asynccontextmanager, nullcontext, suppress
 from datetime import datetime, timedelta, timezone
-from typing import Annotated, Any, AsyncGenerator, ClassVar, Literal
+from typing import Annotated, Any, AsyncGenerator, Awaitable, ClassVar, Literal
 from uuid import uuid4
 
 import fastapi
@@ -45,6 +45,7 @@
     enable_sqlite_wal,
     get_async_database_url,
 )
+from skyrl.tinker.external_future_store import ExternalFutureStore
 from skyrl.tinker.extra import (
     ExternalInferenceClient,
     SkyRLTrainInferenceForwardingClient,
@@ -182,6 +183,56 @@ async def poll_futures(
         await asyncio.sleep(poll_interval_sec)
 
 
+def _finish_forwarding_task(tasks: set[asyncio.Task], task: asyncio.Task) -> None:
+    tasks.discard(task)
+    if task.cancelled():
+        return
+    if error := task.exception():
+        logger.error("Forwarding task failed: %r", error)
+
+
+def _start_forwarding_task(app: FastAPI, operation: Awaitable[None]) -> None:
+    task = asyncio.create_task(operation)
+    app.state.forwarding_tasks.add(task)
+    task.add_done_callback(lambda done: _finish_forwarding_task(app.state.forwarding_tasks, done))
+
+
+async def _close_external_inference(app: FastAPI) -> None:
+    if app.state.forwarding_tasks:
+        await asyncio.gather(*tuple(app.state.forwarding_tasks), return_exceptions=True)
+
+    inference_client = getattr(app.state, "external_inference_client", None)
+    aclose = getattr(inference_client, "aclose", None)
+    if aclose is not None:
+        with suppress(Exception):
+            await aclose()
+
+    if app.state.external_future_store is not None:
+        await app.state.external_future_store.close()
+
+
+async def _close_runtime(app: FastAPI, background_engine: asyncio.subprocess.Process) -> None:
+    try:
+        await _close_external_inference(app)
+    finally:
+        logger.info(f"Stopping background engine (PID {background_engine.pid})")
+        with suppress(ProcessLookupError):
+            background_engine.terminate()
+            try:
+                await asyncio.wait_for(background_engine.wait(), timeout=5)
+            except asyncio.TimeoutError:
+                logger.warning(f"Background engine (PID {background_engine.pid}) did not terminate gracefully, killing")
+                background_engine.kill()
+                await background_engine.wait()
+        logger.info("Background engine stopped")
+
+
+def _get_db_write_context(db_engine):
+    if db_engine.dialect.name == "sqlite":
+        return asyncio.Lock()
+    return nullcontext()
+
+
 def _get_parent_uv_run_args(parent_cmd: list[str]) -> list[str]:
     """Extract parent `uv run <uv run args>` flags for the engine launch given the parent process's startup command
 
@@ -241,6 +292,13 @@ async def lifespan(app: FastAPI):
 
     app.state.future_waiters = {}
     app.state.future_poller = asyncio.create_task(poll_futures(app.state.db_engine, app.state.future_waiters))
+    app.state.forwarding_tasks = set()
+    app.state.external_future_store = None
+    app.state.db_write_lock = _get_db_write_context(app.state.db_engine)
+    app.state.sampling_model_cache = {}
+    app.state.sampling_model_cache_lock = asyncio.Lock()
+    app.state.validated_sampler_checkpoints = set()
+    app.state.sampler_checkpoint_validation_lock = asyncio.Lock()
 
     # Setup external inference client if configured.
     #
@@ -265,8 +323,10 @@ async def lifespan(app: FastAPI):
         app.state.external_inference_client = ExternalInferenceClient(app.state.engine_config, app.state.db_engine)
         logger.info(f"External engine configured: {app.state.engine_config.external_inference_url}")
     elif backend_name in ("megatron", "fsdp") and not is_colocated:
+        app.state.external_future_store = ExternalFutureStore(app.state.db_engine, app.state.db_write_lock)
+        await app.state.external_future_store.start()
         app.state.external_inference_client = SkyRLTrainInferenceForwardingClient(
-            app.state.engine_config, app.state.db_engine
+            app.state.engine_config, app.state.db_engine, app.state.external_future_store
         )
         logger.info(
             "SkyRL-Train inference forwarding client enabled for non-colocated backend=%s",
@@ -319,25 +379,7 @@ def force_exit():
     with suppress(asyncio.CancelledError):
         await app.state.future_poller
 
-    # Close the forwarding client's persistent httpx connection pool if we
-    # installed one. Cheap no-op when external_inference_client doesn't own
-    # an httpx client (ExternalInferenceClient creates one per call).
-    inference_client = getattr(app.state, "external_inference_client", None)
-    aclose = getattr(inference_client, "aclose", None)
-    if aclose is not None:
-        with suppress(Exception):
-            await aclose()
-
-    logger.info(f"Stopping background engine (PID {app.state.background_engine.pid})")
-    with suppress(ProcessLookupError):
-        background_engine.terminate()
-        try:
-            await asyncio.wait_for(background_engine.wait(), timeout=5)
-        except asyncio.TimeoutError:
-            logger.warning(f"Background engine (PID {background_engine.pid}) did not terminate gracefully, killing")
-            background_engine.kill()
-            await background_engine.wait()
-    logger.info("Background engine stopped")
+    await _close_runtime(app, background_engine)
 
 
 app = FastAPI(title="Tinker API Mock", version="0.0.1", lifespan=lifespan)
@@ -925,14 +967,19 @@ async def create_session(request: CreateSessionRequest, session: AsyncSession =
 
 
 @app.post("/api/v1/session_heartbeat", response_model=SessionHeartbeatResponse)
-async def session_heartbeat(request: SessionHeartbeatRequest, session: AsyncSession = Depends(get_session)):
+async def session_heartbeat(
+    request: SessionHeartbeatRequest,
+    raw_request: Request,
+    session: AsyncSession = Depends(get_session),
+):
     """Heartbeat for an active session to keep it alive."""
-    session_db = await session.get(SessionDB, request.session_id)
-    if session_db is None:
-        raise HTTPException(status_code=404, detail="Session not found")
-    session_db.last_heartbeat_at = datetime.now(timezone.utc)
-    session_db.heartbeat_count += 1
-    await session.commit()
+    async with raw_request.app.state.db_write_lock:
+        session_db = await session.get(SessionDB, request.session_id)
+        if session_db is None:
+            raise HTTPException(status_code=404, detail="Session not found")
+        session_db.last_heartbeat_at = datetime.now(timezone.utc)
+        session_db.heartbeat_count += 1
+        await session.commit()
     return SessionHeartbeatResponse()
 
 
@@ -1128,17 +1175,16 @@ async def _read_forward_backward_request(request: Request) -> tuple[ForwardBackw
 async def forward_backward(request: Request, session: AsyncSession = Depends(get_session)):
     """Compute and accumulate gradients (or run forward-only when the proto body asks for it)."""
     req, forward_only = await _read_forward_backward_request(request)
-    await get_model(session, req.model_id)
-
-    request_id = await create_future(
-        session=session,
-        request_type=types.RequestType.FORWARD if forward_only else types.RequestType.FORWARD_BACKWARD,
-        model_id=req.model_id,
-        request_data=req.forward_backward_input.to_types(),
-        seq_id=req.seq_id,
-    )
-
-    await session.commit()
+    async with request.app.state.db_write_lock:
+        await get_model(session, req.model_id)
+        request_id = await create_future(
+            session=session,
+            request_type=types.RequestType.FORWARD if forward_only else types.RequestType.FORWARD_BACKWARD,
+            model_id=req.model_id,
+            request_data=req.forward_backward_input.to_types(),
+            seq_id=req.seq_id,
+        )
+        await session.commit()
 
     return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id))
 
@@ -1298,15 +1344,60 @@ async def save_weights_for_sampler(request: SaveWeightsForSamplerRequest, sessio
     return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id))
 
 
-async def get_sampling_model(request: SampleRequest, session: AsyncSession) -> (str | None, str | None):
+async def get_sampling_model(
+    request: SampleRequest,
+    req: Request,
+    session: AsyncSession,
+) -> tuple[str | None, str | None]:
     """Return (base_model, model_path) for a sampling request."""
-    # Resolve model/base from sampling_session_id if provided
-    if request.sampling_session_id is not None:
-        sampling_session = await session.get(SamplingSessionDB, request.sampling_session_id)
+    sampling_session_id = request.sampling_session_id
+    if sampling_session_id is None:
+        return (request.base_model, request.model_path)
+
+    cache = req.app.state.sampling_model_cache
+    cached = cache.get(sampling_session_id)
+    if cached is not None:
+        return cached
+
+    async with req.app.state.sampling_model_cache_lock:
+        cached = cache.get(sampling_session_id)
+        if cached is not None:
+            return cached
+        sampling_session = await session.get(SamplingSessionDB, sampling_session_id)
         if sampling_session is None:
             raise HTTPException(status_code=404, detail="Sampling session not found")
-        return (sampling_session.base_model, sampling_session.model_path)
-    return (request.base_model, request.model_path)
+        sampling_model = (
+            sampling_session.base_model,
+            sampling_session.model_path,
+        )
+        cache[sampling_session_id] = sampling_model
+        return sampling_model
+
+
+async def validate_sampler_checkpoint_once(
+    request: Request,
+    model_id: str,
+    checkpoint_id: str,
+    session: AsyncSession,
+) -> None:
+    """Validate an immutable sampler checkpoint once before serving it."""
+    key = (model_id, checkpoint_id)
+    validated = request.app.state.validated_sampler_checkpoints
+    if key in validated:
+        return
+
+    async with request.app.state.sampler_checkpoint_validation_lock:
+        if key in validated:
+            return
+        await get_model(session, model_id)
+        await validate_checkpoint(
+            request,
+            model_id,
+            checkpoint_id,
+            types.CheckpointType.SAMPLER,
+            session,
+        )
+        validated.add(key)
 
 
 @app.post("/api/v1/asample", response_model=FutureResponse)
@@ -1318,7 +1409,7 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession =
             detail="sampling_session_id must not contain ':' (the routing-key delimiter)",
         )
 
-    base_model, model_path = await get_sampling_model(request, session)
+    base_model, model_path = await get_sampling_model(request, req, session)
 
     if base_model:
         model_id = checkpoint_id = ""
@@ -1336,38 +1427,40 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession =
                 status_code=400,
                 detail="model_path must be tinker://model_id/checkpoint_id or tinker://model_id/sampler_weights/checkpoint_id",
             )
-        await get_model(session, model_id)
-        # Validate that the checkpoint exists and is ready
-        await validate_checkpoint(req, model_id, checkpoint_id, types.CheckpointType.SAMPLER, session)
+        await validate_sampler_checkpoint_once(req, model_id, checkpoint_id, session)
 
-    request_id = await create_future(
-        session=session,
-        request_type=(
-            types.RequestType.EXTERNAL if req.app.state.external_inference_client else types.RequestType.SAMPLE
-        ),
-        model_id=model_id,
-        request_data=types.SampleInput(
-            base_model=base_model,
-            prompt=request.prompt.to_types(),
-            sampling_params=request.sampling_params.to_types(),
-            num_samples=request.num_samples,
-            checkpoint_id=checkpoint_id,
-            # A positive topk implies prompt logprobs: both are read off the same
-            # prompt forward pass, so asking for one asks for the other.
-            prompt_logprobs=bool(request.prompt_logprobs) or request.topk_prompt_logprobs > 0,
-            topk_prompt_logprobs=request.topk_prompt_logprobs,
-            seq_id=request.seq_id,
-            sampling_session_id=request.sampling_session_id,
-        ),
+    sample_input = types.SampleInput(
+        base_model=base_model,
+        prompt=request.prompt.to_types(),
+        sampling_params=request.sampling_params.to_types(),
+        num_samples=request.num_samples,
+        checkpoint_id=checkpoint_id,
+        # A positive topk implies prompt logprobs: both are read off the same
+        # prompt forward pass, so asking for one asks for the other.
+        prompt_logprobs=bool(request.prompt_logprobs) or request.topk_prompt_logprobs > 0,
+        topk_prompt_logprobs=request.topk_prompt_logprobs,
+        seq_id=request.seq_id,
+        sampling_session_id=request.sampling_session_id,
     )
-
-    await session.commit()
+    if req.app.state.external_future_store is not None:
+        request_id = req.app.state.external_future_store.create(model_id, sample_input)
+    else:
+        request_id = await create_future(
+            session=session,
+            request_type=(
+                types.RequestType.EXTERNAL if req.app.state.external_inference_client else types.RequestType.SAMPLE
+            ),
+            model_id=model_id,
+            request_data=sample_input,
+        )
+        await session.commit()
 
     if req.app.state.external_inference_client:
-        asyncio.create_task(
+        _start_forwarding_task(
+            req.app,
             req.app.state.external_inference_client.call_and_store_result(
                 request_id, request, model_id, checkpoint_id, base_model=base_model
-            )
+            ),
         )
 
     return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id))
@@ -1391,10 +1484,19 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request):
     """Retrieve the result of an async operation, waiting until it's available."""
     request_id = int(request.request_id)
 
-    try:
-        row = await wait_for_future(req.app.state.future_waiters, request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
-    except KeyError:
-        raise HTTPException(status_code=404, detail="Future not found")
+    found_in_memory = False
+    external_future_store = req.app.state.external_future_store
+    if external_future_store is not None:
+        try:
+            row = await external_future_store.wait(request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
+            found_in_memory = True
+        except KeyError:
+            pass
+    if not found_in_memory:
+        try:
+            row = await wait_for_future(req.app.state.future_waiters, request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
+        except KeyError:
+            raise HTTPException(status_code=404, detail="Future not found")
 
     if row is None:
         raise HTTPException(status_code=408, detail="Timeout waiting for result")
@@ -1604,19 +1706,24 @@ async def delete_checkpoint(
             ),
         )
 
-    checkpoint_db = await session.get(CheckpointDB, (unique_id, checkpoint_id, resolved_checkpoint_type))
-    if not checkpoint_db:
-        raise HTTPException(status_code=404, detail=f"Checkpoint not found: {unique_id}/{checkpoint_id}")
-
-    if checkpoint_db.status == CheckpointStatus.PENDING:
-        raise HTTPException(status_code=425, detail="Checkpoint is still being created")
-
-    # Commit the row deletion before unlinking the artifact. If the commit fails we
-    # leave an orphaned file (GC-able) rather than a row that lists a checkpoint whose
-    # archive is gone, which would make every subsequent download 500.
-    path = checkpoint_file_path(request, unique_id, checkpoint_id, resolved_checkpoint_type)
-    await session.delete(checkpoint_db)
-    await session.commit()
+    sampler_checkpoint = resolved_checkpoint_type == types.CheckpointType.SAMPLER
+    validation_context = request.app.state.sampler_checkpoint_validation_lock if sampler_checkpoint else nullcontext()
+    async with validation_context:
+        checkpoint_db = await session.get(CheckpointDB, (unique_id, checkpoint_id, resolved_checkpoint_type))
+        if not checkpoint_db:
+            raise HTTPException(status_code=404, detail=f"Checkpoint not found: {unique_id}/{checkpoint_id}")
+
+        if checkpoint_db.status == CheckpointStatus.PENDING:
+            raise HTTPException(status_code=425, detail="Checkpoint is still being created")
+
+        # Commit the row deletion before unlinking the artifact. If the commit fails we
+        # leave an orphaned file (GC-able) rather than a row that lists a checkpoint whose
+        # archive is gone, which would make every subsequent download 500.
+        path = checkpoint_file_path(request, unique_id, checkpoint_id, resolved_checkpoint_type)
+        await session.delete(checkpoint_db)
+        await session.commit()
+        if sampler_checkpoint:
+            request.app.state.validated_sampler_checkpoints.discard((unique_id, checkpoint_id))
     await asyncio.to_thread(delete_checkpoint_file, path)