Merge pull request #48 from candiceT233/master

Add nodelist parameter for slurm; Fix ppn to ntasks-per-node for slurm
grc-iit · Jan 24, 2024 · 9e2587c · 9e2587c
2 parents 72fcd4a + 6d70bad
commit 9e2587c
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/jarvis_util/shell/slurm_exec.py b/jarvis_util/shell/slurm_exec.py
@@ -36,6 +36,7 @@ def __init__(self, cmd, exec_info):
         self.gres = exec_info.gres
         self.exclusive = exec_info.exclusive
         self.host_suffix = exec_info.host_suffix
+        self.nodelist = exec_info.nodelist
 
         super().__init__(self.slurmcmd(),
                          exec_info.mod(env=exec_info.basic_env))
@@ -47,7 +48,7 @@ def generate_sbatch_command(self):
         options_map = {
             'job_name': 'job-name',
             'num_nodes': 'nodes',
-            'ppn': 'ntasks',
+            'ppn': 'ntasks-per-node',
             'cpus_per_task': 'cpus-per-task',
             'time': 'time',
             'partition': 'partition',
@@ -56,7 +57,8 @@ def generate_sbatch_command(self):
             'error': 'error',
             'mem': 'mem',
             'gres': 'gres',
-            'exclusive': 'exclusive'
+            'exclusive': 'exclusive',
+            'nodelist': 'nodelist',
         }
 
         for attr, option in options_map.items():
@@ -82,7 +84,7 @@ class SlurmExecInfo(ExecInfo):
     def __init__(self, job_name=None, num_nodes=1, **kwargs):
         super().__init__(exec_type=ExecType.SLURM, **kwargs)
         allowed_options = ['job_name', 'num_nodes', 'cpus_per_task', 'time', 'partition', 'mail_type',
-                           'mail_user', 'mem', 'gres', 'exclusive', 'host_suffix']
+                           'mail_user', 'mem', 'gres', 'exclusive', 'host_suffix', 'nodelist']
         self.keys += allowed_options
         # We use ppn, and the output and error file from the base Exec Info
         for key in allowed_options: