fix: write lock for warmup

siibra/retrieval/cache.py

@@ -23,6 +23,7 @@
 from typing import Callable, List, NamedTuple, Union
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
+from filelock import FileLock as Lock
 
 from ..commons import logger, SIIBRA_CACHEDIR, SKIP_CACHEINIT_MAINTENANCE, siibra_tqdm
 from ..exceptions import WarmupRegException
@@ -146,6 +147,7 @@ def build_filename(self, str_rep: str, suffix=None):
 
 
 class WarmupLevel(int, Enum):
+    TEST = -1000
     INSTANCE = 1
     DATA = 5
 
@@ -202,16 +204,17 @@ def call_fn(fn: WarmupParam):
             for f in return_val:
                 f()
 
-        with ThreadPoolExecutor(max_workers=max_workers) as ex:
-            for _ in siibra_tqdm(
-                ex.map(
-                    call_fn,
-                    all_fns
-                ),
-                desc="Warming cache",
-                total=len(all_fns),
-            ):
-                ...
+        with Lock(CACHE.build_filename("lockfile", ".warmup")):
+            with ThreadPoolExecutor(max_workers=max_workers) as ex:
+                for _ in siibra_tqdm(
+                    ex.map(
+                        call_fn,
+                        all_fns
+                    ),
+                    desc="Warming cache",
+                    total=len(all_fns),
+                ):
+                    ...
 
 
 try:

siibra/retrieval/requests.py

@@ -42,13 +42,7 @@
 from functools import wraps
 from time import sleep
 import sys
-import platform
-
-if platform.system() == "Linux":
-    from filelock import FileLock as Lock
-else:
-    from filelock import SoftFileLock as Lock
-
+from filelock import FileLock as Lock
 if TYPE_CHECKING:
     from .repositories import GitlabConnector
 

test/retrieval/test_cache.py

@@ -1,5 +1,6 @@
 import pytest
 from unittest.mock import MagicMock
+from multiprocessing import Pool
 
 from siibra.retrieval.cache import WarmupLevel, Warmup, WarmupRegException
 
@@ -126,3 +127,22 @@ def test_register_warmup_called_level_high(register_all):
 
     dummy_data1.assert_called_once()
     dummy_data2.assert_called_once()
+
+
+def test_write_lock():
+    import time
+
+    sleep_time = 1
+
+    @Warmup.register_warmup_fn(WarmupLevel.TEST)
+    def sleep_for_x():
+        time.sleep(sleep_time)
+
+    tstart_s = time.time()
+    with Pool(2) as p:
+        p.map(Warmup.warmup, (WarmupLevel.TEST, WarmupLevel.TEST))
+    tend_s = time.time()
+    expected_baseline = sleep_time * 2
+
+    time_perf = tend_s - tstart_s
+    assert time_perf > expected_baseline, "Expect second call to block"