Use filelocks for queues in MapInfra

exca/map.py

@@ -11,19 +11,16 @@
 import itertools
 import logging
 import os
-import pickle
 import typing as tp
-import uuid
 from concurrent import futures
 from pathlib import Path
 
 import numpy as np
 import pydantic
-import submitit
-from submitit.core import utils
 
 from . import base, slurm
 from .cachedict import CacheDict
+from .utils import LockManager
 
 MapFunc = tp.Callable[[tp.Sequence[tp.Any]], tp.Iterator[tp.Any]]
 X = tp.TypeVar("X")
@@ -65,45 +62,6 @@ def __call__(self, items: tp.Sequence[tp.Any]) -> tp.Iterator[tp.Any]:
         return self.infra._method_override(items)
 
 
-class JobChecker:
-    """Keeps a record of running jobs in a folder
-    and enables waiting for them to complete.
-    """
-
-    def __init__(self, folder: Path | str) -> None:
-        basefolder = utils.JobPaths.get_first_id_independent_folder(folder)
-        self.folder = basefolder / "running-jobs"
-
-    def add(self, jobs: tp.Iterable[tp.Any]) -> None:
-        """Add jobs to the list of running jobs"""
-        self.folder.mkdir(exist_ok=True, parents=True)
-        for job in jobs:
-            if not job.done():
-                job_path = self.folder / (uuid.uuid4().hex[:8] + ".pkl")
-                with job_path.open("wb") as f:
-                    pickle.dump(job, f)
-
-    def wait(self) -> bool:
-        """Wait for completion of running jobs"""
-        waited = False
-        for fp in self.folder.glob("*.pkl"):
-            try:  # avoid concurrency issues with deleted items
-                with fp.open("rb") as f:
-                    job: tp.Any = pickle.load(f)
-            except Exception:  # pylint: disable=broad-except
-                continue
-            if not job.done():
-                msg = "Waiting for completion of pre-existing map job: %s\nin '%s'"
-                logger.info(msg, job, self.folder)
-                job.wait()
-                waited = True
-            # delete the file as it is not needed anymore
-            fp.unlink(missing_ok=True)
-        if waited:
-            logger.info("Waiting is over")
-        return waited
-
-
 def to_chunks(
     items: tp.List[X], *, max_chunks: int | None, min_items_per_chunk: int = 1
 ) -> tp.Iterator[tp.List[X]]:
@@ -188,6 +146,7 @@ class MapInfra(base.BaseInfra, slurm.SubmititMixin):
     # - force: cache is ignored, and result is (re)computed (and cached)
     # - read-only: never compute anything
     mode: Mode = "cached"
+    lock_timeout: int = 3600
 
     # internals
     _recomputed: tp.Set[str] = set()  # for mode="force"
@@ -333,13 +292,6 @@ def _find_missing(self, items: tp.Dict[str, tp.Any]) -> tp.Dict[str, tp.Any]:
         if missing:
             if self.mode == "read-only":
                 raise RuntimeError(f"{self.mode=} but found {len(missing)} missing items")
-            executor: submitit.Executor | None = self.executor()
-            if executor is not None:  # wait for items being computed
-                jcheck = JobChecker(folder=executor.folder)
-                jcheck.wait()
-                # update cache dict and recheck as actual checking for keys updates the dict
-                keys = set(self.cache_dict)  # update cache dict
-                missing = {k: item for k, item in missing.items() if k not in keys}
         if len(items) == len(missing) == 1 and self.forbid_single_item_computation:
             key, item = next(iter(missing.items()))
             raise RuntimeError(
@@ -398,8 +350,6 @@ def _method_override(self, *args: tp.Any, **kwargs: tp.Any) -> tp.Iterator[tp.An
                     # select a batch/chunk of samples_per_job items to send to a job
                     j = executor.submit(self._call_and_store, chunk, use_cache_dict=True)
                     jobs.append(j)
-            jcheck = JobChecker(folder=executor.folder)
-            jcheck.add(jobs)
             # pylint: disable=expression-not-assigned
             uid = self.uid()
             msg = "Sent %s samples for %s into %s jobs on cluster '%s' (eg: %s)"
@@ -495,14 +445,46 @@ def _call_and_store(
         if imethod is None:
             raise RuntimeError(f"Infra was not applied: {self!r}")
         item_uid = imethod.item_uid
-        if items:  # make sure some overlapping job did not already run stuff
-            keys = set(d)  # update cache dict
+
+        # Filter out already cached items
+        if items:
+            keys = set(d)
             items = [item for item in items if item_uid(item) not in keys]
+
+        # No locking needed
+        if not items or not use_cache_dict or self.cache_dict.folder is None:
+            return self._process_items(items, d, item_uid)
+
+        # With locking: delegate to LockManager
+        lock_manager = LockManager(
+            lock_dir=self.cache_dict.folder / "locks",
+            item_uid=item_uid,
+            cache_contains=lambda uid: uid in self.cache_dict,
+            lock_timeout=self.lock_timeout,
+        )
+        result = lock_manager.process_with_locks(
+            items, lambda itms: self._process_items(itms, d, item_uid)
+        )
+        return {} if isinstance(d, CacheDict) else result
+
+    def _process_items(
+        self,
+        items: tp.Sequence[tp.Any],
+        d: dict[str, tp.Any],
+        item_uid: tp.Callable[[tp.Any], str],
+    ) -> dict[str, tp.Any]:
+        """Process items and store results"""
+        if not items:
+            return {}
+
         if isinstance(self, slurm.SubmititMixin):  # dependence to mixin
-            if self.workdir is not None and self.cluster is not None and items:
+            if self.workdir is not None and self.cluster is not None:
                 logger.info("Running from working directory: '%s'", os.getcwd())
+
         outputs = self._run_method(items)
         sentinel = base.Sentinel()
+        result = {}
+
         with contextlib.ExitStack() as estack:
             writer = d
             if isinstance(d, CacheDict):
@@ -512,9 +494,13 @@ def _call_and_store(
                 if item is sentinel or output is sentinel:
                     msg = f"Cached function did not yield exactly once per item: {item=!r}, {output=!r}"
                     raise RuntimeError(msg)
-                writer[item_uid(item)] = output
-        # don't return the whole cache dict if data is cached
-        return {} if use_cache_dict else d
+                uid = item_uid(item)
+                writer[uid] = output
+                if not isinstance(d, CacheDict):
+                    # don't return the whole output dict if data is cached
+                    result[uid] = output
+
+        return result
 
 
 @dataclasses.dataclass

exca/test_locking.py

@@ -0,0 +1,178 @@
+import os
+import threading
+import time
+import typing as tp
+from pathlib import Path
+
+import filelock
+import pydantic
+
+from .map import MapInfra
+
+
+class Worker(pydantic.BaseModel):
+    infra: MapInfra = MapInfra()
+    computed: list[int] = []
+
+    @infra.apply(item_uid=str)
+    def process(self, items: tp.Sequence[int]) -> tp.Iterator[int]:
+        for item in items:
+            self.computed.append(item)
+            time.sleep(0.1)  # Simulate work
+            yield item * 2
+
+
+def test_locking_wait(tmp_path: Path) -> None:
+    """Test that a worker waits if item is locked"""
+    worker = Worker(infra={"folder": tmp_path})
+
+    # Get the actual cache folder (triggers creation)
+    cache_folder = worker.infra.cache_dict.folder
+    assert cache_folder is not None
+    (cache_folder / "locks").mkdir(parents=True, exist_ok=True)
+
+    # Lock item "2" manually
+    lock_path = cache_folder / "locks" / "2.lock"
+    lock = filelock.FileLock(lock_path)
+    lock.acquire()
+
+    # Function to run worker in a thread
+    def run_worker():
+        # This should block on item 2 until we release it
+        list(worker.process([2]))
+
+    t = threading.Thread(target=run_worker)
+    t.start()
+
+    # Give it a moment to start and block
+    time.sleep(0.5)
+    assert t.is_alive()  # Should be still waiting
+
+    # Release lock
+    lock.release()
+
+    # Should finish now
+    t.join(timeout=2.0)
+    assert not t.is_alive()
+    assert 2 in worker.computed
+
+
+def test_locking_completed(tmp_path: Path) -> None:
+    """Test that if locked item is completed, second worker skips it"""
+    worker1 = Worker(infra={"folder": tmp_path})
+    # Use same config to get same UID folder
+    worker2 = Worker(infra={"folder": tmp_path})
+
+    cache_folder = worker1.infra.cache_dict.folder
+    assert cache_folder is not None
+    (cache_folder / "locks").mkdir(parents=True, exist_ok=True)
+
+    # Lock item "2"
+    lock_path = cache_folder / "locks" / "2.lock"
+    lock = filelock.FileLock(lock_path)
+    lock.acquire()
+
+    def run_worker2():
+        # Should wait then skip
+        list(worker2.process([2]))
+
+    t = threading.Thread(target=run_worker2)
+    t.start()
+
+    time.sleep(0.5)
+    assert t.is_alive()
+
+    # Pretend we computed it and wrote to cache
+    with worker1.infra.cache_dict.writer() as w:
+        w["2"] = 4
+
+    # Release lock
+    lock.release()
+
+    t.join(timeout=2.0)
+    assert not t.is_alive()
+    # Worker 2 should NOT have computed it
+    assert 2 not in worker2.computed
+
+
+def test_incremental_processing(tmp_path: Path) -> None:
+    """Test that workers can process non-overlapping items in parallel"""
+    worker1 = Worker(infra={"folder": tmp_path})
+    worker2 = Worker(infra={"folder": tmp_path})
+
+    start_time = time.time()
+
+    def run_worker1():
+        # Process items [1, 2, 3, 4]
+        list(worker1.process([1, 2, 3, 4]))
+
+    def run_worker2():
+        # Process items [3, 4, 5, 6] - overlaps with worker1 on [3, 4]
+        time.sleep(0.05)  # Start slightly after worker1
+        list(worker2.process([3, 4, 5, 6]))
+
+    t1 = threading.Thread(target=run_worker1)
+    t2 = threading.Thread(target=run_worker2)
+
+    t1.start()
+    t2.start()
+
+    t1.join()
+    t2.join()
+
+    elapsed = time.time() - start_time
+
+    # Worker1 processes [1, 2, 3, 4]
+    # Worker2 should process [5, 6] while worker1 works on [1, 2]
+    # Then worker2 gets [3, 4] from cache (worker1 already computed them)
+    # So worker2 only computes [5, 6]
+
+    # Check that worker2 only computed non-overlapping items
+    assert 5 in worker2.computed
+    assert 6 in worker2.computed
+    # Worker2 should NOT have computed 3 or 4 (worker1 did)
+    assert 3 not in worker2.computed
+    assert 4 not in worker2.computed
+
+    # Verify all items were processed by someone
+    all_results = set(worker1.computed) | set(worker2.computed)
+    assert all_results >= {1, 2, 3, 4, 5, 6}
+
+    # Parallel execution should be faster than sequential
+    # Each worker processes 0.1s per item
+    # Worker1: 4 items = 0.4s
+    # Worker2: 2 items (5, 6) = 0.2s, runs in parallel with worker1
+    # Total should be ~0.5s if incremental, ~0.6s if sequential
+    assert elapsed < 1.0, f"Took {elapsed}s, should be < 1.0s if processing incrementally"
+
+
+def test_stale_lock_detection(tmp_path: Path) -> None:
+    """Test that stale locks from crashed workers are detected and removed"""
+    worker = Worker(
+        infra={"folder": tmp_path, "lock_timeout": 2}
+    )  # Short timeout for test
+
+    cache_folder = worker.infra.cache_dict.folder
+    assert cache_folder is not None
+    lock_dir = cache_folder / "locks"
+    lock_dir.mkdir(parents=True, exist_ok=True)
+
+    # Create a "stale" lock by creating the lock file and making it old
+    stale_lock_path = lock_dir / "5.lock"
+    stale_lock_path.touch()
+
+    # Make the lock file old (older than lock_timeout)
+    old_time = time.time() - 5  # 5 seconds ago
+    os.utime(stale_lock_path, (old_time, old_time))
+
+    # Worker should detect stale lock, remove it, and process the item
+    start = time.time()
+    result = list(worker.process([5]))
+    elapsed = time.time() - start
+
+    # Should complete quickly (not wait for full timeout)
+    assert elapsed < 1.0, f"Took {elapsed}s, should not wait for stale lock"
+
+    # Item should have been processed
+    assert 5 in worker.computed
+    assert result == [10]  # 5 * 2