feat: support for custom metrics

In this example, the user is allowed to provide a custom script
that will be used against the log, and it needs to return a dictionary
of values (the custom metrics). These are passed back to the manager
from the state machine step and can influence workflow behavior (e.g.,
stop early, grow, or shrink.

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

Author: vsoch

examples/test/lammps-custom-metrics.yaml

@@ -5,22 +5,17 @@ metadata:
 spec:
   manager:
     pullPolicy: Never
-    interactive: true
   workflow:
     completed: 10
-
     events:
-      # Custom metric - derived from parsing lammps log
+       # Custom metric - derived from parsing lammps log
+        # Max decreases down to 1 (default), with 3 breaks in between, and 3 times
       - metric: mean.lammps.lammps-walltime
-        when: ">= 12"
-        action: grow
-        # Size 4 cluster, max decreases down to 1
+        when: "<= 10"
+        action: shrink
         repetitions: 3
-        # Require checks between before doing again
-        # TODO check if this is working correctly.
-        # also check the .get()
         backoff: 3
-        maxSize: 10
+        minSize: 1
 
   cluster:
     maxSize: 2

python/state_machine_operator/config/config.py

@@ -1,6 +1,4 @@
-import importlib
 import os
-import shutil
 import sys
 
 import jsonschema
@@ -179,31 +177,8 @@ def load_jobs(self):
             job["config"]["name"] = job["name"]
 
             # Parse custom event functions on job
-            self.add_custom_events(job)
             self.jobs[job["name"]] = job
 
-    def add_custom_events(self, job):
-        """
-        Add (parse) custom job events.
-        """
-        tmpdir = utils.get_tmpdir()
-        script_path = os.path.join(tmpdir, job["name"] + ".py")
-
-        # Parse custom job functions.
-        event = job.get("events") or {}
-        script = event.get("script")
-        if not script:
-            return
-
-        utils.write_file(script, script_path)
-
-        # module will have custom functions
-        spec = importlib.util.spec_from_file_location(job["name"], script_path)
-        module = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(module)
-        job["events"]["module"] = module
-        shutil.rmtree(tmpdir)
-
     def validate(self):
         jsonschema.validate(self.cfg, schema=schema.state_machine_config_schema)
 

python/state_machine_operator/machine/machine.py

@@ -157,14 +157,15 @@ def on_change(self):
     tracker = self.trackers[self.current_state.id]
     tracker.submit_job(self.jobid)
 
+
 def metrics(self):
     """
     Yield (pop) current metrics.
     """
-    tracker = self.trackers[self.current_state.id]
-    while tracker.metrics:
-        yield tracker.metrics.pop(0)
-    
+    for step_name, step in self.trackers.items():
+        while step.metrics:
+            yield step.metrics.pop(0)
+
 
 def cleanup(self):
     """

python/state_machine_operator/manager/manager.py

@@ -553,12 +553,12 @@ def load_custom_metrics(self, state_machine):
         """
         # This takes the current step
         # {"job_name": job_name, "step_name": pod.metadata.labels["app"], "metrics": events}
-        for message in state_machine.metrics():
+        for m in state_machine.metrics():
+            print(f"Loading custom metric {m}")
             try:
-                m = json.loads(message)
                 self.metrics.add_custom_metric(m["metrics"], m["job_name"], m["step_name"])
             except Exception as e:
-                print(f"Issue parsing custom metric {message}: {e}")
+                print(f"Issue parsing custom metric {m}: {e}")
 
     def add_timestamp_first_seen(self, label):
         """

python/state_machine_operator/tracker/kubernetes/event.py

@@ -7,7 +7,6 @@
 
 from kubernetes import client, config, watch
 
-import state_machine_operator.defaults as defaults
 import state_machine_operator.utils as utils
 
 from .job import Job
@@ -29,11 +28,6 @@ def stream_events():
     for event in w.stream(batch_v1.list_namespaced_job, namespace=get_namespace()):
         job = event["object"]
         event = Job(job)
-
-        # This is a job updated with annotated metrics to inform the manager
-        if job.metadata.annotations and defaults.metrics_key in job.metadata.annotations:
-            event._is_event = True
-
         yield event
 
 
@@ -67,7 +61,7 @@ def prepare_watchers(self):
         """
         Prepare watchers for pods and nodes.
         """
-        for function in ["watch_nodes", "receive_metrics"]:
+        for function in ["watch_nodes"]:
             thread = threading.Thread(target=getattr(self, function))
             thread.daemon = True
             self.threads[function] = thread
@@ -158,6 +152,7 @@ def receive_metrics(self):
         for event in w.stream(v1.list_namespaced_event, namespace=get_namespace()):
             e = event["object"]
             if e.reason == "CustomMetric":
+                print(f"Found custom metric event {e.message}")
                 # The message has the metric, the job name, and step name
                 self.metrics.append(e.message)
 

python/state_machine_operator/tracker/kubernetes/job.py

@@ -1,5 +1,3 @@
-import os
-
 import state_machine_operator.defaults as defaults
 import state_machine_operator.utils as utils
 from state_machine_operator.tracker.job import BaseJob

python/state_machine_operator/tracker/kubernetes/tracker.py

@@ -6,7 +6,6 @@
 
 from jinja2 import Template
 from kubernetes import client, config
-from .utils import get_manager_pod
 
 import state_machine_operator.defaults as defaults
 import state_machine_operator.utils as utils
@@ -403,13 +402,49 @@ def get_metric_events(self, pod, log):
         except Exception as e:
             print(f"Error parsing custom metric for {pod.metadata.name}: {e}")
             return
-        
+
         # Metadata about the Job to associate to
         job_name = pod.metadata.labels["job-name"]
 
         # In the message we send the metrics, job name, and step name
         return {"job_name": job_name, "step_name": self.job_desc["name"], "metrics": events}
 
+    def send_kubernetes_event(self, pod, metrics):
+        """
+        Send a kubernetes event instead of metrics via the state machine.
+
+        This function is currently not used.
+        """
+        v1 = client.CoreV1Api()
+        now = datetime.datetime.utcnow()
+        job_name = pod.metadata.labels["job-name"]
+        job_uid = pod.metadata.labels["controller-uid"]
+
+        event = client.CoreV1Event(
+            metadata=client.V1ObjectMeta(
+                generate_name=job_name,
+            ),
+            involved_object=client.V1ObjectReference(
+                kind="Job",
+                api_version="batch/v1",
+                namespace=pod.metadata.namespace,
+                name=job_name,
+                uid=job_uid,
+            ),
+            reason="CustomMetric",
+            message=json.dumps(metrics),
+            first_timestamp=now.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            last_timestamp=now.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            type="Normal",
+            source=client.V1EventSource(component="state-machine"),
+        )
+
+        try:
+            v1.create_namespaced_event(pod.metadata.namespace, event)
+        except Exception as e:
+            print(f"Error creating event: {e}")
+        return metrics
+
     def cancel_jobs(self, joblist):
         """
         For the given job list, cancel each job. This is not currently use,
@@ -466,7 +501,7 @@ def save_log(self, job=None):
         Save a log identifier for a finished pod (job)
         """
         # No job, no purpose to save
-        if not Job:
+        if not job:
             return
         api = client.CoreV1Api()
 
@@ -486,7 +521,6 @@ def save_log(self, job=None):
         # For metrics, we assume logs coming from main (index 0) pod
         # This can change if needed
         for i, pod in enumerate(pods):
-            print(f"Saving log for {pod.metadata.name}")
             try:
                 log = api.read_namespaced_pod_log(
                     name=pod.metadata.name,
@@ -495,6 +529,8 @@ def save_log(self, job=None):
                     timestamps=True,
                 )
                 if i == 0:
+                    print(f"Saving log for {pod.metadata.name}")
+
                     # We assume lead pod (index 0) is of interest
                     metrics = self.adapter.get_metric_events(pod, log)
                     if metrics:

python/state_machine_operator/tracker/kubernetes/utils.py

@@ -1,6 +1,7 @@
+import os
 
 from kubernetes import client
-import os
+
 
 def get_namespace():
     """
@@ -10,7 +11,8 @@ def get_namespace():
     if os.path.exists(ns_path):
         with open(ns_path) as f:
             return f.read().strip()
-    
+
+
 def get_manager_pod():
     """
     Get the currently running manager pod.
@@ -22,4 +24,3 @@ def get_manager_pod():
     pods = v1.list_namespaced_pod(namespace=get_namespace(), label_selector=label_selector)
     assert len(pods.items) == 1
     return pods.items[0]
-

python/state_machine_operator/tracker/tracker.py

@@ -1,6 +1,8 @@
+import importlib
 import json
 import logging
 import os
+import shutil
 
 import state_machine_operator.utils as utils
 from state_machine_operator.tracker.types import SubmissionCode
@@ -16,15 +18,12 @@ class Job:
     """
 
     def __init__(self, job_desc, workflow, **kwargs):
+        # The copy is important since we modify the dict
         self.job_desc = job_desc
         self.workflow = workflow
 
-        # Remove the module from the job descr
-        # TODO check if this was deleted and only working once?
-        self.module = None
-        if "module" in self.job_desc.get("events", {}):
-            self.module = self.job_desc["events"]["module"]
-            del self.job_desc["events"]["module"]
+        # Load custom module
+        self.load_custom_events()
 
         # Allow for arbitrary extra key value arguments
         for key, value in kwargs.items():
@@ -45,6 +44,29 @@ def extra_environment(self):
             environ.append({"name": key, "value": value})
         return environ
 
+    def load_custom_events(self):
+        """
+        Add (parse) custom job events.
+        """
+        tmpdir = utils.get_tmpdir()
+        script_path = os.path.join(tmpdir, self.job_desc["name"] + ".py")
+
+        # Parse custom job functions.
+        event = self.job_desc.get("events") or {}
+        script = event.get("script")
+        self.module = None
+        if not script:
+            return
+
+        utils.write_file(script, script_path)
+
+        # module will have custom functions
+        spec = importlib.util.spec_from_file_location(self.job_desc["name"], script_path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        self.module = module
+        shutil.rmtree(tmpdir)
+
     @property
     def properties(self):
         """
@@ -80,12 +102,6 @@ def __init__(self, job_name, workflow):
         # We retrieve custom metrics from the log and deliver to the manager
         self.metrics = []
 
-        # TODO this envrionment variable has the max nodes we will allow to autoscale to
-        # We can use this later...
-        self.max_nodes_autoscale = (
-            os.environ.get("STATE_MACHINE_MAX_NODES", self.total_nodes) or self.total_nodes
-        )
-
     @property
     def total_nodes(self):
         return self.workflow.get("cluster", {}).get("max_nodes") or 1