From e805b20537362e895824edc3cd18bc12f6be6d69 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jord=C3=A0=20Polo?= <jorda.polo@amd.com>
Date: Mon, 8 Jul 2024 15:50:42 -0500
Subject: [PATCH] Fix slurm_env execution as a script

---
 omnistat/slurm_env.py | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/omnistat/slurm_env.py b/omnistat/slurm_env.py
index 66b0c74e..61449263 100755
--- a/omnistat/slurm_env.py
+++ b/omnistat/slurm_env.py
@@ -38,22 +38,25 @@
 jobData = {}
 jobFile = "/tmp/omni_slurmjobinfo"
 
-if len(sys.argv) > 1:
-    jobFile = sys.argv[1]
+def main():
+    if len(sys.argv) > 1:
+        jobFile = sys.argv[1]
 
-if "SLURM_JOB_ID" in os.environ:
-    jobData["SLURM_JOB_ID"] = os.getenv("SLURM_JOB_ID")
-    jobData["SLURM_JOB_USER"] = os.getenv("SLURM_JOB_USER")
-    jobData["SLURM_JOB_PARTITION"] = os.getenv("SLURM_JOB_PARTITION")
-    jobData["SLURM_JOB_NUM_NODES"] = os.getenv("SLURM_JOB_NUM_NODES")
-    if "SLURM_PTY_PORT" in os.environ:
-        jobData["SLURM_JOB_BATCHMODE"] = 0
-    else:
-        jobData["SLURM_JOB_BATCHMODE"] = 1
+    if "SLURM_JOB_ID" in os.environ:
+        jobData["SLURM_JOB_ID"] = os.getenv("SLURM_JOB_ID")
+        jobData["SLURM_JOB_USER"] = os.getenv("SLURM_JOB_USER")
+        jobData["SLURM_JOB_PARTITION"] = os.getenv("SLURM_JOB_PARTITION")
+        jobData["SLURM_JOB_NUM_NODES"] = os.getenv("SLURM_JOB_NUM_NODES")
+        if "SLURM_PTY_PORT" in os.environ:
+            jobData["SLURM_JOB_BATCHMODE"] = 0
+        else:
+            jobData["SLURM_JOB_BATCHMODE"] = 1
 
-    json.dump(jobData,open(jobFile,"w"),indent=4)
+        json.dump(jobData,open(jobFile,"w"),indent=4)
 
-else:
-    print("ERROR: SLURM settings not visible in current environment. Verify running in active job")
-    sys.exit(1)
+    else:
+        print("ERROR: SLURM settings not visible in current environment. Verify running in active job")
+        sys.exit(1)
     
+if __name__ == '__main__':
+    main()