pytorch/torch/utils/file_baton.py
Yangqing Jia 852d6e8b65 Fix python2 and python 3 compatibility found by lint. (#13140)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/13140

This is an example about the benefit of proper facebook linter. The old code
was not python 2.x (actually, pre-python 3.3) compatible. Note that FileExistsError
is added in Python 3.3:

https://stackoverflow.com/questions/20790580/python-specifically-handle-file-exists-exception

Reviewed By: mingzhe09088

Differential Revision: D10858804

fbshipit-source-id: a4c995aef9f720cb8b0ce463f0a51db667fc42f2
2018-10-25 17:20:11 -07:00

55 lines
1.6 KiB
Python

from __future__ import absolute_import, division, print_function, unicode_literals
import os
import sys
import time
if sys.version < '3.3':
# Note(jiayq): in Python 2, FileExistsError is not defined and the
# error manifests it as OSError.
FileExistsError = OSError
class FileBaton:
'''A primitive, file-based synchronization utility.'''
def __init__(self, lock_file_path, wait_seconds=0.1):
'''
Creates a new :class:`FileBaton`.
Args:
lock_file_path: The path to the file used for locking.
wait_seconds: The seconds to periorically sleep (spin) when
calling ``wait()``.
'''
self.lock_file_path = lock_file_path
self.wait_seconds = wait_seconds
self.fd = None
def try_acquire(self):
'''
Tries to atomically create a file under exclusive access.
Returns:
True if the file could be created, else False.
'''
try:
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.
'''
while os.path.exists(self.lock_file_path):
time.sleep(self.wait_seconds)
def release(self):
'''Releaes the baton and removes its file.'''
os.close(self.fd)
os.remove(self.lock_file_path)