[tinker] 15/n towards Kimi K2.6: continuous sampling for the Tinker API engine under colocate_all

[casper-hansen]

Part of the Kimi K2.6/K2.7 series. Stacked on #2031 (9/n) — its commits appear in this diff; review the last commit ("Continuous sampling: per-request future completion") for this PR's own changes.

What

Replaces the Tinker API engine's batch-completion sampling with per-request future completion under colocate_all, keeping the engine-lifecycle safety property via admission control instead of serialization:

• _ContinuousSampler (engine): a daemon thread running an asyncio loop. Each admitted sample request executes 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.
• The engine main loop admits pending samples whenever no blocking work exists, waking/bringing up engines on the main thread first (backend.prepare_for_sampling, so the sampler coroutines never touch engine lifecycle). Blocking work (forward/forward_backward, single requests, session cleanup that actually unloads models) drains in-flight samples first, then runs with the engines quiesced exactly as before. Admission filters stale requests first (serial-path parity), and no-op cleanup sweeps don't drain.
• Backend: prepare_for_sampling + sample_batch_async split out of the serial sample() path (shared validation; the async core reuses the loop-aware HTTP client with a persistent session on the sampler loop). RemoteInferenceClient._get_session closes the stale loop's connector synchronously on loop handoff.
• SKYRL_TINKER_CONTINUOUS_SAMPLING=0 restores the serial loop; backends without the async surface (jax) keep it automatically via hasattr feature detection.

CPU unit tests in tests/tinker/test_engine.py cover admission/drain ordering, stale-request filtering, and the env-var opt-out.

Why

Measured on our Kimi K2.7-Code agentic RL runs through the Tinker API (2x8xB300, colocated vLLM, ~64 lockstep rollouts): with batch-completion semantics each wave's turn latency was max(generation in wave) — a single 200s generation stalled 63 finished agents for its full tail, and GPU utilization between waves collapsed. The batching existed only to serialize sampling against ops that sleep the colocated engines; admission control preserves that property while letting every finished agent proceed immediately. The stale-request admission filter fixes a real bring-up failure we hit after a server restart: leftover sample requests from the previous server drove prepare_for_sampling into an engine build with no model registered, instead of failing fast with "model not loaded".

Made with Cursor

[gemini-code-assist]

Code Review

This pull request introduces continuous sampling support in the Tinker training engine, allowing individual sample requests to complete independently on a background daemon thread instead of waiting for a full batch wave. It also manages GPU memory more efficiently by offloading trainer models before waking or initializing inference engines. The review feedback highlights a critical thread-safety and loop-safety issue with the shared _session in RemoteInferenceClient, recommending a loop-safe dictionary approach. Additionally, the reviewer noted potential AttributeError crashes in skyrl_train_backend.py due to missing None checks on self._dispatch.

[gemini-code-assist]

Thread-Safety and Loop-Safety Issue with Shared _session

The RemoteInferenceClient is a shared instance across threads (the main thread and the background tinker-sampler thread). Accessing and mutating self._session concurrently without any synchronization is highly prone to race conditions.

Furthermore, closing and recreating self._session on loop handoff can disrupt active requests on other loops or cause RuntimeError: Task got Future attached to a different loop if any delayed callback or background task attempts to use the session.

Solution

Instead of sharing a single _session and constantly closing/recreating it, we can maintain a loop-safe dictionary of sessions using weakref.WeakKeyDictionary. This ensures:

1. Each event loop gets its own dedicated ClientSession.
2. Complete thread-safety and loop-safety without any race conditions.
3. Automatic cleanup of sessions when transient loops are garbage collected.
4. No need for fragile synchronous connector closing.

Note: A similar loop-safe pattern should also be applied to _get_semaphores to prevent similar event loop mismatch errors.

        if "_sessions" not in self.__dict__:
            import weakref
            self.__dict__["_sessions"] = weakref.WeakKeyDictionary()
        session = self._sessions.get(current_loop)
        if session is None or session.closed:
            connector = aiohttp.TCPConnector(
                limit=SKYRL_HTTP_CONNECTION_LIMIT,
                keepalive_timeout=2,
            )
            session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=None))
            self._sessions[current_loop] = session
        self._session = session

[gemini-code-assist]

Potential AttributeError when self._dispatch is None

At line 430, a defensive None check is performed on self._dispatch before calling offload_for_sampling(). However, at line 435, self._dispatch.set_inference_engine_client(...) is called without a None check. If self._dispatch is indeed None, this will raise an AttributeError and crash the engine initialization.

Please wrap this call in a None check as well.

Suggested change

	self._dispatch.set_inference_engine_client(self._inference_engine_client)
	if self._dispatch is not None:
	self._dispatch.set_inference_engine_client(self._inference_engine_client)

[gemini-code-assist]

Potential AttributeError when self._dispatch is None

In _wake_inference_engines_for_sampling, self._dispatch.offload_for_sampling() is called directly. To prevent potential AttributeError crashes (consistent with the defensive checks in _ensure_inference_engines), please add a None check for self._dispatch before calling offload_for_sampling().

Suggested change

	self._dispatch.offload_for_sampling()
	if self._dispatch is not None:
	self._dispatch.offload_for_sampling()

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 58dfc1c. Configure here.}

[cursor]

Session handoff does not close connectors

Medium Severity

Loop handoff in _get_session calls TCPConnector.close() without awaiting it. In aiohttp 3.14 that method is a coroutine, so the old connector is not closed; the exception handler also swallows failures from closing a live sampler-loop connector on the engine thread's throwaway loop. Continuous sampling leaves the session open (close_client=False) and sleep/wake plus the two asyncio.run wake calls hit this path every sample/train cycle, leaking sockets and still producing unclosed-session warnings.

Additional Locations (1)

• skyrl/backends/skyrl_train_backend.py#L825-L828

 

^{Reviewed by Cursor Bugbot for commit 58dfc1c. Configure here.}