xref: /aosp_15_r20/external/pytorch/torch/utils/file_baton.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# mypy: allow-untyped-defs
2import os
3import time
4
5
6class FileBaton:
7    """A primitive, file-based synchronization utility."""
8
9    def __init__(self, lock_file_path, wait_seconds=0.1):
10        """
11        Create a new :class:`FileBaton`.
12
13        Args:
14            lock_file_path: The path to the file used for locking.
15            wait_seconds: The seconds to periodically sleep (spin) when
16                calling ``wait()``.
17        """
18        self.lock_file_path = lock_file_path
19        self.wait_seconds = wait_seconds
20        self.fd = None
21
22    def try_acquire(self):
23        """
24        Try to atomically create a file under exclusive access.
25
26        Returns:
27            True if the file could be created, else False.
28        """
29        try:
30            self.fd = os.open(self.lock_file_path, os.O_CREAT | os.O_EXCL)
31            return True
32        except FileExistsError:
33            return False
34
35    def wait(self):
36        """
37        Periodically sleeps for a certain amount until the baton is released.
38
39        The amount of time slept depends on the ``wait_seconds`` parameter
40        passed to the constructor.
41        """
42        while os.path.exists(self.lock_file_path):
43            time.sleep(self.wait_seconds)
44
45    def release(self):
46        """Release the baton and removes its file."""
47        if self.fd is not None:
48            os.close(self.fd)
49
50        os.remove(self.lock_file_path)
51