Skip to content

Commit f15c899

Browse files
committed
feat(evaluation): run start, dispatch and cancellation
The run half of `EvaluationService`, dispatched through the `EvaluationRunner` port so the orchestrator stays Ray-free. Setup and teardown of a run's identity live here rather than in the worker, because creating users and partitions is orchestration the API layer already owns. The worker receives a partition it may write to and a token it may use, and nothing else about the system. Ordering in `start_run` is deliberate and each step is regression-tested: 1. Ping the runner first. Dispatch is fire-and-forget, so an unreachable worker would otherwise strand the run in QUEUED with a partition and a token provisioned for nobody. It surfaces as 503. 2. Insert the run row before provisioning anything. The partial unique index is the mutual exclusion; a read-then-insert would let two racing requests both regenerate the shared eval user's token, the second revoking the credentials the first is still indexing with. The loser gets a 409 before touching anything. 3. On a failed provision, release the row and drop the partition. A run left in an active status holds the lock and would block every later run. Runs authenticate as one long-lived non-admin service user (`__openrag_eval__`) whose token is regenerated at the start of every run, so no usable plaintext token is ever stored at rest. Cancellation asks the worker first. A worker that owns the run writes its own terminal status, including metrics. `False` means nobody owns it — the run was orphaned by an actor restart — so the row is reaped here instead, otherwise it would block every subsequent run forever.
1 parent 03026f3 commit f15c899

2 files changed

Lines changed: 476 additions & 13 deletions

File tree

openrag/services/orchestrators/evaluation_service.py

Lines changed: 227 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,19 @@
11
"""EvaluationService — datasets on disk, runs dispatched to the worker layer.
22
3-
This slice covers dataset storage: an admin uploads a corpus plus a test set,
4-
both land under ``<data_dir>/eval/<dataset_id>/``, and a row records what is
5-
there. Run dispatch follows.
3+
Setup and teardown of a run's *identity* live here rather than in the worker,
4+
because creating users and partitions is orchestration the API layer already
5+
owns. The worker receives a partition it may write to and a token it may use,
6+
and nothing else about the system.
7+
8+
Dispatch goes through the :class:`~core.evaluation.runner.EvaluationRunner`
9+
port, so this orchestrator stays Ray-free; the Ray actor lives behind the
10+
adapter in ``services/workers/eval_dispatcher.py``.
11+
12+
The bearer token handed to the worker belongs to a single long-lived service
13+
user (``__openrag_eval__``) whose token is **regenerated at the start of every
14+
run**. That keeps exactly one non-admin service account in the database while
15+
ensuring no usable plaintext token is ever stored at rest — the previous one
16+
stops working the moment a new run starts.
617
"""
718

819
from __future__ import annotations
@@ -14,15 +25,34 @@
1425
from typing import TYPE_CHECKING
1526

1627
from core.evaluation import parse_testset
17-
from core.models.evaluation import EvalDataset
18-
from core.utils.exceptions import ConflictError, NotFoundError, ValidationError
28+
from core.models.evaluation import (
29+
EVAL_PARTITION_PREFIX,
30+
EvalDataset,
31+
EvalRun,
32+
EvalRunStatus,
33+
EvalTestCase,
34+
is_eval_partition,
35+
)
36+
from core.models.user import UserCreate
37+
from core.utils.exceptions import ConflictError, NotFoundError, OpenRAGError, ValidationError
38+
from core.utils.logging import get_logger
1939

2040
if TYPE_CHECKING:
2141
from collections.abc import Sequence
2242
from typing import IO
2343

2444
from core.config.root import Settings
45+
from core.evaluation.runner import EvaluationRunner
2546
from core.ports.evaluation_repo import EvaluationRepository
47+
from core.ports.user_repo import UserRepository
48+
from services.orchestrators.partition_service import PartitionService
49+
from services.orchestrators.user_service import UserService
50+
51+
logger = get_logger()
52+
53+
#: Stable identity of the service account runs authenticate as.
54+
EVAL_USER_EXTERNAL_ID = "__openrag_eval__"
55+
EVAL_USER_DISPLAY_NAME = "OpenRAG Evaluation"
2656

2757
TESTSET_FILENAME = "testset.csv"
2858
CORPUS_DIRNAME = "corpus"
@@ -31,16 +61,35 @@
3161
_COPY_CHUNK_BYTES = 1024 * 1024
3262

3363

64+
def eval_partition_name(run_id: str) -> str:
65+
return f"{EVAL_PARTITION_PREFIX}{run_id}"
66+
67+
68+
class EvaluationRunnerUnavailableError(OpenRAGError):
69+
"""The runner actor could not be reached. Maps to HTTP 503."""
70+
71+
def __init__(self, message: str) -> None:
72+
super().__init__(message, code="EVAL_RUNNER_UNAVAILABLE", status_code=503)
73+
74+
3475
class EvaluationService:
3576
"""Dataset storage plus run dispatch for the admin evaluation page."""
3677

3778
def __init__(
3879
self,
3980
*,
4081
repo: EvaluationRepository,
82+
runner: EvaluationRunner,
83+
user_service: UserService,
84+
user_repo: UserRepository,
85+
partition_service: PartitionService,
4186
config: Settings,
4287
) -> None:
4388
self._repo = repo
89+
self._runner = runner
90+
self._user_service = user_service
91+
self._user_repo = user_repo
92+
self._partition_service = partition_service
4493
self._config = config
4594
self._settings = config.evaluation
4695
self._root = Path(config.paths.data_dir) / "eval"
@@ -170,5 +219,177 @@ async def delete_dataset(self, dataset_id: str) -> None:
170219
raise NotFoundError(f"Evaluation dataset '{dataset_id}' not found")
171220
await asyncio.to_thread(shutil.rmtree, self._dataset_dir(dataset_id), True)
172221

222+
# ── runs ─────────────────────────────────────────────────────────
223+
224+
async def list_runs(self, limit: int = 50) -> list[EvalRun]:
225+
return await self._repo.list_runs(limit)
226+
227+
async def get_run(self, run_id: str) -> EvalRun:
228+
run = await self._repo.get_run(run_id)
229+
if run is None:
230+
raise NotFoundError(f"Evaluation run '{run_id}' not found")
231+
return run
232+
233+
async def start_run(self, dataset_id: str, user_id: int | None) -> EvalRun:
234+
"""Provision a run's partition and token, then dispatch it.
235+
236+
The run row is inserted before anything is provisioned: the partial
237+
unique index ``ux_eval_runs_single_active`` makes that insert the mutual
238+
exclusion between concurrent starts. A read-then-insert would let two
239+
racing requests both regenerate the shared eval user's token, the second
240+
revoking the credentials the first is still indexing with.
241+
242+
Raises:
243+
NotFoundError: The dataset does not exist.
244+
ConflictError: Another run is already in flight — the runner
245+
executes one at a time so timings stay comparable.
246+
"""
247+
dataset = await self._repo.get_dataset(dataset_id)
248+
if dataset is None:
249+
raise NotFoundError(f"Evaluation dataset '{dataset_id}' not found")
250+
251+
directory = self._dataset_dir(dataset_id)
252+
testset_path = directory / TESTSET_FILENAME
253+
if not testset_path.exists():
254+
raise NotFoundError(f"Test set for dataset '{dataset_id}' is missing on disk")
255+
cases = parse_testset(testset_path.read_bytes(), max_rows=self._settings.max_testset_rows)
256+
257+
# Reach the runner before claiming the slot: dispatch is
258+
# fire-and-forget, so an unreachable worker would otherwise strand the
259+
# run in QUEUED with a partition and token provisioned for nobody.
260+
await self._ping_runner()
261+
262+
run_id = uuid.uuid4().hex
263+
partition = eval_partition_name(run_id)
264+
run = await self._repo.create_run(
265+
EvalRun(
266+
id=run_id,
267+
dataset_id=dataset_id,
268+
status=EvalRunStatus.QUEUED,
269+
created_by=user_id,
270+
)
271+
)
272+
273+
try:
274+
eval_user_id = await self._ensure_eval_user()
275+
token = (await self._user_service.regenerate_token(eval_user_id))["token"]
276+
await self._partition_service.create_partition(partition, user_id=eval_user_id)
277+
await self._dispatch(run_id, partition, token, directory, cases)
278+
except Exception as exc:
279+
# The run row is the lock; leaving it active would block every
280+
# later run.
281+
logger.exception(f"Could not start evaluation run {run_id}: {exc}")
282+
await self._repo.update_run_status(
283+
run_id,
284+
EvalRunStatus.FAILED,
285+
error=f"Could not start the run: {exc}",
286+
)
287+
await self._drop_orphaned_partition(run_id)
288+
raise
289+
290+
logger.bind(run_id=run_id, dataset_id=dataset_id).info("Dispatched evaluation run")
291+
return run
292+
293+
async def _dispatch(
294+
self,
295+
run_id: str,
296+
partition: str,
297+
token: str,
298+
directory: Path,
299+
cases: Sequence[EvalTestCase],
300+
) -> None:
301+
"""Hand the run to the worker.
302+
303+
Fire and forget: the worker owns the run from here and records its own
304+
outcome.
305+
"""
306+
await self._runner.dispatch(
307+
run_id=run_id,
308+
partition=partition,
309+
token=token,
310+
api_base_url=self._config.server.internal_url,
311+
corpus_dir=str(directory / CORPUS_DIRNAME),
312+
cases=[
313+
{
314+
"query": case.query,
315+
"expected_answer": case.expected_answer,
316+
"expected_file_ids": list(case.expected_file_ids),
317+
}
318+
for case in cases
319+
],
320+
)
321+
322+
async def cancel_run(self, run_id: str) -> EvalRun:
323+
"""Ask the worker to stop, or reap the run if no worker owns it.
324+
325+
The worker writes the terminal status for a run it is executing. When
326+
it disowns the run — it restarted, or died before picking the run up —
327+
nothing else would ever move that row out of an active status, and it
328+
would block every subsequent run. Cancelling reaps it instead.
329+
"""
330+
run = await self.get_run(run_id)
331+
if run.status.is_terminal:
332+
raise ConflictError(f"Evaluation run '{run_id}' has already finished.")
333+
334+
owned = False
335+
try:
336+
owned = await self._runner.cancel(run_id)
337+
except Exception as exc: # noqa: BLE001 — an unreachable runner still has to be reaped
338+
logger.warning(f"Evaluation runner unreachable while cancelling {run_id}: {exc}")
339+
340+
if not owned:
341+
await self._repo.update_run_status(
342+
run_id,
343+
EvalRunStatus.CANCELLED,
344+
error="No runner owns this run — it was orphaned and has been reaped.",
345+
)
346+
await self._drop_orphaned_partition(run_id)
347+
return await self.get_run(run_id)
348+
349+
async def _drop_orphaned_partition(self, run_id: str) -> None:
350+
"""Best-effort cleanup of the throwaway partition of a reaped run."""
351+
try:
352+
await self._partition_service.delete_partition(eval_partition_name(run_id))
353+
except Exception as exc: # noqa: BLE001 — it may never have been created
354+
logger.debug(f"No eval partition to drop for run {run_id}: {exc}")
355+
356+
# ── internals ────────────────────────────────────────────────────
357+
358+
async def _ping_runner(self) -> None:
359+
"""Fail fast when the runner cannot be reached.
360+
361+
Raises:
362+
OpenRAGError: The worker is unreachable — surfaced to the caller
363+
instead of being discovered as a run that never leaves QUEUED.
364+
"""
365+
try:
366+
await self._runner.is_busy()
367+
except Exception as exc:
368+
logger.exception(f"Evaluation runner is unavailable: {exc}")
369+
raise EvaluationRunnerUnavailableError(f"The evaluation runner could not be reached: {exc}") from exc
370+
371+
async def _ensure_eval_user(self) -> int:
372+
"""Get-or-create the non-admin service user runs authenticate as."""
373+
existing = await self._user_repo.get_user_by_external_id(EVAL_USER_EXTERNAL_ID)
374+
if existing is not None:
375+
return int(existing.id)
376+
created = await self._user_service.create_user(
377+
UserCreate(
378+
display_name=EVAL_USER_DISPLAY_NAME,
379+
external_user_id=EVAL_USER_EXTERNAL_ID,
380+
is_admin=False,
381+
# A corpus is uploaded on every run, so a quota would fail the
382+
# second one for reasons unrelated to the eval.
383+
file_quota=-1,
384+
)
385+
)
386+
return int(created["id"])
387+
173388

174-
__all__ = ["EvaluationService"]
389+
__all__ = [
390+
"EVAL_PARTITION_PREFIX",
391+
"EVAL_USER_EXTERNAL_ID",
392+
"EvaluationService",
393+
"eval_partition_name",
394+
"is_eval_partition",
395+
]

0 commit comments

Comments
 (0)