Make argo events sensor namespace configurable

metaflow/metaflow_config.py

@@ -310,6 +310,9 @@
     "ARGO_EVENTS_INTERNAL_WEBHOOK_URL", ARGO_EVENTS_WEBHOOK_URL
 )
 ARGO_EVENTS_WEBHOOK_AUTH = from_conf("ARGO_EVENTS_WEBHOOK_AUTH", "none")
+ARGO_EVENTS_SENSOR_NAMESPACE = from_conf(
+    "ARGO_EVENTS_SENSOR_NAMESPACE", KUBERNETES_NAMESPACE
+)
 
 ARGO_WORKFLOWS_UI_URL = from_conf("ARGO_WORKFLOWS_UI_URL")
 

metaflow/plugins/argo/argo_client.py

@@ -2,6 +2,7 @@
 import os
 import sys
 
+from metaflow.metaflow_config import ARGO_EVENTS_SENSOR_NAMESPACE
 from metaflow.exception import MetaflowException
 from metaflow.plugins.kubernetes.kubernetes_client import KubernetesClient
 
@@ -349,12 +350,15 @@ def schedule_workflow_template(self, name, schedule=None, timezone=None):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def register_sensor(self, name, sensor=None):
+    def register_sensor(
+        self, name, sensor=None, sensor_namespace=ARGO_EVENTS_SENSOR_NAMESPACE
+    ):
         if sensor is None:
             sensor = {}
         # Unfortunately, Kubernetes client does not handle optimistic
         # concurrency control by itself unlike kubectl
         client = self._client.get()
+
         if not sensor:
             sensor["metadata"] = {}
 
@@ -364,7 +368,7 @@ def register_sensor(self, name, sensor=None):
             ] = client.CustomObjectsApi().get_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 name=name,
             )[
@@ -379,7 +383,7 @@ def register_sensor(self, name, sensor=None):
                     return client.CustomObjectsApi().create_namespaced_custom_object(
                         group=self._group,
                         version=self._version,
-                        namespace=self._namespace,
+                        namespace=sensor_namespace,
                         plural="sensors",
                         body=sensor,
                     )
@@ -397,7 +401,7 @@ def register_sensor(self, name, sensor=None):
             return client.CustomObjectsApi().replace_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 body=sensor,
                 name=name,
@@ -407,7 +411,7 @@ def register_sensor(self, name, sensor=None):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def delete_sensor(self, name):
+    def delete_sensor(self, name, sensor_namespace):
         """
         Issues an API call for deleting a sensor
 
@@ -419,7 +423,7 @@ def delete_sensor(self, name):
             return client.CustomObjectsApi().delete_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 name=name,
             )

metaflow/plugins/argo/argo_workflows.py

@@ -16,6 +16,7 @@
     ARGO_EVENTS_EVENT_BUS,
     ARGO_EVENTS_EVENT_SOURCE,
     ARGO_EVENTS_INTERNAL_WEBHOOK_URL,
+    ARGO_EVENTS_SENSOR_NAMESPACE,
     ARGO_EVENTS_SERVICE_ACCOUNT,
     ARGO_EVENTS_WEBHOOK_AUTH,
     ARGO_WORKFLOWS_ENV_VARS_TO_SKIP,
@@ -66,6 +67,10 @@ class ArgoWorkflowsException(MetaflowException):
     headline = "Argo Workflows error"
 
 
+class ArgoWorkflowsCleanupException(MetaflowException):
+    headline = "Argo Workflows clean up error"
+
+
 class ArgoWorkflowsSchedulingException(MetaflowException):
     headline = "Argo Workflows scheduling error"
 
@@ -175,6 +180,44 @@ def deploy(self):
         except Exception as e:
             raise ArgoWorkflowsException(str(e))
 
+    def cleanup_previous(self):
+        try:
+            client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
+            # Check for existing deployment and do cleanup
+            old_template = client.get_workflow_template(self.name)
+            if not old_template:
+                return None
+            # Clean up old sensors
+            old_has_sensor = old_template["metadata"]["annotations"].get(
+                "metaflow/has_sensor"
+            )
+            if old_has_sensor is None:
+                # This workflow was created before sensor annotations
+                # and may have a sensor in the default namespace
+                # we will delete it and it'll get recreated if need be
+                client.delete_sensor(self.name.replace(".", "-"), client._namespace)
+            elif old_has_sensor == "True":
+                # delete old sensor only if it was somewhere else, otherwise it'll get replaced
+                old_sensor_name = old_template["metadata"]["annotations"][
+                    "metaflow/sensor_name"
+                ]
+                old_sensor_namespace = old_template["metadata"]["annotations"][
+                    "metaflow/sensor_namespace"
+                ]
+                if (
+                    not self._sensor
+                    or old_sensor_namespace != ARGO_EVENTS_SENSOR_NAMESPACE
+                ):
+                    client.delete_sensor(old_sensor_name, old_sensor_namespace)
+            # Clean up old schedules/cronworkflow
+            old_has_schedule = old_template["metadata"]["annotations"].get(
+                "metaflow/has_schedule"
+            )
+            if not self._schedule and (old_has_schedule is None or old_has_schedule):
+                client.delete_cronworkflow(self.name)
+        except Exception as e:
+            raise ArgoWorkflowsCleanupException(str(e))
+
     @staticmethod
     def _sanitize(name):
         # Metaflow allows underscores in node names, which are disallowed in Argo
@@ -205,6 +248,14 @@ def list_templates(flow_name, all=False):
     def delete(name):
         client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
 
+        workflow_template = client.get_workflow_template(name)
+        sensor_name = workflow_template["metadata"]["annotations"].get(
+            "metaflow/sensor_name", name.replace(".", "-")
+        )
+        sensor_namespace = workflow_template["metadata"]["annotations"].get(
+            "metaflow/sensor_namespace", ARGO_EVENTS_SENSOR_NAMESPACE
+        )
+
         # Always try to delete the schedule. Failure in deleting the schedule should not
         # be treated as an error, due to any of the following reasons
         # - there might not have been a schedule, or it was deleted by some other means
@@ -214,7 +265,7 @@ def delete(name):
 
         # The workflow might have sensors attached to it, which consume actual resources.
         # Try to delete these as well.
-        sensor_deleted = client.delete_sensor(name)
+        sensor_deleted = client.delete_sensor(sensor_name, sensor_namespace)
 
         # After cleaning up related resources, delete the workflow in question.
         # Failure in deleting is treated as critical and will be made visible to the user
@@ -327,20 +378,21 @@ def _get_schedule(self):
 
     def schedule(self):
         try:
-            argo_client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
-            argo_client.schedule_workflow_template(
-                self.name, self._schedule, self._timezone
-            )
-            # Register sensor. Unfortunately, Argo Events Sensor names don't allow for
-            # dots (sensors run into an error) which rules out self.name :(
-            # Metaflow will overwrite any existing sensor.
-            sensor_name = self.name.replace(".", "-")
+            client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
+
+            if self._schedule:
+                client.schedule_workflow_template(
+                    self.name, self._schedule, self._timezone
+                )
             if self._sensor:
-                argo_client.register_sensor(sensor_name, self._sensor.to_json())
-            else:
-                # Since sensors occupy real resources, delete existing sensor if needed
-                # Deregister sensors that might have existed before this deployment
-                argo_client.delete_sensor(sensor_name)
+                # Register sensor. Unfortunately, Argo Events Sensor names don't allow for
+                # dots (sensors run into an error) which rules out self.name :(
+                # Metaflow will overwrite any existing sensor.
+                sensor_name = self.name.replace(".", "-")
+                # The new sensor will go into the sensor namespace specified
+                client.register_sensor(
+                    sensor_name, self._sensor.to_json(), ARGO_EVENTS_SENSOR_NAMESPACE
+                )
         except Exception as e:
             raise ArgoWorkflowsSchedulingException(str(e))
 
@@ -620,6 +672,8 @@ def _compile_workflow_template(self):
             "metaflow/owner": self.username,
             "metaflow/user": "argo-workflows",
             "metaflow/flow_name": self.flow.name,
+            "metaflow/has_schedule": str(self._schedule is not None),
+            "metaflow/has_sensor": str(bool(self.triggers)),
         }
 
         if self.parameters:
@@ -645,7 +699,9 @@ def _compile_workflow_template(self):
                             {key: trigger.get(key) for key in ["name", "type"]}
                             for trigger in self.triggers
                         ]
-                    )
+                    ),
+                    "metaflow/sensor_name": self.name.replace(".", "-"),
+                    "metaflow/sensor_namespace": ARGO_EVENTS_SENSOR_NAMESPACE,
                 }
             )
         if self.notify_on_error:
@@ -1793,7 +1849,7 @@ def _compile_sensor(self):
                 # Sensor metadata.
                 ObjectMeta()
                 .name(self.name.replace(".", "-"))
-                .namespace(KUBERNETES_NAMESPACE)
+                .namespace(ARGO_EVENTS_SENSOR_NAMESPACE)
                 .label("app.kubernetes.io/name", "metaflow-sensor")
                 .label("app.kubernetes.io/part-of", "metaflow")
                 .labels(self.kubernetes_labels)

metaflow/plugins/argo/argo_workflows_cli.py

@@ -224,6 +224,7 @@ def create(
         obj.echo_always(str(flow), err=False, no_bold=True)
         # TODO: Support echo-ing Argo Events Sensor template
     else:
+        flow.cleanup_previous()
         flow.deploy()
         obj.echo(
             "Workflow *{workflow_name}* "