perf: run native sampling on managed loop - #2
Conversation
f4e2ea7 to
972da36
Compare
| await async_client.aclose() | ||
| self.deployment_sampler._sync_client.close() | ||
|
|
||
| async def _shutdown_loop(self) -> None: |
There was a problem hiding this comment.
This is reasonable for general hardening, but perhaps a bit complicated for our use case given that the exit stack calls this after the run_training() function has unwound. The modal process will exit shortly after. The important cleanup for us service.close(), which cleans up remote Fireworks resources which is handled separately.
| self._loop: asyncio.AbstractEventLoop | None = None | ||
| self._loop_thread: threading.Thread | None = None | ||
| self._loop_lock = threading.Lock() | ||
| self._loop_lock = threading.RLock() |
There was a problem hiding this comment.
nit: arguably using an RLock here is a smell as it obscures the critical sections. I think the code can be refactored to keep a standard lock so each function has to explicitly acquire the lock to document that part of the function must be atomic.
|
|
||
| try: | ||
| result = asyncio.run(_run()) | ||
| assert captured["thread_name"] == "fireworks-sampling-client" |
There was a problem hiding this comment.
is this a redundant assertion to 1717?
| "logprobs": True, | ||
| "request_marker": marker, | ||
| } | ||
| assert result == [completion] |
There was a problem hiding this comment.
is the assertion between 1730 and 1731 redundant?
Summary
sample_with_prompt_tokensthroughFiretitanSamplingClientWhy
The controlled APEX ablation identified the dedicated sampling loop as the load-bearing intervention. This keeps Fireworks request and response processing off the rollout coordinator loop without adding custom transport or parser behavior.
Scope
This PR is based directly on SDK v1.2.9. It does not change
DeploymentSamplerparsing, SSE transport, trainer or scheduler behavior, model parameters, loss, retries, routing, or raw-sampler behavior.Verification
sampling.pyand_sse.pyare byte-identical to v1.2.9