feat: support for flux minicluster

Flux running in the MiniCluster contexts adds the cool
features to grow/shrink (autoscaling), which can be
paired with an actual cluster autoscaler (or not) if
the mini cluster is not using all the nodes available.
To support this we need a few things - first a heartbeat,
and one that runs at a user specified increment (and likely
this should also be exposed in the ensemble config) because
it is very likely the case that the triggers are not linked
to jobs (for example, "run this when the queue wait time is
over X for this job group"). We then need the MiniCluster
member, which is exactly the same as flux, but instead has
added the grow/shrink. I think likely what I want to do is
have the heartbeat disabled unless the setting is found in
the ensemble config OR the action to grow/shrink is found
(and we would use a default heartbeat seconds of 60). This
time should obviously be tested for different applications.

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

Author: vsoch

Makefile

@@ -17,6 +17,10 @@ python: python ## Generate python proto files in python
 docker-build:
 	docker build -t ${IMG} .
 
+.PHONY: kind-load
+kind-load: docker-build # For local development with kind
+	kind load docker-image ${IMG}
+
 .PHONY: arm-build
 arm-build:
 	docker buildx build --platform linux/arm64 --build-arg ARCH=aarch_64 -t ${ARMIMG} .

README.md

@@ -65,15 +65,27 @@ Note that yes, this means "submit" is both an action and an event.
 
 The design of a rule is to have an action, and the action is something your ensemble can be tasked to do when an event is triggered. Right now we support the following:
 
-- **submit**: submit a job
-- **terminate**: terminate the member. This is usually what you want to call to have the entire thing exit on 0
-- **custom**: run a custom function that will receive known kwargs, and then should return an action (read more below)
+**Flux and MiniCluster**
 
+- *submit*: submit a job
+- *terminate*: terminate the member. This is usually what you want to call to have the entire thing exit on 0
+- *custom*: run a custom function that will receive known kwargs, and then should return an action (read more below)
+
+**MiniCluster Only**
+
+- *grow*: grow (or request) the cluster to scale up
+- *shrink*: shrink (or request) the cluster to scale down
+
+For the scale operations, since this brings in the issue of resource contention between different ensembles, we have to assume to start that a request to scale:
+
+1. Should only happen once. If it's granted, great, if not, we aren't going to ask again.
+2. Does not need to consider another level of scheduler (e.g., fair shaire)
+
+I started thinking about putting a second scheduler (or fair share algorithm) in the grpc service, but realized this was another level of complexity that although we might work on it later, is not warranted yet.
 We see "submit" as two examples in the above, which is a common thing you'd want to do! For each action, you should minimally define the "name" and a "label" that typically corresponds to a job group.
 You can also optionally define "repetitions," which are the number of times the action should be run before expiring. If you want a backoff period between repetitions, set "backoff" to a non zero value.
 By default, when no repetitions or backoff are set, the action is assumed to have a repetition of 1. It will be run once! Let's now look at a custom action. Here is what your function should look like in your `ensemble.yaml`
 
-
 ##### Actions
 
 

ensemble/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.15"
+__version__ = "0.0.16"

ensemble/heartbeat.py

@@ -0,0 +1,53 @@
+import signal
+import threading
+
+
+class GracefulExit(Exception):
+    """
+    Ensure if we press control+C, it doesn't throw up.
+    """
+
+    pass
+
+
+class QueueHeartbeat(threading.Thread):
+    """
+    The Queue Heartbeat triggers at a user specified interval, with the
+    intention to be able to run events that might not be linked to jobs.
+    """
+
+    def __init__(self, interval_seconds, callback, **kwargs):
+        super().__init__()
+        self.stop_event = threading.Event()
+        self.interval_seconds = interval_seconds
+        self.callback = callback
+        self.kwargs = kwargs
+
+    def run(self):
+        """
+        Run the heartbeat function, and exit gracefully
+        """
+        while not self.stop_event.wait(self.interval_seconds):
+            try:
+                self.callback(**self.kwargs)
+            except (KeyboardInterrupt, GracefulExit):
+                break
+
+    def stop(self):
+        """
+        Explicit stop of the thread, like that will happen :)
+        """
+        self.stop_event.set()
+
+
+def signal_handler(signum, frame):
+    """
+    The signal handler follows an expected pattern, but
+    just calls a graceful exit
+    """
+    raise GracefulExit()
+
+
+# Signals for our heartbeat
+signal.signal(signal.SIGINT, signal_handler)
+signal.signal(signal.SIGTERM, signal_handler)

ensemble/members/__init__.py

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

ensemble/members/auth.py

@@ -0,0 +1,18 @@
+from contextlib import contextmanager
+
+import grpc
+
+
+@contextmanager
+def grpc_channel(host, use_ssl=False):
+    """
+    Yield a channel, either with or without ssl, and close properly.
+    """
+    if use_ssl:
+        channel = grpc.secure_channel(host, grpc.ssl_channel_credentials())
+    else:
+        channel = grpc.insecure_channel(host)
+    try:
+        yield channel
+    finally:
+        channel.close()

ensemble/members/base.py

@@ -87,7 +87,9 @@ def announce(self, message, meta=None, ljust=15, color="cyan"):
     def execute_action(self, rule, record=None):
         """
         Execute a general action, it's assumed that
-        the trigger was called if we get here.
+        the trigger was called if we get here. Those supported by
+        Flux and the MiniCluster: submit, custom, terminate.
+        MiniCluster Only: shrink and grow.
         """
         # This function checks for repetitions and backoff
         # periods, and determines if we should continue (to run)
@@ -101,9 +103,15 @@ def execute_action(self, rule, record=None):
             self.announce(f"   submit {rule.action.label} ", self.cfg.pretty_job(rule.action.label))
             return self.submit(rule, record)
 
-        if rule.action.name == "custom":
-            self.announce(f"   custom {rule.action.label}")
-            return self.custom_action(rule, record)
+        # These all have the same logic to call the function of the same name
+        if rule.action.name in ["custom", "grow", "shrink"]:
+            run_action = getattr(self, rule.action.name, 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}")
+            run_action(rule, record)
 
         # Note that terminate exits but does not otherwise touch
         # the queue, etc. Given a reactor, we should just stop it
@@ -153,9 +161,6 @@ def start(self, *args, **kwargs):
         """
         raise NotImplementedError
 
-    def custom_action(self, rule, record=None):
-        raise NotImplementedError
-
     def submit(self, *args, **kwargs):
         """
         Submit a job

ensemble/members/client.py

@@ -0,0 +1,36 @@
+import ensemble.members.auth as auth
+from ensemble.protos import ensemble_service_pb2, ensemble_service_pb2_grpc
+
+
+class EnsembleClient:
+    """
+    The EnsembleClient is used by members to communicate with the grpc service.
+
+    The grpc service will receive requests to grow/shrink, etc (requests that fall outside
+    of the scope of the ensemble and require changing the member) and issue some kind of event
+    that can be listened to by an entity to do it (e.g., Kubernetes). Right now we support
+    update requests (to scale up and down) and status requests (to check on state that
+    some grpc endpoint sees).
+    """
+
+    def __init__(self, host="localhost:50051", use_ssl=False):
+        self.host = host
+        self.use_ssl = use_ssl
+
+    def update_request(self, member, action, payload):
+        """
+        Send an update request to the grpc server.
+        """
+        # These are submit variables. A more substantial submit script would have argparse, etc.
+        request = ensemble_service_pb2.UpdateRequest(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")
+        return response

ensemble/members/flux/minicluster.py

@@ -0,0 +1,35 @@
+from ensemble.members.flux import FluxQueue as MemberBase
+
+# These are triggers supported for rules
+rules = ["start", "metric"]
+
+
+class FluxQueue(MemberBase):
+    """
+    The Flux Queue MiniCluster member type
+
+    It uses the FluxQueue as a base. The main difference is that this
+    member supports scale up and scale down for actions.
+    """
+
+    def grow(self, rule, record=None):
+        """
+        Request to the API to grow the MiniCluster
+        """
+        print("GROW Vanessa implement me")
+        import IPython
+
+        IPython.embed()
+
+    def shrink(self, rule, record=None):
+        """
+        Request to the API to shrink the MiniCluster
+        """
+        print("SHRINK Vanessa implement me")
+        import IPython
+
+        IPython.embed()
+
+    @property
+    def name(self):
+        return "MiniCluster"

ensemble/members/flux/queue.py

@@ -2,12 +2,14 @@
 import sys
 from datetime import datetime
 
+from ensemble.heartbeat import QueueHeartbeat
 from ensemble.members.base import MemberBase
 
 try:
     import flux
     import flux.constants
     import flux.job
+    import flux.message
 except ImportError:
     sys.exit("flux python is required to use the flux queue member")
 
@@ -61,6 +63,7 @@ def __init__(
         include_inactive=False,
         refresh_interval=0,
         summary_frequency=10,
+        heartbeat_seconds=60,
         **kwargs,
     ):
         """
@@ -70,6 +73,8 @@ def __init__(
         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
         refresh_interval (int) : number of seconds to refresh all metrics (0 indicates no refresh)
+        heartbeat_seconds (int): number of seconds to heartbeat (for grow/shrink). This defaults
+        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 {}
@@ -80,6 +85,7 @@ def __init__(
         # How often on job completions to summarize?
         self.summary_freqency = summary_frequency
         self.completion_counter = 0
+        self.heartbeat_seconds = heartbeat_seconds
 
         # The flux sentinel tells us when we have finished with the backlog
         # https://github.com/flux-framework/flux-core/blob/master/src/modules/job-manager/journal.c#L28-L33
@@ -98,6 +104,8 @@ def terminate(self):
         Custom termination function for flux.
         """
         self.handle.reactor_stop()
+        if self.heartbeat_seconds:
+            self.heartbeat.stop()
 
     def record_metrics(self, record):
         """
@@ -258,11 +266,65 @@ def event_callback(response):
             flags=flux.constants.FLUX_RPC_STREAMING,
         )
         events.then(event_callback)
-        self.handle.reactor_run()
+        self.setup_heartbeat()
+        self.reactor_start()
 
-    def custom_action(self, rule, record=None):
+    def reactor_start(self):
         """
-        Custom action runs the action (and runs another action, if returned)
+        Courtesy function to start the reactor and more
+        gracefully handle keyboard interrupts.
+        """
+        try:
+            self.handle.reactor_run()
+        except KeyboardInterrupt:
+            self.terminate()
+
+    def setup_heartbeat(self):
+        """
+        Setup the heartbeat - a threading.Thread
+        """
+        if not self.heartbeat_seconds:
+            return
+
+        def heartbeat_callback(cls):
+            print("💗 HEARTBEAT")
+            print(cls)
+
+        # Instead we are using threading, which works!
+        self.heartbeat = QueueHeartbeat(self.heartbeat_seconds, heartbeat_callback, cls=self)
+        self.heartbeat.start()
+
+    def cron_heartbeat(self):
+        """
+        cron heartbeat provided by flux (does not work)
+        """
+
+        def heartbeat_callback(response):
+            print("💗 HEARTBEAT")
+            print(response)
+
+        # Create a cron heartbeat every N seconds, only if we have a heartbeat set
+        # This is intended for grow/shrink actions that might need a regular check
+        print(f"  💗 Creating flux heartbeat every {self.heartbeat_seconds} seconds")
+        heartbeat = self.handle.rpc(
+            "cron.create",
+            {
+                "type": "interval",
+                "name": "heartbeat",
+                "command": "sleep 0",
+                "args": {"interval": self.heartbeat_seconds},
+            },
+            flux.constants.FLUX_NODEID_ANY,
+            flags=flux.constants.FLUX_RPC_STREAMING,
+        )
+
+        self.handle.flux_event_subscribe("cron.*")
+        self.handle.event_subscribe("cron.*")
+        heartbeat.then(heartbeat_callback)
+
+    def custom(self, rule, record=None):
+        """
+        Custom runs a custom action (and runs another action, if returned)
         and passes forward the flux handle and other metadata.
         """
         kwargs = {