Memory support (Issue #42)

scripts/smart-dispatch

@@ -146,18 +146,22 @@ def main():
                                            log_folder=path_job_logs, worker_call_suffix=worker_call_suffix)
     commands = [COMMAND_STRING.format(ID=i) for i in range(args.pool)]
 
-    # TODO: use args.memPerNode instead of args.memPerNode
-    queue = Queue(args.queueName, CLUSTER_NAME, args.walltime, args.coresPerNode, args.gpusPerNode, float('inf'), args.modules)
+    queue = Queue(args.queueName, CLUSTER_NAME, args.walltime, args.coresPerNode, args.gpusPerNode, args.memPerNode, args.modules)
 
-    # Check that requested core number does not exceed node total
+    # Check that requested per command resources do not exceed node total
     if args.coresPerCommand > queue.nb_cores_per_node:
         sys.stderr.write("smart-dispatch: error: coresPerCommand exceeds nodes total: asked {req_cores} cores, nodes have {node_cores}\n"
                          .format(req_cores=args.coresPerCommand, node_cores=queue.nb_cores_per_node))
         sys.exit(2)
 
+    if args.memPerCommand > queue.mem_per_node:
+        sys.stderr.write("smart-dispatch: error: memPerCommand exceeds nodes total: asked {req_mem} Gb, nodes have {node_mem} Gb\n"
+                         .format(req_mem=args.memPerCommand, node_mem=queue.mem_per_node))
+        sys.exit(2)
+
     command_params = {'nb_cores_per_command': args.coresPerCommand,
                       'nb_gpus_per_command': args.gpusPerCommand,
-                      'mem_per_command': None  # args.memPerCommand
+                      'mem_per_command': args.memPerCommand
                       }
 
     prolog = []
@@ -172,7 +176,7 @@ def main():
     for pbs_id, pbs in enumerate(job_generator.pbs_list):
         proper_size_name = utils.jobname_generator(jobname, pbs_id)
         pbs.add_options(N=proper_size_name)
-    
+
     if args.pbsFlags is not None:
         job_generator.add_pbs_flags(args.pbsFlags.split(' '))
     pbs_filenames = job_generator.write_pbs_files(path_job_commands)
@@ -193,11 +197,11 @@ def parse_arguments():
     parser.add_argument('-L', '--launcher', choices=['qsub', 'msub'], required=False, help='Which launcher to use. Default: qsub')
     parser.add_argument('-C', '--coresPerNode', type=int, required=False, help='How many cores there are per node.')
     parser.add_argument('-G', '--gpusPerNode', type=int, required=False, help='How many gpus there are per node.')
-    # parser.add_argument('-M', '--memPerNode', type=int, required=False, help='How much memory there are per node (in Gb).')
+    parser.add_argument('-M', '--memPerNode', type=float, required=False, help='How much memory there are per node (in Gb).')
 
     parser.add_argument('-c', '--coresPerCommand', type=int, required=False, help='How many cores a command needs.', default=1)
     parser.add_argument('-g', '--gpusPerCommand', type=int, required=False, help='How many gpus a command needs.', default=1)
-    # parser.add_argument('-m', '--memPerCommand', type=float, required=False, help='How much memory a command needs (in Gb).')
+    parser.add_argument('-m', '--memPerCommand', type=float, required=False, help='How much memory a command needs (in Gb).')
     parser.add_argument('-f', '--commandsFile', type=file, required=False, help='File containing commands to launch. Each command must be on a seperate line. (Replaces commandAndOptions)')
 
     parser.add_argument('-l', '--modules', type=str, required=False, help='List of additional modules to load.', nargs='+')
@@ -224,7 +228,9 @@ def parse_arguments():
         if args.queueName not in AVAILABLE_QUEUES and ((args.coresPerNode is None and args.gpusPerNode is None) or args.walltime is None):
             parser.error("Unknown queue, --coresPerNode/--gpusPerNode and --walltime must be set.")
         if args.coresPerCommand < 1:
-            parser.error("coresPerNode must be at least 1")
+            parser.error("coresPerCommand must be at least 1")
+        if args.memPerCommand is not None and args.memPerCommand <= 0:
+            parser.error("memPerCommand must be positive")
 
     return args
 

smartdispatch/job_generator.py

@@ -46,7 +46,7 @@ def __init__(self, queue, commands, prolog=[], epilog=[], command_params={}, bas
 
         self.nb_cores_per_command = command_params.get('nb_cores_per_command', 1)
         self.nb_gpus_per_command = command_params.get('nb_gpus_per_command', 1)
-        #self.mem_per_command = command_params.get('mem_per_command', 0.0)
+        self.mem_per_command = command_params.get('mem_per_command', None)
 
         self.pbs_list = self._generate_base_pbs()
         self._add_cluster_specific_rules()
@@ -80,6 +80,10 @@ def _generate_base_pbs(self):
         if self.queue.nb_gpus_per_node > 0 and self.nb_gpus_per_command > 0:
             nb_commands_per_node = min(nb_commands_per_node, self.queue.nb_gpus_per_node // self.nb_gpus_per_command)
 
+        # Limit number of running commands by the amount of available memory on the node.
+        if self.mem_per_command is not None:
+            nb_commands_per_node = min(nb_commands_per_node, self.queue.mem_per_node // self.mem_per_command)
+
         pbs_files = []
         # Distribute equally the jobs among the PBS files and generate those files
         for i, commands in enumerate(utils.chunks(self.commands, n=nb_commands_per_node)):
@@ -92,9 +96,12 @@ def _generate_base_pbs(self):
             resource = "1:ppn={ppn}".format(ppn=len(commands) * self.nb_cores_per_command)
             if self.queue.nb_gpus_per_node > 0:
                 resource += ":gpus={gpus}".format(gpus=len(commands) * self.nb_gpus_per_command)
-
             pbs.add_resources(nodes=resource)
 
+            if self.mem_per_command is not None:
+                resource = "{mem}Gb".format(mem=len(commands) * self.mem_per_command)
+                pbs.add_resources(mem=resource)
+
             pbs.add_modules_to_load(*self.queue.modules)
             pbs.add_to_prolog(*self.prolog)
             pbs.add_commands(*commands)

smartdispatch/tests/test_job_generator.py

@@ -43,15 +43,15 @@ def test_generate_pbs(self):
         assert_equal(job_generator.pbs_list[0].epilog, self.epilog)
 
     def test_generate_pbs2_cpu(self):
-        # Should needs two PBS file
+        # Should need two PBS files
         command_params = {'nb_cores_per_command': self.cores // 2}
         job_generator = JobGenerator(self.queue, self.commands, command_params=command_params)
         assert_equal(len(job_generator.pbs_list), 2)
         assert_equal(job_generator.pbs_list[0].commands, self.commands[:2])
         assert_equal(job_generator.pbs_list[1].commands, self.commands[2:])
 
     def test_generate_pbs4_cpu(self):
-        # Should needs four PBS file
+        # Should need four PBS files
         command_params = {'nb_cores_per_command': self.cores}
         job_generator = JobGenerator(self.queue, self.commands, command_params=command_params)
         assert_equal(len(job_generator.pbs_list), 4)
@@ -64,17 +64,39 @@ def test_generate_pbs4_cpu(self):
         # Check if needed modules for this queue are included in the PBS file
         assert_equal(job_generator.pbs_list[0].modules, self.modules)
 
+    def test_generate_pbs2_mem(self):
+        # Should need two PBS files
+        command_params = {'mem_per_command': self.mem_per_node // 2}
+        job_generator = JobGenerator(self.queue, self.commands, command_params=command_params)
+        assert_equal(len(job_generator.pbs_list), 2)
+        assert_equal(job_generator.pbs_list[0].commands, self.commands[:2])
+        assert_equal(job_generator.pbs_list[1].commands, self.commands[2:])
+
+    def test_generate_pbs4_mem(self):
+        # Should need four PBS files
+        command_params = {'mem_per_command': self.mem_per_node}
+        job_generator = JobGenerator(self.queue, self.commands, command_params=command_params)
+        assert_equal(len(job_generator.pbs_list), 4)
+        assert_equal([pbs.commands[0] for pbs in job_generator.pbs_list], self.commands)
+
+        # Since queue has no gpus it should not be specified in PBS resource `nodes`
+        assert_true('gpus' not in job_generator.pbs_list[0].resources['nodes'])
+
+        # Test modules to load
+        # Check if needed modules for this queue are included in the PBS file
+        assert_equal(job_generator.pbs_list[0].modules, self.modules)
+
     def test_generate_pbs2_gpu(self):
         # Test nb_gpus_per_command argument
-        # Should needs two PBS file
+        # Should need two PBS files
         command_params = {'nb_gpus_per_command': self.gpus // 2}
         job_generator = JobGenerator(self.queue_gpu, self.commands, command_params=command_params)
         assert_equal(len(job_generator.pbs_list), 2)
         assert_equal(job_generator.pbs_list[0].commands, self.commands[:2])
         assert_equal(job_generator.pbs_list[1].commands, self.commands[2:])
 
     def test_generate_pbs4_gpu(self):
-        # Should needs four PBS files
+        # Should need four PBS files
         command_params = {'nb_gpus_per_command': self.gpus}
         job_generator = JobGenerator(self.queue_gpu, self.commands, command_params=command_params)
         assert_equal(len(job_generator.pbs_list), 4)

tests/test_smart_dispatch.py

@@ -23,17 +23,20 @@ def setUp(self):
         self.nb_commands = len(self.commands)
 
         scripts_path = abspath(pjoin(os.path.dirname(__file__), os.pardir, "scripts"))
-        self.smart_dispatch_command = '{} -C 1 -q test -t 5:00 -x'.format(pjoin(scripts_path, 'smart-dispatch'))
+        self.smart_dispatch_command = '{} -C 1 -M 1 -q test -t 5:00 -x'.format(pjoin(scripts_path, 'smart-dispatch'))
         self.launch_command = "{0} launch {1}".format(self.smart_dispatch_command, self.folded_commands)
         self.resume_command = "{0} resume {{0}}".format(self.smart_dispatch_command)
 
-        smart_dispatch_command_with_pool = '{} --pool 10 -C 1 -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
+        smart_dispatch_command_with_pool = '{} --pool 10 -C 1 -M 1 -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
         self.launch_command_with_pool = smart_dispatch_command_with_pool.format('launch ' + self.folded_commands)
         self.nb_workers = 10
 
-        smart_dispatch_command_with_cores = '{} -C 1 -c {{cores}} -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
+        smart_dispatch_command_with_cores = '{} -C 1 -M 1 -c {{cores}} -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
         self.launch_command_with_cores = smart_dispatch_command_with_cores.format('launch ' + self.folded_commands, cores='{cores}')
 
+        smart_dispatch_command_with_memory = '{} -C 1 -M 1 -m {{memory}} -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
+        self.launch_command_with_memory = smart_dispatch_command_with_memory.format('launch ' + self.folded_commands, memory='{memory}')
+
         self._cwd = os.getcwd()
         os.chdir(self.testing_dir)
 
@@ -95,6 +98,18 @@ def test_main_launch_with_cores_command(self):
         assert_equal(exit_status_100, 2)        
         assert_true(os.path.isdir(self.logs_dir))
 
+    def test_main_launch_with_memory_command(self):
+        # Actual test
+        exit_status_0 = call(self.launch_command_with_memory.format(memory=0), shell=True)
+        exit_status_05 = call(self.launch_command_with_memory.format(memory=0.5), shell=True)
+        exit_status_100 = call(self.launch_command_with_memory.format(memory=100), shell=True)
+
+        # Test validation
+        assert_equal(exit_status_0, 2)
+        assert_equal(exit_status_05, 0)
+        assert_equal(exit_status_100, 2)        
+        assert_true(os.path.isdir(self.logs_dir))
+
     def test_main_resume(self):
         # Setup
         call(self.launch_command, shell=True)