Add list_nodes script

Author: milliams

citc/list_nodes.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+from typing import List, Optional
+
+from tabulate import tabulate
+
+from . import slurm, cloud, utils
+
+
+def create_table(
+    slurm_nodes: List[slurm.SlurmNode], cloud_nodes: List[cloud.CloudNode]
+):
+    table = []
+    for slurm_node in slurm_nodes:
+        matches = [node for node in cloud_nodes if node.name == slurm_node.name]
+        if len(matches) == 1:
+            cloud_state = matches[0].state  # type: Optional[cloud.NodeState]
+        else:
+            cloud_state = None
+        table.append(
+            [
+                slurm_node.name,
+                slurm_node.state,
+                slurm.NODE_STATE_FLAGS.get(slurm_node.state_flag or "", ""),
+                cloud_state,
+            ]
+        )
+
+    return table
+
+
+def main():
+    slurm_conf = Path("/mnt/shared/etc/slurm/slurm.conf")
+
+    cloud_nodes = utils.get_cloud_nodes()
+    slurm_nodes = slurm.all_nodes(slurm_conf)
+
+    table_data = create_table(slurm_nodes, cloud_nodes)
+    table = tabulate(table_data, headers=["Name", "Slurm state", "Flag", "Cloud state"])
+    print(table)
+
+
+if __name__ == "__main__":
+    main()

citc/slurm.py

@@ -20,12 +20,12 @@ def node_list(slurm_conf: pathlib.Path) -> Iterator[str]:
 
 
 NODE_STATE_FLAGS = {
-    "*": "not_responding",
-    "~": "power_save",
-    "#": "powering_up",
-    "%": "powering_down",
-    "$": "main_reservation",
-    "@": "pending_reboot",
+    "*": "not responding",
+    "~": "power save",
+    "#": "powering up",
+    "%": "powering down",
+    "$": "main reservation",
+    "@": "pending reboot",
 }
 
 

pyproject.toml

@@ -10,6 +10,7 @@ python = "^3.6"
 boto3 = "^1.10.46"
 PyYaml = "^5.2"
 oci = "^2.8.0"
+tabulate = "^0.8.6"
 
 [tool.poetry.dev-dependencies]
 pytest = "^5.3"
@@ -21,6 +22,7 @@ moto = "^1.3.14"
 
 [tool.poetry.scripts]
 watchdog = 'citc.watchdog:main'
+list_nodes = 'citc.list_nodes:main'
 
 [build-system]
 requires = ["poetry>=0.12"]

tests/test_list_nodes.py

@@ -0,0 +1,26 @@
+from tabulate import tabulate
+
+from citc.aws import AwsNode
+from citc.list_nodes import create_table
+from citc.cloud import NodeState
+from citc.slurm import SlurmNode
+
+
+def test_create_table():
+    slurm_nodes = [
+        SlurmNode(name="foo-1", state="idle", state_flag=None, features={}),
+        SlurmNode(name="foo-2", state="idle", state_flag="~", features={}),
+    ]
+    cloud_nodes = [AwsNode(name="foo-1", state=NodeState.RUNNING)]
+    table = create_table(slurm_nodes, cloud_nodes)
+    assert table == [
+        ["foo-1", "idle", "", NodeState.RUNNING],
+        ["foo-2", "idle", "power save", None],
+    ]
+
+
+def test_print_table():
+    slurm_nodes = [SlurmNode(name="foo-1", state="idle", state_flag=None, features={})]
+    cloud_nodes = [AwsNode(name="foo-1", state=NodeState.RUNNING)]
+    table_data = create_table(slurm_nodes, cloud_nodes)
+    print(tabulate(table_data))