pytorch

1
# mypy: allow-untyped-defs
2
import os
3
import time
4

5

6
class 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