Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e
This commit is contained in:
Public Source Publisher committed
commit
7ca726caf2
659 files changed
+98171
No files matched your search
@@ -0,0 +1,161 @@
|
||||
"""JobWorker: ingestion job outcomes and claim age are really recorded."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from opentelemetry.sdk.metrics import MeterProvider
|
||||
from opentelemetry.sdk.metrics.export import InMemoryMetricReader
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
|
||||
from ragcore.application.jobs.models import (
|
||||
Job,
|
||||
JobClaim,
|
||||
JobContext,
|
||||
JobError,
|
||||
RetryableJobError,
|
||||
TerminalJobError,
|
||||
)
|
||||
from ragcore.application.jobs.worker import JobWorker
|
||||
from ragcore.platform.telemetry import TelemetryMetrics
|
||||
|
||||
|
||||
def _telemetry() -> tuple[TelemetryMetrics, InMemoryMetricReader]:
|
||||
reader = InMemoryMetricReader()
|
||||
provider = MeterProvider(
|
||||
resource=Resource.create({"service.name": "test"}), metric_readers=[reader]
|
||||
)
|
||||
return TelemetryMetrics(meter_provider=provider), reader
|
||||
|
||||
|
||||
def _metric_points(reader: InMemoryMetricReader, name: str) -> list[object]:
|
||||
data = reader.get_metrics_data()
|
||||
if data is None:
|
||||
return []
|
||||
for resource_metric in data.resource_metrics:
|
||||
for scope_metric in resource_metric.scope_metrics:
|
||||
for metric in scope_metric.metrics:
|
||||
if metric.name == name:
|
||||
return list(metric.data.data_points)
|
||||
return []
|
||||
|
||||
|
||||
class _FakeJobRepository:
|
||||
"""Enough of ``JobRepository`` for one claim/complete or claim/fail cycle."""
|
||||
|
||||
def __init__(self, job: Job) -> None:
|
||||
self._job = job
|
||||
self._claimed = False
|
||||
self.completed: list[JobClaim] = []
|
||||
self.failed: list[tuple[JobClaim, JobError]] = []
|
||||
|
||||
async def enqueue(self, job: Job) -> Job:
|
||||
raise NotImplementedError
|
||||
|
||||
async def claim_next(
|
||||
self, worker_id: str, *, now: datetime, lease_duration: timedelta
|
||||
) -> JobClaim | None:
|
||||
if self._claimed:
|
||||
return None
|
||||
self._claimed = True
|
||||
return JobClaim(
|
||||
job=self._job, worker_id=worker_id, attempt_number=1, lease_until=now + lease_duration
|
||||
)
|
||||
|
||||
async def heartbeat(
|
||||
self, claim: JobClaim, *, now: datetime, lease_duration: timedelta
|
||||
) -> bool:
|
||||
return True
|
||||
|
||||
async def complete(self, claim: JobClaim, *, now: datetime) -> None:
|
||||
self.completed.append(claim)
|
||||
|
||||
async def fail(
|
||||
self,
|
||||
claim: JobClaim,
|
||||
error: JobError,
|
||||
*,
|
||||
now: datetime,
|
||||
retry_after: timedelta = timedelta(0),
|
||||
) -> None:
|
||||
self.failed.append((claim, error))
|
||||
|
||||
|
||||
def _job(*, available_at: datetime) -> Job:
|
||||
return Job(
|
||||
id=uuid4(),
|
||||
job_key=f"job-{uuid4()}",
|
||||
operation="ingest_upload",
|
||||
target_key="target",
|
||||
profile_version="v1",
|
||||
available_at=available_at,
|
||||
)
|
||||
|
||||
|
||||
async def test_successful_job_records_claimed_age_and_succeeded_result() -> None:
|
||||
telemetry, reader = _telemetry()
|
||||
available_at = datetime.now(UTC) - timedelta(seconds=5)
|
||||
repository = _FakeJobRepository(_job(available_at=available_at))
|
||||
|
||||
async def handler(ctx: JobContext) -> None:
|
||||
return None
|
||||
|
||||
worker = JobWorker(
|
||||
repository, {"ingest_upload": handler}, worker_id="worker-1", metrics=telemetry
|
||||
)
|
||||
assert await worker.run_once() is True
|
||||
|
||||
age_points = _metric_points(reader, "ragcore_job_age_seconds")
|
||||
assert len(age_points) == 1
|
||||
assert age_points[0].sum == pytest.approx(5.0, abs=1.0)
|
||||
|
||||
result_points = _metric_points(reader, "ragcore_ingestion_job_result_total")
|
||||
assert len(result_points) == 1
|
||||
assert dict(result_points[0].attributes) == {
|
||||
"operation": "ingest_upload",
|
||||
"result": "succeeded",
|
||||
}
|
||||
assert repository.completed
|
||||
|
||||
|
||||
async def test_retryable_failure_records_retryable_failed_result() -> None:
|
||||
telemetry, reader = _telemetry()
|
||||
repository = _FakeJobRepository(_job(available_at=datetime.now(UTC)))
|
||||
|
||||
async def handler(ctx: JobContext) -> None:
|
||||
raise RetryableJobError("transient", retry_after_seconds=1.0)
|
||||
|
||||
worker = JobWorker(
|
||||
repository, {"ingest_upload": handler}, worker_id="worker-1", metrics=telemetry
|
||||
)
|
||||
assert await worker.run_once() is True
|
||||
|
||||
result_points = _metric_points(reader, "ragcore_ingestion_job_result_total")
|
||||
assert len(result_points) == 1
|
||||
assert dict(result_points[0].attributes) == {
|
||||
"operation": "ingest_upload",
|
||||
"result": "retryable_failed",
|
||||
}
|
||||
assert repository.failed
|
||||
|
||||
|
||||
async def test_terminal_failure_records_terminal_failed_result() -> None:
|
||||
telemetry, reader = _telemetry()
|
||||
repository = _FakeJobRepository(_job(available_at=datetime.now(UTC)))
|
||||
|
||||
async def handler(ctx: JobContext) -> None:
|
||||
raise TerminalJobError("unrecoverable")
|
||||
|
||||
worker = JobWorker(
|
||||
repository, {"ingest_upload": handler}, worker_id="worker-1", metrics=telemetry
|
||||
)
|
||||
assert await worker.run_once() is True
|
||||
|
||||
result_points = _metric_points(reader, "ragcore_ingestion_job_result_total")
|
||||
assert len(result_points) == 1
|
||||
assert dict(result_points[0].attributes) == {
|
||||
"operation": "ingest_upload",
|
||||
"result": "terminal_failed",
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
"""A worker process runs more than one job at a time, or it does not scale.
|
||||
|
||||
``claim_next`` has always used ``FOR UPDATE SKIP LOCKED``, so several claim
|
||||
loops can safely draw from the same queue -- but nothing ever ran more than
|
||||
one. A folder source with a few thousand files was therefore ingested strictly
|
||||
one document after another.
|
||||
|
||||
The test below cannot pass on a serial pool: each handler waits until every
|
||||
other one has started before it returns, so if the loops take turns the whole
|
||||
thing deadlocks instead of finishing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from ragcore.application.jobs.models import Job, JobClaim, JobContext, JobError
|
||||
from ragcore.application.jobs.worker import run_worker_pool
|
||||
|
||||
CONCURRENCY = 3
|
||||
|
||||
|
||||
def _job(index: int) -> Job:
|
||||
return Job(
|
||||
id=uuid4(),
|
||||
workspace_id=uuid4(),
|
||||
job_key=f"job-{index}",
|
||||
operation="ingest_upload",
|
||||
target_key="blob:abc",
|
||||
profile_version="upload-v1",
|
||||
available_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
|
||||
class _QueueRepository:
|
||||
"""Hands every job out exactly once, to whoever asks first."""
|
||||
|
||||
def __init__(self, jobs: list[Job]) -> None:
|
||||
self._pending = list(jobs)
|
||||
self._expected = len(jobs)
|
||||
self.completed: list[Job] = []
|
||||
self.claimed_by: list[str] = []
|
||||
self.drained = asyncio.Event()
|
||||
|
||||
async def claim_next(
|
||||
self, worker_id: str, *, now: datetime, lease_duration: timedelta
|
||||
) -> JobClaim | None:
|
||||
if not self._pending:
|
||||
return None
|
||||
job = self._pending.pop(0)
|
||||
self.claimed_by.append(worker_id)
|
||||
return JobClaim(job=job, worker_id=worker_id, attempt_number=1, lease_until=now)
|
||||
|
||||
async def heartbeat(self, claim: JobClaim, **kwargs: object) -> bool:
|
||||
return True
|
||||
|
||||
async def complete(self, claim: JobClaim, *, now: datetime) -> None:
|
||||
self.completed.append(claim.job)
|
||||
if len(self.completed) == self._expected:
|
||||
self.drained.set()
|
||||
|
||||
async def fail(self, claim: JobClaim, error: JobError, **kwargs: object) -> None:
|
||||
raise AssertionError(f"no job should fail here: {error}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_worker_pool_runs_jobs_at_the_same_time() -> None:
|
||||
repository = _QueueRepository([_job(index) for index in range(CONCURRENCY)])
|
||||
started = asyncio.Barrier(CONCURRENCY)
|
||||
stop = asyncio.Event()
|
||||
|
||||
async def handler(context: JobContext) -> None:
|
||||
# Returns only once every other handler has also started. A pool that
|
||||
# runs one job at a time can never get past this.
|
||||
await started.wait()
|
||||
|
||||
async def drain() -> None:
|
||||
await repository.drained.wait()
|
||||
stop.set()
|
||||
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(
|
||||
run_worker_pool(
|
||||
repository, # type: ignore[arg-type]
|
||||
{"ingest_upload": handler},
|
||||
stop=stop,
|
||||
concurrency=CONCURRENCY,
|
||||
worker_id_prefix="pool-test",
|
||||
poll_interval=0.01,
|
||||
),
|
||||
drain(),
|
||||
),
|
||||
timeout=5.0,
|
||||
)
|
||||
|
||||
assert len(repository.completed) == CONCURRENCY
|
||||
# Every loop identifies itself: "which worker holds this lease" is the
|
||||
# first question an operator asks about a stuck job.
|
||||
assert len(set(repository.claimed_by)) == CONCURRENCY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pool_of_one_is_the_previous_behaviour() -> None:
|
||||
repository = _QueueRepository([_job(0)])
|
||||
stop = asyncio.Event()
|
||||
|
||||
async def handler(context: JobContext) -> None:
|
||||
stop.set()
|
||||
|
||||
await asyncio.wait_for(
|
||||
run_worker_pool(
|
||||
repository, # type: ignore[arg-type]
|
||||
{"ingest_upload": handler},
|
||||
stop=stop,
|
||||
concurrency=1,
|
||||
worker_id_prefix="pool-test",
|
||||
poll_interval=0.01,
|
||||
),
|
||||
timeout=5.0,
|
||||
)
|
||||
|
||||
assert len(repository.completed) == 1
|
||||
assert repository.claimed_by == ["pool-test-0"]
|
||||
Reference in new issue
Block a user