Files
pytorch/torch/utils/file_baton.py
Maggie Moss 086dec3235 Pyrefly suppressions 6/n (#164877)
Adds suppressions to pyrefly will typecheck clean: https://github.com/pytorch/pytorch/issues/163283

Almost there!

Test plan:
dmypy restart && python3 scripts/lintrunner.py -a
pyrefly check

step 1: delete lines in the pyrefly.toml file from the project-excludes field
step 2: run pyrefly check
step 3: add suppressions, clean up unused suppressions
before: https://gist.github.com/maggiemoss/4b3bf2037014e116bc00706a16aef199

after:

INFO 0 errors (5,064 ignored)

Only four directories left to enable

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164877
Approved by: https://github.com/oulgen
2025-10-08 02:30:57 +00:00

65 lines
2.0 KiB
Python

# mypy: allow-untyped-defs
import os
import time
import warnings
class FileBaton:
"""A primitive, file-based synchronization utility."""
def __init__(self, lock_file_path, wait_seconds=0.1, warn_after_seconds=None):
"""
Create a new :class:`FileBaton`.
Args:
lock_file_path: The path to the file used for locking.
wait_seconds: The seconds to periodically sleep (spin) when
calling ``wait()``.
warn_after_seconds: The seconds to wait before showing
lock file path to warn existing lock file.
"""
self.lock_file_path = lock_file_path
self.wait_seconds = wait_seconds
self.fd = None
self.warn_after_seconds = warn_after_seconds
def try_acquire(self):
"""
Try to atomically create a file under exclusive access.
Returns:
True if the file could be created, else False.
"""
try:
# pyrefly: ignore # bad-assignment
self.fd = os.open(self.lock_file_path, os.O_CREAT | os.O_EXCL)
return True
except FileExistsError:
return False
def wait(self):
"""
Periodically sleeps for a certain amount until the baton is released.
The amount of time slept depends on the ``wait_seconds`` parameter
passed to the constructor.
"""
has_warned = False
start_time = time.time()
while os.path.exists(self.lock_file_path):
time.sleep(self.wait_seconds)
if self.warn_after_seconds is not None:
if time.time() - start_time > self.warn_after_seconds and not has_warned:
warnings.warn(f'Waited on lock file "{self.lock_file_path}" for '
f'{self.warn_after_seconds} seconds.')
has_warned = True
def release(self):
"""Release the baton and removes its file."""
if self.fd is not None:
os.close(self.fd)
os.remove(self.lock_file_path)