|
1 | 1 | """EvaluationService — datasets on disk, runs dispatched to the worker layer. |
2 | 2 |
|
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. |
6 | 17 | """ |
7 | 18 |
|
8 | 19 | from __future__ import annotations |
|
14 | 25 | from typing import TYPE_CHECKING |
15 | 26 |
|
16 | 27 | 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 |
19 | 39 |
|
20 | 40 | if TYPE_CHECKING: |
21 | 41 | from collections.abc import Sequence |
22 | 42 | from typing import IO |
23 | 43 |
|
24 | 44 | from core.config.root import Settings |
| 45 | + from core.evaluation.runner import EvaluationRunner |
25 | 46 | 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" |
26 | 56 |
|
27 | 57 | TESTSET_FILENAME = "testset.csv" |
28 | 58 | CORPUS_DIRNAME = "corpus" |
|
31 | 61 | _COPY_CHUNK_BYTES = 1024 * 1024 |
32 | 62 |
|
33 | 63 |
|
| 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 | + |
34 | 75 | class EvaluationService: |
35 | 76 | """Dataset storage plus run dispatch for the admin evaluation page.""" |
36 | 77 |
|
37 | 78 | def __init__( |
38 | 79 | self, |
39 | 80 | *, |
40 | 81 | repo: EvaluationRepository, |
| 82 | + runner: EvaluationRunner, |
| 83 | + user_service: UserService, |
| 84 | + user_repo: UserRepository, |
| 85 | + partition_service: PartitionService, |
41 | 86 | config: Settings, |
42 | 87 | ) -> None: |
43 | 88 | 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 |
44 | 93 | self._config = config |
45 | 94 | self._settings = config.evaluation |
46 | 95 | self._root = Path(config.paths.data_dir) / "eval" |
@@ -170,5 +219,177 @@ async def delete_dataset(self, dataset_id: str) -> None: |
170 | 219 | raise NotFoundError(f"Evaluation dataset '{dataset_id}' not found") |
171 | 220 | await asyncio.to_thread(shutil.rmtree, self._dataset_dir(dataset_id), True) |
172 | 221 |
|
| 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 | + |
173 | 388 |
|
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