feat: grow/shrink requests are being hit

I need to put this into the ensemble operator next to have the request
actually do something, like request the minicluster to scale up or
down. I will also need to have a way to communicate the member name
and namespace. This could either be done via discovery (requiring the
kubernetes API within the ensemble python and the rbac to use it),
or more simply done, just put the member name that is expected in
the same namespace. More ideally there can be a registration step at
the onset that generates a random name and sends it over to the grpc
service to associate.

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

Author: vsoch

README.md

@@ -208,6 +208,9 @@ 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
+
+# Here is how to add on the fly debug (logging->debug true)
+ensemble run --debug --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

@@ -46,6 +46,12 @@ def get_parser():
         choices=defaults.supported_members,
         default="flux",
     )
+    run.add_argument(
+        "--debug",
+        help="Enable debug logging for the config",
+        action="store_true",
+        default=False,
+    )
 
     for command in [run]:
         command.add_argument(

ensemble/client/run.py

@@ -4,5 +4,5 @@
 def main(args, parser, extra, subparser):
     # This will raise an error if the member type (e.g., minicluster) is not known
     member = members.get_member(args.executor)
-    member.load(args.config)
+    member.load(args.config, args.debug)
     member.start()

ensemble/config/config.py

@@ -19,11 +19,18 @@
 heartbeat_actions = {"grow", "shrink"}
 
 
-def load_config(config_path):
+def load_config(config_path, debug=False):
     """
     Load the config path, validating with the schema
     """
     cfg = utils.read_yaml(config_path)
+
+    # On the fly debugging
+    if debug:
+        if "logging" not in cfg:
+            cfg["logging"] = {}
+        cfg["logging"]["debug"] = True
+
     jsonschema.validate(cfg, schema=schema.ensemble_config_schema)
     return EnsembleConfig(cfg)
 

ensemble/config/types.py

@@ -32,11 +32,14 @@ def run_when(self, value):
             return True
 
         # If we have a direct value, we check for equality
-        if isinstance(value, int) and value != self.when:
+        number = (int, float)
+        if isinstance(self.when, number) and value != self.when:
             return False
+        if isinstance(self.when, number) and value == self.when:
+            return True
 
         # Otherwise, parse for inequality
-        match = re.search(r"(?P<inequality>[<>]=?)\s*(?P<comparator>\w+)", value).groupdict()
+        match = re.search(r"(?P<inequality>[<>]=?)\s*(?P<comparator>\w+)", self.when).groupdict()
 
         # This could technically be a float value
         comparator = float(match["comparator"])

ensemble/members/base.py

@@ -110,7 +110,7 @@ def execute_action(self, rule, record=None):
             # Not supported if the function does not exist
             if not run_action:
                 raise NotImplementedError("Action {rule.action.name} is not supported.")
-            self.announce(f"   {rule.action.name} {rule.action.label}")
+            self.announce(f"   {rule.action.name} {rule.action.label or ''}")
             run_action(rule, record)
 
         # Note that terminate exits but does not otherwise touch
@@ -146,11 +146,11 @@ def validate_rules(self):
         """
         self.cfg.check_supported(self.rules_supported)
 
-    def load(self, config_path):
+    def load(self, config_path, debug=False):
         """
         Load and validate the config path
         """
-        self.cfg = cfg.load_config(config_path)
+        self.cfg = cfg.load_config(config_path, debug)
         # All rules that the ensemble provides must be
         # supported by the queue executor
         self.validate_rules()

ensemble/members/client.py

@@ -1,3 +1,5 @@
+import json
+
 import ensemble.members.auth as auth
 from ensemble.protos import ensemble_service_pb2, ensemble_service_pb2_grpc
 
@@ -17,20 +19,17 @@ def __init__(self, host="localhost:50051", use_ssl=False):
         self.host = host
         self.use_ssl = use_ssl
 
-    def update_request(self, member, action, payload):
+    def action_request(self, member, action, payload):
         """
-        Send an update request to the grpc server.
+        Send an action request to the grpc server.
         """
+        payload = json.dumps(payload)
+
         # These are submit variables. A more substantial submit script would have argparse, etc.
-        request = ensemble_service_pb2.UpdateRequest(member=member, action=action, payload=payload)
+        request = ensemble_service_pb2.ActionRequest(member=member, action=action, payload=payload)
 
         with auth.grpc_channel(self.host, self.use_ssl) as channel:
             stub = ensemble_service_pb2_grpc.EnsembleOperatorStub(channel)
-            response = stub.ReceiveJobs(request)
-
-            # Case 1: SUCCESS
-            if response.status == 1:
-                print("Successful update.")
-            else:
-                print("Issue with requesting update")
+            response = stub.RequestAction(request)
+            print(f"Action request: {response.status}")
         return response

ensemble/members/flux/minicluster.py

@@ -1,3 +1,4 @@
+from ensemble.members.client import EnsembleClient
 from ensemble.members.flux.queue import FluxQueue as MemberBase
 
 # These are triggers supported for rules
@@ -12,24 +13,40 @@ class FluxMiniClusterQueue(MemberBase):
     member supports scale up and scale down for actions.
     """
 
+    # TODO for ensemble operator - we need the ensemble name to identify it
+    # perhaps generate randomly or similar?
+
+    @property
+    def client(self):
+        """
+        Ensure we have a connection to the service client.
+        """
+        # Do we have the ensemble client?
+        # TODO allow to customize host, etc.
+        if hasattr(self, "_client"):
+            return self._client
+        self._client = EnsembleClient()
+        return self._client
+
     def grow(self, rule, record=None):
         """
         Request to the API to grow the MiniCluster
         """
-        print("GROW Vanessa implement me")
-        import IPython
-
-        IPython.embed()
+        # For now use the ensemble type as the name
+        # TODO this needs to be caught and decided upon - retry?
+        response = self.client.action_request(self.name, "grow", {})
+        print(response)
 
     def shrink(self, rule, record=None):
         """
         Request to the API to shrink the MiniCluster
         """
-        print("SHRINK Vanessa implement me")
-        import IPython
-
-        IPython.embed()
+        response = self.client.action_request(self.name, "shrink", {})
+        print(response)
 
     @property
     def name(self):
-        return "MiniCluster"
+        """
+        Name is used to identify the ensemble member type.
+        """
+        return "minicluster"

ensemble/members/flux/queue.py

@@ -59,8 +59,6 @@ class FluxQueue(MemberBase):
 
     def __init__(
         self,
-        job_filters=None,
-        include_inactive=False,
         refresh_interval=0,
         summary_frequency=10,
         **kwargs,
@@ -69,15 +67,11 @@ def __init__(
         Create a new flux handle to monitor a queue.
 
         Parameters:
-        job_filters (dict)     : key value pairs for attributes or states to filter jobs to
-        include_inactive (bool): consider and include inactive jobs when doing initial parse
+        summary_frequency (int): how often (events) to show summary
         refresh_interval (int) : number of seconds to refresh all metrics (0 indicates no refresh)
         to 60 seconds. If you set to 0, it will not be set.
         """
-        # Filter jobs to any attributes, states, or similar.
-        self.filters = job_filters or {}
         self.handle = flux.Flux()
-        self.include_inactive = include_inactive
         self.refresh_interval = refresh_interval
 
         # How often on job completions to summarize?
@@ -88,13 +82,19 @@ def __init__(
         # https://github.com/flux-framework/flux-core/blob/master/src/modules/job-manager/journal.c#L28-L33
         self.seen_sentinel = False
 
+        # Have we run the start events?
+        self.started = False
+
         # We store the job id associated with a group until it's cleaned up
         self.jobids = {}
         super().__init__(**kwargs)
 
     @property
     def name(self):
-        return "FluxQueue"
+        """
+        Name is used to identify the ensemble member type.
+        """
+        return "flux"
 
     def terminate(self):
         """
@@ -148,8 +148,8 @@ def record_heartbeat_metrics(self):
         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 = set()
+        for _, group in self.jobids.items():
             groups.add(group["name"])
 
             # If we have a submit but not a start, we haven't included
@@ -231,12 +231,11 @@ def record_event(self, record):
         Record the event. This needs to be specific to the workload manager.
         """
         # The sentinel tells us when the "backlog" is finished
-        if not self.seen_sentinel:
-            if record["id"] == -1:
-                if self.cfg.debug_logging:
-                    print("Sentinel is seen, starting event monitoring.")
-                self.seen_sentinel = True
-                return
+        if record["id"] == -1:
+            if self.cfg.debug_logging:
+                print("Sentinel is seen, starting event monitoring.")
+            self.seen_sentinel = True
+            return
 
         # Record metrics for the event
         self.record_metrics(record)
@@ -260,21 +259,33 @@ def check_event(self, record, event):
         for rule in self.iter_rules(job_event):
             self.execute_rule(rule, record)
 
+    def ensure_started(self):
+        """
+        Run start rules.
+        """
+        if self.started:
+            return
+
+        # Iterate through actions provided by the rule
+        for rule in self.iter_rules("start"):
+            self.execute_rule(rule)
+        self.started = True
+
     def start(self):
         """
         Init the events subscriber (no longer pub sub but a callback)
         to the flux queue. See:
 
         https://github.com/flux-framework/flux-core/blob/master/src/modules/job-manager/journal.c#L11-L41
         """
-        # Iterate through actions provided by the rule
-        for rule in self.iter_rules("start"):
-            self.execute_rule(rule)
 
         def event_callback(response):
             """
             Receive callback when a flux job posts an event.
             """
+            # This needs to happen when the event callback is running
+            # Otherwise we get a race and can miss start events
+            self.ensure_started()
             payload = response.get()
 
             # Only print if the config has logging->debug set to true
@@ -312,7 +323,8 @@ def setup_heartbeat(self):
 
         def heartbeat_callback(cls):
             print("💗 HEARTBEAT")
-            print(cls)
+            self.summarize()
+            self.record_heartbeat_metrics()
 
         # Instead we are using threading, which works!
         self.heartbeat = QueueHeartbeat(self.cfg.heartbeat, heartbeat_callback, cls=self)
@@ -374,6 +386,11 @@ def submit(self, rule, record=None):
             jobset = self.extract_jobs(group)
 
             for job in jobset:
+                # If the number of tasks < node count we get an error
+                # assume the user wants one task per node
+                if job["tasks"] < job["nodes"]:
+                    job["tasks"] = job["nodes"]
+
                 jobspec = flux.job.JobspecV1.from_command(
                     command=job["command"], num_nodes=job["nodes"], num_tasks=job["tasks"]
                 )

ensemble/server.py

@@ -91,8 +91,9 @@ def RequestAction(self, request, context):
         response = ensemble_service_pb2.Response()
 
         # The member primarily is directed to take the action
-        member = members.get_member(request.member)
-        print(member)
+        # member = members.get_member(request.member)
+        # print(member)
+        print(response)
 
         if request.action == "grow":
             print("REQUEST TO GROW")

examples/grow-shrink-example.yaml

@@ -1,10 +1,13 @@
-# 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)
+logging:
+  heartbeat: 10
+
 jobs:
+
+  # Note that echo seems to lose events here, not great
+  # not reliable for submit at the onset at least
   - name: echo
     command: echo hello world
-    count: 5
+    count: 10
     nodes: 1
 
   # Note that this takes up the entirety of our 4 node faux
@@ -15,13 +18,6 @@ jobs:
     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
@@ -39,7 +35,7 @@ rules:
 
   # When the pending time of sleep is > 5 seconds, grow the cluster.
   - trigger: metric
-    name: pending-duration.mean
+    name: mean.sleep-pending
     when: "> 5"
     action:
       name: grow