Skip to content

fix(pipelines): lock datasets across workers - #4167

Draft
GautamSharma99 wants to merge 1 commit into
topoteretes:mainfrom
GautamSharma99:fix/cross-worker-dataset-pipeline-lock
Draft

fix(pipelines): lock datasets across workers#4167
GautamSharma99 wants to merge 1 commit into
topoteretes:mainfrom
GautamSharma99:fix/cross-worker-dataset-pipeline-lock

Conversation

@GautamSharma99

Copy link
Copy Markdown

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:

  • duplicate extraction and processing;
  • concurrent graph and vector writes;
  • inconsistent pipeline run statuses;
  • database write conflicts.

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.

    • Lock keys are deterministic, signed 64-bit values derived from dataset UUIDs.
    • The advisory lock uses a dedicated database connection for its lifetime.
    • The transaction opened while acquiring the lock is committed immediately.
    • The lock is released in finally, including when pipeline execution fails.
  • Added host-wide locking for file-backed SQLite and Turso using OS advisory file locks.

    • Lock acquisition uses non-blocking polling so the event loop is not blocked.
    • Lock paths use 256 deterministic stripes, preventing an unbounded number of lock files.
    • Stripe collisions may serialize unrelated datasets but cannot weaken same-dataset exclusion.
  • 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.

    • Entries are removed after the final holder or waiter exits.
    • Cancelled waiters do not leave stale entries.
  • 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:

  • serialization of concurrent local callers;
  • cleanup of local registry entries;
  • cleanup when a waiting task is cancelled;
  • SQLite advisory-lock contention;
  • cancellation without orphaning a SQLite lock;
  • bounded SQLite lock-file stripes;
  • PostgreSQL advisory-lock release after pipeline failure;
  • lock lifetime around the complete pipeline generator;
  • nested same-dataset pipeline execution without lock reacquisition.

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 plotly and neo4j
packages are not installed. The failure occurs during test collection and is unrelated to these changes.

Impact

  • Prevents simultaneous same-dataset pipeline execution across FastAPI workers.
  • Reduces the risk of duplicate processing and conflicting graph/vector writes.
  • Prevents unbounded growth of the process-local dataset lock registry.
  • Does not change the public API, MCP server, or frontend.
  • Different datasets can still execute concurrently, except for occasional SQLite lock-stripe collisions.

Checklist

  • Code formatted with Ruff
  • Ruff checks pass
  • Targeted regression tests added
  • Relevant pipeline, cognify, and memify tests pass
  • Existing nested pipeline behavior preserved
  • Commit signed
  • DCO sign-off included

Signed-off-by: GautamSharma99 <gautamsharma99067@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant