Method handle_single

tests/benchmark/benchmark_store.py:307–337

(task_index: int)

            await tracker.handle_progress(progress_made=False, pending_rollout_ids=pending_ids, store=store)

        async def handle_single(task_index: int) -> None:
            task_name = f"task-{task_index}"
            async with semaphore:
                console.print(f"Submitting task {task_index} of {total_tasks}")
                await store.add_resources(
                    {
                        "llm": agl.LLM(
                            endpoint=f"http://localhost:{task_index}/v1",
                            model=f"test-model-{task_index}",
                        )
                    }
                )
                rollout = await store.enqueue_rollout(input=task_name, mode="train")
                rollout_id = rollout.rollout_id
                async with active_lock:
                    active_rollouts.add(rollout_id)
                try:
                    while True:
                        current = await store.get_rollout_by_id(rollout_id)
                        if current is not None and current.status in ("failed", "succeeded", "cancelled"):
                            if current.status != "succeeded":
                                raise RuntimeError(f"Rollout {rollout_id} finished with status {current.status}")
                            break
                        await emit_progress(progress_made=False)
                        await asyncio.sleep(5.0)
                    spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
                    check_spans(spans, task_name)
                    await emit_progress(progress_made=True)
                finally:
                    async with active_lock:
                        active_rollouts.discard(rollout_id)

        all_tasks = [handle_single(i) for i in range(total_tasks)]
        await asyncio.gather(*all_tasks)