diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index 608ffda9f8..8237d15ef7 100644 --- a/skyrl/tinker/api.py +++ b/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 ` 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) diff --git a/skyrl/tinker/external_future_store.py b/skyrl/tinker/external_future_store.py new file mode 100644 index 0000000000..4a450f48c0 --- /dev/null +++ b/skyrl/tinker/external_future_store.py @@ -0,0 +1,141 @@ +import asyncio +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass, field +from datetime import datetime, timezone + +from pydantic import BaseModel +from sqlmodel import func, select +from sqlmodel.ext.asyncio.session import AsyncSession + +from skyrl.tinker import types +from skyrl.tinker.db_models import FutureDB, RequestStatus +from skyrl.utils.log import logger + + +@dataclass +class ExternalFuture: + request_id: int + model_id: str | None + request_data: dict + status: RequestStatus = RequestStatus.PENDING + result_data: str | None = None + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + completed_at: datetime | None = None + event: asyncio.Event = field(default_factory=asyncio.Event) + persistence_error: Exception | None = None + + +class ExternalFutureStore: + """Keeps forwarded sample futures off the database hot path.""" + + _PERSIST_BATCH_SIZE = 64 + _PERSIST_QUEUE_SIZE = 2048 + + def __init__(self, db_engine, db_write_lock: AbstractAsyncContextManager): + self.db_engine = db_engine + self.db_write_lock = db_write_lock + self._entries: dict[int, ExternalFuture] = {} + self._persist_queue: asyncio.Queue[ExternalFuture] = asyncio.Queue(maxsize=self._PERSIST_QUEUE_SIZE) + self._persist_worker: asyncio.Task | None = None + self._persist_error: Exception | None = None + self._next_request_id = -1 + + async def start(self) -> None: + async with AsyncSession(self.db_engine) as session: + statement = select(func.min(FutureDB.request_id)).where(FutureDB.request_id < 0) + minimum_request_id = (await session.exec(statement)).one() + if minimum_request_id is not None: + self._next_request_id = minimum_request_id - 1 + self._persist_worker = asyncio.create_task(self._persist_loop()) + + def create(self, model_id: str | None, request_data: BaseModel) -> int: + request_id = self._next_request_id + self._next_request_id -= 1 + self._entries[request_id] = ExternalFuture( + request_id=request_id, + model_id=model_id, + request_data=request_data.model_dump(mode="json"), + ) + return request_id + + async def wait(self, request_id: int, timeout: float) -> tuple[RequestStatus, types.RequestType, str | None] | None: + entry = self._entries.get(request_id) + if entry is None: + raise KeyError(request_id) + try: + await asyncio.wait_for(entry.event.wait(), timeout) + except asyncio.TimeoutError: + return None + if entry.persistence_error is not None: + self._entries.pop(request_id, None) + raise RuntimeError(f"Failed to persist external future {request_id}") from entry.persistence_error + return entry.status, types.RequestType.EXTERNAL, entry.result_data + + async def complete(self, request_id: int, result_data: BaseModel, status: RequestStatus) -> None: + entry = self._entries[request_id] + entry.result_data = result_data.model_dump_json() + entry.status = status + entry.completed_at = datetime.now(timezone.utc) + await self._persist_queue.put(entry) + + async def flush(self) -> None: + await self._persist_queue.join() + if self._persist_error is not None: + error, self._persist_error = self._persist_error, None + raise RuntimeError("External future persistence failed") from error + + async def close(self) -> None: + try: + await self.flush() + finally: + if self._persist_worker is not None: + self._persist_worker.cancel() + await asyncio.gather(self._persist_worker, return_exceptions=True) + + async def _persist_loop(self) -> None: + while True: + entries = [await self._persist_queue.get()] + while len(entries) < self._PERSIST_BATCH_SIZE: + try: + entries.append(self._persist_queue.get_nowait()) + except asyncio.QueueEmpty: + break + try: + await self._persist(entries) + except Exception as error: + self._persist_error = error + logger.exception( + "External future persistence failed request_ids=%s..%s", + entries[0].request_id, + entries[-1].request_id, + ) + for entry in entries: + entry.persistence_error = error + entry.event.set() + else: + for entry in entries: + entry.event.set() + self._entries.pop(entry.request_id, None) + finally: + for _ in entries: + self._persist_queue.task_done() + + async def _persist(self, entries: list[ExternalFuture]) -> None: + async with self.db_write_lock: + async with AsyncSession(self.db_engine) as session: + session.add_all( + [ + FutureDB( + request_id=entry.request_id, + request_type=types.RequestType.EXTERNAL, + model_id=entry.model_id, + request_data=entry.request_data, + result_data=entry.result_data, + status=entry.status, + created_at=entry.created_at, + completed_at=entry.completed_at, + ) + for entry in entries + ] + ) + await session.commit() diff --git a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py index 227db2de92..d6d32dae10 100644 --- a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py +++ b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py @@ -15,15 +15,22 @@ from skyrl.tinker import types from skyrl.tinker.config import EngineConfig from skyrl.tinker.db_models import EngineStateDB, FutureDB, RequestStatus +from skyrl.tinker.external_future_store import ExternalFutureStore from skyrl.utils.log import logger class SkyRLTrainInferenceForwardingClient: """Forwards EXTERNAL sample requests to the SkyRL-Train-managed vLLM.""" - def __init__(self, engine_config: EngineConfig, db_engine): + def __init__( + self, + engine_config: EngineConfig, + db_engine, + external_future_store: ExternalFutureStore | None = None, + ): self.engine_config = engine_config self.db_engine = db_engine + self.external_future_store = external_future_store self._cached_proxy_url: str | None = None self._cache_lock = asyncio.Lock() # Backpressure layered: httpx pool -> vllm-router -> vLLM max_num_seqs. @@ -80,6 +87,10 @@ async def call_and_store_result( result = types.ErrorResponse(error=str(e), status="failed") status = RequestStatus.FAILED + if self.external_future_store is not None: + await self.external_future_store.complete(request_id, result, status) + return + async with AsyncSession(self.db_engine) as session: future = await session.get(FutureDB, request_id) if future is None: diff --git a/tests/tinker/skyrl_train/test_async_sample_routing.py b/tests/tinker/skyrl_train/test_async_sample_routing.py index 72f98751c8..b6b92314f0 100644 --- a/tests/tinker/skyrl_train/test_async_sample_routing.py +++ b/tests/tinker/skyrl_train/test_async_sample_routing.py @@ -169,17 +169,18 @@ def _read_engine_state(db_path: str): engine.dispose() -def _read_future_request_type(db_path: str, request_id: int) -> str: - """Read the request_type of a single future from the test server's DB.""" - from sqlmodel import Session, create_engine +def _read_external_future_ids(db_path: str) -> set[int]: + """Read the IDs of terminal EXTERNAL futures from the test server's DB.""" + from sqlmodel import Session, create_engine, select + from skyrl.tinker import types as skyrl_types from skyrl.tinker.db_models import FutureDB engine = create_engine(f"sqlite:///{db_path}", echo=False) try: with Session(engine) as session: - row = session.get(FutureDB, request_id) - return None if row is None else str(row.request_type) + statement = select(FutureDB.request_id).where(FutureDB.request_type == skyrl_types.RequestType.EXTERNAL) + return set(session.exec(statement).all()) finally: engine.dispose() @@ -216,11 +217,6 @@ def test_sample_uses_external_path(server_db_path): This is the "test" half of the design: the API hoists the sample off the engine's serial loop and into the API process's asyncio loop. """ - from sqlmodel import Session, create_engine, func, select - - from skyrl.tinker import types as skyrl_types - from skyrl.tinker.db_models import FutureDB - proc, db_path, _ = server_db_path sc = tinker.ServiceClient(base_url=f"http://0.0.0.0:{TEST_PORT}/", api_key=TINKER_API_KEY) tc = sc.create_lora_training_client(base_model=BASE_MODEL, rank=8) @@ -229,14 +225,7 @@ def test_sample_uses_external_path(server_db_path): _train_one_step(tc, tok) sampler = tc.save_weights_and_get_sampling_client(name="external_path_a") - # Snapshot the max future_id before submitting our sample so we can - # filter out any EXTERNAL futures from earlier tests. - eng = create_engine(f"sqlite:///{db_path}", echo=False) - try: - with Session(eng) as s: - max_before = s.exec(select(func.max(FutureDB.request_id))).one() or 0 - finally: - eng.dispose() + future_ids_before = _read_external_future_ids(db_path) out = sampler.sample( prompt=tinker_types.ModelInput.from_ints(tok.encode("Hi", add_special_tokens=True)), @@ -245,24 +234,16 @@ def test_sample_uses_external_path(server_db_path): ).result() assert len(out.sequences) == 1 - # Look for an EXTERNAL future with id > max_before. If async routing - # is on, every sample creates exactly one such row. - eng = create_engine(f"sqlite:///{db_path}", echo=False) - try: - with Session(eng) as s: - stmt = ( - select(FutureDB.request_id, FutureDB.request_type) - .where(FutureDB.request_id > max_before) - .where(FutureDB.request_type == skyrl_types.RequestType.EXTERNAL) - ) - rows = s.exec(stmt).all() - finally: - eng.dispose() - - assert len(rows) >= 1, ( - f"expected at least one EXTERNAL future to be created by the sample call, " - f"found {len(rows)}; async sample routing may not be active" + # Terminal rows flush asynchronously, and in-memory futures use negative IDs + # so they cannot collide with the database's positive autoincrement sequence. + found_new_future = wait_for_condition( + lambda: bool(_read_external_future_ids(db_path) - future_ids_before), + timeout_sec=10, + poll_interval_sec=0.1, ) + assert ( + found_new_future + ), "expected a new EXTERNAL future after the sample call; async sample routing may not be active" def test_sample_concurrent_with_training_is_fast(server_db_path): diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py new file mode 100644 index 0000000000..be88b1b0dc --- /dev/null +++ b/tests/tinker/test_external_future_store.py @@ -0,0 +1,300 @@ +import asyncio +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import create_async_engine +from sqlmodel import SQLModel +from sqlmodel.ext.asyncio.session import AsyncSession + +from skyrl.tinker import api, types +from skyrl.tinker.config import EngineConfig +from skyrl.tinker.db_models import ( + CheckpointDB, + CheckpointStatus, + ModelDB, + RequestStatus, + SessionDB, + enable_sqlite_wal, + get_async_database_url, +) +from skyrl.tinker.external_future_store import ExternalFutureStore +from skyrl.tinker.extra.skyrl_train_inference_forwarding import ( + SkyRLTrainInferenceForwardingClient, +) + + +def _sample_input(seq_id: int) -> types.SampleInput: + return types.SampleInput( + base_model="model_a", + prompt=types.ModelInput(chunks=[types.EncodedTextChunk(tokens=[seq_id])]), + sampling_params=types.SamplingParams(temperature=0.0, max_tokens=1, seed=seq_id), + num_samples=1, + checkpoint_id="", + prompt_logprobs=False, + seq_id=seq_id, + ) + + +@pytest_asyncio.fixture() +async def future_store(tmp_path): + db_url = get_async_database_url(f"sqlite:///{tmp_path / 'tinker.db'}") + engine = create_async_engine(db_url, pool_size=5, max_overflow=10, pool_timeout=0.1) + enable_sqlite_wal(engine.sync_engine) + async with engine.begin() as connection: + await connection.run_sync(SQLModel.metadata.create_all) + + db_write_lock = asyncio.Lock() + store = ExternalFutureStore(engine, db_write_lock) + await store.start() + yield store, engine, db_write_lock + await store.close() + await engine.dispose() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_forwarding_tasks_before_closing_store(): + release_forwarding = asyncio.Event() + events = [] + + class ClosingClient: + async def aclose(self) -> None: + events.append("client_closed") + + class ClosingStore: + async def close(self) -> None: + events.append("store_closed") + + app = SimpleNamespace( + state=SimpleNamespace( + external_inference_client=ClosingClient(), + external_future_store=ClosingStore(), + forwarding_tasks=set(), + ) + ) + + async def finish_forwarding() -> None: + await release_forwarding.wait() + events.append("future_completed") + + api._start_forwarding_task(app, finish_forwarding()) + shutdown = asyncio.create_task(api._close_external_inference(app)) + await asyncio.sleep(0) + assert not shutdown.done() + + release_forwarding.set() + await shutdown + + assert events == ["future_completed", "client_closed", "store_closed"] + assert not app.state.forwarding_tasks + + +@pytest.mark.asyncio +async def test_shutdown_stops_engine_when_future_persistence_failed(monkeypatch): + events = [] + + class BackgroundEngine: + pid = 123 + + def terminate(self) -> None: + events.append("engine_terminated") + + async def wait(self) -> int: + events.append("engine_waited") + return 0 + + async def fail_external_close(_app) -> None: + events.append("external_close_failed") + raise RuntimeError("persistence failed") + + monkeypatch.setattr(api, "_close_external_inference", fail_external_close) + + with pytest.raises(RuntimeError, match="persistence failed"): + await api._close_runtime(SimpleNamespace(), BackgroundEngine()) + + assert events == ["external_close_failed", "engine_terminated", "engine_waited"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("dialect", "serializes"), [("sqlite", True), ("postgresql", False)]) +async def test_db_write_context_serializes_only_sqlite(dialect, serializes): + context = api._get_db_write_context(SimpleNamespace(dialect=SimpleNamespace(name=dialect))) + first_entered = asyncio.Event() + release_first = asyncio.Event() + second_entered = asyncio.Event() + + async def first_writer() -> None: + async with context: + first_entered.set() + await release_first.wait() + + async def second_writer() -> None: + await first_entered.wait() + async with context: + second_entered.set() + + first = asyncio.create_task(first_writer()) + second = asyncio.create_task(second_writer()) + await first_entered.wait() + await asyncio.sleep(0) + assert second_entered.is_set() is not serializes + + release_first.set() + await asyncio.gather(first, second) + assert second_entered.is_set() + + +@pytest.mark.asyncio +async def test_sampler_checkpoint_delete_waits_for_validation_and_invalidates_cache(future_store, monkeypatch): + _, engine, _ = future_store + validation_started = asyncio.Event() + release_validation = asyncio.Event() + request = SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + engine_config=EngineConfig(base_model="model_a"), + sampler_checkpoint_validation_lock=asyncio.Lock(), + validated_sampler_checkpoints=set(), + ) + ) + ) + + async with AsyncSession(engine) as session: + session.add( + SessionDB( + session_id="session_a", + tags=[], + user_metadata={}, + sdk_version="test", + ) + ) + session.add( + ModelDB( + model_id="model_a", + base_model="model_a", + lora_config={}, + status="ready", + request_id=0, + session_id="session_a", + ) + ) + session.add( + CheckpointDB( + model_id="model_a", + checkpoint_id="weights_a", + checkpoint_type=types.CheckpointType.SAMPLER, + status=CheckpointStatus.COMPLETED, + ) + ) + await session.commit() + + async def hold_validation(*args) -> None: + validation_started.set() + await release_validation.wait() + + monkeypatch.setattr(api, "validate_checkpoint", hold_validation) + async with AsyncSession(engine) as validation_session, AsyncSession(engine) as deletion_session: + validation = asyncio.create_task( + api.validate_sampler_checkpoint_once( + request, + "model_a", + "weights_a", + validation_session, + ) + ) + await validation_started.wait() + deletion = asyncio.create_task( + api.delete_checkpoint( + request, + "model_a", + "weights_a", + types.CheckpointType.SAMPLER, + deletion_session, + ) + ) + await asyncio.sleep(0) + assert not deletion.done() + + release_validation.set() + await asyncio.gather(validation, deletion) + + assert not request.app.state.validated_sampler_checkpoints + async with AsyncSession(engine) as session: + assert ( + await session.get( + CheckpointDB, + ("model_a", "weights_a", types.CheckpointType.SAMPLER), + ) + is None + ) + + +@pytest.mark.asyncio +async def test_forwarding_client_completes_in_memory_future(future_store, monkeypatch): + store, engine, _ = future_store + request_id = store.create("model_a", _sample_input(1)) + result = types.SampleOutput( + sequences=[types.GeneratedSequence(stop_reason="stop", tokens=[1, 2], logprobs=[-0.5, -1.0])] + ) + client = SkyRLTrainInferenceForwardingClient(EngineConfig(base_model="model_a"), engine, store) + + async def forward(*args, **kwargs): + return result + + monkeypatch.setattr(client, "_forward_with_retry", forward) + try: + await client.call_and_store_result( + request_id, + SimpleNamespace(), + model_id="model_a", + checkpoint_id="", + ) + completed = await store.wait(request_id, timeout=1) + finally: + await client.aclose() + + assert completed == ( + RequestStatus.COMPLETED, + types.RequestType.EXTERNAL, + result.model_dump_json(), + ) + + +@pytest.mark.asyncio +async def test_persistence_failure_is_reported_to_waiter(future_store, monkeypatch): + store, _, _ = future_store + request_id = store.create("model_a", _sample_input(1)) + + async def fail_persistence(entries): + raise RuntimeError("database unavailable") + + monkeypatch.setattr(store, "_persist", fail_persistence) + await store.complete(request_id, types.SampleOutput(sequences=[]), RequestStatus.COMPLETED) + + with pytest.raises(RuntimeError, match=f"Failed to persist external future {request_id}"): + await store.wait(request_id, timeout=1) + with pytest.raises(RuntimeError, match="External future persistence failed"): + await store.flush() + + +@pytest.mark.asyncio +async def test_retrieve_future_serializes_in_memory_result_as_proto(future_store): + from tinker import SampleResponse + from tinker.proto.response_conv import deserialize_proto_response + + store, engine, _ = future_store + request_id = store.create("model_a", _sample_input(1)) + result = types.SampleOutput( + sequences=[types.GeneratedSequence(stop_reason="stop", tokens=[1, 2], logprobs=[-0.5, -1.0])] + ) + await store.complete(request_id, result, RequestStatus.COMPLETED) + + request = SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(db_engine=engine, external_future_store=store, future_waiters={})), + headers={"accept": "application/x-protobuf, application/json"}, + ) + response = await api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_id)), request) + + assert response.media_type == "application/x-protobuf" + result = deserialize_proto_response(response.body, SampleResponse) + assert result.sequences[0].tokens == [1, 2] diff --git a/tests/tinker/test_future_waiting.py b/tests/tinker/test_future_waiting.py index 748f43973f..f0aecf0dd0 100644 --- a/tests/tinker/test_future_waiting.py +++ b/tests/tinker/test_future_waiting.py @@ -191,7 +191,13 @@ def _stub_request(async_engine, waiters, headers: dict | None = None): from types import SimpleNamespace return SimpleNamespace( - app=SimpleNamespace(state=SimpleNamespace(db_engine=async_engine, future_waiters=waiters)), + app=SimpleNamespace( + state=SimpleNamespace( + db_engine=async_engine, + external_future_store=None, + future_waiters=waiters, + ) + ), headers=headers or {}, )