test gpu jobs

gpu_jobs.py

@@ -0,0 +1,23 @@
+import os
+from dotenv import load_dotenv
+from utils.subprocess_operations import get_piped_stdout, stdout_to_gpu_records
+from hooks.slack import send_slack_message
+
+load_dotenv()
+SLACK_WEBHOOK = os.getenv("SLACK_JOB_WEBHOOK")
+
+
+def get_gpu_jobs() -> str:
+    main_cmd = "squeue -t RUNNING"
+    piped_cmd = "grep gpu_cuda"
+    return get_piped_stdout(main_cmd, piped_cmd)
+
+
+def monitor():
+    gpu_jobs = get_gpu_jobs()
+    return stdout_to_gpu_records(gpu_jobs)
+
+
+if __name__ == "__main__":
+    print(monitor())
+

utils/subprocess_operations.py

@@ -4,7 +4,7 @@
 JOB_RECORD = dict[str, str]
 JOB_RECORDS = list[JOB_RECORD]
 QUOTA_RECORD = dict[str, str]
-QUOTA_RECORDS = list(QUOTA_RECORD)
+QUOTA_RECORDS = list[QUOTA_RECORD]
 
 
 def get_cmd_stdout(command: str) -> str:
@@ -77,3 +77,17 @@ def stdout_to_quota_records(s: str) -> QUOTA_RECORD:
     for header, quota in zip(headers, s_list):
         quota_record[header] = quota
     return quota_record
+
+
+def stdout_to_gpu_records(s: str) -> JOB_RECORDS:
+    s = s.strip()
+    s_list = s.split("\n")
+    headers = ["JOBID",  "PARTITION", "NAME", "USER", "ST", "TIME", "NODES"]
+    data = [strip_spaces(element.split(" ")) for element in s_list]
+    job_records = []
+    for l in data:
+        d = {}
+        for key, value in zip(headers, l):
+            d[key] = value
+        job_records.append(d)
+    return job_records