adding gpu_jobs

gpu_jobs.py

@@ -0,0 +1,22 @@
+import os
+from dotenv import load_dotenv
+from utils.subprocess_operations import get_piped_stdout, stdout_to_quota_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())

hooks/slack.py

@@ -1,6 +1,7 @@
-import requests
 import json
 
+import requests
+
 
 def send_slack_message(message: str, webhook: str) -> str:
     payload = {"text": message}

quota.py

@@ -13,8 +13,7 @@ def get_fileset_quota(project_id: str) -> str:
     main_cmd = "rquota"
     piped_cmd = f"grep {project_id}"
 
-    stdout: str = get_piped_stdout(main_command=main_cmd, piped_command=piped_cmd)
-    return stdout
+    return get_piped_stdout(main_command=main_cmd, piped_command=piped_cmd)
 
 
 def monitor():

utils/subprocess_operations.py

@@ -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