[ENH] Addressed @mgermain comments. Added unit tests.

Author: MarcCote

scripts/smart_dispatch.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-import re
 import os
 import sys
 import argparse
 import time as t
 import numpy as np
+from os.path import join as pjoin
 from subprocess import check_output
 from textwrap import dedent
 
@@ -55,7 +55,11 @@ def main():
     else:
         raise ValueError("Unknown subcommand!")
 
-    command_manager = CommandManager(os.path.join(path_job_commands, "commands.txt"))
+    # Keep a log of the command line in the job folder.
+    command_line = " ".join(sys.argv)
+    smartdispatch.log_command_line(path_job, command_line)
+
+    command_manager = CommandManager(pjoin(path_job_commands, "commands.txt"))
 
     # If resume mode, reset running jobs
     if args.mode == "launch":
@@ -84,7 +88,7 @@ def main():
 
     # If no pool size is specified the number of commands is taken
     if args.pool is None:
-        args.pool = command_manager.get_nb_commands_to_run()
+        args.pool = command_manager.get_nb_commands_to_run(command_line)
 
     # Generating all the worker commands
     COMMAND_STRING = 'cd "{cwd}"; smart_worker.py "{commands_file}" "{log_folder}" '\
@@ -104,14 +108,6 @@ def main():
     job_generator = job_generator_factory(queue, commands, command_params, CLUSTER_NAME, path_job)
     pbs_filenames = job_generator.write_pbs_files(path_job_commands)
 
-    # Keep an history of command lines i.e. smart_dispatch.py commands.
-    with utils.open_with_lock(os.path.join(path_job, "command_line_history.txt"), 'a') as command_line_history_file:
-        command_line_history_file.write(t.strftime("## %Y-%m-%d %H:%M:%S ##\n"))
-        command_line = " ".join(sys.argv)
-        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
-        command_line_history_file.write(command_line + "\n\n")
-
     # Launch the jobs
     print "## {nb_commands} command(s) will be executed in {nb_jobs} job(s) ##".format(nb_commands=nb_commands, nb_jobs=len(pbs_filenames))
     print "Batch UID:\n{batch_uid}".format(batch_uid=jobname)
@@ -121,7 +117,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(os.path.join(path_job, "jobs_id.txt"), 'a') as jobs_id_file:
+        with utils.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))
@@ -167,10 +163,10 @@ def parse_arguments():
 
 
 def _gen_job_paths(jobname):
-    path_smartdispatch_logs = os.path.join(os.getcwd(), LOGS_FOLDERNAME)
-    path_job = os.path.join(path_smartdispatch_logs, jobname)
-    path_job_logs = os.path.join(path_job, 'logs')
-    path_job_commands = os.path.join(path_job, 'commands')
+    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
 
@@ -183,10 +179,10 @@ def get_job_folders(jobname):
 
     if not os.path.exists(path_job_logs):
         os.makedirs(path_job_logs)
-    if not os.path.exists(os.path.join(path_job_logs, "worker")):
-        os.makedirs(os.path.join(path_job_logs, "worker"))
-    if not os.path.exists(os.path.join(path_job_logs, "job")):
-        os.makedirs(os.path.join(path_job_logs, "job"))
+    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
 
@@ -200,8 +196,8 @@ def create_job_folders(jobname):
 
     if not os.path.exists(path_job_logs):
         os.makedirs(path_job_logs)
-        os.makedirs(os.path.join(path_job_logs, "worker"))
-        os.makedirs(os.path.join(path_job_logs, "job"))
+        os.makedirs(pjoin(path_job_logs, "worker"))
+        os.makedirs(pjoin(path_job_logs, "job"))
 
     return path_job, path_job_logs, path_job_commands
 

smartdispatch/smartdispatch.py

@@ -4,6 +4,7 @@
 import re
 import itertools
 import time as t
+from os.path import join as pjoin
 
 import smartdispatch
 from smartdispatch import utils
@@ -119,3 +120,23 @@ def get_available_queues(cluster_name=utils.detect_cluster()):
 
     queues_infos = utils.load_dict_from_json_file(config_filepath)
     return queues_infos
+
+
+def log_command_line(path_job, command_line):
+    """ Logs a command line in a job folder.
+
+    The command line is append to a file named 'command_line.log' that resides
+    in the given job folder. The current date and time is also added along
+    each command line logged.
+
+    Notes
+    -----
+    Commands save in log file might differ from sys.argv since we want to make sure
+    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:
+        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
+        command_line_log.write(command_line + "\n\n")

smartdispatch/tests/test_smartdispatch.py

@@ -1,8 +1,15 @@
-import smartdispatch
+import os
+import re
+import shutil
+import time as t
+from os.path import join as pjoin
 from StringIO import StringIO
 
+import tempfile
 from nose.tools import assert_true, assert_equal
 from numpy.testing import assert_array_equal
+
+import smartdispatch
 from smartdispatch import utils
 
 
@@ -137,3 +144,37 @@ def test_get_available_queues():
 
     queues_infos = smartdispatch.get_available_queues(cluster_name="mammouth")
     assert_true(len(queues_infos) > 0)
+
+
+def test_log_command_line():
+    temp_dir = tempfile.mkdtemp()
+    command_line_log_file = pjoin(temp_dir, "command_line.log")
+
+    command_1 = "echo 1 2 3"
+    smartdispatch.log_command_line(temp_dir, command_1)
+    assert_true(os.path.isfile(command_line_log_file))
+
+    lines = open(command_line_log_file).read().strip().split("\n")
+    assert_equal(len(lines), 2)  # Datetime, the command line.
+    assert_true(t.strftime("## %Y-%m-%d %H:%M:") in lines[0])  # Don't check second.
+    assert_equal(lines[1], command_1)
+
+    command_2 = "echo \"bob\""  # With quotes.
+    smartdispatch.log_command_line(temp_dir, command_2)
+    assert_true(os.path.isfile(command_line_log_file))
+
+    lines = open(command_line_log_file).read().strip().split("\n")
+    assert_equal(len(lines), 5)
+    assert_true(t.strftime("## %Y-%m-%d %H:%M:") in lines[3])  # Don't check second.
+    assert_equal(lines[4], command_2.replace('"', r'\"'))
+
+    command_3 = "echo [asd]"  # With square brackets.
+    smartdispatch.log_command_line(temp_dir, command_3)
+    assert_true(os.path.isfile(command_line_log_file))
+
+    lines = open(command_line_log_file).read().strip().split("\n")
+    assert_equal(len(lines), 8)
+    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)