fix(pipelines): lock datasets across workers - #4167
Draft
GautamSharma99 wants to merge 1 commit into
Draft
Conversation
Signed-off-by: GautamSharma99 <gautamsharma99067@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Pipeline execution currently uses an asyncio.Lock keyed by dataset ID. Because this registry exists only within one Python
process, separate FastAPI workers can run pipelines against the same dataset concurrently.
This PR introduces cross-worker locking for per-dataset pipeline execution while retaining an in-process lock as a lightweight
first layer.
Problem
Concurrent pipeline runs for the same dataset can currently be started by different FastAPI workers. The process-local lock
does not coordinate those workers, which may lead to:
The existing _dataset_locks registry also retains one asyncio.Lock for every dataset processed and never removes unused
entries.
Changes
Added a database-aware dataset_pipeline_lock abstraction.
Added PostgreSQL cross-worker locking using session-level advisory locks.
Added host-wide locking for file-backed SQLite and Turso using OS advisory file locks.
Retained an in-process asyncio.Lock as an optimization so only one task per worker contends for the cross-worker resource.
Made the local lock registry reference-counted.
Preserved the existing re-entrant behavior for nested pipelines operating on the same dataset.
Updated the concurrent cognify test documentation to reflect cross-worker locking.
Backend Behavior
Relational backend Coordination scope
━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
PostgreSQL Across workers and hosts sharing the same database
────────────────────────── ─────────────────────────────────────────────────────────────────
File-backed SQLite/Turso Across workers on the same host
────────────────────────── ─────────────────────────────────────────────────────────────────
In-memory SQLite Current process only, because the database itself is not shared
Testing
Added tests covering:
Commands run:
uv run ruff check cognee/infrastructure/locks/dataset_pipeline_lock.py
cognee/modules/pipelines/operations/pipeline.py
cognee/tests/unit/infrastructure/locks/test_dataset_pipeline_lock.py
cognee/tests/unit/modules/pipelines/test_pipeline_dataset_lock.py
cognee/tests/test_concurrent_cognify_e2e.py
uv run ruff format --check cognee/infrastructure/locks/dataset_pipeline_lock.py
cognee/modules/pipelines/operations/pipeline.py
cognee/tests/unit/infrastructure/locks/test_dataset_pipeline_lock.py
cognee/tests/unit/modules/pipelines/test_pipeline_dataset_lock.py
cognee/tests/test_concurrent_cognify_e2e.py
uv run pytest
cognee/tests/unit/modules/pipelines
cognee/tests/unit/modules/cognify
cognee/tests/unit/modules/memify_tasks
cognee/tests/unit/memify_pipelines
-q --timeout=60
Result:
132 passed.The complete unit suite could not be collected in the current local environment because the optional
plotlyandneo4jpackages are not installed. The failure occurs during test collection and is unrelated to these changes.
Impact
Checklist