feat: support for inequality in rule->when

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

Author: vsoch

README.md

@@ -204,6 +204,10 @@ ensemble run examples/terminate-example.yaml
 # Run a heartbeat every 3 seconds.
 # This will trigger a check to see if actions need to be performed
 ensemble run examples/heartbeat-example.yaml
+
+# Grow/shrink requires a minicluster (flux doesn't support it) but we can mock it here
+# TODO - add greater than parsing, then run example here, then in operator
+ensemble run --executor minicluster examples/grow-shrink-example.yaml
 ```
 
 Right now, this will run any rules with "start" triggers, which for this hello world example includes a few hello world jobs! You'll then be able to watch and see flux events coming in!

ensemble/client/__init__.py

@@ -7,6 +7,7 @@
 import sys
 
 import ensemble
+import ensemble.defaults as defaults
 
 
 def get_parser():
@@ -42,7 +43,7 @@ def get_parser():
     run.add_argument(
         "--executor",
         help="Executor to use (defaults to flux)",
-        choices=["flux"],
+        choices=defaults.supported_members,
         default="flux",
     )
 

ensemble/config/types.py

@@ -1,3 +1,5 @@
+import re
+
 import ensemble.defaults as defaults
 
 
@@ -20,9 +22,57 @@ def name(self):
     def trigger(self):
         return self._rule["trigger"]
 
+    def run_when(self, value):
+        """
+        Check if "when" is relevant to be run now, return True/False
+        to say to run or not.
+        """
+        # No when is set, so we just continue assuming there is no when
+        if self.when is None:
+            return True
+
+        # If we have a direct value, we check for equality
+        if isinstance(value, int) and value != self.when:
+            return False
+
+        # Otherwise, parse for inequality
+        match = re.search(r"(?P<inequality>[<>]=?)\s*(?P<comparator>\w+)", value).groupdict()
+
+        # This could technically be a float value
+        comparator = float(match["comparator"])
+        inequality = match["inequality"]
+        assert inequality in {"<", ">", "<=", ">=", "==", "="}
+
+        # Evaluate! Not sure there is a better way than this :)
+        if inequality == "<":
+            return value < comparator
+        if inequality == "<=":
+            return value <= comparator
+        if inequality == ">":
+            return value > comparator
+        if inequality == ">=":
+            return value >= comparator
+        if inequality in ["==", "="]:
+            return value == comparator
+        raise ValueError(f"Invalid comparator {comparator} for rule when")
+
+    def check_when(self):
+        """
+        Ensure we have a valid inequality before running anything!
+        """
+        # If a number, we require greater than == 0
+        if isinstance(self.when, int) and self.when >= 0:
+            return
+
+        # Check running the function with a value, should not raise error
+        try:
+            self.run_when(10)
+        except Exception as err:
+            raise ValueError(f"when: for rule {self} is not valid: {err}")
+
     @property
     def when(self):
-        return self._rule["when"]
+        return self._rule.get("when")
 
     def validate(self):
         """
@@ -33,6 +83,8 @@ def validate(self):
             raise ValueError(
                 f"Rule trigger {self.trigger} has invalid action name {self.action.name}"
             )
+        # Ensure we have a valid number or inequality
+        self.check_when()
 
 
 class Action:

ensemble/defaults.py

@@ -1,7 +1,8 @@
 workers = 10
 port = 50051
 
-valid_actions = ["submit", "custom", "terminate"]
+supported_members = ["flux", "minicluster"]
+valid_actions = ["submit", "custom", "terminate", "grow", "shrink"]
 heartbeat_seconds = 60
 
 job_events = [

ensemble/members/__init__.py

@@ -9,7 +9,7 @@ def get_member(name, options=None):
 
         return FluxQueue(**options)
     if name == "minicluster":
-        from ensemble.members.flux.minicluster import FluxQueue
+        from ensemble.members.flux.minicluster import FluxMiniClusterQueue
 
-        return FluxQueue(**options)
+        return FluxMiniClusterQueue(**options)
     raise ValueError(f"Member type {name} is not known")

ensemble/members/base.py

@@ -132,7 +132,7 @@ def execute_metric_action(self, rule, record=None):
             item = item[path]
 
         # The user set a "when" and it must match exactly.
-        if rule.when is not None and item.get() != rule.when:
+        if not rule.run_when(item.get()):
             return
         if self.cfg.debug_logging:
             print(self.metrics.models)

ensemble/members/flux/minicluster.py

@@ -1,10 +1,10 @@
-from ensemble.members.flux import FluxQueue as MemberBase
+from ensemble.members.flux.queue import FluxQueue as MemberBase
 
 # These are triggers supported for rules
 rules = ["start", "metric"]
 
 
-class FluxQueue(MemberBase):
+class FluxMiniClusterQueue(MemberBase):
     """
     The Flux Queue MiniCluster member type
 

ensemble/members/flux/queue.py

@@ -1,6 +1,6 @@
 import shlex
 import sys
-from datetime import datetime
+import time
 
 from ensemble.heartbeat import QueueHeartbeat
 from ensemble.members.base import MemberBase
@@ -138,6 +138,33 @@ def record_metrics(self, record):
         for rule in self.iter_rules("metric"):
             self.execute_rule(rule, record)
 
+    def record_heartbeat_metrics(self):
+        """
+        Heartbeat metrics cannot rely on an event, but need
+        to update metrics about the queue and then check to
+        see if any actions should be triggered related to that.
+
+        Note: this doesn't reset anything from previous pending, it
+        will be more like a moving average where the same jobs (if they
+        are still pending) get counted again, possibly increasing time.
+        """
+        groups = {}
+        for _, group in self.jobids:
+            groups.add(group["name"])
+
+            # If we have a submit but not a start, we haven't included
+            # pending yet
+            if "submit" in group and "start" not in group:
+                time_in_queue = time.time() - group["submit"]
+                group_name = f"{group['name']}-pending"
+                self.metrics.record_datum(group_name, time_in_queue)
+
+        print(f"Found active groups {groups}")
+
+        # Now execute metric rules that might be impacted
+        for rule in self.iter_rules("metric"):
+            self.execute_rule(rule)
+
     def record_start_metrics(self, event, record):
         """
         We typically want to keep the job start time for the
@@ -364,7 +391,7 @@ def submit(self, rule, record=None):
                 jobid = flux.job.submit(self.handle, jobspec)
 
                 # Don't rely on an event here, this is when the user (us) submits
-                submit_time = datetime.now().timestamp()
+                submit_time = time.time()
 
                 # This is the job id that will show up in events
                 numerical = jobid.as_integer_ratio()[0]

examples/grow-shrink-example.yaml

@@ -0,0 +1,45 @@
+# Note that we don't need to "turn on" the heartbeat because grow is present.
+# I made these the size of the cluster so we trigger scaling at least once
+# We currently assume these are uniform (so no order enforced beyond the listing here)
+jobs:
+  - name: echo
+    command: echo hello world
+    count: 5
+    nodes: 1
+
+  # Note that this takes up the entirety of our 4 node faux
+  # test cluster, so the grow will trigger when the queue time
+  # stacks up
+  - name: sleep
+    command: sleep 100
+    count: 10
+    nodes: 4
+
+  # TODO need a test case for queue time waiting...
+  # Also need to implement when with < > case.
+  - name: big-sleep
+    command: sleep 10
+    count: 5
+    nodes: 10
+
+rules:
+  # 1. This rule says to submit the sleep jobs when we start
+  - trigger: start
+    action:
+      name: submit
+      label: echo
+
+  # This says to submit the sleep jobs when we have 5 successful echo
+  - trigger: metric
+    name: count.echo.success
+    when: 5
+    action:
+      name: submit
+      label: sleep
+
+  # When the pending time of sleep is > 5 seconds, grow the cluster.
+  - trigger: metric
+    name: pending-duration.mean
+    when: "> 5"
+    action:
+      name: grow

examples/heartbeat-example.yaml

@@ -15,6 +15,13 @@ jobs:
     count: 5
     nodes: 1
 
+  # TODO need a test case for queue time waiting...
+  # Also need to implement when with < > case.
+  - name: big-sleep
+    command: sleep 10
+    count: 5
+    nodes: 10
+
 rules:
   # 1. This rule says to submit the sleep jobs when we start
   - trigger: start
@@ -31,3 +38,11 @@ rules:
     action:
       name: submit
       label: echo
+
+  # This says to submit
+  - trigger: metric
+    name: count.sleep.success
+    when: 3
+    action:
+      name: submit
+      label: echo