metaflow 2.15.3__py2.py3-none-any.whl → 2.15.5__py2.py3-none-any.whl

metaflow/plugins/argo/argo_client.py

@@ -256,12 +256,19 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def trigger_workflow_template(self, name, parameters={}):
+    def trigger_workflow_template(self, name, usertype, username, parameters={}):
         client = self._client.get()
         body = {
             "apiVersion": "argoproj.io/v1alpha1",
             "kind": "Workflow",
-            "metadata": {"generateName": name + "-"},
+            "metadata": {
+                "generateName": name + "-",
+                "annotations": {
+                    "metaflow/triggered_by_user": json.dumps(
+                        {"type": usertype, "name": username}
+                    )
+                },
+            },
             "spec": {
                 "workflowTemplateRef": {"name": name},
                 "arguments": {

metaflow/plugins/argo/argo_workflows.py

@@ -64,6 +64,7 @@ from metaflow.util import (
 )
 
 from .argo_client import ArgoClient
+from metaflow.util import resolve_identity
 
 
 class ArgoWorkflowsException(MetaflowException):
@@ -108,8 +109,7 @@ class ArgoWorkflows(object):
         notify_slack_webhook_url=None,
         notify_pager_duty_integration_key=None,
         notify_incident_io_api_key=None,
-        incident_io_success_severity_id=None,
-        incident_io_error_severity_id=None,
+        incident_io_alert_source_config_id=None,
         enable_heartbeat_daemon=True,
         enable_error_msg_capture=False,
     ):
@@ -160,8 +160,7 @@ class ArgoWorkflows(object):
         self.notify_slack_webhook_url = notify_slack_webhook_url
         self.notify_pager_duty_integration_key = notify_pager_duty_integration_key
         self.notify_incident_io_api_key = notify_incident_io_api_key
-        self.incident_io_success_severity_id = incident_io_success_severity_id
-        self.incident_io_error_severity_id = incident_io_error_severity_id
+        self.incident_io_alert_source_config_id = incident_io_alert_source_config_id
         self.enable_heartbeat_daemon = enable_heartbeat_daemon
         self.enable_error_msg_capture = enable_error_msg_capture
         self.parameters = self._process_parameters()
@@ -315,8 +314,16 @@ class ArgoWorkflows(object):
                     "Workflows before proceeding." % name
                 )
         try:
+            id_parts = resolve_identity().split(":")
+            parts_size = len(id_parts)
+            usertype = id_parts[0] if parts_size > 0 else "unknown"
+            username = id_parts[1] if parts_size > 1 else "unknown"
+
             return ArgoClient(namespace=KUBERNETES_NAMESPACE).trigger_workflow_template(
-                name, parameters
+                name,
+                usertype,
+                username,
+                parameters,
             )
         except Exception as e:
             raise ArgoWorkflowsException(str(e))
@@ -2505,25 +2512,49 @@ class ArgoWorkflows(object):
     def _incident_io_alert_template(self):
         if self.notify_incident_io_api_key is None:
             return None
-        if self.incident_io_error_severity_id is None:
+        if self.incident_io_alert_source_config_id is None:
             raise MetaflowException(
-                "Creating incidents for errors requires a severity id."
+                "Creating alerts for errors requires a alert source config ID."
             )
+        ui_links = self._incident_io_ui_urls_for_run()
         return Template("notify-incident-io-on-error").http(
             Http("POST")
-            .url("https://api.incident.io/v2/incidents")
+            .url(
+                "https://api.incident.io/v2/alert_events/http/%s"
+                % self.incident_io_alert_source_config_id
+            )
             .header("Content-Type", "application/json")
             .header("Authorization", "Bearer %s" % self.notify_incident_io_api_key)
             .body(
                 json.dumps(
                     {
                         "idempotency_key": "argo-{{workflow.name}}",  # use run id to deduplicate alerts.
-                        "visibility": "public",
-                        "severity_id": self.incident_io_error_severity_id,
-                        "name": "Flow %s has failed." % self.flow.name,
-                        "summary": "Metaflow run %s/argo-{{workflow.name}} failed! %s"
-                        % (self.flow.name, self._incident_io_ui_urls_for_run()),
-                        # TODO: Add support for custom field entries.
+                        "status": "firing",
+                        "title": "Flow %s has failed." % self.flow.name,
+                        "description": "Metaflow run {run_pathspec} failed!{urls}".format(
+                            run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
+                            urls=(
+                                "\n\nSee details for the run at:\n\n"
+                                + "\n\n".join(ui_links)
+                                if ui_links
+                                else ""
+                            ),
+                        ),
+                        "source_url": (
+                            "%s/%s/%s"
+                            % (
+                                UI_URL.rstrip("/"),
+                                self.flow.name,
+                                "argo-{{workflow.name}}",
+                            )
+                            if UI_URL
+                            else None
+                        ),
+                        "metadata": {
+                            "run_status": "failed",
+                            "flow_name": self.flow.name,
+                            "run_id": "argo-{{workflow.name}}",
+                        },
                     }
                 )
             )
@@ -2532,27 +2563,49 @@ class ArgoWorkflows(object):
     def _incident_io_change_template(self):
         if self.notify_incident_io_api_key is None:
             return None
-        if self.incident_io_success_severity_id is None:
+        if self.incident_io_alert_source_config_id is None:
             raise MetaflowException(
-                "Creating incidents for successes requires a severity id."
+                "Creating alerts for successes requires an alert source config ID."
             )
+        ui_links = self._incident_io_ui_urls_for_run()
         return Template("notify-incident-io-on-success").http(
             Http("POST")
-            .url("https://api.incident.io/v2/incidents")
+            .url(
+                "https://api.incident.io/v2/alert_events/http/%s"
+                % self.incident_io_alert_source_config_id
+            )
             .header("Content-Type", "application/json")
             .header("Authorization", "Bearer %s" % self.notify_incident_io_api_key)
             .body(
                 json.dumps(
                     {
                         "idempotency_key": "argo-{{workflow.name}}",  # use run id to deduplicate alerts.
-                        "visibility": "public",
-                        "severity_id": self.incident_io_success_severity_id,
-                        # TODO: Do we need to make incident type configurable for successes? otherwise they are created as 'investigating'
-                        # "incident_type_id": ""
-                        "name": "Flow %s has succeeded." % self.flow.name,
-                        "summary": "Metaflow run %s/argo-{{workflow.name}} succeeded!%s"
-                        % (self.flow.name, self._incident_io_ui_urls_for_run()),
-                        # TODO: Add support for custom field entries.
+                        "status": "firing",
+                        "title": "Flow %s has succeeded." % self.flow.name,
+                        "description": "Metaflow run {run_pathspec} succeeded!{urls}".format(
+                            run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
+                            urls=(
+                                "\n\nSee details for the run at:\n\n"
+                                + "\n\n".join(ui_links)
+                                if ui_links
+                                else ""
+                            ),
+                        ),
+                        "source_url": (
+                            "%s/%s/%s"
+                            % (
+                                UI_URL.rstrip("/"),
+                                self.flow.name,
+                                "argo-{{workflow.name}}",
+                            )
+                            if UI_URL
+                            else None
+                        ),
+                        "metadata": {
+                            "run_status": "succeeded",
+                            "flow_name": self.flow.name,
+                            "run_id": "argo-{{workflow.name}}",
+                        },
                     }
                 )
             )
@@ -2574,9 +2627,7 @@ class ArgoWorkflows(object):
                 "{{workflow.name}}",
             )
             links.append(url)
-        if links:
-            links = ["See details for the run at: ", *links]
-        return "\n\n".join(links)
+        return links
 
     def _pager_duty_change_template(self):
         # https://developer.pagerduty.com/docs/ZG9jOjExMDI5NTgy-send-a-change-event

metaflow/plugins/argo/argo_workflows_cli.py

@@ -130,6 +130,7 @@ def argo_workflows(obj, name=None):
     is_flag=True,
     default=False,
     help="Only print out JSON sent to Argo Workflows. Do not deploy anything.",
+    hidden=True,
 )
 @click.option(
     "--max-workers",
@@ -182,14 +183,9 @@ def argo_workflows(obj, name=None):
     help="Incident.io API V2 key for workflow success/failure notifications.",
 )
 @click.option(
-    "--incident-io-success-severity-id",
-    default=None,
-    help="Incident.io severity id for success alerts.",
-)
-@click.option(
-    "--incident-io-error-severity-id",
+    "--incident-io-alert-source-config-id",
     default=None,
-    help="Incident.io severity id for error alerts.",
+    help="Incident.io Alert source config ID. Example '01GW2G3V0S59R238FAHPDS1R66'",
 )
 @click.option(
     "--enable-heartbeat-daemon/--no-enable-heartbeat-daemon",
@@ -229,8 +225,7 @@ def create(
     notify_slack_webhook_url=None,
     notify_pager_duty_integration_key=None,
     notify_incident_io_api_key=None,
-    incident_io_success_severity_id=None,
-    incident_io_error_severity_id=None,
+    incident_io_alert_source_config_id=None,
     enable_heartbeat_daemon=True,
     deployer_attribute_file=None,
     enable_error_msg_capture=False,
@@ -287,8 +282,7 @@ def create(
         notify_slack_webhook_url,
         notify_pager_duty_integration_key,
         notify_incident_io_api_key,
-        incident_io_success_severity_id,
-        incident_io_error_severity_id,
+        incident_io_alert_source_config_id,
         enable_heartbeat_daemon,
         enable_error_msg_capture,
     )
@@ -464,8 +458,7 @@ def make_flow(
     notify_slack_webhook_url,
     notify_pager_duty_integration_key,
     notify_incident_io_api_key,
-    incident_io_success_severity_id,
-    incident_io_error_severity_id,
+    incident_io_alert_source_config_id,
     enable_heartbeat_daemon,
     enable_error_msg_capture,
 ):
@@ -488,19 +481,18 @@ def make_flow(
             "https://api.slack.com/messaging/webhooks to generate a webhook url.\n"
             " For notifications through PagerDuty, generate an integration key by following the instructions at "
             "https://support.pagerduty.com/docs/services-and-integrations#create-a-generic-events-api-integration\n"
-            " For notifications through Incident.io, generate an API key with a permission to create incidents."
+            " For notifications through Incident.io, generate an alert source config."
         )
 
-    if notify_incident_io_api_key:
-        if notify_on_error and incident_io_error_severity_id is None:
-            raise MetaflowException(
-                "Incident.io error notifications require a severity id. Please set one with --incident-io-error-severity-id"
-            )
+    if (
+        (notify_on_error or notify_on_success)
+        and notify_incident_io_api_key
+        and incident_io_alert_source_config_id is None
+    ):
+        raise MetaflowException(
+            "Incident.io alerts require an alert source configuration ID. Please set one with --incident-io-alert-source-config-id"
+        )
 
-        if notify_on_success and incident_io_success_severity_id is None:
-            raise MetaflowException(
-                "Incident.io success notifications require a severity id. Please set one with --incident-io-success-severity-id"
-            )
     # Attach @kubernetes and @environment decorator to the flow to
     # ensure that the related decorator hooks are invoked.
     decorators._attach_decorators(
@@ -545,8 +537,7 @@ def make_flow(
         notify_slack_webhook_url=notify_slack_webhook_url,
         notify_pager_duty_integration_key=notify_pager_duty_integration_key,
         notify_incident_io_api_key=notify_incident_io_api_key,
-        incident_io_success_severity_id=incident_io_success_severity_id,
-        incident_io_error_severity_id=incident_io_error_severity_id,
+        incident_io_alert_source_config_id=incident_io_alert_source_config_id,
         enable_heartbeat_daemon=enable_heartbeat_daemon,
         enable_error_msg_capture=enable_error_msg_capture,
     )

metaflow/plugins/argo/argo_workflows_deployer_objects.py

@@ -171,12 +171,16 @@ class ArgoWorkflowsTriggeredRun(TriggeredRun):
         command_obj.sync_wait()
         return command_obj.process.returncode == 0
 
-    def wait_for_completion(self, timeout: Optional[int] = None):
+    def wait_for_completion(
+        self, check_interval: int = 5, timeout: Optional[int] = None
+    ):
         """
         Wait for the workflow to complete or timeout.
 
         Parameters
         ----------
+        check_interval: int, default: 5
+            Frequency of checking for workflow completion, in seconds.
         timeout : int, optional, default None
             Maximum time in seconds to wait for workflow completion.
             If None, waits indefinitely.
@@ -187,7 +191,6 @@ class ArgoWorkflowsTriggeredRun(TriggeredRun):
             If the workflow does not complete within the specified timeout period.
         """
         start_time = time.time()
-        check_interval = 5
         while self.is_running:
             if timeout is not None and (time.time() - start_time) > timeout:
                 raise TimeoutError(