Clustermgtd refactor

Signed-off-by: chenwany <[email protected]>

Author: chenwany

src/common/schedulers/slurm_commands.py

@@ -9,14 +9,21 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-import collections
 import logging
 import re
-from enum import Enum
 
-from common.utils import check_command_output, convert_range_to_list, grouper, run_command
+from common.utils import check_command_output, grouper, run_command
 from retrying import retry
+from slurm_plugin.slurm_resources import (
+    DynamicNode,
+    InvalidNodenameError,
+    PartitionStatus,
+    SlurmPartition,
+    StaticNode,
+    parse_nodename,
+)
+
+log = logging.getLogger(__name__)
 
 PENDING_RESOURCES_REASONS = [
     "Resources",
@@ -46,144 +53,12 @@
 SCONTROL = "sudo /opt/slurm/bin/scontrol"
 SINFO = "/opt/slurm/bin/sinfo"
 
-SlurmPartition = collections.namedtuple("SlurmPartition", ["name", "nodes", "state"])
-NodeType = collections.namedtuple("NodeType", ["instance_type", "partition"])
-
 # Set default timeouts for running different slurm commands.
 # These timeouts might be needed when running on large scale
 DEFAULT_GET_INFO_COMMAND_TIMEOUT = 30
 DEFAULT_UPDATE_COMMAND_TIMEOUT = 60
 
 
-class PartitionStatus(Enum):
-    UP = "UP"
-    DOWN = "DOWN"
-    INACTIVE = "INACTIVE"
-    DRAIN = "DRAIN"
-
-    def __str__(self):
-        return str(self.value)
-
-
-class SlurmNode:
-    SLURM_SCONTROL_BUSY_STATES = {"MIXED", "ALLOCATED", "COMPLETING"}
-    SLURM_SCONTROL_IDLE_STATE = "IDLE"
-    SLURM_SCONTROL_DOWN_STATE = "DOWN"
-    SLURM_SCONTROL_DRAIN_STATE = "DRAIN"
-    SLURM_SCONTROL_POWERING_DOWN_STATE = "POWERING_DOWN"
-    SLURM_SCONTROL_POWER_STATE = "IDLE+CLOUD+POWER"
-    SLURM_SCONTROL_POWER_UP_STATE = "#"
-    SLURM_SCONTROL_ONLINE_STATES = {"IDLE+CLOUD", "MIXED+CLOUD", "ALLOCATED+CLOUD", "COMPLETING+CLOUD"}
-    SLURM_SCONTROL_POWER_WITH_JOB_STATE = "MIXED+CLOUD+POWER"
-    SLURM_SCONTROL_RESUME_FAILED_STATE = "DOWN*+CLOUD+POWER"
-
-    def __init__(self, name, nodeaddr, nodehostname, state, partitions=None):
-        """Initialize slurm node with attributes."""
-        self.name = name
-        self.is_static = is_static_node(name)
-        self.nodeaddr = nodeaddr
-        self.nodehostname = nodehostname
-        self.state = state
-        self.partitions = partitions.strip().split(",") if partitions else None
-
-    def is_nodeaddr_set(self):
-        """Check if nodeaddr(private ip) for the node is set."""
-        return self.nodeaddr != self.name
-
-    def has_job(self):
-        """Check if slurm node is in a working state."""
-        return any(working_state in self.state for working_state in self.SLURM_SCONTROL_BUSY_STATES)
-
-    def _is_drain(self):
-        """Check if slurm node is in any drain(draining, drained) states."""
-        return self.SLURM_SCONTROL_DRAIN_STATE in self.state
-
-    def is_drained(self):
-        """
-        Check if slurm node is in drained state.
-
-        drained(sinfo) is equivalent to IDLE+DRAIN(scontrol) or DOWN+DRAIN(scontrol)
-        """
-        return self._is_drain() and (
-            self.SLURM_SCONTROL_IDLE_STATE in self.state or self.SLURM_SCONTROL_DOWN_STATE in self.state
-        )
-
-    def is_powering_down(self):
-        """Check if slurm node is in powering down state."""
-        return self.SLURM_SCONTROL_POWERING_DOWN_STATE in self.state
-
-    def is_power(self):
-        """Check if slurm node is in power state."""
-        return self.SLURM_SCONTROL_POWER_STATE == self.state
-
-    def is_down(self):
-        """Check if slurm node is in a down state."""
-        return self.SLURM_SCONTROL_DOWN_STATE in self.state and not self.is_powering_down()
-
-    def is_up(self):
-        """Check if slurm node is in a healthy state."""
-        return not self._is_drain() and not self.is_down() and not self.is_powering_down()
-
-    def is_powering_up(self):
-        """Check if slurm node is in powering up state."""
-        return self.SLURM_SCONTROL_POWER_UP_STATE in self.state
-
-    def is_online(self):
-        """Check if slurm node is online with backing instance."""
-        return self.state in self.SLURM_SCONTROL_ONLINE_STATES
-
-    def is_configuring_job(self):
-        """Check if slurm node is configuring with job and haven't begun to run a job."""
-        return self.is_powering_up() and self.has_job()
-
-    def is_power_with_job(self):
-        """Dynamic nodes allocated a job but power up process has not started yet."""
-        return self.state == self.SLURM_SCONTROL_POWER_WITH_JOB_STATE
-
-    def is_running_job(self):
-        """Check if slurm node is running a job but not in configuring job state."""
-        return not self.is_powering_up() and self.has_job() and not self.is_power_with_job()
-
-    def is_resume_failed(self):
-        """Check if node resume timeout expires."""
-        return self.state == self.SLURM_SCONTROL_RESUME_FAILED_STATE
-
-    def __eq__(self, other):
-        """Compare 2 SlurmNode objects."""
-        if isinstance(other, SlurmNode):
-            return self.__dict__ == other.__dict__
-        return False
-
-    def __repr__(self):
-        attrs = ", ".join(["{key}={value}".format(key=key, value=repr(value)) for key, value in self.__dict__.items()])
-        return "{class_name}({attrs})".format(class_name=self.__class__.__name__, attrs=attrs)
-
-    def __str__(self):
-        return f"{self.name}({self.nodeaddr})"
-
-
-class InvalidNodenameError(ValueError):
-    r"""
-    Exception raised when encountering a NodeName that is invalid/incorrectly formatted.
-
-    Valid NodeName format: {queue-name}-{st/dy}-{instancetype}-{number}
-    And match: ^([a-z0-9\-]+)-(st|dy)-([a-z0-9-]+)-\d+$
-    Sample NodeName: queue-1-st-c5xlarge-2
-    """
-
-    pass
-
-
-def parse_nodename(nodename):
-    """Parse queue_name, node_type (st vs dy) and instance_type from nodename."""
-    nodename_capture = re.match(r"^([a-z0-9\-]+)-(st|dy)-([a-z0-9]+)-\d+$", nodename)
-    if not nodename_capture:
-        raise InvalidNodenameError
-
-    queue_name, node_type, instance_name = nodename_capture.groups()
-    return queue_name, node_type, instance_name
-
-
 def is_static_node(nodename):
     """
     Check if the node is static or dynamic.
@@ -227,7 +102,6 @@ def update_nodes(
         update_cmd += f" state={state}"
     if reason:
         update_cmd += f' reason="{reason}"'
-
     for nodenames, addrs, hostnames in batched_node_info:
         node_info = f"nodename={nodenames}"
         if addrs:
@@ -248,7 +122,7 @@ def update_partitions(partitions, state):
             )
             succeeded_partitions.append(partition)
         except Exception as e:
-            logging.error("Failed when setting partition %s to %s with error %s", partition, state, e)
+            log.error("Failed when setting partition %s to %s with error %s", partition, state, e)
 
     return succeeded_partitions
 
@@ -261,15 +135,15 @@ def update_all_partitions(state, reset_node_addrs_hostname):
         partition_to_update = []
         for part in partitions:
             if PartitionStatus(part.state) != PartitionStatus(state):
-                logging.info(f"Setting partition {part.name} state from {part.state} to {state}")
+                log.info(f"Setting partition {part.name} state from {part.state} to {state}")
                 if reset_node_addrs_hostname:
-                    logging.info(f"Resetting partition nodes {part.nodes}")
-                    reset_nodes(part.nodes, state="power_down", reason="stopping cluster")
+                    log.info(f"Resetting partition nodes {part.nodenames}")
+                    reset_nodes(part.nodenames, state="power_down", reason="stopping cluster")
                 partition_to_update.append(part.name)
         succeeded_partitions = update_partitions(partition_to_update, state)
         return succeeded_partitions == partition_to_update
     except Exception as e:
-        logging.error("Failed when updating partitions with error %s", e)
+        log.error("Failed when updating partitions with error %s", e)
         return False
 
 
@@ -296,15 +170,13 @@ def _batch_node_info(nodenames, nodeaddrs, nodehostnames, batch_size):
         try:
             nodeaddrs_batch = _batch_attribute(nodeaddrs, batch_size, expected_length=len(nodenames))
         except ValueError:
-            logging.error("Nodename %s and NodeAddr %s contain different number of entries", nodenames, nodeaddrs)
+            log.error("Nodename %s and NodeAddr %s contain different number of entries", nodenames, nodeaddrs)
             raise
     if nodehostnames:
         try:
             nodehostnames_batch = _batch_attribute(nodehostnames, batch_size, expected_length=len(nodenames))
         except ValueError:
-            logging.error(
-                "Nodename %s and NodeHostname %s contain different number of entries", nodenames, nodehostnames
-            )
+            log.error("Nodename %s and NodeHostname %s contain different number of entries", nodenames, nodehostnames)
             raise
 
     return zip(nodename_batch, nodeaddrs_batch, nodehostnames_batch)
@@ -467,53 +339,14 @@ def _parse_nodes_info(slurm_node_info):
     for node in node_info:
         lines = node.strip().splitlines()
         if lines:
-            slurm_nodes.append(SlurmNode(*lines))
-    return slurm_nodes
-
-
-def get_nodes_type(nodes):
-    """Get node type given nodename. Node type format: NodeType(instancetype, queue_name)."""
-    nodes_types = set()
-    for node in nodes:
-        queue_name, _, instance_name = parse_nodename(node.name)
-        nodes_types.add(NodeType(instance_name, queue_name))
-    return nodes_types
+            try:
+                if is_static_node(lines[0]):
+                    node = StaticNode(*lines)
+                    slurm_nodes.append(node)
+                else:
+                    node = DynamicNode(*lines)
+                    slurm_nodes.append(node)
+            except InvalidNodenameError:
+                log.warning("Ignoring node %s because it has an invalid name", lines[0])
 
-
-def get_inactive_partition_names(partitions):
-    """Get the name of all inactive partitions from partitions and name mapping."""
-    inactive_partitions_name = []
-    for partition_name, partition in partitions.items():
-        if partition.state == "INACTIVE":
-            inactive_partitions_name.append(partition_name)
-    return inactive_partitions_name
-
-
-def retrieve_partitions_from_node_types(node_types):
-    """Retrieve partitions for node types."""
-    return {node_type.partition for node_type in node_types}
-
-
-def get_partition_node_names(nodenames):
-    """
-    Convert partition nodenames read from partition to a list of nodes.
-
-    Example input nodenames: "queue1-st-c5xlarge-[1,3,4-5],queue1-st-c5large-20"
-    Example output [queue1-st-c5xlarge-1, queue1-st-c5xlarge-3, queue1-st-c5xlarge-4, queue1-st-c5xlarge-5,
-    queue1-st-c5large-20]
-    """
-    if type(nodenames) is str:
-        matches = re.findall(r"((([a-z0-9\-]+)-(st|dy)-([a-z0-9]+)-)(\[[\d+,-]+\]|\d+))", nodenames)
-        # [('queue1-st-c5xlarge-[1,3,4-5]', 'queue1-st-c5xlarge-', 'queue1', 'st', 'c5xlarge', '[1,3,4-5]'),
-        # ('queue1-st-c5large-20', 'queue1-st-c5large-', 'queue1', 'st', 'c5large', '20')]
-    partition_node_names = []
-    for match in matches:
-        nodename, prefix, _, _, _, nodes = match
-        if "[" not in nodes:
-            # Single nodename
-            partition_node_names.append(nodename)
-        else:
-            # Multiple nodenames
-            node_list = convert_range_to_list(nodes.strip("[]"))
-            partition_node_names += [prefix + str(n) for n in node_list]
-    return partition_node_names
+    return slurm_nodes

src/common/utils.py

@@ -259,3 +259,22 @@ def convert_range_to_list(node_range):
         ),
         [],
     )
+
+
+def time_is_up(initial_time, current_time, grace_time):
+    """Check if timeout is exceeded."""
+    # Localize datetime objects to UTC if not previously localized
+    # All timestamps used in this function should be already localized
+    # Assume timestamp was taken from system local time if there is no localization info
+    if not initial_time.tzinfo:
+        logging.warning(
+            "Timestamp %s is not localized. Please double check that this is expected, localizing to UTC.", initial_time
+        )
+        initial_time = initial_time.astimezone(tz=timezone.utc)
+    if not current_time.tzinfo:
+        logging.warning(
+            "Timestamp %s is not localized. Please double check that this is expected, localizing to UTC", current_time
+        )
+        current_time = current_time.astimezone(tz=timezone.utc)
+    time_diff = (current_time - initial_time).total_seconds()
+    return time_diff >= grace_time