From 6134c7ce63be662a7dff01aca5ecb7f8036663ac Mon Sep 17 00:00:00 2001 From: soufianeamini Date: Thu, 27 Aug 2026 10:04:19 +0100 Subject: [PATCH] Computer improvements --- .../_computers/_async_callables_computers.py | 91 +++++++------------ .../_computers/test_evaluator_concurrency.py | 40 -------- 2 files changed, 32 insertions(+), 99 deletions(-) diff --git a/pyagentspec/src/pyagentspec/evaluation/_computers/_async_callables_computers.py b/pyagentspec/src/pyagentspec/evaluation/_computers/_async_callables_computers.py index f583aef8..f6354b23 100644 --- a/pyagentspec/src/pyagentspec/evaluation/_computers/_async_callables_computers.py +++ b/pyagentspec/src/pyagentspec/evaluation/_computers/_async_callables_computers.py @@ -12,9 +12,21 @@ """ import json -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Generic, Hashable, Tuple, TypeVar +from contextlib import nullcontext +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Dict, + Generic, + Hashable, + Tuple, + TypeVar, +) import anyio +from anyio.abc import TaskStatus from pyagentspec._lazy_loader import LazyLoader from pyagentspec.evaluation.datasets.dataset import Dataset @@ -60,26 +72,25 @@ def __init__( self.dataset = dataset self.callables = callables self.max_concurrency = max_concurrency - if max_concurrency == -1: - self.semaphore = None - else: - self.semaphore = anyio.Semaphore(max_concurrency) + self.limiter = ( + anyio.CapacityLimiter(max_concurrency) if max_concurrency != -1 else nullcontext() + ) self._registry = _AsyncRegistry[Tuple[Any, str], T]() - async def _compute(self, sample_id: Any, callable_id: str) -> None: + async def _compute( + self, + sample_id: Any, + callable_id: str, + task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED, + ) -> None: """Run a single callable against a dataset sample and store the result.""" - # Fetch the sample lazily so IO is naturally parallelised by the caller. - sample = await self.dataset.get_sample(sample_id) - result = await self.callables[callable_id](**sample) - await self._registry.register((sample_id, callable_id), result) - - async def _queue(self, sample_id: Any, callable_id: str) -> None: - """Wrapper that honours the semaphore before delegating to ``_compute``.""" - if self.semaphore is not None: - async with self.semaphore: - await self._compute(sample_id, callable_id) - else: - await self._compute(sample_id, callable_id) + async with self.limiter: + task_status.started() + + # Fetch the sample lazily so IO is naturally parallelised by the caller. + sample = await self.dataset.get_sample(sample_id) + result = await self.callables[callable_id](**sample) + await self._registry.register((sample_id, callable_id), result) async def run(self) -> Dict[Tuple[Hashable, str], T]: """Kick off all pending computations and return the populated registry.""" @@ -88,48 +99,10 @@ async def run(self) -> Dict[Tuple[Hashable, str], T]: if not metrics_names: return {} - # For "unlimited" concurrency we still spawn one task per work item since callers - # explicitly opted out of concurrency caps. The producer/worker pattern below - # is primarily meant to prevent memory blow-ups when a bounded concurrency limit is used. - if self.semaphore is None: - async with anyio.create_task_group() as tg: - async for sample_id in self.dataset.ids(): - for metric_name in metrics_names: - tg.start_soon(self._queue, sample_id, metric_name) - return self._registry.store - - # Avoid spawning one task per (sample, metric) pair: for large datasets - # that can create millions of tasks and consume large amounts of memory. - # - # Instead, use a producer/worker pattern: - # - one producer enumerates dataset sample ids and enqueues work items - # - N workers consume items from the queue and run computations - - num_workers = max(1, self.max_concurrency) - queue_max_size = max(1, num_workers * self._QUEUE_BUFFER_FACTOR) - work_queue: anyio.abc.ObjectSendStream[Tuple[Any, str]] - receive_stream: anyio.abc.ObjectReceiveStream[Tuple[Any, str]] - work_queue, receive_stream = anyio.create_memory_object_stream(queue_max_size) - - async def producer() -> None: - async with work_queue: - async for sample_id in self.dataset.ids(): - for metric_name in metrics_names: - await work_queue.send((sample_id, metric_name)) - - async def worker(worker_id: int) -> None: - del worker_id - while True: - try: - sample_id, metric_name = await receive_stream.receive() - except anyio.EndOfStream: - return - await self._queue(sample_id, metric_name) - async with anyio.create_task_group() as tg: - tg.start_soon(producer) - for i in range(num_workers): - tg.start_soon(worker, i) + async for sample_id in self.dataset.ids(): + for metric_name in metrics_names: + await tg.start(self._compute, sample_id, metric_name) return self._registry.store diff --git a/pyagentspec/tests/evaluation/_computers/test_evaluator_concurrency.py b/pyagentspec/tests/evaluation/_computers/test_evaluator_concurrency.py index 8b7a7cb3..0a3d21ba 100644 --- a/pyagentspec/tests/evaluation/_computers/test_evaluator_concurrency.py +++ b/pyagentspec/tests/evaluation/_computers/test_evaluator_concurrency.py @@ -121,46 +121,6 @@ async def test_unlimited_concurrency() -> None: assert num_runnings_sequence[-i - 1] == i -@pytest.mark.anyio -async def test_run_does_not_spawn_one_task_per_item() -> None: - """ - Ensure ``_AsyncCallablesComputer.run`` does not create O(N) tasks. - This is a regression test for memory blow-ups when datasets are large. - """ - - class CountingTaskGroup: - def __init__(self, max_allowed: int) -> None: - self.max_allowed = max_allowed - self.started = 0 - - async def __aenter__(self) -> "CountingTaskGroup": - return self - - async def __aexit__(self, exc_type, exc, tb) -> None: - return None - - def start_soon(self, func, *args) -> None: - self.started += 1 - assert self.started <= self.max_allowed - - dataset = Dataset.from_dict([{"dummy_arg": i} for i in range(10000)]) - callables = {"dummy_callable": (lambda **kwargs: asyncio.sleep(0))} - computer = _AsyncCallablesComputer( - dataset=dataset, - callables=callables, - max_concurrency=10, - ) - - import anyio # imported here to keep the patch localized to this test - - original = anyio.create_task_group - try: - anyio.create_task_group = lambda: CountingTaskGroup(max_allowed=1 + 10) - await computer.run() - finally: - anyio.create_task_group = original - - @pytest.mark.anyio @pytest.mark.parametrize("max_concurrency", [5, 10, 20]) async def test_max_concurrency_is_respected(max_concurrency: int) -> None: