🐛 bug fixes. Arrays should now show up correctly

Author: mil-ad

stui/backend.py

@@ -102,11 +102,13 @@ def _thread_fn(self):
 
                 sleep(1)
 
-            except (
-                EOFError,
-                OSError,
-                socket.error,
-            ):
+            except:
+                # except (
+                #     EOFError,
+                #     OSError,
+                #     socket.error,
+                # ):
+                # TODO implement a proper reconnect state with timeout?
                 # TODO:: Where's the best place to do error handing
                 with self.lock:
                     self.fabric_connection = fabric.Connection(self.remote)
@@ -150,13 +152,23 @@ def _get_partition_info(self):
         return my_p, all_p
 
     def _get_jobs(self):
-        cmd = 'squeue --all --format="%A|%C|%b|%F|%K|%j|%P|%r|%u|%y|%T|%M|%b|%N"'
+        # %A : Job id. This will have a unique value for each element of job arrays. (Valid for jobs only)
+        # %F : Job array's job ID. This is the base job ID. For non-array jobs, this is the job ID. (Valid for jobs only)
+        # %i : Job or job step id. In the case of job arrays, the job ID format will be of the form "<base_job_id>_<index>"
+        # %K : Job array index
+        cmd = r'squeue --all --format="%A|%i|%K|%F|%C|%b|%j|%P|%r|%u|%y|%T|%M|%b|%N"'
         o = self._run_command(cmd)
 
         jobs = []
         fields = o[0].split("|")
+
+        # squeue gives the same column name (JOBID) to %A and %i so we'd have to
+        # manually differentiate them here
+        fields[1] = "JOB_ID_COMBINED"
+
         for line in o[1:]:
             job = {k: v for k, v in zip(fields, line.split("|"))}
+            job["whole_line"] = line
             jobs.append(Job(job))
 
         return jobs
@@ -194,36 +206,65 @@ def cancel_my_oldest_job(self):
 
 
 class Job(object):
-    def __init__(self, string):
+    def __init__(self, d: dict):
         super().__init__()
 
-        self.job_id = string["JOBID"]
-        self.nodes = string["NODELIST"].split(",")
-        self.partition = string["PARTITION"]
-        self.name = string["NAME"]
-        self.user = string["USER"]
-        self.state = string["STATE"]
-        self.time = string["TIME"]
-        self.nice = string["NICE"]
-        self.cpus = string["CPUS"]
-        self.gres = string["GRES"] if "GRES" in string else None
+        self.job_id = d["JOBID"]
+        self.job_id_combined = d["JOB_ID_COMBINED"]
+        self.nodes = d["NODELIST"].split(",")
+        self.partition = d["PARTITION"]
+        self.name = d["NAME"]
+        self.user = d["USER"]
+        self.state = d["STATE"]
+        self.time = d["TIME"]
+        self.nice = d["NICE"]
+        self.cpus = d["CPUS"]
+        self.gres = d["GRES"] if "GRES" in d else None
 
-        self.array_base_id = string["ARRAY_JOB_ID"]
-        self.array_task_id = string["ARRAY_TASK_ID"]
+        self.whole_line = d["whole_line"]
+
+        self.array_base_id = d["ARRAY_JOB_ID"]
+        self.array_task_id = d["ARRAY_TASK_ID"]
 
         self.is_array_job = False if self.array_task_id == "N/A" else True
 
-        if self.is_array_job and "%" in self.array_task_id:
-            match = re.search(r"\d+%(\d+)", self.array_task_id)
-            self.array_throttle = match.group(1)
-        else:
-            self.array_throttle = None
+        if self.is_array_job:
+            if self.is_pending():
+                if "%" in self.array_task_id:
+                    match = re.search(r"-(\d+)%(\d+)$", self.array_task_id)
+                    self.array_total_jobs = match.group(2)
+                    self.array_throttle = match.group(2)
+                else:
+                    match = re.search(r"-(\d+)$", self.array_task_id)
+                    self.array_total_jobs = match.group(1)
+                    self.array_throttle = 0  # TODO: 0 means unlimited. Is this good?
+            else:
+                self.array_throttle = None
+                self.array_total_jobs = None
+
+        # if self.is_array_job and "%" in self.array_task_id:
+        #     match = re.search(r"\d+%(\d+)", self.array_task_id)
+        #     self.array_throttle = match.group(1)
+        # else:
+        #     self.array_throttle = None
 
     def __repr__(self):
         return f"Job {self.job_id} - State{self.state}"
 
     def is_running(self):
         return self.state == "RUNNING"
 
+    def is_pending(self):
+        return self.state == "PENDING"
+
     def uses_gpu(self):
         return "gpu" in self.gres
+
+    def is_array_job_f(self):
+        return self.is_array_job  # TODO: use property?
+
+    def array_str(self):
+        if not self.is_array_job:
+            return ""
+        else:
+            return self.array_task_id

stui/stui.py

@@ -19,7 +19,7 @@ class JobWidget(urwid.WidgetWrap):
         "CANCELLED": {None: ""},
         "COMPLETED": {None: ""},
         "CONFIGURING": {None: ""},
-        "COMPLETING": {None: ""},
+        "COMPLETING": {None: "job_state_completeing"},
         "DEADLINE": {None: ""},
         "FAILED": {None: ""},
         "NODE FAIL": {None: ""},
@@ -46,6 +46,7 @@ def __init__(self, job):
         self.columns = OrderedDict()
         self.columns["selected"] = urwid.Text("", wrap="ellipsis")
         self.columns["job_id"] = urwid.Text("", wrap="ellipsis")
+        self.columns["array"] = urwid.Text("", wrap="ellipsis")
         self.columns["user"] = urwid.Text("", wrap="ellipsis")
         self.columns["name"] = urwid.Text("", wrap="ellipsis")
         self.columns["state"] = urwid.AttrMap(urwid.Text("", wrap="ellipsis"), None)
@@ -73,6 +74,7 @@ def __init__(self, job):
 
     def update_values(self, job):
         self.columns["job_id"].set_text(job.job_id)
+        self.columns["array"].set_text(job.array_str())
         self.columns["user"].set_text(job.user)
         self.columns["name"].set_text(job.name)
         self.columns["state"]._original_widget.set_text(job.state.title())
@@ -91,7 +93,12 @@ def set_selected_attr(self, in_focus):
             attr = "highlight_out_of_focus"
 
         self._w.set_focus_map(
-            {None: attr, "job_state_running": attr, "job_state_pending": attr}
+            {
+                None: attr,
+                "job_state_running": attr,
+                "job_state_pending": attr,
+                "job_state_completeing": attr,
+            }
         )
 
     def keypress(self, size, key):
@@ -118,6 +125,7 @@ class JobQueueWidget(urwid.WidgetWrap):
     column_widths = [
         (2,),
         (10,),
+        (10,),
         ("weight", 1),
         ("weight", 2),
         (14,),
@@ -133,6 +141,7 @@ def __init__(self):
         column_labels = [
             "",
             "Job ID",
+            "Job Array",
             "User",
             "Name",
             "State",
@@ -298,9 +307,15 @@ def __init__(self):
 
         super().__init__(w)
 
-    def set_nice(self, value):
+    def set_nice_value(self, value):
         self.nice_spinbutton.set_value(value)
 
+    def enable_nice(self):
+        self.nice_spinbutton.enable()
+
+    def disable_nice(self):
+        self.nice_spinbutton.disable()
+
     def enable_throttle(self):
         self.throttle_spinbutton.enable()
 
@@ -384,13 +399,16 @@ def on_jobs_focus_changed(self):
         if job is None:
             return
 
-        self.apanel.set_nice(job.nice)
+        self.apanel.set_nice_value(job.nice)
 
-        if job.array_throttle is not None:
-            self.apanel.enable_throttle()
-            self.apanel.set_throttle_value(job.array_throttle)
-        else:
+        if job.is_running():
+            self.apanel.disable_nice()
             self.apanel.disable_throttle()
+        else:
+            self.apanel.enable_nice()
+            if job.is_array_job_f():
+                self.apanel.enable_throttle()
+                self.apanel.set_throttle_value(job.array_throttle)
 
     def filter_jobs(self, jobs):
 
@@ -706,6 +724,7 @@ def __init__(self, args):
         self.palette = [
             ("job_state_running", "light cyan", ""),
             ("job_state_pending", "yellow", ""),
+            ("job_state_completeing", "light magenta", ""),
             ("active_tab_label", "yellow", ""),
             ("focus_and_active_tab_label", "yellow,underline", ""),
             ("focus_and_inactive_tab_label", "underline", ""),
@@ -748,13 +767,13 @@ def exit(*args, **kwargs):
 
         elif message == b"connection established":
             self.fd = None  # TODO: delattr?
-        self.register_refresh()
-        self.w.cluster_connected_callback()
+            self.register_refresh()
+            self.w.cluster_connected_callback()
 
             # Return False so that the watch is removed from main loop and its read-end
             # of the pipe is closed. The write-end of the pipe will be closed on the
             # backend.
-        return False
+            return False
 
     def run(self):
         # self.loop.screen.set_terminal_properties(bright_is_bold=False)