feat: shrink with flux minicluster example working.

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

Author: vsoch

api/v1alpha1/statemachine_types.go

@@ -88,6 +88,9 @@ type WorkflowEvent struct {
 	// Backoff and repetitions to respond to event
 	Backoff     int32 `json:"backoff,omitempty"`
 	Repetitions int32 `json:"repetitions,omitempty"`
+	MinCompletions int32 `json:"minCompletions,omitempty"`
+	MaxSize int32 `json:"maxSize,omitempty"`
+	MinSize int32 `json:"minSize,omitempty"`
 }
 
 // A JobSequence is a list of JobSteps

config/crd/bases/state-machine.converged-computing.org_statemachines.yaml

@@ -250,10 +250,19 @@ spec:
                           description: Backoff and repetitions to respond to event
                           format: int32
                           type: integer
+                        maxSize:
+                          format: int32
+                          type: integer
                         metric:
                           description: Name of metric, indexed into model lookup (e.g.,
                             count.job_b.failed)
                           type: string
+                        minCompletions:
+                          format: int32
+                          type: integer
+                        minSize:
+                          format: int32
+                          type: integer
                         repetitions:
                           format: int32
                           type: integer

examples/dist/state-machine-operator-dev.yaml

@@ -258,10 +258,19 @@ spec:
                           description: Backoff and repetitions to respond to event
                           format: int32
                           type: integer
+                        maxSize:
+                          format: int32
+                          type: integer
                         metric:
                           description: Name of metric, indexed into model lookup (e.g.,
                             count.job_b.failed)
                           type: string
+                        minCompletions:
+                          format: int32
+                          type: integer
+                        minSize:
+                          format: int32
+                          type: integer
                         repetitions:
                           format: int32
                           type: integer

examples/test/flux-minicluster.yaml

@@ -7,6 +7,17 @@ spec:
     pullPolicy: Never
   workflow:
     completed: 10
+    events:
+      # When more than half of job b's fail, finish the entire workflow
+      - metric: mean.job_b.duration
+        when: "<= 60"
+        action: shrink
+        # Size 4 cluster, max decreases down to 1
+        repetitions: 3
+        # Require checks between before doing again
+        # TODO check if this is working correctly.
+        # also check the .get()
+        backoff: 3
 
   cluster:
     maxSize: 2
@@ -22,7 +33,7 @@ spec:
   # Note that this step always fails and we never make it to C
   # This should end the workflow early
   - name: job_b
-    properties: 
+    properties:
       minicluster: "yes"
     config:
       nodes: 4

internal/controller/manager/templates/entrypoint.sh

@@ -12,7 +12,12 @@ workflow:
   events: {{ range .Spec.Workflow.Events }}
     - action: {{ .Action }}
       when: "{{ if .When }}{{ .When }}{{ else }}> 0{{ end }}"
+      minCompletions: {{ if .MinCompletions }}{{ .MinCompetions }}{{ else }}1{{ end }}
+      minSize: {{ if .MinSize }}{{ .MinSize }}{{ else }}1{{ end }}
+      maxSize: {{ if .MaxSize }}{{ .MaxSize }}{{ else }}null{{ end }}
+      repetitions: {{ if .Repetitions }}{{ .Repetitions }}{{ else }}1{{ end }}
       metric: {{ .Metric }}
+      backoff: {{ if .Backoff }}{{ .Backoff }}{{ else }}null{{ end }}
 {{ end }}{{ end }}
 cluster:
   max_size: {{ .Spec.Cluster.MaxSize }}

python/state_machine_operator/config/types.py

@@ -23,6 +23,10 @@ def should_trigger(self, value):
         if self.when is None:
             return True
 
+        # This ensures we check backoff, etc.
+        if not self.action.perform():
+            return False
+
         # If we have a direct value, we check for equality
         number = (int, float)
         if isinstance(self.when, number) and value != self.when:
@@ -156,6 +160,18 @@ def perform_backoff(self):
     def name(self):
         return self._action["action"]
 
+    @property
+    def min_completions(self):
+        return self._action.get("minCompletions")
+
+    @property
+    def max_size(self):
+        return self._action.get("maxSize")
+
+    @property
+    def min_size(self):
+        return self._action.get("minSize")
+
     @property
     def metric(self):
         return self._action["metric"]

python/state_machine_operator/manager/manager.py

@@ -469,18 +469,66 @@ def check_metrics(self, job):
 
                 # And check against the condition
                 if trigger.should_trigger(value):
-                    self.trigger_workflow_action(trigger.action.name)
+                    self.trigger_workflow_action(trigger, step_name, value)
 
-    def trigger_workflow_action(self, action_name):
+    def trigger_grow(self, trigger, step_name, value):
+        """
+        Trigger the job to grow
+        """
+        previous = self.workflow.jobs[step_name]["config"]["nnodes"]
+        max_size = trigger.action.max_size
+        if max_size >= previous + 1:
+            LOGGER.info(
+                f"Grow triggered: {trigger.action.metric} {trigger.when} ({value}), already >= max size {max_size}"
+            )
+            return
+
+        self.workflow.jobs[step_name]["config"]["nnodes"] += 1
+        updated = self.workflow.jobs[step_name]["config"]["nnodes"]
+        LOGGER.info(
+            f"Grow triggered: {trigger.action.metric} {trigger.when} ({value}), nodes {previous}=>{updated}"
+        )
+
+    def trigger_shrink(self, trigger, step_name, value):
+        """
+        Trigger the job to shrink, down to a min size of 1
+        """
+        previous = self.workflow.jobs[step_name]["config"]["nnodes"]
+        min_size = trigger.action.min_size or 1
+        if previous <= min_size:
+            LOGGER.info(
+                f"Shrink triggered: {trigger.action.metric} {trigger.when} ({value}), already at min size {min_size}"
+            )
+            return
+        self.workflow.jobs[step_name]["config"]["nnodes"] -= 1
+        updated = self.workflow.jobs[step_name]["config"]["nnodes"]
+        LOGGER.info(
+            f"Shrink triggered: {trigger.action.metric} {trigger.when} ({value}), nodes {previous}=>{updated}"
+        )
+
+    def trigger_workflow_action(self, trigger, step_name, value):
         """
         Given an action name, issue it for the workflow
         """
-        if action_name == "finish-workflow":
+        # This action has a minimum number of total completions
+        if trigger.action.min_completions:
+            completions = len(self.get_current_state()["completed"])
+            if completions < trigger.action.min_completions:
+                return
+
+        # Check if we have enough completions
+        if trigger.action.name == "finish-workflow":
             # TODO: add more detail in calling function to event trigger
-            LOGGER.info("Workflow completion triggered, ending workflow.")
+            LOGGER.info(
+                f"Workflow completion triggered: {trigger.action.metric} {trigger.when} ({value})"
+            )
             self.complete_workflow()
 
-        # TODO add support for grow, shrink
+        if trigger.action.name == "grow":
+            self.trigger_grow(trigger, step_name, value)
+
+        if trigger.action.name == "shrink":
+            self.trigger_shrink(trigger, step_name, value)
 
     def update_metrics(self, job):
         """
@@ -495,7 +543,7 @@ def update_metrics(self, job):
         # Add job duration if supported by tracker
         duration = job.duration()
         if job.is_completed() and duration is not None:
-            self.metrics.add_model_entry("duration", duration)
+            self.metrics.add_model_entry("duration", duration, step=job.step_name)
 
     def add_timestamp_first_seen(self, label):
         """

python/state_machine_operator/manager/metrics.py

@@ -1,6 +1,6 @@
-from river import stats
+import json
 
-import state_machine_operator.utils as utils
+from river import stats
 
 model_inits = {
     "variance": stats.Var,
@@ -28,29 +28,21 @@ def __init__(self):
         # Counters are separate
         self.models["count"] = {}
 
-        # Cache for all group model keys
-        self.keys = set()
-
     def summarize_all(self):
         """
         Summarize all models
         """
-        for key in self.keys:
-            self.summary(key)
-
-    def summary(self, key):
-        """
-        Summarize currently known models under a key
-        """
-        print(f"🌊 Streaming ML Model Summary {key}: ", end="")
+        print("🌊 Streaming ML Model Summary: ", end="")
         items = {}
-        for model_name in model_inits:
-            models = self.models[model_name]
-            if key not in models:
-                continue
-            model = models[key]
-            items[model_name] = model.get()
-        print(utils.pretty_print_list(items))
+        for model_name, models in self.models.items():
+            if model_name not in items:
+                items[model_name] = {}
+            for step_name, keys in models.items():
+                if step_name not in items[model_name]:
+                    items[model_name][step_name] = {}
+                for key, model in keys.items():
+                    items[model_name][step_name] = round(model.get(), 3)
+        print(json.dumps(items))
         return items
 
     def increment_counter(self, key, step=None):
@@ -68,20 +60,22 @@ def increment_counter(self, key, step=None):
             self.models["count"][step][key] = stats.Count()
         self.models["count"][step][key].update()
 
-    def add_model_entry(self, key, value, model_name=None):
+    def add_model_entry(self, key, value, step=None, model_name=None):
         """
         Record a datum for one or more models.
 
         If model_name is not set, add to add models.
         """
-        self.keys.add(key)
+        step = step or "global"
 
         # This should be all models except for counts
         model_names = list(model_inits)
         if model_name is not None:
             model_names = [model_name]
 
         for model_name in model_names:
-            if key not in self.models[model_name]:
-                self.models[model_name][key] = model_inits[model_name]()
-            self.models[model_name][key].update(value)
+            if step not in self.models[model_name]:
+                self.models[model_name][step] = {}
+            if key not in self.models[model_name][step]:
+                self.models[model_name][step][key] = model_inits[model_name]()
+            self.models[model_name][step][key].update(value)

python/state_machine_operator/schema.py

@@ -30,6 +30,11 @@
                             "action": {"type": "string"},
                             "when": {"type": "string"},
                             "metric": {"type": "string"},
+                            "minCompletions": {"type": ["number", "null"]},
+                            "maxSize": {"type": ["number", "null"]},
+                            "minSize": {"type": ["number", "null"]},
+                            "repetitions": {"type": ["number", "null"]},
+                            "backoff": {"type": ["number", "null"]},
                         },
                     },
                 },

python/state_machine_operator/tracker/kubernetes/tracker.py

@@ -342,6 +342,9 @@ def submit_minicluster_job(self, step, jobid):
             "maxSize": step.nodes + 100,
             "size": step.nodes,
             "tasks": step.cores_per_task,
+            "network": {
+                "headlessName": step.name,
+            },
         }
         if walltime:
             spec["deadlineSeconds"] = int(walltime)