Fix dead-worker condition and test for run_batch abort
The condition only checked claimed jobs, but queued jobs with no live workers also hang forever. Check live_workers == 0 regardless of whether outstanding work is queued or claimed. Rewrite the test to actually crash workers: patch _process to raise a bare Exception (which _worker_loop doesn't catch), use worker_count=1 so the single crash leaves live_workers == 0, and assert the abort log message fires.
This commit is contained in:
parent
da8e7dc568
commit
b144e620de
2 changed files with 27 additions and 35 deletions
|
|
@ -204,12 +204,13 @@ class IngesterApp:
|
||||||
counts = await self._jobs.counts_by_status()
|
counts = await self._jobs.counts_by_status()
|
||||||
if not counts.get("queued") and not counts.get("claimed"):
|
if not counts.get("queued") and not counts.get("claimed"):
|
||||||
break
|
break
|
||||||
if counts.get("claimed") and self._pool.live_workers == 0:
|
if self._pool.live_workers == 0:
|
||||||
|
outstanding = counts.get("queued", 0) + counts.get("claimed", 0)
|
||||||
logger.error(
|
logger.error(
|
||||||
"All workers have died with %d claimed job(s) — "
|
"All workers have died with %d outstanding job(s) "
|
||||||
"aborting batch; stranded jobs will be reaped on "
|
"— aborting batch; stranded jobs will be reaped "
|
||||||
"next start",
|
"on next start",
|
||||||
counts["claimed"],
|
outstanding,
|
||||||
)
|
)
|
||||||
break
|
break
|
||||||
await asyncio.sleep(0.1)
|
await asyncio.sleep(0.1)
|
||||||
|
|
|
||||||
|
|
@ -235,45 +235,36 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_run_batch_aborts_when_all_workers_die(tmp_path, use_client, monkeypatch):
|
async def test_run_batch_aborts_when_all_workers_die(
|
||||||
"""If all workers crash with claimed jobs still outstanding, run_batch
|
tmp_path, use_client, monkeypatch, caplog
|
||||||
should break out of the drain loop instead of spinning forever."""
|
):
|
||||||
|
"""If all workers crash with outstanding jobs, run_batch should break
|
||||||
|
out of the drain loop instead of spinning forever."""
|
||||||
(tmp_path / "a.md").write_text("hello")
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
|
||||||
client = _mock_client()
|
client = _mock_client()
|
||||||
# Block the worker forever so the job stays claimed until we kill it.
|
|
||||||
stall = asyncio.Event()
|
|
||||||
client.create_document_from_source.side_effect = lambda *a, **k: stall.wait()
|
|
||||||
|
|
||||||
use_client(client)
|
use_client(client)
|
||||||
config = _config(tmp_path)
|
config = _config(tmp_path)
|
||||||
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb")
|
|
||||||
|
|
||||||
async def _kill_workers_after_claim():
|
config.ingester.workers.worker_count = 1
|
||||||
"""Wait until at least one job is claimed, then kill all workers."""
|
|
||||||
pool = app._pool
|
|
||||||
assert pool is not None
|
|
||||||
for _ in range(200):
|
|
||||||
counts = await app._jobs.counts_by_status()
|
|
||||||
if counts.get("claimed"):
|
|
||||||
break
|
|
||||||
await asyncio.sleep(0.05)
|
|
||||||
for task in pool._workers:
|
|
||||||
task.cancel()
|
|
||||||
await asyncio.gather(*pool._workers, return_exceptions=True)
|
|
||||||
|
|
||||||
# Run the killer concurrently with run_batch.
|
# Patch _process to raise an unhandled exception, simulating a hard
|
||||||
batch_task = asyncio.create_task(app.run_batch())
|
# worker crash. _process only catches CancelledError, PermanentError,
|
||||||
# Give run_batch a moment to start, then schedule the killer.
|
# and TransientError — anything else propagates and kills the task.
|
||||||
await asyncio.sleep(0.1)
|
monkeypatch.setattr(
|
||||||
killer_task = asyncio.create_task(_kill_workers_after_claim())
|
WorkerPool,
|
||||||
|
"_process",
|
||||||
|
AsyncMock(side_effect=Exception("worker crash")),
|
||||||
|
)
|
||||||
|
|
||||||
|
with caplog.at_level("ERROR", logger="haiku.rag.ingester.app"):
|
||||||
|
report = await asyncio.wait_for(
|
||||||
|
IngesterApp(config=config, db_path=tmp_path / "db.lancedb").run_batch(),
|
||||||
|
timeout=10.0,
|
||||||
|
)
|
||||||
|
|
||||||
report = await asyncio.wait_for(batch_task, timeout=10.0)
|
|
||||||
# Killer may still be running against a closed DB — suppress errors.
|
|
||||||
killer_task.cancel()
|
|
||||||
await asyncio.gather(killer_task, return_exceptions=True)
|
|
||||||
# The batch should have exited without hanging.
|
|
||||||
assert report is not None
|
assert report is not None
|
||||||
|
assert "All workers have died" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
async def _wait_until(predicate, *, timeout: float = 5.0):
|
async def _wait_until(predicate, *, timeout: float = 5.0):
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue