feat(kg-extract): concurrent backfill via streaming task pool - #217
Conversation
Five tightly-coupled changes that together unstick KG backfill throughput after PR #208 removed the asyncio.Lock bottleneck: 1. Semaphore narrowed to the LLM call only. Extracted ``_extract_under_sem`` so the ``asyncio.Semaphore`` slot is held during ``extract_triples`` and released before any AGE writes begin. Previously a slow DB write held an LLM slot, starving llama server's ``--parallel`` capacity. 2. Streaming task pool replaces the gather-over-batch barrier. The old loop did claim → ``asyncio.gather(*batch)`` → next claim, so the slowest drawer in each batch stalled the next claim. New model: one producer task tops up an ``asyncio.Queue`` whenever it dips below ``batch_size // 2``; ``max_concurrency`` persistent consumer tasks pull and run ``_process_one``. No batch boundaries. 3. Default knobs bumped for the new floor: - ``DEFAULT_CONCURRENCY`` 8 → 24 (matches llama-server --parallel) - ``_SyncConnPool.max_size`` 10 → 32 (covers 24 LLM slots + slack) 4. New ``--db-pool-size`` CLI flag (env: ``MEMPALACE_KG_DB_POOL_SIZE``). Validated ``db_pool_size >= max_concurrency`` so every in-flight LLM call is guaranteed a write conn. Default is ``max_concurrency + 8``. 5. ``DEFAULT_ENDPOINT`` switched from ``http://familiar.jphe.in:11436`` to ``http://familiar:11436``. The FQDN was stale; bare ``familiar`` resolves via Tailscale for inter-host calls. Tests: 23/23 in tests/test_kg_triple_worker.py — including new coverage for streaming-pool no-barrier behaviour, consumer count = max_concurrency, sem-release-before-DB-write, db-pool-size validation, and the bare-host endpoint default.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request optimizes the KG backfill throughput by decoupling LLM inference from database write operations. By moving to a streaming producer-consumer model and narrowing the semaphore scope, the worker can maintain higher concurrency without being bottlenecked by individual slow database transactions. These changes are supported by updated default resource limits and a new database pool configuration flag. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the worker concurrency model in kg_triple_worker.py to use a producer/consumer pattern with an asyncio.Queue. This decoupling prevents slow database writes from stalling the LLM processing pipeline. Additionally, the semaphore scope is narrowed to only wrap the LLM extraction step. The review feedback suggests a highly efficient consumer shutdown mechanism by pushing None sentinels to the queue upon producer completion, avoiding busy-waiting.
| finally: | ||
| producer_done.set() |
There was a problem hiding this comment.
To support a non-polling, highly efficient consumer shutdown, update the producer's finally block to push a None sentinel for each consumer task. This avoids busy-waiting in the consumers when the queue is empty.
| finally: | |
| producer_done.set() | |
| finally: | |
| producer_done.set() | |
| for _ in range(max_concurrency): | |
| await work_queue.put(None) |
1 participant
Summary
Five tightly-coupled changes that together unstick KG backfill throughput
after PR #208 removed the
asyncio.Lockbottleneck._extract_under_semso the slot is held duringextract_triplesand released before AGE writes begin. Previously slow DB writes held LLM slots, starving llama-server's--parallelcapacity.asyncio.gather(*batch)→ next claim — slowest drawer stalled the next claim cycle. New model: one producer task tops up anasyncio.Queuewhenever it dips belowbatch_size // 2;max_concurrencypersistent consumer tasks pull and run_process_one. No batch boundaries.DEFAULT_CONCURRENCY8 → 24 (matchesllama-server --parallel),_SyncConnPool.max_size10 → 32 (covers 24 LLM slots + slack).--db-pool-sizeCLI flag (env:MEMPALACE_KG_DB_POOL_SIZE). Validateddb_pool_size >= max_concurrencyso every in-flight LLM call is guaranteed a write conn. Default ismax_concurrency + 8.DEFAULT_ENDPOINTswitched fromhttp://familiar.jphe.in:11436tohttp://familiar:11436— bare Tailscale hostname for inter-host calls.Why these belong in one PR
Changes #1–#4 are interlocking: narrowing the sem only helps if the producer/consumer model doesn't serialize on batch boundaries; bumping concurrency without bumping the DB pool would starve writes; the
--db-pool-sizeflag exists precisely so #3 can be retuned at deploy time. #5 is a one-line stale-URL fix bundled here because it would otherwise need its own PR for a single literal change.Test plan
ruff check mempalace/kg_triple_worker.py tests/test_kg_triple_worker.py— cleanruff format --check— cleanpytest tests/test_kg_triple_worker.py— 23/23 passing, including new coverage for:max_concurrencydb_pool_size < max_concurrencyraisesValueError--db-pool-sizeflag threads through topool_factorydb_pool_sizeismax_concurrency + 8Related
AsyncConnectionPool)run_workerthe new bottleneckfamiliar(Tailscale)