Merge pull request #92 from felix5572/master

dpdispatcher/submission.py

@@ -646,7 +646,7 @@ def __init__(self,
         self.envs = envs
         # self.if_cuda_multi_devices = if_cuda_multi_devices
 
-        self.kwargs = kwargs
+        self.kwargs = kwargs.get('kwargs', kwargs)
 
         self.gpu_in_use = 0
         self.task_in_para = 0

tests/jsons/machine_lazy_local_slurm.json

@@ -0,0 +1,24 @@
+{
+    "machine":{
+        "batch_type": "Slurm",
+        "context_type": "LazyLocalContext",
+        "local_root": "./test_slurm_dir"
+        },
+    "resources":{
+        "number_node": 1,
+        "cpu_per_node": 4,
+        "gpu_per_node": 2,
+        "queue_name": "GPU_2080Ti",
+        "group_size": 4,
+        "custom_flags": ["#SBATCH --nice=100", "#SBATCH --time=24:00:00"],
+        "strategy": {
+            "if_cuda_multi_devices": true
+        },
+        "para_deg": 2,
+        "module_unload_list": ["singularity"],
+        "module_list": ["singularity/3.0.0"],
+        "source_list": ["./slurm_test.env"],
+        "envs": {"DP_DISPATCHER_EXPORT_VAR": "test_foo_bar_baz"},
+        "custom_gpu_line":"#SBATCH --gres=gpu:2080Ti:2"
+    }
+}

tests/test_slurm_script_generation.py

@@ -0,0 +1,49 @@
+# from dpdispatcher.batch_object import BatchObject
+# from dpdispatcher.batch import Batch
+import sys, os, shutil
+import textwrap
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+__package__ = 'tests'
+from .context import Submission, Job, Task, Resources
+from .context import Shell
+from .context import LocalContext
+from .context import get_file_md5
+from .context import Machine
+
+import unittest
+import json
+
+
+class TestSlurmScriptGeneration(unittest.TestCase):
+    def setUp(self):
+        self.maxDiff = None
+
+    def test_shell_trival(self):
+        with open('jsons/machine_lazy_local_slurm.json', 'r') as f:
+            machine_dict = json.load(f)
+
+        machine = Machine(**machine_dict['machine'])
+        resources = Resources(**machine_dict['resources'])
+
+        task1 = Task(command='cat example.txt', task_work_path='dir1/', forward_files=['example.txt'], backward_files=['out.txt'], outlog='out.txt')
+        task2 = Task(command='cat example.txt', task_work_path='dir2/', forward_files=['example.txt'], backward_files=['out.txt'], outlog='out.txt')
+        task3 = Task(command='cat example.txt', task_work_path='dir3/', forward_files=['example.txt'], backward_files=['out.txt'], outlog='out.txt')
+        task4 = Task(command='cat example.txt', task_work_path='dir4/', forward_files=['example.txt'], backward_files=['out.txt'], outlog='out.txt')
+        task_list = [task1, task2, task3, task4]
+
+        submission = Submission(work_base='parent_dir/',
+            machine=machine,
+            resources=resources,
+            forward_common_files=['graph.pb'],
+            backward_common_files=[],
+            task_list=task_list
+        )
+        submission.generate_jobs()
+        str = machine.gen_script_header(submission.belonging_jobs[0])
+        benchmark_str = textwrap.dedent("""\
+            #!/bin/bash -l
+            #SBATCH --nodes 1
+            #SBATCH --ntasks-per-node 4
+            #SBATCH --gres=gpu:2080Ti:2
+            #SBATCH --partition GPU_2080Ti""")
+        self.assertEqual(str, benchmark_str)