ob-metaflow 2.15.17.1__py2.py3-none-any.whl → 2.15.21.1__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of ob-metaflow might be problematic. Click here for more details.

Files changed (39) hide show
  1. metaflow/_vendor/imghdr/__init__.py +180 -0
  2. metaflow/cli.py +12 -0
  3. metaflow/cmd/develop/stub_generator.py +19 -2
  4. metaflow/metaflow_config.py +0 -2
  5. metaflow/plugins/__init__.py +3 -0
  6. metaflow/plugins/airflow/airflow.py +6 -0
  7. metaflow/plugins/argo/argo_workflows.py +331 -297
  8. metaflow/plugins/argo/exit_hooks.py +209 -0
  9. metaflow/plugins/aws/aws_utils.py +1 -1
  10. metaflow/plugins/aws/step_functions/step_functions.py +6 -0
  11. metaflow/plugins/cards/card_cli.py +20 -1
  12. metaflow/plugins/cards/card_creator.py +24 -1
  13. metaflow/plugins/cards/card_datastore.py +8 -36
  14. metaflow/plugins/cards/card_decorator.py +57 -1
  15. metaflow/plugins/cards/card_modules/convert_to_native_type.py +5 -2
  16. metaflow/plugins/cards/card_modules/test_cards.py +16 -0
  17. metaflow/plugins/cards/metadata.py +22 -0
  18. metaflow/plugins/exit_hook/__init__.py +0 -0
  19. metaflow/plugins/exit_hook/exit_hook_decorator.py +46 -0
  20. metaflow/plugins/exit_hook/exit_hook_script.py +52 -0
  21. metaflow/plugins/pypi/conda_environment.py +8 -4
  22. metaflow/plugins/pypi/micromamba.py +9 -1
  23. metaflow/plugins/secrets/__init__.py +3 -0
  24. metaflow/plugins/secrets/secrets_decorator.py +9 -173
  25. metaflow/plugins/secrets/secrets_func.py +49 -0
  26. metaflow/plugins/secrets/secrets_spec.py +101 -0
  27. metaflow/plugins/secrets/utils.py +74 -0
  28. metaflow/runner/metaflow_runner.py +16 -1
  29. metaflow/runtime.py +45 -0
  30. metaflow/version.py +1 -1
  31. {ob_metaflow-2.15.17.1.data → ob_metaflow-2.15.21.1.data}/data/share/metaflow/devtools/Tiltfile +27 -2
  32. {ob_metaflow-2.15.17.1.dist-info → ob_metaflow-2.15.21.1.dist-info}/METADATA +2 -2
  33. {ob_metaflow-2.15.17.1.dist-info → ob_metaflow-2.15.21.1.dist-info}/RECORD +39 -30
  34. {ob_metaflow-2.15.17.1.data → ob_metaflow-2.15.21.1.data}/data/share/metaflow/devtools/Makefile +0 -0
  35. {ob_metaflow-2.15.17.1.data → ob_metaflow-2.15.21.1.data}/data/share/metaflow/devtools/pick_services.sh +0 -0
  36. {ob_metaflow-2.15.17.1.dist-info → ob_metaflow-2.15.21.1.dist-info}/WHEEL +0 -0
  37. {ob_metaflow-2.15.17.1.dist-info → ob_metaflow-2.15.21.1.dist-info}/entry_points.txt +0 -0
  38. {ob_metaflow-2.15.17.1.dist-info → ob_metaflow-2.15.21.1.dist-info}/licenses/LICENSE +0 -0
  39. {ob_metaflow-2.15.17.1.dist-info → ob_metaflow-2.15.21.1.dist-info}/top_level.txt +0 -0
@@ -66,6 +66,7 @@ from metaflow.util import (
66
66
  )
67
67
 
68
68
  from .argo_client import ArgoClient
69
+ from .exit_hooks import ExitHookHack, HttpExitHook, ContainerHook
69
70
  from metaflow.util import resolve_identity
70
71
 
71
72
 
@@ -140,6 +141,14 @@ class ArgoWorkflows(object):
140
141
  # ensure that your Argo Workflows controller doesn't restrict
141
142
  # templateReferencing.
142
143
 
144
+ # get initial configs
145
+ self.initial_configs = init_config()
146
+ for entry in ["OBP_PERIMETER", "OBP_INTEGRATIONS_URL"]:
147
+ if entry not in self.initial_configs:
148
+ raise ArgoWorkflowsException(
149
+ f"{entry} was not found in metaflow config. Please make sure to run `outerbounds configure <...>` command which can be found on the Outerbounds UI or reach out to your Outerbounds support team."
150
+ )
151
+
143
152
  self.name = name
144
153
  self.graph = graph
145
154
  self.flow = flow
@@ -796,6 +805,7 @@ class ArgoWorkflows(object):
796
805
 
797
806
  dag_annotation = {"metaflow/dag": json.dumps(graph_info)}
798
807
 
808
+ lifecycle_hooks = self._lifecycle_hooks()
799
809
  return (
800
810
  WorkflowTemplate()
801
811
  .metadata(
@@ -904,97 +914,20 @@ class ArgoWorkflows(object):
904
914
  if self.enable_error_msg_capture
905
915
  else None
906
916
  )
907
- # Set exit hook handlers if notifications are enabled
917
+ # Set lifecycle hooks if notifications are enabled
908
918
  .hooks(
909
919
  {
910
- **(
911
- {
912
- # workflow status maps to Completed
913
- "notify-slack-on-success": LifecycleHook()
914
- .expression("workflow.status == 'Succeeded'")
915
- .template("notify-slack-on-success"),
916
- }
917
- if self.notify_on_success and self.notify_slack_webhook_url
918
- else {}
919
- ),
920
- **(
921
- {
922
- # workflow status maps to Completed
923
- "notify-pager-duty-on-success": LifecycleHook()
924
- .expression("workflow.status == 'Succeeded'")
925
- .template("notify-pager-duty-on-success"),
926
- }
927
- if self.notify_on_success
928
- and self.notify_pager_duty_integration_key
929
- else {}
930
- ),
931
- **(
932
- {
933
- # workflow status maps to Completed
934
- "notify-incident-io-on-success": LifecycleHook()
935
- .expression("workflow.status == 'Succeeded'")
936
- .template("notify-incident-io-on-success"),
937
- }
938
- if self.notify_on_success
939
- and self.notify_incident_io_api_key
940
- else {}
941
- ),
942
- **(
943
- {
944
- # workflow status maps to Failed or Error
945
- "notify-slack-on-failure": LifecycleHook()
946
- .expression("workflow.status == 'Failed'")
947
- .template("notify-slack-on-error"),
948
- "notify-slack-on-error": LifecycleHook()
949
- .expression("workflow.status == 'Error'")
950
- .template("notify-slack-on-error"),
951
- }
952
- if self.notify_on_error and self.notify_slack_webhook_url
953
- else {}
954
- ),
955
- **(
956
- {
957
- # workflow status maps to Failed or Error
958
- "notify-pager-duty-on-failure": LifecycleHook()
959
- .expression("workflow.status == 'Failed'")
960
- .template("notify-pager-duty-on-error"),
961
- "notify-pager-duty-on-error": LifecycleHook()
962
- .expression("workflow.status == 'Error'")
963
- .template("notify-pager-duty-on-error"),
964
- }
965
- if self.notify_on_error
966
- and self.notify_pager_duty_integration_key
967
- else {}
968
- ),
969
- **(
970
- {
971
- # workflow status maps to Failed or Error
972
- "notify-incident-io-on-failure": LifecycleHook()
973
- .expression("workflow.status == 'Failed'")
974
- .template("notify-incident-io-on-error"),
975
- "notify-incident-io-on-error": LifecycleHook()
976
- .expression("workflow.status == 'Error'")
977
- .template("notify-incident-io-on-error"),
978
- }
979
- if self.notify_on_error and self.notify_incident_io_api_key
980
- else {}
981
- ),
982
- # Warning: terrible hack to workaround a bug in Argo Workflow
983
- # where the hooks listed above do not execute unless
984
- # there is an explicit exit hook. as and when this
985
- # bug is patched, we should remove this effectively
986
- # no-op hook.
987
- **(
988
- {"exit": LifecycleHook().template("exit-hook-hack")}
989
- if self.notify_on_error or self.notify_on_success
990
- else {}
991
- ),
920
+ lifecycle.name: lifecycle
921
+ for hook in lifecycle_hooks
922
+ for lifecycle in hook.lifecycle_hooks
992
923
  }
993
924
  )
994
925
  # Top-level DAG template(s)
995
926
  .templates(self._dag_templates())
996
927
  # Container templates
997
928
  .templates(self._container_templates())
929
+ # Lifecycle hook template(s)
930
+ .templates([hook.template for hook in lifecycle_hooks])
998
931
  # Exit hook template(s)
999
932
  .templates(self._exit_hook_templates())
1000
933
  # Sidecar templates (Daemon Containers)
@@ -1955,17 +1888,10 @@ class ArgoWorkflows(object):
1955
1888
  and k not in set(ARGO_WORKFLOWS_ENV_VARS_TO_SKIP.split(","))
1956
1889
  }
1957
1890
 
1958
- # get initial configs
1959
- initial_configs = init_config()
1960
- for entry in ["OBP_PERIMETER", "OBP_INTEGRATIONS_URL"]:
1961
- if entry not in initial_configs:
1962
- raise ArgoWorkflowsException(
1963
- f"{entry} was not found in metaflow config. Please make sure to run `outerbounds configure <...>` command which can be found on the Ourebounds UI or reach out to your Outerbounds support team."
1964
- )
1965
-
1891
+ # OBP configs
1966
1892
  additional_obp_configs = {
1967
- "OBP_PERIMETER": initial_configs["OBP_PERIMETER"],
1968
- "OBP_INTEGRATIONS_URL": initial_configs["OBP_INTEGRATIONS_URL"],
1893
+ "OBP_PERIMETER": self.initial_configs["OBP_PERIMETER"],
1894
+ "OBP_INTEGRATIONS_URL": self.initial_configs["OBP_INTEGRATIONS_URL"],
1969
1895
  }
1970
1896
 
1971
1897
  # Tmpfs variables
@@ -2359,40 +2285,190 @@ class ArgoWorkflows(object):
2359
2285
  templates.append(self._heartbeat_daemon_template())
2360
2286
  return templates
2361
2287
 
2362
- # Return exit hook templates for workflow execution notifications.
2363
- def _exit_hook_templates(self):
2364
- templates = []
2288
+ # Return lifecycle hooks for workflow execution notifications.
2289
+ def _lifecycle_hooks(self):
2290
+ hooks = []
2365
2291
  if self.notify_on_error:
2366
- templates.append(self._slack_error_template())
2367
- templates.append(self._pager_duty_alert_template())
2368
- templates.append(self._incident_io_alert_template())
2292
+ hooks.append(self._slack_error_template())
2293
+ hooks.append(self._pager_duty_alert_template())
2294
+ hooks.append(self._incident_io_alert_template())
2369
2295
  if self.notify_on_success:
2370
- templates.append(self._slack_success_template())
2371
- templates.append(self._pager_duty_change_template())
2372
- templates.append(self._incident_io_change_template())
2296
+ hooks.append(self._slack_success_template())
2297
+ hooks.append(self._pager_duty_change_template())
2298
+ hooks.append(self._incident_io_change_template())
2299
+
2300
+ exit_hook_decos = self.flow._flow_decorators.get("exit_hook", [])
2301
+
2302
+ for deco in exit_hook_decos:
2303
+ hooks.extend(self._lifecycle_hook_from_deco(deco))
2373
2304
 
2374
2305
  # Clean up None values from templates.
2375
- templates = list(filter(None, templates))
2376
-
2377
- if self.notify_on_error or self.notify_on_success:
2378
- # Warning: terrible hack to workaround a bug in Argo Workflow where the
2379
- # templates listed above do not execute unless there is an
2380
- # explicit exit hook. as and when this bug is patched, we should
2381
- # remove this effectively no-op template.
2382
- # Note: We use the Http template because changing this to an actual no-op container had the side-effect of
2383
- # leaving LifecycleHooks in a pending state even when they have finished execution.
2384
- templates.append(
2385
- Template("exit-hook-hack").http(
2386
- Http("GET")
2387
- .url(
2306
+ hooks = list(filter(None, hooks))
2307
+
2308
+ if hooks:
2309
+ hooks.append(
2310
+ ExitHookHack(
2311
+ url=(
2388
2312
  self.notify_slack_webhook_url
2389
2313
  or "https://events.pagerduty.com/v2/enqueue"
2390
2314
  )
2391
- .success_condition("true == true")
2392
2315
  )
2393
2316
  )
2317
+ return hooks
2318
+
2319
+ def _lifecycle_hook_from_deco(self, deco):
2320
+ from kubernetes import client as kubernetes_sdk
2321
+
2322
+ start_step = [step for step in self.graph if step.name == "start"][0]
2323
+ # We want to grab the base image used by the start step, as this is known to be pullable from within the cluster,
2324
+ # and it might contain the required libraries, allowing us to start up faster.
2325
+ start_kube_deco = [
2326
+ deco for deco in start_step.decorators if deco.name == "kubernetes"
2327
+ ][0]
2328
+ resources = dict(start_kube_deco.attributes)
2329
+ kube_defaults = dict(start_kube_deco.defaults)
2330
+
2331
+ # OBP Configs
2332
+ additional_obp_configs = {
2333
+ "OBP_PERIMETER": self.initial_configs["OBP_PERIMETER"],
2334
+ "OBP_INTEGRATIONS_URL": self.initial_configs["OBP_INTEGRATIONS_URL"],
2335
+ }
2336
+
2337
+ run_id_template = "argo-{{workflow.name}}"
2338
+ metaflow_version = self.environment.get_environment_info()
2339
+ metaflow_version["flow_name"] = self.graph.name
2340
+ metaflow_version["production_token"] = self.production_token
2341
+ env = {
2342
+ # These values are needed by Metaflow to set it's internal
2343
+ # state appropriately.
2344
+ "METAFLOW_CODE_URL": self.code_package_url,
2345
+ "METAFLOW_CODE_SHA": self.code_package_sha,
2346
+ "METAFLOW_CODE_DS": self.flow_datastore.TYPE,
2347
+ "METAFLOW_SERVICE_URL": SERVICE_INTERNAL_URL,
2348
+ "METAFLOW_SERVICE_HEADERS": json.dumps(SERVICE_HEADERS),
2349
+ "METAFLOW_USER": "argo-workflows",
2350
+ "METAFLOW_DEFAULT_DATASTORE": self.flow_datastore.TYPE,
2351
+ "METAFLOW_DEFAULT_METADATA": DEFAULT_METADATA,
2352
+ "METAFLOW_OWNER": self.username,
2353
+ }
2354
+ # pass on the Run pathspec for script
2355
+ env["RUN_PATHSPEC"] = f"{self.graph.name}/{run_id_template}"
2356
+
2357
+ # support Metaflow sandboxes
2358
+ env["METAFLOW_INIT_SCRIPT"] = KUBERNETES_SANDBOX_INIT_SCRIPT
2359
+
2360
+ # support fetching secrets
2361
+ env.update(additional_obp_configs)
2362
+
2363
+ env["METAFLOW_WORKFLOW_NAME"] = "{{workflow.name}}"
2364
+ env["METAFLOW_WORKFLOW_NAMESPACE"] = "{{workflow.namespace}}"
2365
+ env = {
2366
+ k: v
2367
+ for k, v in env.items()
2368
+ if v is not None
2369
+ and k not in set(ARGO_WORKFLOWS_ENV_VARS_TO_SKIP.split(","))
2370
+ }
2371
+
2372
+ def _cmd(fn_name):
2373
+ mflog_expr = export_mflog_env_vars(
2374
+ datastore_type=self.flow_datastore.TYPE,
2375
+ stdout_path="$PWD/.logs/mflog_stdout",
2376
+ stderr_path="$PWD/.logs/mflog_stderr",
2377
+ flow_name=self.flow.name,
2378
+ run_id=run_id_template,
2379
+ step_name=f"_hook_{fn_name}",
2380
+ task_id="1",
2381
+ retry_count="0",
2382
+ )
2383
+ cmds = " && ".join(
2384
+ [
2385
+ # For supporting sandboxes, ensure that a custom script is executed
2386
+ # before anything else is executed. The script is passed in as an
2387
+ # env var.
2388
+ '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"}',
2389
+ "mkdir -p $PWD/.logs",
2390
+ mflog_expr,
2391
+ ]
2392
+ + self.environment.get_package_commands(
2393
+ self.code_package_url, self.flow_datastore.TYPE
2394
+ )[:-1]
2395
+ # Replace the line 'Task in starting'
2396
+ + [f"mflog 'Lifecycle hook {fn_name} is starting.'"]
2397
+ + [
2398
+ f"python -m metaflow.plugins.exit_hook.exit_hook_script {metaflow_version['script']} {fn_name} $RUN_PATHSPEC"
2399
+ ]
2400
+ )
2401
+
2402
+ cmds = shlex.split('bash -c "%s"' % cmds)
2403
+ return cmds
2404
+
2405
+ def _container(cmds):
2406
+ return to_camelcase(
2407
+ kubernetes_sdk.V1Container(
2408
+ name="main",
2409
+ command=cmds,
2410
+ image=deco.attributes["options"].get("image", None)
2411
+ or resources["image"],
2412
+ env=[
2413
+ kubernetes_sdk.V1EnvVar(name=k, value=str(v))
2414
+ for k, v in env.items()
2415
+ ],
2416
+ env_from=[
2417
+ kubernetes_sdk.V1EnvFromSource(
2418
+ secret_ref=kubernetes_sdk.V1SecretEnvSource(
2419
+ name=str(k),
2420
+ # optional=True
2421
+ )
2422
+ )
2423
+ for k in list(
2424
+ []
2425
+ if not resources.get("secrets")
2426
+ else (
2427
+ [resources.get("secrets")]
2428
+ if isinstance(resources.get("secrets"), str)
2429
+ else resources.get("secrets")
2430
+ )
2431
+ )
2432
+ + KUBERNETES_SECRETS.split(",")
2433
+ + ARGO_WORKFLOWS_KUBERNETES_SECRETS.split(",")
2434
+ if k
2435
+ ],
2436
+ resources=kubernetes_sdk.V1ResourceRequirements(
2437
+ requests={
2438
+ "cpu": str(kube_defaults["cpu"]),
2439
+ "memory": "%sM" % str(kube_defaults["memory"]),
2440
+ }
2441
+ ),
2442
+ ).to_dict()
2443
+ )
2444
+
2445
+ # create lifecycle hooks from deco
2446
+ hooks = []
2447
+ for success_fn_name in deco.success_hooks:
2448
+ hook = ContainerHook(
2449
+ name=f"success-{success_fn_name.replace('_', '-')}",
2450
+ container=_container(cmds=_cmd(success_fn_name)),
2451
+ service_account_name=resources["service_account"],
2452
+ on_success=True,
2453
+ )
2454
+ hooks.append(hook)
2455
+
2456
+ for error_fn_name in deco.error_hooks:
2457
+ hook = ContainerHook(
2458
+ name=f"error-{error_fn_name.replace('_', '-')}",
2459
+ service_account_name=resources["service_account"],
2460
+ container=_container(cmds=_cmd(error_fn_name)),
2461
+ on_error=True,
2462
+ )
2463
+ hooks.append(hook)
2464
+
2465
+ return hooks
2466
+
2467
+ def _exit_hook_templates(self):
2468
+ templates = []
2394
2469
  if self.enable_error_msg_capture:
2395
2470
  templates.extend(self._error_msg_capture_hook_templates())
2471
+
2396
2472
  return templates
2397
2473
 
2398
2474
  def _error_msg_capture_hook_templates(self):
@@ -2541,30 +2617,30 @@ class ArgoWorkflows(object):
2541
2617
  # https://developer.pagerduty.com/docs/ZG9jOjExMDI5NTgx-send-an-alert-event
2542
2618
  if self.notify_pager_duty_integration_key is None:
2543
2619
  return None
2544
- return Template("notify-pager-duty-on-error").http(
2545
- Http("POST")
2546
- .url("https://events.pagerduty.com/v2/enqueue")
2547
- .header("Content-Type", "application/json")
2548
- .body(
2549
- json.dumps(
2550
- {
2551
- "event_action": "trigger",
2552
- "routing_key": self.notify_pager_duty_integration_key,
2553
- # "dedup_key": self.flow.name, # TODO: Do we need deduplication?
2554
- "payload": {
2555
- "source": "{{workflow.name}}",
2556
- "severity": "info",
2557
- "summary": "Metaflow run %s/argo-{{workflow.name}} failed!"
2558
- % self.flow.name,
2559
- "custom_details": {
2560
- "Flow": self.flow.name,
2561
- "Run ID": "argo-{{workflow.name}}",
2562
- },
2620
+ return HttpExitHook(
2621
+ name="notify-pager-duty-on-error",
2622
+ method="POST",
2623
+ url="https://events.pagerduty.com/v2/enqueue",
2624
+ headers={"Content-Type": "application/json"},
2625
+ body=json.dumps(
2626
+ {
2627
+ "event_action": "trigger",
2628
+ "routing_key": self.notify_pager_duty_integration_key,
2629
+ # "dedup_key": self.flow.name, # TODO: Do we need deduplication?
2630
+ "payload": {
2631
+ "source": "{{workflow.name}}",
2632
+ "severity": "info",
2633
+ "summary": "Metaflow run %s/argo-{{workflow.name}} failed!"
2634
+ % self.flow.name,
2635
+ "custom_details": {
2636
+ "Flow": self.flow.name,
2637
+ "Run ID": "argo-{{workflow.name}}",
2563
2638
  },
2564
- "links": self._pager_duty_notification_links(),
2565
- }
2566
- )
2567
- )
2639
+ },
2640
+ "links": self._pager_duty_notification_links(),
2641
+ }
2642
+ ),
2643
+ on_error=True,
2568
2644
  )
2569
2645
 
2570
2646
  def _incident_io_alert_template(self):
@@ -2575,50 +2651,52 @@ class ArgoWorkflows(object):
2575
2651
  "Creating alerts for errors requires a alert source config ID."
2576
2652
  )
2577
2653
  ui_links = self._incident_io_ui_urls_for_run()
2578
- return Template("notify-incident-io-on-error").http(
2579
- Http("POST")
2580
- .url(
2654
+ return HttpExitHook(
2655
+ name="notify-incident-io-on-error",
2656
+ method="POST",
2657
+ url=(
2581
2658
  "https://api.incident.io/v2/alert_events/http/%s"
2582
2659
  % self.incident_io_alert_source_config_id
2583
- )
2584
- .header("Content-Type", "application/json")
2585
- .header("Authorization", "Bearer %s" % self.notify_incident_io_api_key)
2586
- .body(
2587
- json.dumps(
2588
- {
2589
- "idempotency_key": "argo-{{workflow.name}}", # use run id to deduplicate alerts.
2590
- "status": "firing",
2591
- "title": "Flow %s has failed." % self.flow.name,
2592
- "description": "Metaflow run {run_pathspec} failed!{urls}".format(
2593
- run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
2594
- urls=(
2595
- "\n\nSee details for the run at:\n\n"
2596
- + "\n\n".join(ui_links)
2597
- if ui_links
2598
- else ""
2599
- ),
2600
- ),
2601
- "source_url": (
2602
- "%s/%s/%s"
2603
- % (
2604
- UI_URL.rstrip("/"),
2605
- self.flow.name,
2606
- "argo-{{workflow.name}}",
2607
- )
2608
- if UI_URL
2609
- else None
2660
+ ),
2661
+ headers={
2662
+ "Content-Type": "application/json",
2663
+ "Authorization": "Bearer %s" % self.notify_incident_io_api_key,
2664
+ },
2665
+ body=json.dumps(
2666
+ {
2667
+ "idempotency_key": "argo-{{workflow.name}}", # use run id to deduplicate alerts.
2668
+ "status": "firing",
2669
+ "title": "Flow %s has failed." % self.flow.name,
2670
+ "description": "Metaflow run {run_pathspec} failed!{urls}".format(
2671
+ run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
2672
+ urls=(
2673
+ "\n\nSee details for the run at:\n\n"
2674
+ + "\n\n".join(ui_links)
2675
+ if ui_links
2676
+ else ""
2610
2677
  ),
2611
- "metadata": {
2612
- **(self.incident_io_metadata or {}),
2613
- **{
2614
- "run_status": "failed",
2615
- "flow_name": self.flow.name,
2616
- "run_id": "argo-{{workflow.name}}",
2617
- },
2678
+ ),
2679
+ "source_url": (
2680
+ "%s/%s/%s"
2681
+ % (
2682
+ UI_URL.rstrip("/"),
2683
+ self.flow.name,
2684
+ "argo-{{workflow.name}}",
2685
+ )
2686
+ if UI_URL
2687
+ else None
2688
+ ),
2689
+ "metadata": {
2690
+ **(self.incident_io_metadata or {}),
2691
+ **{
2692
+ "run_status": "failed",
2693
+ "flow_name": self.flow.name,
2694
+ "run_id": "argo-{{workflow.name}}",
2618
2695
  },
2619
- }
2620
- )
2621
- )
2696
+ },
2697
+ }
2698
+ ),
2699
+ on_error=True,
2622
2700
  )
2623
2701
 
2624
2702
  def _incident_io_change_template(self):
@@ -2629,50 +2707,52 @@ class ArgoWorkflows(object):
2629
2707
  "Creating alerts for successes requires an alert source config ID."
2630
2708
  )
2631
2709
  ui_links = self._incident_io_ui_urls_for_run()
2632
- return Template("notify-incident-io-on-success").http(
2633
- Http("POST")
2634
- .url(
2710
+ return HttpExitHook(
2711
+ name="notify-incident-io-on-success",
2712
+ method="POST",
2713
+ url=(
2635
2714
  "https://api.incident.io/v2/alert_events/http/%s"
2636
2715
  % self.incident_io_alert_source_config_id
2637
- )
2638
- .header("Content-Type", "application/json")
2639
- .header("Authorization", "Bearer %s" % self.notify_incident_io_api_key)
2640
- .body(
2641
- json.dumps(
2642
- {
2643
- "idempotency_key": "argo-{{workflow.name}}", # use run id to deduplicate alerts.
2644
- "status": "firing",
2645
- "title": "Flow %s has succeeded." % self.flow.name,
2646
- "description": "Metaflow run {run_pathspec} succeeded!{urls}".format(
2647
- run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
2648
- urls=(
2649
- "\n\nSee details for the run at:\n\n"
2650
- + "\n\n".join(ui_links)
2651
- if ui_links
2652
- else ""
2653
- ),
2654
- ),
2655
- "source_url": (
2656
- "%s/%s/%s"
2657
- % (
2658
- UI_URL.rstrip("/"),
2659
- self.flow.name,
2660
- "argo-{{workflow.name}}",
2661
- )
2662
- if UI_URL
2663
- else None
2716
+ ),
2717
+ headers={
2718
+ "Content-Type": "application/json",
2719
+ "Authorization": "Bearer %s" % self.notify_incident_io_api_key,
2720
+ },
2721
+ body=json.dumps(
2722
+ {
2723
+ "idempotency_key": "argo-{{workflow.name}}", # use run id to deduplicate alerts.
2724
+ "status": "firing",
2725
+ "title": "Flow %s has succeeded." % self.flow.name,
2726
+ "description": "Metaflow run {run_pathspec} succeeded!{urls}".format(
2727
+ run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
2728
+ urls=(
2729
+ "\n\nSee details for the run at:\n\n"
2730
+ + "\n\n".join(ui_links)
2731
+ if ui_links
2732
+ else ""
2664
2733
  ),
2665
- "metadata": {
2666
- **(self.incident_io_metadata or {}),
2667
- **{
2668
- "run_status": "succeeded",
2669
- "flow_name": self.flow.name,
2670
- "run_id": "argo-{{workflow.name}}",
2671
- },
2734
+ ),
2735
+ "source_url": (
2736
+ "%s/%s/%s"
2737
+ % (
2738
+ UI_URL.rstrip("/"),
2739
+ self.flow.name,
2740
+ "argo-{{workflow.name}}",
2741
+ )
2742
+ if UI_URL
2743
+ else None
2744
+ ),
2745
+ "metadata": {
2746
+ **(self.incident_io_metadata or {}),
2747
+ **{
2748
+ "run_status": "succeeded",
2749
+ "flow_name": self.flow.name,
2750
+ "run_id": "argo-{{workflow.name}}",
2672
2751
  },
2673
- }
2674
- )
2675
- )
2752
+ },
2753
+ }
2754
+ ),
2755
+ on_success=True,
2676
2756
  )
2677
2757
 
2678
2758
  def _incident_io_ui_urls_for_run(self):
@@ -2697,27 +2777,27 @@ class ArgoWorkflows(object):
2697
2777
  # https://developer.pagerduty.com/docs/ZG9jOjExMDI5NTgy-send-a-change-event
2698
2778
  if self.notify_pager_duty_integration_key is None:
2699
2779
  return None
2700
- return Template("notify-pager-duty-on-success").http(
2701
- Http("POST")
2702
- .url("https://events.pagerduty.com/v2/change/enqueue")
2703
- .header("Content-Type", "application/json")
2704
- .body(
2705
- json.dumps(
2706
- {
2707
- "routing_key": self.notify_pager_duty_integration_key,
2708
- "payload": {
2709
- "summary": "Metaflow run %s/argo-{{workflow.name}} Succeeded"
2710
- % self.flow.name,
2711
- "source": "{{workflow.name}}",
2712
- "custom_details": {
2713
- "Flow": self.flow.name,
2714
- "Run ID": "argo-{{workflow.name}}",
2715
- },
2780
+ return HttpExitHook(
2781
+ name="notify-pager-duty-on-success",
2782
+ method="POST",
2783
+ url="https://events.pagerduty.com/v2/change/enqueue",
2784
+ headers={"Content-Type": "application/json"},
2785
+ body=json.dumps(
2786
+ {
2787
+ "routing_key": self.notify_pager_duty_integration_key,
2788
+ "payload": {
2789
+ "summary": "Metaflow run %s/argo-{{workflow.name}} Succeeded"
2790
+ % self.flow.name,
2791
+ "source": "{{workflow.name}}",
2792
+ "custom_details": {
2793
+ "Flow": self.flow.name,
2794
+ "Run ID": "argo-{{workflow.name}}",
2716
2795
  },
2717
- "links": self._pager_duty_notification_links(),
2718
- }
2719
- )
2720
- )
2796
+ },
2797
+ "links": self._pager_duty_notification_links(),
2798
+ }
2799
+ ),
2800
+ on_success=True,
2721
2801
  )
2722
2802
 
2723
2803
  def _pager_duty_notification_links(self):
@@ -2839,8 +2919,12 @@ class ArgoWorkflows(object):
2839
2919
  blocks = self._get_slack_blocks(message)
2840
2920
  payload = {"text": message, "blocks": blocks}
2841
2921
 
2842
- return Template("notify-slack-on-error").http(
2843
- Http("POST").url(self.notify_slack_webhook_url).body(json.dumps(payload))
2922
+ return HttpExitHook(
2923
+ name="notify-slack-on-error",
2924
+ method="POST",
2925
+ url=self.notify_slack_webhook_url,
2926
+ body=json.dumps(payload),
2927
+ on_error=True,
2844
2928
  )
2845
2929
 
2846
2930
  def _slack_success_template(self):
@@ -2855,8 +2939,12 @@ class ArgoWorkflows(object):
2855
2939
  blocks = self._get_slack_blocks(message)
2856
2940
  payload = {"text": message, "blocks": blocks}
2857
2941
 
2858
- return Template("notify-slack-on-success").http(
2859
- Http("POST").url(self.notify_slack_webhook_url).body(json.dumps(payload))
2942
+ return HttpExitHook(
2943
+ name="notify-slack-on-success",
2944
+ method="POST",
2945
+ url=self.notify_slack_webhook_url,
2946
+ body=json.dumps(payload),
2947
+ on_success=True,
2860
2948
  )
2861
2949
 
2862
2950
  def _heartbeat_daemon_template(self):
@@ -4227,57 +4315,3 @@ class TriggerParameter(object):
4227
4315
 
4228
4316
  def __str__(self):
4229
4317
  return json.dumps(self.payload, indent=4)
4230
-
4231
-
4232
- class Http(object):
4233
- # https://argoproj.github.io/argo-workflows/fields/#http
4234
-
4235
- def __init__(self, method):
4236
- tree = lambda: defaultdict(tree)
4237
- self.payload = tree()
4238
- self.payload["method"] = method
4239
- self.payload["headers"] = []
4240
-
4241
- def header(self, header, value):
4242
- self.payload["headers"].append({"name": header, "value": value})
4243
- return self
4244
-
4245
- def body(self, body):
4246
- self.payload["body"] = str(body)
4247
- return self
4248
-
4249
- def url(self, url):
4250
- self.payload["url"] = url
4251
- return self
4252
-
4253
- def success_condition(self, success_condition):
4254
- self.payload["successCondition"] = success_condition
4255
- return self
4256
-
4257
- def to_json(self):
4258
- return self.payload
4259
-
4260
- def __str__(self):
4261
- return json.dumps(self.payload, indent=4)
4262
-
4263
-
4264
- class LifecycleHook(object):
4265
- # https://argoproj.github.io/argo-workflows/fields/#lifecyclehook
4266
-
4267
- def __init__(self):
4268
- tree = lambda: defaultdict(tree)
4269
- self.payload = tree()
4270
-
4271
- def expression(self, expression):
4272
- self.payload["expression"] = str(expression)
4273
- return self
4274
-
4275
- def template(self, template):
4276
- self.payload["template"] = template
4277
- return self
4278
-
4279
- def to_json(self):
4280
- return self.payload
4281
-
4282
- def __str__(self):
4283
- return json.dumps(self.payload, indent=4)