Change asyncio.run_all() --> semaphore queue for sub-calls (alexzhang13#131)

* cap batching to 16

* add mock LM for tests, add test for semaphore

* remove mock lm in rlm/ and update into tests

* linting

* more linting

---------

Co-authored-by: Ubuntu <az@primeintellect.ai>

Author: alexzhang13

rlm/core/lm_handler.py

@@ -85,8 +85,14 @@ def _handle_batched(self, request: LMRequest, handler: "LMHandler") -> LMRespons
 
         start_time = time.perf_counter()
 
+        sem = asyncio.Semaphore(handler.batch_max_concurrent)
+
+        async def run_one(prompt: str):
+            async with sem:
+                return await client.acompletion(prompt)
+
         async def run_all():
-            tasks = [client.acompletion(prompt) for prompt in request.prompts]
+            tasks = [run_one(prompt) for prompt in request.prompts]
             return await asyncio.gather(*tasks)
 
         results = asyncio.run(run_all())
@@ -132,6 +138,7 @@ def __init__(
         host: str = "127.0.0.1",
         port: int = 0,  # auto-assign available port
         other_backend_client: BaseLM | None = None,
+        batch_max_concurrent: int = 16,
     ):
         self.default_client = client
         self.other_backend_client = other_backend_client
@@ -140,6 +147,7 @@ def __init__(
         self._server: ThreadingLMServer | None = None
         self._thread: Thread | None = None
         self._port = port
+        self.batch_max_concurrent = batch_max_concurrent
 
         self.register_client(client.model_name, client)
 

tests/mock_lm.py

@@ -1,27 +1,61 @@
+"""Mock LM for testing without a real model."""
+
+from collections.abc import Callable
+from typing import Any
+
 from rlm.clients.base_lm import BaseLM
 from rlm.core.types import ModelUsageSummary, UsageSummary
 
 
 class MockLM(BaseLM):
-    """Simple mock LM that echoes prompts."""
+    """
+    In-memory mock LM that implements BaseLM for testing.
 
-    def __init__(self):
-        super().__init__(model_name="mock-model")
+    - completion() / acompletion(): echo-style response by default, or use
+      optional responses list / response_fn if provided.
+    - get_usage_summary() / get_last_usage(): return usage suitable for tests.
+    """
 
-    def completion(self, prompt):
-        return f"Mock response to: {prompt[:50]}"
+    def __init__(
+        self,
+        model_name: str = "mock-model",
+        responses: list[str] | None = None,
+        response_fn: Callable[[str | dict[str, Any]], str] | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+        self._responses = list(responses) if responses is not None else None
+        self._response_fn = response_fn
+        self._call_count = 0
 
-    async def acompletion(self, prompt):
+    def completion(self, prompt: str | dict[str, Any]) -> str:
+        self._call_count += 1
+        if self._responses is not None:
+            if not self._responses:
+                raise IndexError("MockLM: no more responses in list")
+            return self._responses.pop(0)
+        if self._response_fn is not None:
+            return self._response_fn(prompt)
+        prompt_str = prompt if isinstance(prompt, str) else str(prompt)[:80]
+        return f"Mock response to: {prompt_str}"
+
+    async def acompletion(self, prompt: str | dict[str, Any]) -> str:
         return self.completion(prompt)
 
-    def get_usage_summary(self):
+    def get_usage_summary(self) -> UsageSummary:
         return UsageSummary(
             model_usage_summaries={
-                "mock-model": ModelUsageSummary(
-                    total_calls=1, total_input_tokens=10, total_output_tokens=10
+                self.model_name: ModelUsageSummary(
+                    total_calls=self._call_count,
+                    total_input_tokens=self._call_count * 10,
+                    total_output_tokens=self._call_count * 10,
                 )
             }
         )
 
-    def get_last_usage(self):
-        return self.get_usage_summary()
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=10,
+            total_output_tokens=10,
+        )

tests/test_lm_handler.py

@@ -0,0 +1,45 @@
+"""Tests for LMHandler using MockLM (no real LM required)."""
+
+from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
+from rlm.core.lm_handler import LMHandler
+from tests.mock_lm import MockLM
+
+
+def test_lm_handler_single_request():
+    """Single prompt request returns success and echo-style content."""
+    mock = MockLM(responses=["hello back"])
+    with LMHandler(client=mock) as handler:
+        request = LMRequest(prompt="hello")
+        response = send_lm_request(handler.address, request)
+    assert response.success
+    assert response.chat_completion is not None
+    assert response.chat_completion.response == "hello back"
+
+
+def test_lm_handler_batched_request():
+    """Batched prompts return one response per prompt in order."""
+    responses = [f"r{i}" for i in range(5)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=3) as handler:
+        prompts = [f"prompt-{i}" for i in range(5)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == 5
+    for i, resp in enumerate(result):
+        assert resp.success, resp.error
+        assert resp.chat_completion is not None
+        assert resp.chat_completion.response == f"r{i}"
+
+
+def test_lm_handler_batched_many_prompts_semaphore_cap():
+    """Many prompts complete successfully with semaphore limiting concurrency."""
+    # 50 prompts, max 4 concurrent: should still all complete
+    count = 50
+    responses = [f"resp-{i}" for i in range(count)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=4) as handler:
+        prompts = [f"p-{i}" for i in range(count)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == count
+    for i, resp in enumerate(result):
+        assert resp.success, (i, resp.error)
+        assert resp.chat_completion.response == f"resp-{i}"