提交 8c276acb authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Lazy filelock import

上级 175b67b9
...@@ -8,8 +8,6 @@ import threading ...@@ -8,8 +8,6 @@ import threading
from contextlib import contextmanager from contextlib import contextmanager
from pathlib import Path from pathlib import Path
import filelock
from pytensor.configdefaults import config from pytensor.configdefaults import config
...@@ -35,8 +33,9 @@ def force_unlock(lock_dir: os.PathLike): ...@@ -35,8 +33,9 @@ def force_unlock(lock_dir: os.PathLike):
lock_dir : os.PathLike lock_dir : os.PathLike
Path to a directory that was locked with `lock_ctx`. Path to a directory that was locked with `lock_ctx`.
""" """
from filelock import FileLock
fl = filelock.FileLock(Path(lock_dir) / ".lock") fl = FileLock(Path(lock_dir) / ".lock")
fl.release(force=True) fl.release(force=True)
dir_key = f"{lock_dir}-{os.getpid()}" dir_key = f"{lock_dir}-{os.getpid()}"
...@@ -62,6 +61,8 @@ def lock_ctx( ...@@ -62,6 +61,8 @@ def lock_ctx(
Timeout in seconds for waiting in lock acquisition. Timeout in seconds for waiting in lock acquisition.
Defaults to `pytensor.config.compile__timeout`. Defaults to `pytensor.config.compile__timeout`.
""" """
from filelock import FileLock
if lock_dir is None: if lock_dir is None:
lock_dir = config.compiledir lock_dir = config.compiledir
...@@ -73,7 +74,7 @@ def lock_ctx( ...@@ -73,7 +74,7 @@ def lock_ctx(
if dir_key not in local_mem._locks: if dir_key not in local_mem._locks:
local_mem._locks[dir_key] = True local_mem._locks[dir_key] = True
fl = filelock.FileLock(Path(lock_dir) / ".lock") fl = FileLock(Path(lock_dir) / ".lock")
fl.acquire(timeout=timeout) fl.acquire(timeout=timeout)
try: try:
yield yield
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论