File size: 1,567 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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
|