Add timeout to WeakFileLock (#2751)

* Add timeout to WeakFileLock

* mypy

Author: Wauplin

src/huggingface_hub/_local_folder.py

@@ -413,9 +413,14 @@ def _huggingface_dir(local_dir: Path) -> Path:
     gitignore_lock = path / ".gitignore.lock"
     if not gitignore.exists():
         try:
-            with WeakFileLock(gitignore_lock):
+            with WeakFileLock(gitignore_lock, timeout=0.1):
                 gitignore.write_text("*")
+        except IndexError:
+            pass
+        except OSError:  # TimeoutError, FileNotFoundError, PermissionError, etc.
+            pass
+        try:
             gitignore_lock.unlink()
-        except OSError:  # FileNotFoundError, PermissionError, etc.
+        except OSError:
             pass
     return path

src/huggingface_hub/utils/_fixes.py

@@ -13,6 +13,7 @@
 import shutil
 import stat
 import tempfile
+import time
 from functools import partial
 from pathlib import Path
 from typing import Callable, Generator, Optional, Union
@@ -83,39 +84,50 @@ def _set_write_permission_and_retry(func, path, excinfo):
 
 
 @contextlib.contextmanager
-def WeakFileLock(lock_file: Union[str, Path]) -> Generator[BaseFileLock, None, None]:
+def WeakFileLock(
+    lock_file: Union[str, Path], *, timeout: Optional[float] = None
+) -> Generator[BaseFileLock, None, None]:
     """A filelock with some custom logic.
 
     This filelock is weaker than the default filelock in that:
     1. It won't raise an exception if release fails.
     2. It will default to a SoftFileLock if the filesystem does not support flock.
 
     An INFO log message is emitted every 10 seconds if the lock is not acquired immediately.
+    If a timeout is provided, a `filelock.Timeout` exception is raised if the lock is not acquired within the timeout.
     """
-    lock = FileLock(lock_file, timeout=constants.FILELOCK_LOG_EVERY_SECONDS)
+    log_interval = constants.FILELOCK_LOG_EVERY_SECONDS
+    lock = FileLock(lock_file, timeout=log_interval)
+    start_time = time.time()
+
     while True:
+        elapsed_time = time.time() - start_time
+        if timeout is not None and elapsed_time >= timeout:
+            raise Timeout(str(lock_file))
+
         try:
-            lock.acquire()
+            lock.acquire(timeout=min(log_interval, timeout - elapsed_time) if timeout else log_interval)
         except Timeout:
-            logger.info("still waiting to acquire lock on %s", lock_file)
+            logger.info(
+                f"Still waiting to acquire lock on {lock_file} (elapsed: {time.time() - start_time:.1f} seconds)"
+            )
         except NotImplementedError as e:
             if "use SoftFileLock instead" in str(e):
-                # It's possible that the system does support flock, expect for one partition or filesystem.
-                # In this case, let's default to a SoftFileLock.
                 logger.warning(
                     "FileSystem does not appear to support flock. Falling back to SoftFileLock for %s", lock_file
                 )
-                lock = SoftFileLock(lock_file, timeout=constants.FILELOCK_LOG_EVERY_SECONDS)
+                lock = SoftFileLock(lock_file, timeout=log_interval)
                 continue
         else:
             break
 
-    yield lock
-
     try:
-        return lock.release()
-    except OSError:
+        yield lock
+    finally:
         try:
-            Path(lock_file).unlink()
+            lock.release()
         except OSError:
-            pass
+            try:
+                Path(lock_file).unlink()
+            except OSError:
+                pass

tests/test_local_folder.py

@@ -19,6 +19,7 @@
 
 import logging
 import os
+import threading
 import time
 from pathlib import Path, WindowsPath
 
@@ -49,6 +50,18 @@ def test_creates_huggingface_dir_with_gitignore(tmp_path: Path):
     assert (huggingface_dir / ".gitignore").read_text() == "*"
 
 
+def test_gitignore_lock_timeout_is_ignored(tmp_path: Path):
+    local_dir = tmp_path / "path" / "to" / "local"
+
+    threads = [threading.Thread(target=_huggingface_dir, args=(local_dir,)) for _ in range(10)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+    assert (local_dir / ".cache" / "huggingface" / ".gitignore").exists()
+    assert not (local_dir / ".cache" / "huggingface" / ".gitignore.lock").exists()
+
+
 def test_local_download_paths(tmp_path: Path):
     """Test local download paths are valid + usable."""
     paths = get_local_download_paths(tmp_path, "path/in/repo.txt")

tests/test_utils_fixes.py

@@ -1,7 +1,11 @@
+import logging
 import unittest
 from pathlib import Path
 
-from huggingface_hub.utils import SoftTemporaryDirectory, yaml_dump
+import filelock
+import pytest
+
+from huggingface_hub.utils import SoftTemporaryDirectory, WeakFileLock, yaml_dump
 
 
 class TestYamlDump(unittest.TestCase):
@@ -24,3 +28,21 @@ def test_temporary_directory(self) -> None:
             self.assertTrue(path.is_dir())
         # Tmpdir is deleted
         self.assertFalse(path.is_dir())
+
+
+class TestWeakFileLock:
+    def test_lock_log_every(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        monkeypatch.setattr("huggingface_hub.constants.FILELOCK_LOG_EVERY_SECONDS", 0.1)
+        lock_file = tmp_path / ".lock"
+
+        with caplog.at_level(logging.INFO, logger="huggingface_hub.utils._fixes"):
+            with WeakFileLock(lock_file):
+                with pytest.raises(filelock.Timeout) as exc_info:
+                    with WeakFileLock(lock_file, timeout=0.3):
+                        pass
+                assert exc_info.value.lock_file == str(lock_file)
+
+        assert len(caplog.records) >= 3
+        assert caplog.records[0].message.startswith(f"Still waiting to acquire lock on {lock_file}")

utils/check_contrib_list.py

@@ -55,7 +55,7 @@ def check_contrib_list(update: bool) -> NoReturn:
 
     # Check workflow is consistent with list
     workflow_content = WORKFLOW_PATH.read_text()
-    _substitute = "\n".join(f'{" "*10}"{lib}",' for lib in contrib_list)
+    _substitute = "\n".join(f'{" " * 10}"{lib}",' for lib in contrib_list)
     workflow_content_expected = WORKFLOW_REGEX.sub(rf"\g<before>{_substitute}\n\g<after>", workflow_content)
 
     #