From b8f7240597c4484d310b2c2b14737294ab2426ff Mon Sep 17 00:00:00 2001 From: Hersh Godse Date: Thu, 27 Aug 2026 03:49:54 +0000 Subject: [PATCH 1/7] Keep external sample futures off the database hot path Signed-off-by: Hersh Godse --- skyrl/tinker/api.py | 68 ++++-- skyrl/tinker/external_future_store.py | 149 +++++++++++ .../extra/skyrl_train_inference_forwarding.py | 13 +- .../skyrl_train/test_async_sample_routing.py | 51 ++-- tests/tinker/test_external_future_store.py | 231 ++++++++++++++++++ tests/tinker/test_future_waiting.py | 8 +- 6 files changed, 459 insertions(+), 61 deletions(-) create mode 100644 skyrl/tinker/external_future_store.py create mode 100644 tests/tinker/test_external_future_store.py diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index 608ffda9f8..ae542f8d81 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -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, @@ -241,6 +242,7 @@ 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.external_future_store = None # Setup external inference client if configured. # @@ -265,8 +267,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) + 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,6 +323,9 @@ def force_exit(): with suppress(asyncio.CancelledError): await app.state.future_poller + if app.state.external_future_store is not None: + await app.state.external_future_store.close() + # 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). @@ -1340,26 +1347,30 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession = # Validate that the checkpoint exists and is ready await validate_checkpoint(req, model_id, checkpoint_id, types.CheckpointType.SAMPLER, 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, ) + 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() @@ -1391,10 +1402,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") diff --git a/skyrl/tinker/external_future_store.py b/skyrl/tinker/external_future_store.py new file mode 100644 index 0000000000..f8968109b1 --- /dev/null +++ b/skyrl/tinker/external_future_store.py @@ -0,0 +1,149 @@ +import asyncio +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: BaseModel + 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) + persisted: bool = False + retrieved: bool = False + + +class ExternalFutureStore: + """Keeps forwarded sample futures off the database hot path.""" + + _PERSIST_BATCH_SIZE = 64 + _PERSIST_QUEUE_SIZE = 2048 + + def __init__(self, db_engine): + self.db_engine = db_engine + self._entries: dict[int, ExternalFuture] = {} + self._persist_queue: asyncio.Queue[ExternalFuture] = asyncio.Queue( + maxsize=self._PERSIST_QUEUE_SIZE + ) + self._persist_workers: list[asyncio.Task] = [] + self._persist_errors: list[Exception] = [] + 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_workers = [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, + ) + 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 + entry.retrieved = True + self._remove_finished_entry(entry) + 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) + entry.event.set() + + async def flush(self) -> None: + await self._persist_queue.join() + if self._persist_errors: + raise RuntimeError( + "External future persistence failed" + ) from self._persist_errors[0] + + async def close(self) -> None: + try: + await self.flush() + finally: + for worker in self._persist_workers: + worker.cancel() + await asyncio.gather(*self._persist_workers, 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_errors.append(error) + logger.exception( + "External future persistence failed request_ids=%s..%s", + entries[0].request_id, + entries[-1].request_id, + ) + else: + for entry in entries: + entry.persisted = True + self._remove_finished_entry(entry) + finally: + for _ in entries: + self._persist_queue.task_done() + + async def _persist(self, entries: list[ExternalFuture]) -> None: + 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.model_dump(mode="json"), + 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() + + def _remove_finished_entry(self, entry: ExternalFuture) -> None: + if entry.persisted and entry.retrieved: + self._entries.pop(entry.request_id, None) 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..90b3fd5016 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,23 +234,15 @@ 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" ) diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py new file mode 100644 index 0000000000..0018902fda --- /dev/null +++ b/tests/tinker/test_external_future_store.py @@ -0,0 +1,231 @@ +import asyncio +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import create_async_engine +from sqlmodel import SQLModel, func, select +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 ( + FutureDB, + RequestStatus, + SamplingSessionDB, + 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, + ) + + +class _CompletingForwarder: + def __init__(self, store: ExternalFutureStore): + self.store = store + self.gate = asyncio.Event() + + async def call_and_store_result( + self, + request_id: int, + sample_req, + model_id: str, + checkpoint_id: str, + base_model: str | None = None, + ) -> None: + await self.gate.wait() + await self.store.complete( + request_id, + types.SampleOutput(sequences=[]), + RequestStatus.COMPLETED, + ) + + +@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=1, max_overflow=0) + enable_sqlite_wal(engine.sync_engine) + async with engine.begin() as connection: + await connection.run_sync(SQLModel.metadata.create_all) + + store = ExternalFutureStore(engine) + await store.start() + yield store, engine + await store.close() + await engine.dispose() + + +@pytest.mark.asyncio +async def test_two_full_rollout_waves_complete_and_persist(future_store): + store, engine = future_store + forwarder = _CompletingForwarder(store) + request = SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + external_future_store=store, + external_inference_client=forwarder, + ) + ) + ) + result_data = types.SampleOutput(sequences=[]).model_dump_json() + + async with AsyncSession(engine) as session: + session.add( + SamplingSessionDB( + sampling_session_id="session_a", + session_id="session_a", + sampling_session_seq_id=0, + base_model="model_a", + ) + ) + await session.commit() + + for wave in range(2): + forwarder.gate.clear() + + async def create_sample(index: int) -> int: + async with AsyncSession(engine) as session: + response = await api.asample( + api.SampleRequest( + prompt=api.ModelInput( + chunks=[api.EncodedTextChunk(tokens=[index])] + ), + sampling_params=api.SamplingParams( + temperature=0.0, max_tokens=1, seed=index + ), + sampling_session_id="session_a", + seq_id=wave * 512 + index, + ), + request, + session, + ) + return int(response.request_id) + + request_ids = await asyncio.gather( + *(create_sample(index) for index in range(512)) + ) + waiters = [ + asyncio.create_task(store.wait(request_id, timeout=5)) + for request_id in request_ids + ] + + async with AsyncSession(engine) as session: + pending_after_creation = ( + await session.exec( + select(func.count()) + .select_from(FutureDB) + .where(FutureDB.status == RequestStatus.PENDING) + ) + ).one() + assert pending_after_creation == 0 + + forwarder.gate.set() + results = await asyncio.gather(*waiters) + + assert all( + result == (RequestStatus.COMPLETED, types.RequestType.EXTERNAL, result_data) + for result in results + ) + + await store.flush() + async with AsyncSession(engine) as session: + persisted = ( + await session.exec(select(func.count()).select_from(FutureDB)) + ).one() + pending = ( + await session.exec( + select(func.count()) + .select_from(FutureDB) + .where(FutureDB.status == RequestStatus.PENDING) + ) + ).one() + + assert persisted == 1024 + assert pending == 0 + + +@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_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 {}, ) From 625df777b84433fe9ec6d1c9d1fdd7d75b333a3e Mon Sep 17 00:00:00 2001 From: Hersh Godse Date: Thu, 27 Aug 2026 08:22:16 +0000 Subject: [PATCH 2/7] Serialize SQLite writes under rollout load --- skyrl/tinker/api.py | 61 ++++++++++++++++------ skyrl/tinker/external_future_store.py | 38 +++++++------- tests/tinker/test_external_future_store.py | 49 ++++++++++++----- 3 files changed, 102 insertions(+), 46 deletions(-) diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index ae542f8d81..fe9f445826 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -243,6 +243,9 @@ 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.external_future_store = None + app.state.db_write_lock = asyncio.Lock() + app.state.sampling_model_cache = {} + app.state.sampling_model_cache_lock = asyncio.Lock() # Setup external inference client if configured. # @@ -267,7 +270,9 @@ 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.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.external_future_store @@ -932,14 +937,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() @@ -1305,15 +1315,36 @@ 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 @app.post("/api/v1/asample", response_model=FutureResponse) @@ -1325,7 +1356,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 = "" diff --git a/skyrl/tinker/external_future_store.py b/skyrl/tinker/external_future_store.py index f8968109b1..b3ce06aa8f 100644 --- a/skyrl/tinker/external_future_store.py +++ b/skyrl/tinker/external_future_store.py @@ -31,8 +31,9 @@ class ExternalFutureStore: _PERSIST_BATCH_SIZE = 64 _PERSIST_QUEUE_SIZE = 2048 - def __init__(self, db_engine): + def __init__(self, db_engine, db_write_lock: asyncio.Lock): 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 @@ -126,23 +127,24 @@ async def _persist_loop(self) -> None: self._persist_queue.task_done() async def _persist(self, entries: list[ExternalFuture]) -> None: - 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.model_dump(mode="json"), - 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() + 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.model_dump(mode="json"), + 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() def _remove_finished_entry(self, entry: ExternalFuture) -> None: if entry.persisted and entry.retrieved: diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index 0018902fda..0ef6678df7 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -13,6 +13,7 @@ FutureDB, RequestStatus, SamplingSessionDB, + SessionDB, enable_sqlite_wal, get_async_database_url, ) @@ -39,7 +40,6 @@ def _sample_input(seq_id: int) -> types.SampleInput: class _CompletingForwarder: def __init__(self, store: ExternalFutureStore): self.store = store - self.gate = asyncio.Event() async def call_and_store_result( self, @@ -49,7 +49,6 @@ async def call_and_store_result( checkpoint_id: str, base_model: str | None = None, ) -> None: - await self.gate.wait() await self.store.complete( request_id, types.SampleOutput(sequences=[]), @@ -60,33 +59,47 @@ async def call_and_store_result( @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=1, max_overflow=0) + 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) - store = ExternalFutureStore(engine) + db_write_lock = asyncio.Lock() + store = ExternalFutureStore(engine, db_write_lock) await store.start() - yield store, engine + yield store, engine, db_write_lock await store.close() await engine.dispose() @pytest.mark.asyncio async def test_two_full_rollout_waves_complete_and_persist(future_store): - store, engine = future_store + store, engine, db_write_lock = future_store forwarder = _CompletingForwarder(store) request = SimpleNamespace( app=SimpleNamespace( state=SimpleNamespace( external_future_store=store, external_inference_client=forwarder, + db_write_lock=db_write_lock, + sampling_model_cache={}, + sampling_model_cache_lock=asyncio.Lock(), ) ) ) result_data = types.SampleOutput(sequences=[]).model_dump_json() async with AsyncSession(engine) as session: + session.add( + SessionDB( + session_id="session_a", + tags=[], + user_metadata={}, + sdk_version="test", + ) + ) session.add( SamplingSessionDB( sampling_session_id="session_a", @@ -98,8 +111,6 @@ async def test_two_full_rollout_waves_complete_and_persist(future_store): await session.commit() for wave in range(2): - forwarder.gate.clear() - async def create_sample(index: int) -> int: async with AsyncSession(engine) as session: response = await api.asample( @@ -118,9 +129,19 @@ async def create_sample(index: int) -> int: ) return int(response.request_id) - request_ids = await asyncio.gather( - *(create_sample(index) for index in range(512)) + async def heartbeat() -> None: + async with AsyncSession(engine) as session: + await api.session_heartbeat( + api.SessionHeartbeatRequest(session_id="session_a"), + request, + session, + ) + + responses = await asyncio.gather( + *(create_sample(index) for index in range(512)), + *(heartbeat() for _ in range(32)), ) + request_ids = responses[:512] waiters = [ asyncio.create_task(store.wait(request_id, timeout=5)) for request_id in request_ids @@ -136,7 +157,6 @@ async def create_sample(index: int) -> int: ).one() assert pending_after_creation == 0 - forwarder.gate.set() results = await asyncio.gather(*waiters) assert all( @@ -156,14 +176,17 @@ async def create_sample(index: int) -> int: .where(FutureDB.status == RequestStatus.PENDING) ) ).one() + session_db = await session.get(SessionDB, "session_a") assert persisted == 1024 assert pending == 0 + assert session_db is not None + assert session_db.heartbeat_count == 64 @pytest.mark.asyncio async def test_forwarding_client_completes_in_memory_future(future_store, monkeypatch): - store, engine = future_store + store, engine, _ = future_store request_id = store.create("model_a", _sample_input(1)) result = types.SampleOutput( sequences=[ @@ -203,7 +226,7 @@ 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 + store, engine, _ = future_store request_id = store.create("model_a", _sample_input(1)) result = types.SampleOutput( sequences=[ From 4561290ff7205a9c763131f577d0777732e81605 Mon Sep 17 00:00:00 2001 From: Hersh Godse Date: Thu, 27 Aug 2026 10:41:32 +0000 Subject: [PATCH 3/7] Serialize forward backward future writes --- skyrl/tinker/api.py | 21 ++- tests/tinker/test_external_future_store.py | 145 +++++++++++++++++++++ 2 files changed, 155 insertions(+), 11 deletions(-) diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index fe9f445826..869ae0ce06 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -1145,17 +1145,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)) diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index 0ef6678df7..9394533e32 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -6,11 +6,13 @@ from sqlalchemy.ext.asyncio import create_async_engine from sqlmodel import SQLModel, func, select from sqlmodel.ext.asyncio.session import AsyncSession +from starlette.requests import Request from skyrl.tinker import api, types from skyrl.tinker.config import EngineConfig from skyrl.tinker.db_models import ( FutureDB, + ModelDB, RequestStatus, SamplingSessionDB, SessionDB, @@ -56,6 +58,47 @@ async def call_and_store_result( ) +def _forward_backward_request(seq_id: int, db_write_lock: asyncio.Lock) -> Request: + body = api.ForwardBackwardRequest( + model_id="model_a", + seq_id=seq_id, + forward_backward_input=api.ForwardBackwardInput( + data=[ + api.Datum( + model_input=api.ModelInput( + chunks=[api.EncodedTextChunk(tokens=[1, 2])] + ), + loss_fn_inputs={ + "target_tokens": api.TensorData(data=[2, 3]), + "weights": api.TensorData(data=[1.0, 1.0]), + }, + ) + ], + loss_fn="cross_entropy", + ), + ).model_dump_json().encode() + body_sent = False + + async def receive(): + nonlocal body_sent + if body_sent: + return {"type": "http.disconnect"} + body_sent = True + return {"type": "http.request", "body": body, "more_body": False} + + app = SimpleNamespace(state=SimpleNamespace(db_write_lock=db_write_lock)) + return Request( + { + "type": "http", + "method": "POST", + "path": "/api/v1/forward_backward", + "headers": [(b"content-type", b"application/json")], + "app": app, + }, + receive, + ) + + @pytest_asyncio.fixture() async def future_store(tmp_path): db_url = get_async_database_url(f"sqlite:///{tmp_path / 'tinker.db'}") @@ -184,6 +227,108 @@ async def heartbeat() -> None: assert session_db.heartbeat_count == 64 +@pytest.mark.asyncio +async def test_sustained_rollouts_training_futures_and_heartbeats(future_store): + store, engine, db_write_lock = future_store + forwarder = _CompletingForwarder(store) + sample_request = SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + external_future_store=store, + external_inference_client=forwarder, + db_write_lock=db_write_lock, + sampling_model_cache={}, + sampling_model_cache_lock=asyncio.Lock(), + ) + ) + ) + + async with AsyncSession(engine) as session: + session.add( + SessionDB( + session_id="session_a", + tags=[], + user_metadata={}, + sdk_version="test", + ) + ) + session.add( + SamplingSessionDB( + sampling_session_id="session_a", + session_id="session_a", + sampling_session_seq_id=0, + base_model="model_a", + ) + ) + session.add( + ModelDB( + model_id="model_a", + base_model="model_a", + lora_config={}, + status="ready", + request_id=0, + session_id="session_a", + ) + ) + await session.commit() + + for wave in range(4): + async def create_sample(index: int) -> None: + async with AsyncSession(engine) as session: + await api.asample( + api.SampleRequest( + prompt=api.ModelInput( + chunks=[api.EncodedTextChunk(tokens=[index])] + ), + sampling_params=api.SamplingParams( + temperature=0.0, max_tokens=1, seed=index + ), + sampling_session_id="session_a", + seq_id=wave * 512 + index, + ), + sample_request, + session, + ) + + async def create_training_future(index: int) -> None: + async with AsyncSession(engine) as session: + await api.forward_backward( + _forward_backward_request(wave * 512 + index, db_write_lock), + session, + ) + + async def heartbeat() -> None: + async with AsyncSession(engine) as session: + await api.session_heartbeat( + api.SessionHeartbeatRequest(session_id="session_a"), + sample_request, + session, + ) + + await asyncio.gather( + *(create_sample(index) for index in range(512)), + *(create_training_future(index) for index in range(512)), + *(heartbeat() for _ in range(32)), + ) + + await store.flush() + async with AsyncSession(engine) as session: + persisted_by_type = dict( + ( + await session.exec( + select(FutureDB.request_type, func.count()) + .group_by(FutureDB.request_type) + ) + ).all() + ) + session_db = await session.get(SessionDB, "session_a") + + assert persisted_by_type[types.RequestType.EXTERNAL] == 2048 + assert persisted_by_type[types.RequestType.FORWARD_BACKWARD] == 2048 + assert session_db is not None + assert session_db.heartbeat_count == 128 + + @pytest.mark.asyncio async def test_forwarding_client_completes_in_memory_future(future_store, monkeypatch): store, engine, _ = future_store From 8d0fe222624235b453483e352444f359cb7f0e8d Mon Sep 17 00:00:00 2001 From: Hersh Godse Date: Thu, 27 Aug 2026 12:34:56 +0000 Subject: [PATCH 4/7] Cache sampler checkpoint validation under rollout load --- skyrl/tinker/api.py | 35 ++++++++++++++++++---- tests/tinker/test_external_future_store.py | 20 +++++++++++-- 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index 869ae0ce06..006990a6c6 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -246,6 +246,8 @@ async def lifespan(app: FastAPI): app.state.db_write_lock = asyncio.Lock() 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. # @@ -1346,6 +1348,32 @@ async def get_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) async def asample(request: SampleRequest, req: Request, session: AsyncSession = Depends(get_session)): """Generates samples from the model (async version).""" @@ -1373,9 +1401,7 @@ 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) sample_input = types.SampleInput( base_model=base_model, @@ -1401,8 +1427,7 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession = model_id=model_id, request_data=sample_input, ) - - await session.commit() + await session.commit() if req.app.state.external_inference_client: asyncio.create_task( diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index 9394533e32..4b16a88499 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -11,6 +11,8 @@ from skyrl.tinker import api, types from skyrl.tinker.config import EngineConfig from skyrl.tinker.db_models import ( + CheckpointDB, + CheckpointStatus, FutureDB, ModelDB, RequestStatus, @@ -228,7 +230,7 @@ async def heartbeat() -> None: @pytest.mark.asyncio -async def test_sustained_rollouts_training_futures_and_heartbeats(future_store): +async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(future_store): store, engine, db_write_lock = future_store forwarder = _CompletingForwarder(store) sample_request = SimpleNamespace( @@ -236,9 +238,12 @@ async def test_sustained_rollouts_training_futures_and_heartbeats(future_store): state=SimpleNamespace( external_future_store=store, external_inference_client=forwarder, + engine_config=EngineConfig(base_model="model_a"), db_write_lock=db_write_lock, sampling_model_cache={}, sampling_model_cache_lock=asyncio.Lock(), + validated_sampler_checkpoints=set(), + sampler_checkpoint_validation_lock=asyncio.Lock(), ) ) ) @@ -257,7 +262,7 @@ async def test_sustained_rollouts_training_futures_and_heartbeats(future_store): sampling_session_id="session_a", session_id="session_a", sampling_session_seq_id=0, - base_model="model_a", + model_path="tinker://model_a/sampler_weights/weights_a", ) ) session.add( @@ -270,6 +275,14 @@ async def test_sustained_rollouts_training_futures_and_heartbeats(future_store): 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() for wave in range(4): @@ -327,6 +340,9 @@ async def heartbeat() -> None: assert persisted_by_type[types.RequestType.FORWARD_BACKWARD] == 2048 assert session_db is not None assert session_db.heartbeat_count == 128 + assert sample_request.app.state.validated_sampler_checkpoints == { + ("model_a", "weights_a") + } @pytest.mark.asyncio From 79cbe3c642157576f3a760a2d318f123d6fc72e3 Mon Sep 17 00:00:00 2001 From: Hersh Godse Date: Thu, 27 Aug 2026 21:11:45 +0000 Subject: [PATCH 5/7] Fix external future lifecycle --- skyrl/tinker/api.py | 97 ++-- skyrl/tinker/external_future_store.py | 66 ++- .../skyrl_train/test_async_sample_routing.py | 6 +- tests/tinker/test_external_future_store.py | 425 ++++++++++-------- 4 files changed, 334 insertions(+), 260 deletions(-) diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index 006990a6c6..22ee05dd0f 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 @@ -183,6 +183,40 @@ 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() + + +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 @@ -242,8 +276,9 @@ 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 = asyncio.Lock() + 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() @@ -272,9 +307,7 @@ 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 - ) + 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.external_future_store @@ -330,17 +363,7 @@ def force_exit(): with suppress(asyncio.CancelledError): await app.state.future_poller - if app.state.external_future_store is not None: - await app.state.external_future_store.close() - - # 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() + await _close_external_inference(app) logger.info(f"Stopping background engine (PID {app.state.background_engine.pid})") with suppress(ProcessLookupError): @@ -1335,9 +1358,7 @@ async def get_sampling_model( cached = cache.get(sampling_session_id) if cached is not None: return cached - sampling_session = await session.get( - SamplingSessionDB, sampling_session_id - ) + sampling_session = await session.get(SamplingSessionDB, sampling_session_id) if sampling_session is None: raise HTTPException(status_code=404, detail="Sampling session not found") sampling_model = ( @@ -1430,10 +1451,11 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession = 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)) @@ -1679,19 +1701,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 index b3ce06aa8f..4a450f48c0 100644 --- a/skyrl/tinker/external_future_store.py +++ b/skyrl/tinker/external_future_store.py @@ -1,4 +1,5 @@ import asyncio +from contextlib import AbstractAsyncContextManager from dataclasses import dataclass, field from datetime import datetime, timezone @@ -15,14 +16,13 @@ class ExternalFuture: request_id: int model_id: str | None - request_data: BaseModel + 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) - persisted: bool = False - retrieved: bool = False + persistence_error: Exception | None = None class ExternalFutureStore: @@ -31,26 +31,22 @@ class ExternalFutureStore: _PERSIST_BATCH_SIZE = 64 _PERSIST_QUEUE_SIZE = 2048 - def __init__(self, db_engine, db_write_lock: asyncio.Lock): + 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_workers: list[asyncio.Task] = [] - self._persist_errors: list[Exception] = [] + 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 - ) + 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_workers = [asyncio.create_task(self._persist_loop())] + 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 @@ -58,13 +54,11 @@ def create(self, model_id: str | None, request_data: BaseModel) -> int: self._entries[request_id] = ExternalFuture( request_id=request_id, model_id=model_id, - request_data=request_data, + 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: + 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) @@ -72,34 +66,31 @@ async def wait( await asyncio.wait_for(entry.event.wait(), timeout) except asyncio.TimeoutError: return None - entry.retrieved = True - self._remove_finished_entry(entry) + 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: + 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) - entry.event.set() async def flush(self) -> None: await self._persist_queue.join() - if self._persist_errors: - raise RuntimeError( - "External future persistence failed" - ) from self._persist_errors[0] + 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: - for worker in self._persist_workers: - worker.cancel() - await asyncio.gather(*self._persist_workers, return_exceptions=True) + 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: @@ -112,16 +103,19 @@ async def _persist_loop(self) -> None: try: await self._persist(entries) except Exception as error: - self._persist_errors.append(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.persisted = True - self._remove_finished_entry(entry) + entry.event.set() + self._entries.pop(entry.request_id, None) finally: for _ in entries: self._persist_queue.task_done() @@ -135,7 +129,7 @@ async def _persist(self, entries: list[ExternalFuture]) -> None: request_id=entry.request_id, request_type=types.RequestType.EXTERNAL, model_id=entry.model_id, - request_data=entry.request_data.model_dump(mode="json"), + request_data=entry.request_data, result_data=entry.result_data, status=entry.status, created_at=entry.created_at, @@ -145,7 +139,3 @@ async def _persist(self, entries: list[ExternalFuture]) -> None: ] ) await session.commit() - - def _remove_finished_entry(self, entry: ExternalFuture) -> None: - if entry.persisted and entry.retrieved: - self._entries.pop(entry.request_id, None) diff --git a/tests/tinker/skyrl_train/test_async_sample_routing.py b/tests/tinker/skyrl_train/test_async_sample_routing.py index 90b3fd5016..b6b92314f0 100644 --- a/tests/tinker/skyrl_train/test_async_sample_routing.py +++ b/tests/tinker/skyrl_train/test_async_sample_routing.py @@ -241,9 +241,9 @@ def test_sample_uses_external_path(server_db_path): 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" - ) + 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 index 4b16a88499..724e593012 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -31,9 +31,7 @@ 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 - ), + sampling_params=types.SamplingParams(temperature=0.0, max_tokens=1, seed=seq_id), num_samples=1, checkpoint_id="", prompt_logprobs=False, @@ -61,24 +59,26 @@ async def call_and_store_result( def _forward_backward_request(seq_id: int, db_write_lock: asyncio.Lock) -> Request: - body = api.ForwardBackwardRequest( - model_id="model_a", - seq_id=seq_id, - forward_backward_input=api.ForwardBackwardInput( - data=[ - api.Datum( - model_input=api.ModelInput( - chunks=[api.EncodedTextChunk(tokens=[1, 2])] - ), - loss_fn_inputs={ - "target_tokens": api.TensorData(data=[2, 3]), - "weights": api.TensorData(data=[1.0, 1.0]), - }, - ) - ], - loss_fn="cross_entropy", - ), - ).model_dump_json().encode() + body = ( + api.ForwardBackwardRequest( + model_id="model_a", + seq_id=seq_id, + forward_backward_input=api.ForwardBackwardInput( + data=[ + api.Datum( + model_input=api.ModelInput(chunks=[api.EncodedTextChunk(tokens=[1, 2])]), + loss_fn_inputs={ + "target_tokens": api.TensorData(data=[2, 3]), + "weights": api.TensorData(data=[1.0, 1.0]), + }, + ) + ], + loss_fn="cross_entropy", + ), + ) + .model_dump_json() + .encode() + ) body_sent = False async def receive(): @@ -104,9 +104,7 @@ async def receive(): @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 - ) + 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) @@ -120,21 +118,27 @@ async def future_store(tmp_path): @pytest.mark.asyncio -async def test_two_full_rollout_waves_complete_and_persist(future_store): +async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(future_store): store, engine, db_write_lock = future_store forwarder = _CompletingForwarder(store) - request = SimpleNamespace( + sample_request = SimpleNamespace( app=SimpleNamespace( state=SimpleNamespace( + db_engine=engine, external_future_store=store, external_inference_client=forwarder, + forwarding_tasks=set(), + future_waiters={}, + engine_config=EngineConfig(base_model="model_a"), db_write_lock=db_write_lock, sampling_model_cache={}, sampling_model_cache_lock=asyncio.Lock(), + validated_sampler_checkpoints=set(), + sampler_checkpoint_validation_lock=asyncio.Lock(), ) - ) + ), + headers={}, ) - result_data = types.SampleOutput(sequences=[]).model_dump_json() async with AsyncSession(engine) as session: session.add( @@ -150,100 +154,179 @@ async def test_two_full_rollout_waves_complete_and_persist(future_store): sampling_session_id="session_a", session_id="session_a", sampling_session_seq_id=0, + model_path="tinker://model_a/sampler_weights/weights_a", + ) + ) + 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() - for wave in range(2): - async def create_sample(index: int) -> int: - async with AsyncSession(engine) as session: - response = await api.asample( - api.SampleRequest( - prompt=api.ModelInput( - chunks=[api.EncodedTextChunk(tokens=[index])] - ), - sampling_params=api.SamplingParams( - temperature=0.0, max_tokens=1, seed=index + future_poller = asyncio.create_task( + api.poll_futures(engine, sample_request.app.state.future_waiters, poll_interval_sec=0.001) + ) + expected_sample = types.SampleOutput(sequences=[]).model_dump_json().encode() + try: + request_ids = [] + for wave in range(4): + + async def create_sample(index: int) -> int: + async with AsyncSession(engine) as session: + response = await api.asample( + api.SampleRequest( + prompt=api.ModelInput(chunks=[api.EncodedTextChunk(tokens=[index])]), + sampling_params=api.SamplingParams(temperature=0.0, max_tokens=1, seed=index), + sampling_session_id="session_a", + seq_id=wave * 512 + index, ), - sampling_session_id="session_a", - seq_id=wave * 512 + index, - ), - request, - session, - ) - return int(response.request_id) - - async def heartbeat() -> None: - async with AsyncSession(engine) as session: - await api.session_heartbeat( - api.SessionHeartbeatRequest(session_id="session_a"), - request, - session, - ) - - responses = await asyncio.gather( - *(create_sample(index) for index in range(512)), - *(heartbeat() for _ in range(32)), - ) - request_ids = responses[:512] - waiters = [ - asyncio.create_task(store.wait(request_id, timeout=5)) - for request_id in request_ids - ] - - async with AsyncSession(engine) as session: - pending_after_creation = ( - await session.exec( - select(func.count()) - .select_from(FutureDB) - .where(FutureDB.status == RequestStatus.PENDING) + sample_request, + session, + ) + return int(response.request_id) + + async def create_training_future(index: int) -> None: + async with AsyncSession(engine) as session: + await api.forward_backward( + _forward_backward_request(wave * 512 + index, db_write_lock), + session, + ) + + async def heartbeat() -> None: + async with AsyncSession(engine) as session: + await api.session_heartbeat( + api.SessionHeartbeatRequest(session_id="session_a"), + sample_request, + session, + ) + + responses = await asyncio.gather( + *(create_sample(index) for index in range(512)), + *(create_training_future(index) for index in range(512)), + *(heartbeat() for _ in range(32)), + ) + request_ids = responses[:512] + retrievals = await asyncio.gather( + *( + api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_id)), sample_request) + for request_id in request_ids ) - ).one() - assert pending_after_creation == 0 - - results = await asyncio.gather(*waiters) + ) + assert all(response.body == expected_sample for response in retrievals) + assert not store._entries - assert all( - result == (RequestStatus.COMPLETED, types.RequestType.EXTERNAL, result_data) - for result in results - ) + repeated = await api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_ids[-1])), sample_request) + assert repeated.body == expected_sample + finally: + future_poller.cancel() + await asyncio.gather(future_poller, return_exceptions=True) await store.flush() async with AsyncSession(engine) as session: - persisted = ( - await session.exec(select(func.count()).select_from(FutureDB)) - ).one() - pending = ( - await session.exec( - select(func.count()) - .select_from(FutureDB) - .where(FutureDB.status == RequestStatus.PENDING) - ) - ).one() + persisted_by_type = dict( + (await session.exec(select(FutureDB.request_type, func.count()).group_by(FutureDB.request_type))).all() + ) session_db = await session.get(SessionDB, "session_a") - assert persisted == 1024 - assert pending == 0 + assert persisted_by_type[types.RequestType.EXTERNAL] == 2048 + assert persisted_by_type[types.RequestType.FORWARD_BACKWARD] == 2048 assert session_db is not None - assert session_db.heartbeat_count == 64 + assert session_db.heartbeat_count == 128 + assert sample_request.app.state.validated_sampler_checkpoints == {("model_a", "weights_a")} + assert not sample_request.app.state.forwarding_tasks @pytest.mark.asyncio -async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(future_store): - store, engine, db_write_lock = future_store - forwarder = _CompletingForwarder(store) - sample_request = SimpleNamespace( +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 +@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( - external_future_store=store, - external_inference_client=forwarder, engine_config=EngineConfig(base_model="model_a"), - db_write_lock=db_write_lock, - sampling_model_cache={}, - sampling_model_cache_lock=asyncio.Lock(), - validated_sampler_checkpoints=set(), sampler_checkpoint_validation_lock=asyncio.Lock(), + validated_sampler_checkpoints=set(), ) ) ) @@ -257,14 +340,6 @@ async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(fut sdk_version="test", ) ) - session.add( - SamplingSessionDB( - sampling_session_id="session_a", - session_id="session_a", - sampling_session_seq_id=0, - model_path="tinker://model_a/sampler_weights/weights_a", - ) - ) session.add( ModelDB( model_id="model_a", @@ -285,64 +360,45 @@ async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(fut ) await session.commit() - for wave in range(4): - async def create_sample(index: int) -> None: - async with AsyncSession(engine) as session: - await api.asample( - api.SampleRequest( - prompt=api.ModelInput( - chunks=[api.EncodedTextChunk(tokens=[index])] - ), - sampling_params=api.SamplingParams( - temperature=0.0, max_tokens=1, seed=index - ), - sampling_session_id="session_a", - seq_id=wave * 512 + index, - ), - sample_request, - session, - ) - - async def create_training_future(index: int) -> None: - async with AsyncSession(engine) as session: - await api.forward_backward( - _forward_backward_request(wave * 512 + index, db_write_lock), - session, - ) - - async def heartbeat() -> None: - async with AsyncSession(engine) as session: - await api.session_heartbeat( - api.SessionHeartbeatRequest(session_id="session_a"), - sample_request, - session, - ) - - await asyncio.gather( - *(create_sample(index) for index in range(512)), - *(create_training_future(index) for index in range(512)), - *(heartbeat() for _ in range(32)), + 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() - await store.flush() + release_validation.set() + await asyncio.gather(validation, deletion) + + assert not request.app.state.validated_sampler_checkpoints async with AsyncSession(engine) as session: - persisted_by_type = dict( - ( - await session.exec( - select(FutureDB.request_type, func.count()) - .group_by(FutureDB.request_type) - ) - ).all() + assert ( + await session.get( + CheckpointDB, + ("model_a", "weights_a", types.CheckpointType.SAMPLER), + ) + is None ) - session_db = await session.get(SessionDB, "session_a") - - assert persisted_by_type[types.RequestType.EXTERNAL] == 2048 - assert persisted_by_type[types.RequestType.FORWARD_BACKWARD] == 2048 - assert session_db is not None - assert session_db.heartbeat_count == 128 - assert sample_request.app.state.validated_sampler_checkpoints == { - ("model_a", "weights_a") - } @pytest.mark.asyncio @@ -350,15 +406,9 @@ async def test_forwarding_client_completes_in_memory_future(future_store, monkey 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 + 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 @@ -382,6 +432,23 @@ async def forward(*args, **kwargs): ) +@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 @@ -390,25 +457,15 @@ async def test_retrieve_future_serializes_in_memory_result_as_proto(future_store 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] - ) - ] + 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={} - ) - ), + 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 - ) + 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) From 6778f452e4ec5124fa2ea277cbd4a73eb8fc1301 Mon Sep 17 00:00:00 2001 From: Hersh Godse Date: Fri, 28 Aug 2026 01:16:55 +0000 Subject: [PATCH 6/7] Ensure engine stops after persistence failure --- skyrl/tinker/api.py | 29 +++++++++++++--------- tests/tinker/test_external_future_store.py | 26 +++++++++++++++++++ 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index 22ee05dd0f..8237d15ef7 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -211,6 +211,22 @@ async def _close_external_inference(app: FastAPI) -> 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() @@ -363,18 +379,7 @@ def force_exit(): with suppress(asyncio.CancelledError): await app.state.future_poller - await _close_external_inference(app) - - 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) diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index 724e593012..3a58d98bc6 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -287,6 +287,32 @@ async def finish_forwarding() -> None: 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): From 6cae3656ef3afd69d2ba8153cb240d83ef75217c Mon Sep 17 00:00:00 2001 From: j316chuck <13524881+j316chuck@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:44:41 +0000 Subject: [PATCH 7/7] Move SQLite contention repro out of unit coverage --- tests/tinker/test_external_future_store.py | 200 +-------------------- 1 file changed, 1 insertion(+), 199 deletions(-) diff --git a/tests/tinker/test_external_future_store.py b/tests/tinker/test_external_future_store.py index 3a58d98bc6..be88b1b0dc 100644 --- a/tests/tinker/test_external_future_store.py +++ b/tests/tinker/test_external_future_store.py @@ -4,19 +4,16 @@ import pytest import pytest_asyncio from sqlalchemy.ext.asyncio import create_async_engine -from sqlmodel import SQLModel, func, select +from sqlmodel import SQLModel from sqlmodel.ext.asyncio.session import AsyncSession -from starlette.requests import Request from skyrl.tinker import api, types from skyrl.tinker.config import EngineConfig from skyrl.tinker.db_models import ( CheckpointDB, CheckpointStatus, - FutureDB, ModelDB, RequestStatus, - SamplingSessionDB, SessionDB, enable_sqlite_wal, get_async_database_url, @@ -39,68 +36,6 @@ def _sample_input(seq_id: int) -> types.SampleInput: ) -class _CompletingForwarder: - def __init__(self, store: ExternalFutureStore): - self.store = store - - async def call_and_store_result( - self, - request_id: int, - sample_req, - model_id: str, - checkpoint_id: str, - base_model: str | None = None, - ) -> None: - await self.store.complete( - request_id, - types.SampleOutput(sequences=[]), - RequestStatus.COMPLETED, - ) - - -def _forward_backward_request(seq_id: int, db_write_lock: asyncio.Lock) -> Request: - body = ( - api.ForwardBackwardRequest( - model_id="model_a", - seq_id=seq_id, - forward_backward_input=api.ForwardBackwardInput( - data=[ - api.Datum( - model_input=api.ModelInput(chunks=[api.EncodedTextChunk(tokens=[1, 2])]), - loss_fn_inputs={ - "target_tokens": api.TensorData(data=[2, 3]), - "weights": api.TensorData(data=[1.0, 1.0]), - }, - ) - ], - loss_fn="cross_entropy", - ), - ) - .model_dump_json() - .encode() - ) - body_sent = False - - async def receive(): - nonlocal body_sent - if body_sent: - return {"type": "http.disconnect"} - body_sent = True - return {"type": "http.request", "body": body, "more_body": False} - - app = SimpleNamespace(state=SimpleNamespace(db_write_lock=db_write_lock)) - return Request( - { - "type": "http", - "method": "POST", - "path": "/api/v1/forward_backward", - "headers": [(b"content-type", b"application/json")], - "app": app, - }, - receive, - ) - - @pytest_asyncio.fixture() async def future_store(tmp_path): db_url = get_async_database_url(f"sqlite:///{tmp_path / 'tinker.db'}") @@ -117,139 +52,6 @@ async def future_store(tmp_path): await engine.dispose() -@pytest.mark.asyncio -async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(future_store): - store, engine, db_write_lock = future_store - forwarder = _CompletingForwarder(store) - sample_request = SimpleNamespace( - app=SimpleNamespace( - state=SimpleNamespace( - db_engine=engine, - external_future_store=store, - external_inference_client=forwarder, - forwarding_tasks=set(), - future_waiters={}, - engine_config=EngineConfig(base_model="model_a"), - db_write_lock=db_write_lock, - sampling_model_cache={}, - sampling_model_cache_lock=asyncio.Lock(), - validated_sampler_checkpoints=set(), - sampler_checkpoint_validation_lock=asyncio.Lock(), - ) - ), - headers={}, - ) - - async with AsyncSession(engine) as session: - session.add( - SessionDB( - session_id="session_a", - tags=[], - user_metadata={}, - sdk_version="test", - ) - ) - session.add( - SamplingSessionDB( - sampling_session_id="session_a", - session_id="session_a", - sampling_session_seq_id=0, - model_path="tinker://model_a/sampler_weights/weights_a", - ) - ) - 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() - - future_poller = asyncio.create_task( - api.poll_futures(engine, sample_request.app.state.future_waiters, poll_interval_sec=0.001) - ) - expected_sample = types.SampleOutput(sequences=[]).model_dump_json().encode() - try: - request_ids = [] - for wave in range(4): - - async def create_sample(index: int) -> int: - async with AsyncSession(engine) as session: - response = await api.asample( - api.SampleRequest( - prompt=api.ModelInput(chunks=[api.EncodedTextChunk(tokens=[index])]), - sampling_params=api.SamplingParams(temperature=0.0, max_tokens=1, seed=index), - sampling_session_id="session_a", - seq_id=wave * 512 + index, - ), - sample_request, - session, - ) - return int(response.request_id) - - async def create_training_future(index: int) -> None: - async with AsyncSession(engine) as session: - await api.forward_backward( - _forward_backward_request(wave * 512 + index, db_write_lock), - session, - ) - - async def heartbeat() -> None: - async with AsyncSession(engine) as session: - await api.session_heartbeat( - api.SessionHeartbeatRequest(session_id="session_a"), - sample_request, - session, - ) - - responses = await asyncio.gather( - *(create_sample(index) for index in range(512)), - *(create_training_future(index) for index in range(512)), - *(heartbeat() for _ in range(32)), - ) - request_ids = responses[:512] - retrievals = await asyncio.gather( - *( - api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_id)), sample_request) - for request_id in request_ids - ) - ) - assert all(response.body == expected_sample for response in retrievals) - assert not store._entries - - repeated = await api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_ids[-1])), sample_request) - assert repeated.body == expected_sample - finally: - future_poller.cancel() - await asyncio.gather(future_poller, return_exceptions=True) - - await store.flush() - async with AsyncSession(engine) as session: - persisted_by_type = dict( - (await session.exec(select(FutureDB.request_type, func.count()).group_by(FutureDB.request_type))).all() - ) - session_db = await session.get(SessionDB, "session_a") - - assert persisted_by_type[types.RequestType.EXTERNAL] == 2048 - assert persisted_by_type[types.RequestType.FORWARD_BACKWARD] == 2048 - assert session_db is not None - assert session_db.heartbeat_count == 128 - assert sample_request.app.state.validated_sampler_checkpoints == {("model_a", "weights_a")} - assert not sample_request.app.state.forwarding_tasks - - @pytest.mark.asyncio async def test_shutdown_waits_for_forwarding_tasks_before_closing_store(): release_forwarding = asyncio.Event()