Merge pull request #100 from MarcCote/smarter_lock_for_helios

Use the atomicity of folder creation to simulate a locking mechanism

Author: mgermain

scripts/smart_dispatch.py

@@ -16,6 +16,7 @@
 from smartdispatch.job_generator import job_generator_factory
 from smartdispatch import get_available_queues
 from smartdispatch import utils
+from smartdispatch.filelock import open_with_lock
 
 import logging
 import smartdispatch
@@ -126,7 +127,7 @@ def main():
             qsub_output = check_output('{launcher} {pbs_filename}'.format(launcher=LAUNCHER if args.launcher is None else args.launcher, pbs_filename=pbs_filename), shell=True)
             jobs_id += [qsub_output.strip()]
 
-        with utils.open_with_lock(pjoin(path_job, "jobs_id.txt"), 'a') as jobs_id_file:
+        with open_with_lock(pjoin(path_job, "jobs_id.txt"), 'a') as jobs_id_file:
             jobs_id_file.writelines(t.strftime("## %Y-%m-%d %H:%M:%S ##\n"))
             jobs_id_file.writelines("\n".join(jobs_id) + "\n")
         print "\nJobs id:\n{jobs_id}".format(jobs_id=" ".join(jobs_id))

setup.py

@@ -12,6 +12,6 @@
     license='LICENSE.txt',
     description='A batch job launcher for the Mammouth supercomputer.',
     long_description=open('README.txt').read(),
-    install_requires=[],
+    install_requires=['psutil>=1'],
     package_data={'smartdispatch': ['config/*.json']}
 )

smartdispatch/command_manager.py

@@ -1,5 +1,5 @@
 import os
-from smartdispatch import utils
+from .filelock import open_with_lock
 
 
 class CommandManager(object):
@@ -24,13 +24,13 @@ def _move_line_between_files(self, file1, file2, line):
         file2.write(line)
 
     def set_commands_to_run(self, commands):
-        with utils.open_with_lock(self._commands_filename, 'a') as commands_file:
+        with open_with_lock(self._commands_filename, 'a') as commands_file:
             commands = [command + '\n' for command in commands]
             commands_file.writelines(commands)
 
     def get_command_to_run(self):
-        with utils.open_with_lock(self._commands_filename, 'r+') as commands_file:
-            with utils.open_with_lock(self._running_commands_filename, 'a') as running_commands_file:
+        with open_with_lock(self._commands_filename, 'r+') as commands_file:
+            with open_with_lock(self._running_commands_filename, 'a') as running_commands_file:
                 command = commands_file.readline()
                 if command == '':
                     return None
@@ -54,14 +54,14 @@ def set_running_command_as_finished(self, command, error_code=0):
         else:
             file_name = self._failed_commands_filename
 
-        with utils.open_with_lock(self._running_commands_filename, 'r+') as running_commands_file:
-            with utils.open_with_lock(file_name, 'a') as finished_commands_file:
+        with open_with_lock(self._running_commands_filename, 'r+') as running_commands_file:
+            with open_with_lock(file_name, 'a') as finished_commands_file:
                 self._move_line_between_files(running_commands_file, finished_commands_file, command + '\n')
 
     def reset_running_commands(self):
         if os.path.isfile(self._running_commands_filename):
-            with utils.open_with_lock(self._commands_filename, 'r+') as commands_file:
-                with utils.open_with_lock(self._running_commands_filename, 'r+') as running_commands_file:
+            with open_with_lock(self._commands_filename, 'r+') as commands_file:
+                with open_with_lock(self._running_commands_filename, 'r+') as running_commands_file:
                     commands = running_commands_file.readlines()
                     if len(commands) > 0:
                         running_commands_file.seek(0, os.SEEK_SET)

smartdispatch/filelock.py

@@ -0,0 +1,86 @@
+import os
+import time
+import fcntl
+import psutil
+import logging
+
+from contextlib import contextmanager
+
+# Constants needed for `open_with_dirlock` function.
+MAX_ATTEMPTS = 1000  # This would correspond to be blocked for ~15min.
+TIME_BETWEEN_ATTEMPTS = 1  # In seconds
+
+
+def find_mount_point(path='.'):
+    """ Finds the mount point used to access `path`. """
+    path = os.path.abspath(path)
+    while not os.path.ismount(path):
+        path = os.path.dirname(path)
+
+    return path
+
+
+def get_fs(path='.'):
+    """ Gets info about the filesystem on which `path` lives. """
+    mount = find_mount_point(path)
+
+    for fs in psutil.disk_partitions(True):
+        if fs.mountpoint == mount:
+            return fs
+
+
+@contextmanager
+def open_with_flock(*args, **kwargs):
+    """ Context manager for opening file with an exclusive lock. """
+    f = open(*args, **kwargs)
+    try:
+        fcntl.lockf(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    except IOError:
+        logging.info("Can't immediately write-lock the file ({0}), waiting ...".format(f.name))
+        fcntl.lockf(f, fcntl.LOCK_EX)
+
+    yield f
+    fcntl.lockf(f, fcntl.LOCK_UN)
+    f.close()
+
+
+@contextmanager
+def open_with_dirlock(*args, **kwargs):
+    """ Context manager for opening file with an exclusive lock using. """
+    dirname = os.path.dirname(args[0])
+    filename = os.path.basename(args[0])
+    lockfile = os.path.join(dirname, "." + filename)
+
+    no_attempt = 0
+    while no_attempt < MAX_ATTEMPTS:
+        try:
+            os.mkdir(lockfile)  # Atomic operation
+            f = open(*args, **kwargs)
+            yield f
+            f.close()
+            os.rmdir(lockfile)
+            break
+        except OSError:
+            logging.info("Can't immediately write-lock the file ({0}), retrying in {1} sec. ...".format(filename, TIME_BETWEEN_ATTEMPTS))
+            time.sleep(TIME_BETWEEN_ATTEMPTS)
+            no_attempt += 1
+
+
+def _fs_support_globalflock(fs):
+    if fs.fstype == "lustre":
+        return ("flock" in fs.opts) and "localflock" not in fs.opts
+
+    elif fs.fstype == "gpfs":
+        return True
+
+    return False  # We don't know.
+
+
+# Determine if we can rely on the fcntl module for locking files on the cluster.
+# Otherwise, fallback on using the directory creation atomicity as a locking mechanism.
+fs = get_fs('.')
+if _fs_support_globalflock(fs):
+    open_with_lock = open_with_flock
+else:
+    logging.warn("Cluster does not support flock! Falling back to folder lock.")
+    open_with_lock = open_with_dirlock

smartdispatch/smartdispatch.py

@@ -8,6 +8,7 @@
 
 import smartdispatch
 from smartdispatch import utils
+from smartdispatch.filelock import open_with_lock
 from smartdispatch.argument_template import argument_templates
 
 UID_TAG = "{UID}"
@@ -153,7 +154,7 @@ def log_command_line(path_job, command_line):
     we can paste the command line as-is in the terminal. This means that the quotes
     symbole " and the square brackets will be escaped.
     """
-    with utils.open_with_lock(pjoin(path_job, "command_line.log"), 'a') as command_line_log:
+    with open_with_lock(pjoin(path_job, "command_line.log"), 'a') as command_line_log:
         command_line_log.write(t.strftime("## %Y-%m-%d %H:%M:%S ##\n"))
         command_line = command_line.replace('"', r'\"')  # Make sure we can paste the command line as-is
         command_line = re.sub(r'(\[)([^\[\]]*\\ [^\[\]]*)(\])', r'"\1\2\3"', command_line)  # Make sure we can paste the command line as-is

smartdispatch/tests/test_filelock.py

@@ -0,0 +1,73 @@
+import os
+import time
+import tempfile
+import shutil
+
+from subprocess import Popen, PIPE
+from nose.tools import assert_equal, assert_true
+
+from smartdispatch.filelock import open_with_lock, open_with_dirlock, open_with_flock
+from smartdispatch.filelock import find_mount_point, get_fs
+
+
+def _test_open_with_lock(lock_func):
+    """ Test a particular file lock.
+
+    Notes
+    -----
+    This test only checks if the locking mechanism works on a single
+    computer/compute node. There is *no* check for multi-node lock.
+    """
+    temp_dir = tempfile.mkdtemp()
+    filename = os.path.join(temp_dir, "testing.lck")
+
+    python_script = os.path.join(temp_dir, "test_lock.py")
+
+    script = ["import logging",
+              "from smartdispatch.filelock import {}".format(lock_func.__name__),
+              "logging.root.setLevel(logging.INFO)",
+              "with {}('{}', 'r+'): pass".format(lock_func.__name__, filename)]
+
+    open(os.path.join(temp_dir, "test_lock.py"), 'w').write("\n".join(script))
+
+    command = "python " + python_script
+
+    # Lock the commands file before running python command
+    with lock_func(filename, 'w'):
+        process = Popen(command, stdout=PIPE, stderr=PIPE, shell=True)
+        time.sleep(1)
+
+    stdout, stderr = process.communicate()
+    assert_equal(stdout, "")
+    assert_true("Traceback" not in stderr, msg="Unexpected error: '{}'".format(stderr))
+    assert_true("write-lock" in stderr, msg="Forcing a race condition, try increasing sleeping time above.")
+
+    shutil.rmtree(temp_dir)  # Cleaning up.
+
+
+def test_open_with_default_lock():
+    _test_open_with_lock(open_with_lock)
+
+
+def test_open_with_dirlock():
+    _test_open_with_lock(open_with_dirlock)
+
+
+def test_open_with_flock():
+    _test_open_with_lock(open_with_flock)
+
+
+def test_find_mount_point():
+    assert_equal(find_mount_point('/'), '/')
+
+    for d in os.listdir('/mnt'):
+        path = os.path.join('/mnt', d)
+        if os.path.ismount(path):
+            assert_equal(find_mount_point(path), path)
+        else:
+            assert_equal(find_mount_point(path), '/')
+
+
+def test_get_fs():
+    fs = get_fs('/')
+    assert_true(fs is not None)

smartdispatch/tests/test_utils.py

@@ -1,12 +1,6 @@
 # -*- coding: utf-8 -*-
-import os
-import time
-import tempfile
-import shutil
 import unittest
 
-from subprocess import Popen, PIPE
-
 from smartdispatch import utils
 
 from nose.tools import assert_equal, assert_true
@@ -51,31 +45,3 @@ def test_slugify():
 
     for arg, expected in testing_arguments:
         assert_equal(utils.slugify(arg), expected)
-
-
-def test_open_with_lock():
-    temp_dir = tempfile.mkdtemp()
-    filename = os.path.join(temp_dir, "testing.lck")
-
-    python_script = os.path.join(temp_dir, "test_lock.py")
-
-    script = ["import logging",
-              "from smartdispatch.utils import open_with_lock",
-              "logging.root.setLevel(logging.INFO)",
-              "with open_with_lock('{0}', 'r+'): pass".format(filename)]
-
-    open(os.path.join(temp_dir, "test_lock.py"), 'w').write("\n".join(script))
-
-    command = "python " + python_script
-
-    # Lock the commands file before running python command
-    with utils.open_with_lock(filename, 'w'):
-        process = Popen(command, stdout=PIPE, stderr=PIPE, shell=True)
-        time.sleep(1)
-
-    stdout, stderr = process.communicate()
-    assert_equal(stdout, "")
-    assert_true("write-lock" in stderr, msg="Forcing a race condition, try increasing sleeping time above.")
-    assert_true("Traceback" not in stderr, msg="Unexpected error: " + stderr)  # Check that there are no errors.
-
-    shutil.rmtree(temp_dir)

smartdispatch/utils.py

@@ -1,13 +1,10 @@
 import re
-import fcntl
-import logging
 import hashlib
 import unicodedata
 import json
 
 from distutils.util import strtobool
 from subprocess import Popen, PIPE
-from contextlib import contextmanager
 
 
 def print_boxed(string):
@@ -81,20 +78,6 @@ def unhexify(match):
     return re.sub(r"\\x..", unhexify, text)
 
 
-@contextmanager
-def open_with_lock(*args, **kwargs):
-    """ Context manager for opening file with an exclusive lock. """
-    f = open(*args, **kwargs)
-    try:
-        fcntl.lockf(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
-    except IOError:
-        logging.info("Can't immediately write-lock the file ({0}), blocking ...".format(f.name))
-        fcntl.lockf(f, fcntl.LOCK_EX)
-    yield f
-    fcntl.lockf(f, fcntl.LOCK_UN)
-    f.close()
-
-
 def save_dict_to_json_file(path, dictionary):
     with open(path, "w") as json_file:
         json_file.write(json.dumps(dictionary, indent=4, separators=(',', ': ')))

tests/test_smart_worker.py

@@ -5,6 +5,7 @@
 import shutil
 
 from smartdispatch import utils
+from smartdispatch.filelock import open_with_lock
 from smartdispatch.command_manager import CommandManager
 
 from subprocess import Popen, call, PIPE
@@ -106,7 +107,7 @@ def test_lock(self):
         command = ["smart_worker.py", self.command_manager._commands_filename, self.logs_dir]
 
         # Lock the commands file before running 'smart_worker.py'
-        with utils.open_with_lock(self.command_manager._commands_filename, 'r+'):
+        with open_with_lock(self.command_manager._commands_filename, 'r+'):
             process = Popen(command, stdout=PIPE, stderr=PIPE)
             time.sleep(1)