Merge pull request #98 from MarcCote/allow_resuming_from_path

Support both path and Batch UID when resuming jobs.

Author: mgermain

scripts/smart_dispatch.py

@@ -20,7 +20,6 @@
 import logging
 import smartdispatch
 
-
 LOGS_FOLDERNAME = "SMART_DISPATCH_LOGS"
 CLUSTER_NAME = utils.detect_cluster()
 AVAILABLE_QUEUES = get_available_queues(CLUSTER_NAME)
@@ -32,6 +31,7 @@ def main():
     logging.root.setLevel(logging.INFO)
 
     args = parse_arguments()
+    path_smartdispatch_logs = pjoin(os.getcwd(), LOGS_FOLDERNAME)
 
     # Check if RESUME or LAUNCH mode
     if args.mode == "launch":
@@ -48,13 +48,20 @@ def main():
         commands = smartdispatch.replace_uid_tag(commands)
         nb_commands = len(commands)  # For print at the end
 
-        path_job, path_job_logs, path_job_commands = create_job_folders(jobname)
     elif args.mode == "resume":
         jobname = args.batch_uid
-        path_job, path_job_logs, path_job_commands = get_job_folders(args.batch_uid)
+        if os.path.isdir(jobname):
+            # We assume `jobname` is `path_job` repo, we extract the real `jobname`.
+            jobname = os.path.basename(os.path.abspath(jobname))
+
+        if not os.path.isdir(pjoin(path_smartdispatch_logs, jobname)):
+            raise LookupError("Batch UID ({0}) does not exist! Cannot resume.".format(jobname))
     else:
         raise ValueError("Unknown subcommand!")
 
+    job_folders_paths = smartdispatch.get_job_folders(path_smartdispatch_logs, jobname)
+    path_job, path_job_logs, path_job_commands = job_folders_paths
+
     # Keep a log of the command line in the job folder.
     command_line = " ".join(sys.argv)
     smartdispatch.log_command_line(path_job, command_line)
@@ -165,45 +172,5 @@ def parse_arguments():
     return args
 
 
-def _gen_job_paths(jobname):
-    path_smartdispatch_logs = pjoin(os.getcwd(), LOGS_FOLDERNAME)
-    path_job = pjoin(path_smartdispatch_logs, jobname)
-    path_job_logs = pjoin(path_job, 'logs')
-    path_job_commands = pjoin(path_job, 'commands')
-
-    return path_job, path_job_logs, path_job_commands
-
-
-def get_job_folders(jobname):
-    path_job, path_job_logs, path_job_commands = _gen_job_paths(jobname)
-
-    if not os.path.exists(path_job_commands):
-        raise LookupError("Batch UID ({0}) does not exist! Cannot resume.".format(jobname))
-
-    if not os.path.exists(path_job_logs):
-        os.makedirs(path_job_logs)
-    if not os.path.exists(pjoin(path_job_logs, "worker")):
-        os.makedirs(pjoin(path_job_logs, "worker"))
-    if not os.path.exists(pjoin(path_job_logs, "job")):
-        os.makedirs(pjoin(path_job_logs, "job"))
-
-    return path_job, path_job_logs, path_job_commands
-
-
-def create_job_folders(jobname):
-    """Creates the folders where the logs, commands and QSUB files will be saved."""
-    path_job, path_job_logs, path_job_commands = _gen_job_paths(jobname)
-
-    if not os.path.exists(path_job_commands):
-        os.makedirs(path_job_commands)
-
-    if not os.path.exists(path_job_logs):
-        os.makedirs(path_job_logs)
-        os.makedirs(pjoin(path_job_logs, "worker"))
-        os.makedirs(pjoin(path_job_logs, "job"))
-
-    return path_job, path_job_logs, path_job_commands
-
-
 if __name__ == "__main__":
     main()

smartdispatch/smartdispatch.py

@@ -110,10 +110,10 @@ def get_available_queues(cluster_name=utils.detect_cluster()):
         return {}
 
     smartdispatch_dir, _ = os.path.split(smartdispatch.__file__)
-    config_dir = os.path.join(smartdispatch_dir, 'config')
+    config_dir = pjoin(smartdispatch_dir, 'config')
 
     config_filename = cluster_name + ".json"
-    config_filepath = os.path.join(config_dir, config_filename)
+    config_filepath = pjoin(config_dir, config_filename)
 
     if not os.path.isfile(config_filepath):
         return {}  # Unknown cluster
@@ -122,6 +122,24 @@ def get_available_queues(cluster_name=utils.detect_cluster()):
     return queues_infos
 
 
+def get_job_folders(path, jobname, create_if_needed=False):
+    """ Get all folder paths for a specific job (creating them if needed). """
+    path_job = pjoin(path, jobname)
+    path_job_logs = pjoin(path_job, 'logs')
+    path_job_commands = pjoin(path_job, 'commands')
+
+    if not os.path.isdir(path_job_commands):
+        os.makedirs(path_job_commands)
+    if not os.path.isdir(path_job_logs):
+        os.makedirs(path_job_logs)
+    if not os.path.isdir(pjoin(path_job_logs, "worker")):
+        os.makedirs(pjoin(path_job_logs, "worker"))
+    if not os.path.isdir(pjoin(path_job_logs, "job")):
+        os.makedirs(pjoin(path_job_logs, "job"))
+
+    return path_job, path_job_logs, path_job_commands
+
+
 def log_command_line(path_job, command_line):
     """ Logs a command line in a job folder.
 

smartdispatch/tests/test_smartdispatch.py

@@ -146,6 +146,52 @@ def test_get_available_queues():
     assert_true(len(queues_infos) > 0)
 
 
+def test_get_job_folders():
+    temp_dir = tempfile.mkdtemp()
+    jobname = "this_is_the_name_of_my_job"
+    job_folders_paths = smartdispatch.get_job_folders(temp_dir, jobname)
+    path_job, path_job_logs, path_job_commands = job_folders_paths
+
+    assert_true(jobname in path_job)
+    assert_true(os.path.isdir(path_job))
+    assert_equal(os.path.basename(path_job), jobname)
+
+    assert_true(jobname in path_job_logs)
+    assert_true(os.path.isdir(path_job_logs))
+    assert_true(os.path.isdir(pjoin(path_job_logs, 'worker')))
+    assert_true(os.path.isdir(pjoin(path_job_logs, 'job')))
+    assert_true(os.path.isdir(path_job_logs))
+    assert_equal(os.path.basename(path_job_logs), "logs")
+
+    assert_true(jobname in path_job_commands)
+    assert_true(os.path.isdir(path_job_commands))
+    assert_equal(os.path.basename(path_job_commands), "commands")
+
+    # In theory the following should not create new folders.
+    # Insteead it will return the paths to existing folders.
+    jobname += "2"
+    os.rename(path_job, path_job + "2")
+    job_folders_paths = smartdispatch.get_job_folders(temp_dir, jobname)
+    path_job, path_job_logs, path_job_commands = job_folders_paths
+
+    assert_true(jobname in path_job)
+    assert_true(os.path.isdir(path_job))
+    assert_equal(os.path.basename(path_job), jobname)
+
+    assert_true(jobname in path_job_logs)
+    assert_true(os.path.isdir(path_job_logs))
+    assert_true(os.path.isdir(pjoin(path_job_logs, 'worker')))
+    assert_true(os.path.isdir(pjoin(path_job_logs, 'job')))
+    assert_true(os.path.isdir(path_job_logs))
+    assert_equal(os.path.basename(path_job_logs), "logs")
+
+    assert_true(jobname in path_job_commands)
+    assert_true(os.path.isdir(path_job_commands))
+    assert_equal(os.path.basename(path_job_commands), "commands")
+
+    shutil.rmtree(temp_dir)
+
+
 def test_log_command_line():
     temp_dir = tempfile.mkdtemp()
     command_line_log_file = pjoin(temp_dir, "command_line.log")
@@ -177,4 +223,4 @@ def test_log_command_line():
     assert_true(t.strftime("## %Y-%m-%d %H:%M:") in lines[6])  # Don't check second.
     assert_equal(lines[7], re.sub(r'(\[)([^\[\]]*\\ [^\[\]]*)(\])', r'"\1\2\3"', command_3))
 
-    shutil.rmtree(temp_dir)
+    shutil.rmtree(temp_dir)

tests/test_smart_dispatch.py

@@ -23,7 +23,6 @@ def setUp(self):
         os.chdir(self.testing_dir)
 
     def tearDown(self):
-        print "Tear down"
         os.chdir(self._cwd)
         shutil.rmtree(self.testing_dir)
 
@@ -37,7 +36,7 @@ def test_main_launch(self):
         assert_equal(len(os.listdir(self.logs_dir)), 1)
 
     def test_main_resume(self):
-        # SetUp
+        # Setup
         call(self.launch_command, shell=True)
         batch_uid = os.listdir(self.logs_dir)[0]
 
@@ -61,6 +60,32 @@ def test_main_resume(self):
         assert_equal(len(open(running_commands_file).readlines()), 0)
         assert_equal(len(open(pending_commands_file).readlines()), len(commands))
 
+		# Test when batch_uid is a path instead of a jobname.
+		# Setup
+        call(self.launch_command, shell=True)
+        batch_uid = os.path.join(self.logs_dir, os.listdir(self.logs_dir)[0])
+
+		# Simulate that some commands are in the running state.
+        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
+        pending_commands_file = pjoin(path_job_commands, "commands.txt")
+        running_commands_file = pjoin(path_job_commands, "running_commands.txt")
+        commands = open(pending_commands_file).read().strip().split("\n")
+        with open(running_commands_file, 'w') as running_commands:
+            running_commands.write("\n".join(commands[::2]) + "\n")
+        with open(pending_commands_file, 'w') as pending_commands:
+            pending_commands.write("\n".join(commands[1::2]) + "\n")
+
+        # Actual test (should move running commands back to pending).
+        exit_status = call(self.resume_command.format(batch_uid), shell=True)
+
+        # Test validation
+        assert_equal(exit_status, 0)
+        assert_true(os.path.isdir(self.logs_dir))
+        assert_equal(len(os.listdir(self.logs_dir)), 1)
+        assert_equal(len(open(running_commands_file).readlines()), 0)
+        assert_equal(len(open(pending_commands_file).readlines()), len(commands))
+
+
     def test_main_resume_only_pending(self):
         # SetUp
         call(self.launch_command, shell=True)