From 7582fd07294c83d8e3c5369cd064bc48c7d36ba0 Mon Sep 17 00:00:00 2001
From: lukemartinlogan <lukemartinlogan@gmail.com>
Date: Thu, 26 Oct 2023 16:50:11 -0500
Subject: [PATCH] Commit slurm hostfile

---
 jarvis_util/shell/slurm_exec.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/jarvis_util/shell/slurm_exec.py b/jarvis_util/shell/slurm_exec.py
index 6b38c3a..135d140 100644
--- a/jarvis_util/shell/slurm_exec.py
+++ b/jarvis_util/shell/slurm_exec.py
@@ -35,6 +35,7 @@ def __init__(self, cmd, exec_info):
         self.mem = exec_info.mem
         self.gres = exec_info.gres
         self.exclusive = exec_info.exclusive
+        self.host_suffix = exec_info.host_suffix
 
         super().__init__(self.slurmcmd(),
                          exec_info.mod(env=exec_info.basic_env))
@@ -91,3 +92,17 @@ def __init__(self, job_name=None, num_nodes=1, **kwargs):
                 setattr(self, key, None)
         self.job_name = job_name
         self.num_nodes = num_nodes
+
+
+class SlurmHostfile(LocalExec):
+    def __init__(self, file_location, host_suffix=None):
+        self.file_location = file_location
+        self.host_suffix = host_suffix
+        cmd = f'scontrol show hostnames $SLURM_JOB_NODELIST > {file_location}'
+        super().__init__(cmd, LocalExecInfo())
+        if host_suffix is not None:
+            with open(file_location, 'r', encoding='utf-8') as fp:
+                lines = fp.readlines()
+                lines = [f'{line}{suffix}' for line in lines]
+            with open(file_location, 'w', encoding='utf-8') as fp:
+                fp.writelines(lines)