[tinker] Continuous sampling: per-request future completion under colocate_all

The engine's serial loop batched all pending sample requests and completed
every future only when the whole batch finished, so with ~64 lockstep
agentic rollouts each wave's turn latency was max(generation in wave): a
single 200s turn stalled 63 finished agents for its full tail. The batch
semantics existed only to serialize sampling against ops that sleep the
colocated engines.

Keep that safety property with admission control instead of serialization:

- _ContinuousSampler: daemon thread + asyncio loop; each admitted sample
  runs as its own coroutine (backend.sample_batch_async -> per-request
  vLLM HTTP; vLLM continuous-batches internally) and its future is written
  the moment its own generation completes.
- Engine loop admits pending samples whenever no blocking work exists;
  engines are brought up/woken on the main thread first
  (backend.prepare_for_sampling). Blocking work (forward/forward_backward,
  single requests, stale-session model unloads) drains in-flight samples
  first -- generations keep decoding while the blocking op waits, then run
  with the engines quiesced exactly as before. Deferred samples stay
  PENDING and are admitted next iteration, never as a convoy batch.
- Admission runs _filter_valid_requests first (serial-path parity): stale
  sample requests left over from a previous server fail with "model not
  loaded" instead of driving prepare_for_sampling into an engine build
  with no model registered.
- No-op session-cleanup sweeps don't drain (only sweeps that actually
  unload models do), so the periodic reaper can't reintroduce a stall.
- backend: prepare_for_sampling + sample_batch_async (validation shared
  with the serial path; the async core reuses the loop-aware HTTP client,
  keeping the session open on the sampler's persistent loop).
  RemoteInferenceClient._get_session closes the stale loop's connector
  synchronously on loop handoff so sockets release immediately.
- SKYRL_TINKER_CONTINUOUS_SAMPLING=0 restores the serial loop; backends
  without the async surface (jax) keep it automatically.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: casper-hansen

skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py

@@ -289,7 +289,19 @@ async def _get_session(self) -> aiohttp.ClientSession:
         # aiohttp.ClientSession is tied to the event loop.
         current_loop = asyncio.get_running_loop()
         if self._session is not None and not self._session.closed and self._session.loop != current_loop:
-            # Event loop changed - the old session is unusable (bound to a dead loop).
+            # Event loop changed - the old session is unusable (bound to another,
+            # possibly dead, loop) and cannot be aclose()'d from here. Close its
+            # connector synchronously (sync API in aiohttp 3.x) so the sockets are
+            # released now and GC doesn't log "Unclosed client session" errors —
+            # this handoff happens routinely between the engine main thread's
+            # transient loops (wake/sleep/adapter loads) and the continuous
+            # sampler's persistent loop.
+            old_connector = getattr(self._session, "_connector", None)
+            try:
+                if old_connector is not None:
+                    old_connector.close()
+            except Exception:
+                pass
             self._session = None
         if self._session is None or self._session.closed:
             # keepalive_timeout must be shorter than the server's timeout_keep_alive

skyrl/backends/skyrl_train_backend.py

@@ -1049,16 +1049,38 @@ def sample(
         have been updated.
         """
         # 1. Ensure inference engines are initialized and awake
+        self.prepare_for_sampling()
+
+        # 2. Validate model ids, 3. dispatch to the sampling path.
+        errors = self._validate_sample_models(prepared_batch)
+        if errors is not None:
+            return errors
+        return self._sample_with_remote_client(prepared_batch)
+
+    def prepare_for_sampling(self) -> None:
+        """Ensure inference engines exist and are awake for sampling.
+
+        Idempotent and cheap when the engines are already up and awake. The
+        Tinker engine's continuous sampler calls this from its main loop
+        before admitting sample requests to the background executor, so the
+        executor's coroutines never have to touch engine lifecycle (which is
+        not thread-safe) — they only issue data-plane HTTP calls.
+        """
         self._ensure_inference_engines()
         self._wake_inference_engines_for_sampling()
 
-        # 2. Validate every model_id in the batch is a known policy. Multi-LoRA
-        # mixes adapters in one batched sample call (the engine batches across
-        # model_ids in find_batchable_sample); we route each request via the
-        # `model` field in _sample_with_remote_client below. An empty model_id
-        # is base-model sampling (create_sampling_client(base_model=...): the
-        # API maps it to model_id "") and must not be treated as unknown --
-        # _sample_with_remote_client routes it to the served base model name.
+    def _validate_sample_models(
+        self, prepared_batch: types.PreparedSampleBatch
+    ) -> dict[str, types.ErrorResponse] | None:
+        """Validate every model_id in the batch is a known policy; None if OK.
+
+        Multi-LoRA mixes adapters in one batched sample call (the engine
+        batches across model_ids in find_batchable_sample); we route each
+        request via the `model` field in _sample_with_remote_client. An empty
+        model_id is base-model sampling (create_sampling_client(base_model=...):
+        the API maps it to model_id "") and must not be treated as unknown --
+        _sample_with_remote_client routes it to the served base model name.
+        """
         unique_models = set(prepared_batch.all_model_ids)
         unknown = [mid for mid in unique_models if mid and mid not in self._model_ids_to_role]
         if unknown:
@@ -1073,15 +1095,45 @@ def sample(
                 status="error",
             )
             return {req_id: error for req_id, *_ in prepared_batch.request_batch_slices}
+        return None
 
-        # 3. Dispatch to the sampling path
-        return self._sample_with_remote_client(prepared_batch)
+    async def sample_batch_async(
+        self, prepared_batch: types.PreparedSampleBatch
+    ) -> dict[str, types.SampleOutput | types.ErrorResponse]:
+        """Awaitable sample path for the engine's continuous sampler.
+
+        The caller (engine main loop) must have called prepare_for_sampling()
+        before scheduling this, and must not sleep the engines, swap adapters,
+        or tear down the runtime while calls are in flight (the engine drains
+        the sampler before such ops). This coroutine runs on the sampler
+        thread's event loop and touches only read-only backend state plus the
+        data-plane HTTP client, which recreates its session/semaphores per
+        event loop.
+        """
+        errors = self._validate_sample_models(prepared_batch)
+        if errors is not None:
+            return errors
+        return await self._sample_with_remote_client_async(prepared_batch, close_client=False)
 
     def _sample_with_remote_client(
         self,
         prepared_batch: types.PreparedSampleBatch,
     ) -> dict[str, types.SampleOutput | types.ErrorResponse]:
-        """Sample using RemoteInferenceClient, forwarding model input chunks directly."""
+        """Sync wrapper over the async sampling core (serial engine-loop path)."""
+        return asyncio.run(self._sample_with_remote_client_async(prepared_batch, close_client=True))
+
+    async def _sample_with_remote_client_async(
+        self,
+        prepared_batch: types.PreparedSampleBatch,
+        close_client: bool,
+    ) -> dict[str, types.SampleOutput | types.ErrorResponse]:
+        """Sample using RemoteInferenceClient, forwarding model input chunks directly.
+
+        ``close_client`` closes the HTTP session when done: the serial path
+        runs each batch on a throwaway event loop (asyncio.run), so its
+        session must not outlive the loop. The continuous sampler runs on a
+        persistent loop and reuses the session across requests instead.
+        """
 
         # Resolve the inference-engine model name per request. With multi-LoRA
         # the adapter name on vLLM IS the Tinker model_id (registered by
@@ -1141,10 +1193,11 @@ async def sample_all():
             try:
                 return await asyncio.gather(*tasks, return_exceptions=True)
             finally:
-                await self._inference_engine_client.aclose()
+                if close_client:
+                    await self._inference_engine_client.aclose()
 
-        sample_outputs = asyncio.run(sample_all())
-        logger.info(f"Collected {len(sample_outputs)} sample outputs")
+        sample_outputs = await sample_all()
+        logger.debug(f"Collected {len(sample_outputs)} sample outputs")
         return self._aggregate_sample_results(prepared_batch, sample_outputs)
 
     def _aggregate_sample_results(
@@ -1153,7 +1206,7 @@ def _aggregate_sample_results(
         sample_outputs: list,
     ) -> dict[str, types.SampleOutput | types.ErrorResponse]:
         """Convert sample outputs to Tinker format."""
-        logger.info(f"Aggregating sample results for {len(sample_outputs)} samples")
+        logger.debug(f"Aggregating sample results for {len(sample_outputs)} samples")
 
         def _extract_sequences(output):
             """Yield (tokens, logprobs, stop_reason) from a single sample output."""