This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 021d78d [Test] Add test for filelock utility (#216)
021d78d is described below
commit 021d78dbb8c5c5ba9ea07f8035eb9ce01a48bbc7
Author: Yaoyao Ding <[email protected]>
AuthorDate: Tue Nov 4 11:42:34 2025 -0500
[Test] Add test for filelock utility (#216)
This PR adds tests for the filelock utility, to make sure it works well
on multiple platforms.
---
python/tvm_ffi/utils/lockfile.py | 41 +++++++++--
tests/python/utils/filelock_worker.py | 49 +++++++++++++
tests/python/utils/test_filelock.py | 126 ++++++++++++++++++++++++++++++++++
3 files changed, 210 insertions(+), 6 deletions(-)
diff --git a/python/tvm_ffi/utils/lockfile.py b/python/tvm_ffi/utils/lockfile.py
index da98491..5fa69d3 100644
--- a/python/tvm_ffi/utils/lockfile.py
+++ b/python/tvm_ffi/utils/lockfile.py
@@ -34,7 +34,9 @@ class FileLock:
"""Provide a cross-platform file locking mechanism using Python's stdlib.
This class implements an advisory lock, which must be respected by all
- cooperating processes.
+ cooperating processes. Please note that this lock does not prevent the
same process
+ from acquiring the lock multiple times; it is the caller's responsibility
to
+ manage this.
Examples
--------
@@ -71,9 +73,20 @@ class FileLock:
def acquire(self) -> bool:
"""Acquire an exclusive, non-blocking lock on the file.
- Returns True if the lock was acquired, False otherwise.
+ Returns
+ -------
+ ret: bool
+ True if the lock was acquired, False otherwise.
+
+ Raises
+ ------
+ RuntimeError
+ If an unexpected error occurs during lock acquisition.
+
"""
try:
+ if self._file_descriptor is not None:
+ return False # Lock is already held by this instance
if sys.platform == "win32":
self._file_descriptor = os.open(
self.lock_file_path, os.O_RDWR | os.O_CREAT | os.O_BINARY
@@ -97,12 +110,28 @@ class FileLock:
def blocking_acquire(self, timeout: float | None = None, poll_interval:
float = 0.1) -> bool:
"""Wait until an exclusive lock can be acquired, with an optional
timeout.
- Args:
- timeout (float): The maximum time to wait for the lock in seconds.
- A value of None means wait indefinitely.
- poll_interval (float): The time to wait between lock attempts in
seconds.
+ Parameters
+ ----------
+ timeout: float, optional
+ The maximum time to wait for the lock in seconds. A value of None
means wait indefinitely.
+ poll_interval: float
+ The time to wait between lock attempts in seconds.
+
+ Returns
+ -------
+ ret: bool
+ True if the lock was acquired.
+
+ Raises
+ ------
+ TimeoutError
+ If the lock is not acquired within the timeout period.
+ RuntimeError
+ If the lock is already held by this instance.
"""
+ if self._file_descriptor is not None:
+ raise RuntimeError("Lock is already held by this instance.")
start_time = time.time()
while True:
if self.acquire():
diff --git a/tests/python/utils/filelock_worker.py
b/tests/python/utils/filelock_worker.py
new file mode 100644
index 0000000..46ff29b
--- /dev/null
+++ b/tests/python/utils/filelock_worker.py
@@ -0,0 +1,49 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import random
+import sys
+import time
+
+from tvm_ffi.utils.lockfile import FileLock
+
+
+def worker(worker_id: int, lock_path: str, counter_file: str) -> int:
+ """Worker function that tries to acquire lock and increment counter."""
+ try:
+ with FileLock(lock_path):
+ # Critical section - read, increment, write counter
+ with open(counter_file) as f: # noqa: PTH123
+ current_value = int(f.read().strip())
+
+ time.sleep(random.uniform(0.01, 0.1)) # Simulate some work
+
+ with open(counter_file, "w") as f: # noqa: PTH123
+ f.write(str(current_value + 1))
+
+ print(f"Worker {worker_id}: success")
+ return 0
+ except Exception as e:
+ print(f"Worker {worker_id}: error: {e}")
+ return 1
+
+
+if __name__ == "__main__":
+ worker_id = int(sys.argv[1])
+ lock_path = sys.argv[2]
+ counter_file = sys.argv[3]
+ sys.exit(worker(worker_id, lock_path, counter_file))
diff --git a/tests/python/utils/test_filelock.py
b/tests/python/utils/test_filelock.py
new file mode 100644
index 0000000..fa68edf
--- /dev/null
+++ b/tests/python/utils/test_filelock.py
@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for the FileLock utility."""
+
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import pytest
+from tvm_ffi.utils.lockfile import FileLock
+
+
+def test_basic_lock_acquire_and_release() -> None:
+ """Test basic lock acquisition and release."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ lock_path = Path(temp_dir) / "test.lock"
+ lock = FileLock(str(lock_path))
+
+ # Test acquire
+ assert lock.acquire() is True
+ assert lock._file_descriptor is not None
+
+ # Test release
+ lock.release()
+ assert lock._file_descriptor is None
+
+
+def test_context_manager() -> None:
+ """Test FileLock as a context manager."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ lock_path = Path(temp_dir) / "test.lock"
+
+ with FileLock(str(lock_path)) as lock:
+ assert lock._file_descriptor is not None
+ # Lock should be acquired
+ assert lock_path.exists()
+
+ # Lock should be released after exiting context
+ # Note: file may still exist but should be unlocked
+
+
+def test_multiple_acquire_attempts() -> None:
+ """Test multiple acquire attempts on the same lock instance."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ lock_path = Path(temp_dir) / "test.lock"
+ lock = FileLock(str(lock_path))
+
+ # First acquire should succeed
+ assert lock.acquire() is True
+
+ # Second acquire on same instance should fail
+ # (can't acquire same lock twice)
+ assert lock.acquire() is False
+
+ lock.release()
+
+
+def test_exception_in_context_manager() -> None:
+ """Test that lock is properly released even when exception occurs."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ lock_path = Path(temp_dir) / "test.lock"
+
+ # Test that exception is propagated and lock is released
+ with pytest.raises(ValueError, match="test exception"):
+ with FileLock(str(lock_path)) as lock:
+ assert lock._file_descriptor is not None
+ raise ValueError("test exception")
+
+ # Lock should be released, so we can acquire it again
+ lock2 = FileLock(str(lock_path))
+ assert lock2.acquire() is True
+ lock2.release()
+
+
+def test_concurrent_access() -> None:
+ """Test concurrent access from multiple processes."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ lock_path = Path(temp_dir) / "test.lock"
+ counter_file = Path(temp_dir) / "counter.txt"
+
+ # Initialize counter file
+ counter_file.write_text("0")
+
+ # Create worker script content
+
+ # Write worker script to a temporary file
+ worker_script_path = Path(__file__).parent / "filelock_worker.py"
+
+ # Run multiple worker processes concurrently
+ num_workers = 16
+ processes = []
+ for i in range(num_workers):
+ p = subprocess.Popen(
+ [sys.executable, str(worker_script_path), str(i),
str(lock_path), str(counter_file)]
+ )
+ processes.append(p)
+
+ # Wait for all processes to complete
+ for p in processes:
+ p.wait(timeout=60.0)
+ assert p.returncode == 0, f"Worker process failed with return code
{p.returncode}"
+
+ # Check final counter value
+ final_count = int(counter_file.read_text().strip())
+
+ # Counter should equal number of workers (no race conditions)
+ assert final_count == num_workers
+
+
+if __name__ == "__main__":
+ pytest.main([__file__])
