fix(pipelines): lock datasets across workers

cognee/infrastructure/locks/dataset_pipeline_lock.py

@@ -0,0 +1,142 @@
+"""Cross-worker locking for pipeline runs that target the same dataset."""
+
+import asyncio
+import os
+import threading
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from hashlib import sha256
+from pathlib import Path
+from uuid import UUID
+
+from sqlalchemy import text
+
+from cognee.infrastructure.databases.relational import get_relational_engine
+
+
+_SQLITE_LOCK_STRIPES = 256
+_SQLITE_LOCK_POLL_INTERVAL_SECONDS = 0.05
+
+
+@dataclass
+class _LocalLockEntry:
+    lock: asyncio.Lock
+    users: int = 0
+
+
+# Protecting this small registry with a threading lock avoids binding the guard to
+# one event loop. The registry key includes the loop because asyncio locks cannot
+# safely be shared between loops (which is common in SDK users and unit tests).
+_local_locks: dict[tuple[asyncio.AbstractEventLoop, UUID], _LocalLockEntry] = {}
+_local_locks_guard = threading.Lock()
+
+
+def _postgres_lock_key(dataset_id: UUID) -> int:
+    """Return a deterministic signed bigint key namespaced to pipeline locks."""
+    digest = sha256(b"cognee:dataset-pipeline:" + dataset_id.bytes).digest()
+    return int.from_bytes(digest[:8], byteorder="big", signed=True)
+
+
+def _sqlite_lock_path(engine, dataset_id: UUID) -> Path | None:
+    """Return one of a bounded number of lock files next to the SQLite database."""
+    database = engine.url.database
+    if not database or database == ":memory:":
+        return None
+
+    # A fixed number of stripes prevents an unbounded lock-file leak. A hash
+    # collision only serializes two unrelated datasets; it cannot compromise
+    # same-dataset exclusion.
+    stripe = sha256(dataset_id.bytes).digest()[0] % _SQLITE_LOCK_STRIPES
+    database_path = Path(os.path.abspath(database))
+    return database_path.parent / f".cognee-pipeline-{stripe:02x}.lock"
+
+
+@asynccontextmanager
+async def _local_dataset_lock(dataset_id: UUID) -> AsyncGenerator[None, None]:
+    """Serialize same-process callers and remove unused registry entries."""
+    loop = asyncio.get_running_loop()
+    key = (loop, dataset_id)
+    with _local_locks_guard:
+        entry = _local_locks.get(key)
+        if entry is None:
+            entry = _LocalLockEntry(asyncio.Lock())
+            _local_locks[key] = entry
+        entry.users += 1
+
+    try:
+        async with entry.lock:
+            yield
+    finally:
+        with _local_locks_guard:
+            entry.users -= 1
+            if entry.users == 0:
+                _local_locks.pop(key, None)
+
+
+@asynccontextmanager
+async def _sqlite_advisory_lock(engine, dataset_id: UUID) -> AsyncGenerator[None, None]:
+    """Use a host-wide OS advisory lock for a file-backed SQLite database."""
+    lock_path = _sqlite_lock_path(engine, dataset_id)
+    if lock_path is None:
+        # In-memory SQLite databases cannot be shared between worker processes.
+        yield
+        return
+
+    from filelock import FileLock, Timeout
+
+    lock = FileLock(lock_path, thread_local=False)
+    while True:
+        try:
+            # timeout=0 performs one non-blocking OS lock attempt. Polling with
+            # asyncio avoids retaining one worker thread for every long-running
+            # pipeline and makes cancellation immediate and leak-free.
+            lock.acquire(timeout=0)
+            break
+        except Timeout:
+            await asyncio.sleep(_SQLITE_LOCK_POLL_INTERVAL_SECONDS)
+
+    try:
+        yield
+    finally:
+        lock.release()
+
+
+@asynccontextmanager
+async def _cross_worker_dataset_lock(dataset_id: UUID) -> AsyncGenerator[None, None]:
+    """Acquire the database-appropriate cross-worker lock for ``dataset_id``."""
+    engine = get_relational_engine().engine
+    if engine.dialect.name == "postgresql":
+        lock_key = _postgres_lock_key(dataset_id)
+        async with engine.connect() as connection:
+            await connection.execute(text("SELECT pg_advisory_lock(:key)"), {"key": lock_key})
+            # Do not hold an idle transaction for the lifetime of a pipeline run.
+            await connection.commit()
+            try:
+                yield
+            finally:
+                await connection.execute(text("SELECT pg_advisory_unlock(:key)"), {"key": lock_key})
+                await connection.commit()
+        return
+
+    if engine.dialect.name == "sqlite":
+        async with _sqlite_advisory_lock(engine, dataset_id):
+            yield
+        return
+
+    # Unknown relational backends retain process-local safety. SQLAlchemy
+    # adapters currently support PostgreSQL and SQLite, so this is defensive.
+    yield
+
+
+@asynccontextmanager
+async def dataset_pipeline_lock(dataset_id: UUID) -> AsyncGenerator[None, None]:
+    """Serialize pipeline runs for a dataset across tasks and API workers.
+
+    The cheap in-process lock is acquired first, so only one task per worker
+    contends for the cross-worker resource. PostgreSQL advisory locks coordinate
+    across hosts; file-backed SQLite uses host-wide file locks.
+    """
+    async with _local_dataset_lock(dataset_id):
+        async with _cross_worker_dataset_lock(dataset_id):
+            yield

cognee/modules/pipelines/operations/pipeline.py

@@ -1,8 +1,9 @@
 import asyncio
 from contextvars import ContextVar
+from typing import Any, AsyncIterator, Awaitable, Callable, Optional, Union
 from uuid import UUID
-from typing import AsyncIterator, Awaitable, Callable, Optional, Union
 
+from cognee.infrastructure.locks.dataset_pipeline_lock import dataset_pipeline_lock
 from cognee.modules.pipelines.layers.setup_and_check_environment import (
     setup_and_check_environment,
 )
@@ -23,39 +24,19 @@
 from cognee.modules.pipelines.layers.check_pipeline_run_qualification import (
     check_pipeline_run_qualification,
 )
-from typing import Any
 
 logger = get_logger("cognee.pipeline")
 
 update_status_lock = asyncio.Lock()
 
-# Per-dataset locks so concurrent pipeline runs on the SAME dataset are serialized:
-# a run waits until any in-flight run for that dataset finishes, while different
-# datasets still run in parallel.
-# NOTE: process-local only (asyncio) — this does NOT protect against multiple
-# processes/workers running against the same dataset. To be replaced by a
-# cross-process mechanism (e.g. DB-backed lock) later.
-_dataset_locks: dict[UUID, asyncio.Lock] = {}
-_dataset_locks_guard = asyncio.Lock()
-
 # Tracks the dataset ids whose per-dataset lock is already held by the current
 # execution. A pipeline task may legitimately start another pipeline on the same
 # dataset (e.g. cognify_session -> add()/cognify()); without this, re-acquiring the
-# non-reentrant _dataset_locks[dataset_id] from the same execution self-deadlocks.
+# non-reentrant dataset lock from the same execution self-deadlocks.
 # ContextVar propagates into the child tasks run_tasks spawns via asyncio.create_task.
 _held_datasets: ContextVar[frozenset] = ContextVar("_held_datasets", default=frozenset())
 
 
-async def _get_dataset_lock(dataset_id: UUID) -> asyncio.Lock:
-    """Return the asyncio.Lock for a dataset, creating it on first use."""
-    async with _dataset_locks_guard:
-        lock = _dataset_locks.get(dataset_id)
-        if lock is None:
-            lock = asyncio.Lock()
-            _dataset_locks[dataset_id] = lock
-        return lock
-
-
 async def _drive_marking_held(dataset_id: UUID, source: AsyncIterator[Any]) -> AsyncIterator[Any]:
     """Yield from ``source`` while ``dataset_id`` is recorded as locked.
 
@@ -183,6 +164,6 @@ async def _run_body():
 
     # External run: serialize on the per-dataset lock, marking the dataset held so
     # any nested run on it takes the re-entrant path above.
-    async with await _get_dataset_lock(dataset.id):
+    async with dataset_pipeline_lock(dataset.id):
         async for run_info in _drive_marking_held(dataset.id, _run_body()):
             yield run_info

cognee/tests/test_concurrent_cognify_e2e.py

@@ -5,9 +5,10 @@
     add(file_two)   -> cognify(run_in_background=True)   # run #2 starts, returns immediately
     gather both runs to completion
 
-The per-dataset asyncio lock in `run_pipeline_per_dataset` serializes everything
-touching the dataset, so the two runs cannot clobber each other. We assert, purely
-from each run's return value (its PipelineRunCompleted.data_ingestion_info):
+The per-dataset pipeline lock in `run_pipeline_per_dataset` serializes everything
+touching the dataset, including runs started by different worker processes, so the
+two runs cannot clobber each other. We assert, purely from each run's return value
+(its PipelineRunCompleted.data_ingestion_info):
 
     * both runs complete successfully (no errored run),
     * they are two distinct pipeline runs, and