diff --git a/src/olmo_eval/inference/providers/sglang.py b/src/olmo_eval/inference/providers/sglang.py index 71ea10acd..c7cfc8ffe 100644 --- a/src/olmo_eval/inference/providers/sglang.py +++ b/src/olmo_eval/inference/providers/sglang.py @@ -6,8 +6,12 @@ import math import os import sys +import threading +from collections.abc import Coroutine +from concurrent.futures import Future +from dataclasses import dataclass from importlib import import_module -from typing import Any, Literal +from typing import Any, Literal, TypeVar from jinja2 import TemplateError @@ -22,6 +26,8 @@ logger = get_logger(__name__) +_T = TypeVar("_T") + # Keep this allowlist aligned with the pinned olmo-sglang arithmetic options. # Record exact startup values (None means unset), rather than duplicating the # plugin's defaults or claiming that an automatic/per-layer choice is resolved. @@ -63,6 +69,17 @@ def _require_original_logprobs() -> None: ) +@dataclass +class _PreparedBatch: + """A generate call's requests, prepared for one engine call.""" + + params: SamplingParams + grouped: list[list[LMOutput]] + prepared: list[tuple[dict[str, Any], dict[str, Any]]] + positions: list[int] + engine_kwargs: dict[str, Any] + + class SGLangProvider(InferenceProvider): """Run generation with a locally installed SGLang engine. @@ -216,6 +233,37 @@ def __init__( self._closed = False self._capacity_checked = False self._previous_batch: tuple[int, int, int, int] | None = None + self._loop_lock = threading.Lock() + self._loop_thread: threading.Thread | None = None + + def _on_engine_loop(self, coroutine: Coroutine[Any, Any, _T]) -> Future[_T]: + """Run a coroutine on the engine's event loop, from any thread but that loop's. + + SGLang binds its tokenizer manager to ``engine.loop``, and ``Engine.generate`` drives + that loop with ``run_until_complete``, so two threads calling it at once fail with + "this event loop is already running". One daemon thread instead runs the loop for the + provider's lifetime and every later engine call is submitted to it, so concurrent + calls reach the scheduler together and are batched there. Startup calls that ran + before this thread existed (``get_server_info``) used the same loop. + """ + with self._loop_lock: + if self._loop_thread is None: + thread = threading.Thread( + target=self.engine.loop.run_forever, name="sglang-engine-loop", daemon=True + ) + thread.start() + self._loop_thread = thread + return asyncio.run_coroutine_threadsafe(coroutine, self.engine.loop) + + def _server_info(self) -> dict[str, Any]: + """The fields of ``Engine.get_server_info()`` that the capacity hint reads. + + ``Engine.get_server_info`` drives ``engine.loop`` itself, which fails while the + provider's thread runs that loop, so ask the tokenizer manager on the loop instead. + """ + manager = self.engine.tokenizer_manager + states = self._on_engine_loop(manager.get_internal_state()).result() + return {"dp_size": getattr(manager.server_args, "dp_size", 1), "internal_states": states} def get_tokenizer(self) -> Any: """Return the engine's tokenizer.""" @@ -340,7 +388,7 @@ def _consider_larger_batches(self, outputs: list[LMOutput], params: SamplingPara # Telemetry is optional across SGLang versions. A failed diagnostic must # never discard successful generations or turn into an inference retry. try: - info = self.engine.get_server_info() + info = self._server_info() except Exception: logger.debug("SGLang capacity diagnostic unavailable", exc_info=True) return @@ -387,13 +435,23 @@ def generate( sampling_params: SamplingParams | None = None, ) -> list[list[LMOutput]]: """Generate completions and retain sampled-token probabilities and engine metadata.""" + batch = self._prepare_batch(requests, sampling_params) + if batch is None: + return [] + outputs = self._on_engine_loop(self.engine.async_generate(**batch.engine_kwargs)).result() + return self._finish_batch(batch, outputs) + + def _prepare_batch( + self, requests: list[LMRequest], sampling_params: SamplingParams | None + ) -> _PreparedBatch | None: + """Validate and render the requests; None when there is nothing to send.""" params = self._default_sampling_params(sampling_params) if self._closed: raise RuntimeError("SGLangProvider is closed") sampling = self._build_sampling_params(params) _require_original_logprobs() if not requests: - return [] + return None grouped: list[list[LMOutput]] = [[] for _ in requests] prepared = [] positions = [] @@ -419,12 +477,22 @@ def generate( } if "image_data" in inputs and not any(inputs["image_data"]): del inputs["image_data"] - outputs = self.engine.generate( - **inputs, - sampling_params=sampling, - return_logprob=True, - top_logprobs_num=params.logprobs if params.logprobs is not None else 1, + return _PreparedBatch( + params=params, + grouped=grouped, + prepared=prepared, + positions=positions, + engine_kwargs={ + **inputs, + "sampling_params": sampling, + "return_logprob": True, + "top_logprobs_num": params.logprobs if params.logprobs is not None else 1, + }, ) + + def _finish_batch(self, batch: _PreparedBatch, outputs: Any) -> list[list[LMOutput]]: + """Convert the engine's completions and place them back in request order.""" + params, grouped, prepared = batch.params, batch.grouped, batch.prepared if not isinstance(outputs, list) or len(outputs) != len(prepared) * params.num_samples: raise RuntimeError("SGLang returned an unexpected number of completions") converted = [self._convert_output(output) for output in outputs] @@ -432,7 +500,7 @@ def generate( raise RuntimeError( f"SGLang aborted every generation: {converted[0].metadata[REQUEST_ERROR_KEY]}" ) - for offset, position in enumerate(positions): + for offset, position in enumerate(batch.positions): samples = converted[offset * params.num_samples : (offset + 1) * params.num_samples] failures = [ o.metadata[REQUEST_ERROR_KEY] for o in samples if REQUEST_ERROR_KEY in o.metadata @@ -562,13 +630,29 @@ def logprobs( async def agenerate( self, requests: list[LMRequest], sampling_params: SamplingParams | None = None ) -> list[list[LMOutput]]: - """Run sequential engine calls without blocking the harness event loop.""" - return await asyncio.to_thread(self.generate, requests, sampling_params) + """Generate without blocking the harness loop; concurrent calls share engine batches. + + Each call awaits its own engine request instead of holding a thread, so a scaffold's + concurrent conversations are all in the engine at once. Rendering and conversion run + in worker threads, as ``generate`` did here before. + """ + batch = await asyncio.to_thread(self._prepare_batch, requests, sampling_params) + if batch is None: + return [] + outputs = await asyncio.wrap_future( + self._on_engine_loop(self.engine.async_generate(**batch.engine_kwargs)) + ) + return await asyncio.to_thread(self._finish_batch, batch, outputs) def close(self) -> None: """Release engine processes when the evaluation worker finishes.""" if not self._closed: self.engine.shutdown() + with self._loop_lock: + thread, self._loop_thread = self._loop_thread, None + if thread is not None: + self.engine.loop.call_soon_threadsafe(self.engine.loop.stop) + thread.join(timeout=30) self._closed = True def shutdown(self) -> None: diff --git a/tests/providers/test_sglang_provider.py b/tests/providers/test_sglang_provider.py index e4bf4212e..068b437e5 100644 --- a/tests/providers/test_sglang_provider.py +++ b/tests/providers/test_sglang_provider.py @@ -4,6 +4,7 @@ import json import os import sys +import threading from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -38,7 +39,21 @@ def setup_engine(monkeypatch): tokenizer.apply_chat_template.return_value = "question" tokenizer.encode.return_value = [10, 20] tokenizer.decode.side_effect = lambda ids: f"token{ids[0]}" - engine = Mock(tokenizer_manager=SimpleNamespace(tokenizer=tokenizer)) + engine = Mock( + tokenizer_manager=SimpleNamespace( + tokenizer=tokenizer, + server_args=SimpleNamespace(dp_size=1), + get_internal_state=AsyncMock(return_value=None), + ), + loop=asyncio.new_event_loop(), + ) + + async def async_generate(**kwargs): + # The provider calls the engine's native async API. It records through the sync + # mock, so each test reads one mock for what the engine received. + return engine.generate(**kwargs) + + engine.async_generate = async_generate def generate(**kwargs): return [ @@ -58,7 +73,18 @@ def generate(**kwargs): engine.generate.side_effect = generate constructor = Mock(return_value=engine) monkeypatch.setitem(sys.modules, "sglang", SimpleNamespace(Engine=constructor)) - return engine, constructor, tokenizer + yield engine, constructor, tokenizer + _stop_engine_loop(engine.loop) + + +def _stop_engine_loop(loop): + """Stop a provider's engine-loop thread that a test left running, then close the loop.""" + if loop.is_running(): + loop.call_soon_threadsafe(loop.stop) + for thread in threading.enumerate(): + if thread.name == "sglang-engine-loop": + thread.join(timeout=5) + loop.close() def test_provider_factory_and_config(setup_engine, monkeypatch): @@ -312,6 +338,85 @@ def test_bos_policy_multiple_samples_and_async(setup_engine): tokenizer.encode.assert_called_with("abc", add_special_tokens=False) +class LoopBoundEngine: + """An engine with SGLang's event-loop contract. + + Like ``sglang.Engine``, ``generate`` drives ``self.loop`` with ``run_until_complete``, + so two threads calling it at once fail with "this event loop is already running". + ``async_generate`` must run on ``self.loop`` and returns only once ``expected`` + requests are in flight together, so a path that serializes the calls times out. + """ + + def __init__(self, tokenizer, expected): + self.loop = asyncio.new_event_loop() + self.tokenizer_manager = SimpleNamespace(tokenizer=tokenizer) + self.expected = expected + self.in_flight = 0 + self.peak = 0 + self.all_in_flight = asyncio.Event() + self.shutdown = Mock() + + def generate(self, **kwargs): + return self.loop.run_until_complete(self.async_generate(**kwargs)) + + async def async_generate(self, *, prompt, sampling_params, return_logprob, top_logprobs_num): + assert asyncio.get_running_loop() is self.loop + assert return_logprob and top_logprobs_num == 1 + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + if self.in_flight >= self.expected: + self.all_in_flight.set() + try: + await asyncio.wait_for(self.all_in_flight.wait(), timeout=10) + finally: + self.in_flight -= 1 + return [ + { + "text": f"answer to {text}", + "meta_info": { + "output_token_logprobs": [[-0.5, 17, None]], + "output_top_logprobs": [[[-0.5, 17, None]]], + "completion_tokens": 1, + "prompt_tokens": 3, + "finish_reason": {"type": "stop", "matched": 17}, + }, + } + for text in prompt + ] + + +def test_concurrent_agenerate_calls_are_batched_on_one_engine_loop(setup_engine): + _, constructor, tokenizer = setup_engine + engine = LoopBoundEngine(tokenizer, expected=8) + constructor.return_value = engine + provider = SGLangProvider("model") + prompts = [f"question {index}" for index in range(8)] + + async def conversations(): + return await asyncio.gather( + *( + provider.agenerate([LMRequest(RequestType.COMPLETION, prompt=prompt)]) + for prompt in prompts + ) + ) + + outputs = asyncio.run(conversations()) + assert [group[0][0].text for group in outputs] == [f"answer to {p}" for p in prompts] + assert engine.peak == 8 + assert all(group[0][0].metadata["finish_reason"] == "stop" for group in outputs) + + # The sync path shares the running loop rather than driving it a second time. + engine.expected = 1 + assert provider.generate([LMRequest(RequestType.COMPLETION, prompt="sync")])[0][0].text == ( + "answer to sync" + ) + provider.close() + engine.shutdown.assert_called_once_with() + assert not engine.loop.is_running() + assert not any(t.name == "sglang-engine-loop" for t in threading.enumerate()) + engine.loop.close() + + def test_greedy_uncapped_and_empty(setup_engine): engine, _, _ = setup_engine provider = SGLangProvider("model") @@ -620,21 +725,28 @@ def capacity_info(capacity=512, tokens=1000000): } +def serve_capacity(engine, info): + """Answer the capacity hint's engine queries with ``info``, a get_server_info result.""" + info = info if isinstance(info, dict) else {} + engine.tokenizer_manager.server_args.dp_size = info.get("dp_size", 1) + engine.tokenizer_manager.get_internal_state.return_value = info.get("internal_states") + + def test_concurrency_hint_once_without_changing_generation(setup_engine, caplog): engine, _, _ = setup_engine - engine.get_server_info.return_value = capacity_info() + serve_capacity(engine, capacity_info()) provider = SGLangProvider("model") requests = [LMRequest(RequestType.COMPLETION, prompt="test")] * 64 params = SamplingParams(max_tokens=100) first = provider.generate(requests, params) assert "Consider increasing concurrency" not in caplog.text - engine.get_server_info.assert_not_called() + engine.tokenizer_manager.get_internal_state.assert_not_called() assert provider.generate(requests, params) == first assert provider.generate(requests, params) == first assert caplog.text.count("Consider increasing concurrency") == 1 assert "effective engine limit is 512" in caplog.text assert "-o batching.chunk_size=128" in caplog.text - engine.get_server_info.assert_called_once_with() + engine.tokenizer_manager.get_internal_state.assert_awaited_once_with() assert all(call == engine.generate.call_args_list[0] for call in engine.generate.call_args_list) @@ -653,7 +765,7 @@ def test_concurrency_hint_once_without_changing_generation(setup_engine, caplog) ) def test_concurrency_hint_requires_resolved_capacity(setup_engine, caplog, info): engine, _, _ = setup_engine - engine.get_server_info.return_value = info + serve_capacity(engine, info) provider = SGLangProvider("model") requests = [LMRequest(RequestType.COMPLETION, prompt="test")] * 64 for _ in range(2): @@ -663,13 +775,13 @@ def test_concurrency_hint_requires_resolved_capacity(setup_engine, caplog, info) def test_concurrency_hint_counts_samples_and_ignores_small_or_unequal_batches(setup_engine, caplog): engine, _, _ = setup_engine - engine.get_server_info.return_value = capacity_info() + serve_capacity(engine, capacity_info()) provider = SGLangProvider("model") request = LMRequest(RequestType.COMPLETION, prompt="test") params = SamplingParams(max_tokens=100, num_samples=4) for count in (2, 2, 16, 32): provider.generate([request] * count, params) - engine.get_server_info.assert_not_called() + engine.tokenizer_manager.get_internal_state.assert_not_called() provider.generate([request] * 32, params) assert "supplied 128 sequences" in caplog.text assert "-o batching.chunk_size=64" in caplog.text @@ -678,7 +790,7 @@ def test_concurrency_hint_counts_samples_and_ignores_small_or_unequal_batches(se def test_concurrency_diagnostic_failure_preserves_outputs(setup_engine, caplog): engine, _, _ = setup_engine - engine.get_server_info.side_effect = RuntimeError("telemetry unavailable") + engine.tokenizer_manager.get_internal_state.side_effect = RuntimeError("telemetry unavailable") provider = SGLangProvider("model") requests = [LMRequest(RequestType.COMPLETION, prompt="test")] * 64 first = provider.generate(requests, SamplingParams(max_tokens=100)) @@ -696,12 +808,12 @@ def retracted(**kwargs): return outputs engine.generate.side_effect = retracted - engine.get_server_info.return_value = capacity_info() + serve_capacity(engine, capacity_info()) provider = SGLangProvider("model") requests = [LMRequest(RequestType.COMPLETION, prompt="test")] * 64 for _ in range(2): provider.generate(requests, SamplingParams(max_tokens=100)) - engine.get_server_info.assert_not_called() + engine.tokenizer_manager.get_internal_state.assert_not_called() assert "Consider increasing concurrency" not in caplog.text @@ -716,7 +828,7 @@ def with_prompt_length(**kwargs): return outputs engine.generate.side_effect = with_prompt_length - engine.get_server_info.return_value = capacity_info(tokens=1000000) + serve_capacity(engine, capacity_info(tokens=1000000)) provider = SGLangProvider("model") requests = [LMRequest(RequestType.COMPLETION, prompt="test")] * 64 provider.generate(requests, SamplingParams(max_tokens=100)) @@ -727,18 +839,18 @@ def with_prompt_length(**kwargs): def test_small_batch_breaks_consecutive_capacity_observations(setup_engine, caplog): engine, _, _ = setup_engine - engine.get_server_info.return_value = capacity_info() + serve_capacity(engine, capacity_info()) provider = SGLangProvider("model") request = LMRequest(RequestType.COMPLETION, prompt="test") for count in (64, 1, 64): provider.generate([request] * count, SamplingParams(max_tokens=100)) - engine.get_server_info.assert_not_called() + engine.tokenizer_manager.get_internal_state.assert_not_called() assert "Consider increasing concurrency" not in caplog.text def test_uncapped_generation_abstains_from_capacity_estimate(setup_engine, caplog): engine, _, _ = setup_engine - engine.get_server_info.return_value = capacity_info() + serve_capacity(engine, capacity_info()) provider = SGLangProvider("model") requests = [LMRequest(RequestType.COMPLETION, prompt="test")] * 64 # Include a bounded call first to ensure an uncapped call resets history. @@ -748,7 +860,7 @@ def test_uncapped_generation_abstains_from_capacity_estimate(setup_engine, caplo assert len(outputs) == 64 assert engine.generate.call_args.kwargs["sampling_params"]["max_new_tokens"] is None provider.generate(requests, SamplingParams(max_tokens=100)) - engine.get_server_info.assert_not_called() + engine.tokenizer_manager.get_internal_state.assert_not_called() assert "Consider increasing concurrency" not in caplog.text