| 123456789101112131415161718192021222324252627282930313233343536373839404142 |
- from types import TracebackType
- from typing import Optional
- from typing_extensions import Self
- from filelock import FileLock as base_FileLock
- from torch.monitor import _WaitCounter
- class FileLock(base_FileLock):
- """
- This behaves like a normal file lock.
- However, it adds waitcounters for acquiring and releasing the filelock
- as well as for the critical region within it.
- pytorch.filelock.enter - While we're acquiring the filelock.
- pytorch.filelock.region - While we're holding the filelock and doing work.
- pytorch.filelock.exit - While we're releasing the filelock.
- """
- def __enter__(self) -> Self:
- self.region_counter = _WaitCounter("pytorch.filelock.region").guard()
- with _WaitCounter("pytorch.filelock.enter").guard():
- result = super().__enter__()
- self.region_counter.__enter__()
- return result
- def __exit__(
- self,
- exc_type: Optional[type[BaseException]],
- exc_value: Optional[BaseException],
- traceback: Optional[TracebackType],
- ) -> None:
- self.region_counter.__exit__()
- with _WaitCounter("pytorch.filelock.exit").guard():
- # Returns nothing per
- # https://github.com/tox-dev/filelock/blob/57f488ff8fdc2193572efe102408fb63cfefe4e4/src/filelock/_api.py#L379
- super().__exit__(exc_type, exc_value, traceback)
- # Returns nothing per
- # https://github.com/pytorch/pytorch/blob/0f6bfc58a2cfb7a5c052bea618ab62becaf5c912/torch/csrc/monitor/python_init.cpp#L315
- return None
|
