From da8e7dc5681db579b291d1e21aa8e30e29f45176 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 1 Jun 2026 07:52:17 -0400 Subject: [PATCH 1/2] Fix run_batch hanging forever when all workers die MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The drain loop in run_batch() polls counts_by_status() waiting for queued and claimed counts to reach zero. If all worker tasks crash (unhandled exception, OOM), claimed jobs stay claimed forever and the loop never exits — the CLI command hangs. Check live_workers during the drain loop. If claimed jobs exist but no workers are alive to process them, log an error and break out. The stranded jobs will be reaped on the next start. --- haiku_rag_slim/haiku/rag/ingester/app.py | 8 +++++ tests/ingester/test_run_batch.py | 42 ++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 7b64556f..23234582 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -204,6 +204,14 @@ class IngesterApp: counts = await self._jobs.counts_by_status() if not counts.get("queued") and not counts.get("claimed"): break + if counts.get("claimed") and self._pool.live_workers == 0: + logger.error( + "All workers have died with %d claimed job(s) — " + "aborting batch; stranded jobs will be reaped on " + "next start", + counts["claimed"], + ) + break await asyncio.sleep(0.1) completed = await self._jobs.counts_by_status_since(started_at) return BatchReport( diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py index c646d5d6..60a96d8e 100644 --- a/tests/ingester/test_run_batch.py +++ b/tests/ingester/test_run_batch.py @@ -234,6 +234,48 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client): client.create_document_from_source.assert_not_awaited() +@pytest.mark.asyncio +async def test_run_batch_aborts_when_all_workers_die(tmp_path, use_client, monkeypatch): + """If all workers crash with claimed jobs still outstanding, run_batch + should break out of the drain loop instead of spinning forever.""" + (tmp_path / "a.md").write_text("hello") + + 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) + config = _config(tmp_path) + app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") + + async def _kill_workers_after_claim(): + """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. + batch_task = asyncio.create_task(app.run_batch()) + # Give run_batch a moment to start, then schedule the killer. + await asyncio.sleep(0.1) + killer_task = asyncio.create_task(_kill_workers_after_claim()) + + 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 + + async def _wait_until(predicate, *, timeout: float = 5.0): deadline = asyncio.get_running_loop().time() + timeout while not predicate(): From b144e620de8ad9b860b2aba31753e0a836410fd0 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 1 Jun 2026 09:47:10 -0400 Subject: [PATCH 2/2] 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. --- haiku_rag_slim/haiku/rag/ingester/app.py | 11 ++--- tests/ingester/test_run_batch.py | 51 ++++++++++-------------- 2 files changed, 27 insertions(+), 35 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 23234582..2d51dedd 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -204,12 +204,13 @@ class IngesterApp: counts = await self._jobs.counts_by_status() if not counts.get("queued") and not counts.get("claimed"): 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( - "All workers have died with %d claimed job(s) — " - "aborting batch; stranded jobs will be reaped on " - "next start", - counts["claimed"], + "All workers have died with %d outstanding job(s) " + "— aborting batch; stranded jobs will be reaped " + "on next start", + outstanding, ) break await asyncio.sleep(0.1) diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py index 60a96d8e..dbb616c6 100644 --- a/tests/ingester/test_run_batch.py +++ b/tests/ingester/test_run_batch.py @@ -235,45 +235,36 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client): @pytest.mark.asyncio -async def test_run_batch_aborts_when_all_workers_die(tmp_path, use_client, monkeypatch): - """If all workers crash with claimed jobs still outstanding, run_batch - should break out of the drain loop instead of spinning forever.""" +async def test_run_batch_aborts_when_all_workers_die( + tmp_path, use_client, monkeypatch, caplog +): + """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") 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) config = _config(tmp_path) - app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") - async def _kill_workers_after_claim(): - """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) + config.ingester.workers.worker_count = 1 - # Run the killer concurrently with run_batch. - batch_task = asyncio.create_task(app.run_batch()) - # Give run_batch a moment to start, then schedule the killer. - await asyncio.sleep(0.1) - killer_task = asyncio.create_task(_kill_workers_after_claim()) + # Patch _process to raise an unhandled exception, simulating a hard + # worker crash. _process only catches CancelledError, PermanentError, + # and TransientError — anything else propagates and kills the task. + monkeypatch.setattr( + 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 "All workers have died" in caplog.text async def _wait_until(predicate, *, timeout: float = 5.0):