ob-metaflow 2.15.18.1__py2.py3-none-any.whl → 2.15.21.2__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 (34) hide show
  1. metaflow/_vendor/imghdr/__init__.py +180 -0
  2. metaflow/cmd/develop/stub_generator.py +19 -2
  3. metaflow/plugins/__init__.py +3 -0
  4. metaflow/plugins/airflow/airflow.py +6 -0
  5. metaflow/plugins/argo/argo_workflows.py +350 -303
  6. metaflow/plugins/argo/exit_hooks.py +209 -0
  7. metaflow/plugins/aws/aws_utils.py +1 -1
  8. metaflow/plugins/aws/step_functions/step_functions.py +6 -0
  9. metaflow/plugins/cards/card_cli.py +20 -1
  10. metaflow/plugins/cards/card_creator.py +24 -1
  11. metaflow/plugins/cards/card_decorator.py +57 -1
  12. metaflow/plugins/cards/card_modules/convert_to_native_type.py +5 -2
  13. metaflow/plugins/cards/card_modules/test_cards.py +16 -0
  14. metaflow/plugins/cards/metadata.py +22 -0
  15. metaflow/plugins/exit_hook/__init__.py +0 -0
  16. metaflow/plugins/exit_hook/exit_hook_decorator.py +46 -0
  17. metaflow/plugins/exit_hook/exit_hook_script.py +52 -0
  18. metaflow/plugins/secrets/__init__.py +3 -0
  19. metaflow/plugins/secrets/secrets_decorator.py +9 -173
  20. metaflow/plugins/secrets/secrets_func.py +49 -0
  21. metaflow/plugins/secrets/secrets_spec.py +101 -0
  22. metaflow/plugins/secrets/utils.py +74 -0
  23. metaflow/runner/metaflow_runner.py +16 -1
  24. metaflow/runtime.py +45 -0
  25. metaflow/version.py +1 -1
  26. {ob_metaflow-2.15.18.1.data → ob_metaflow-2.15.21.2.data}/data/share/metaflow/devtools/Tiltfile +27 -2
  27. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.15.21.2.dist-info}/METADATA +2 -2
  28. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.15.21.2.dist-info}/RECORD +34 -25
  29. {ob_metaflow-2.15.18.1.data → ob_metaflow-2.15.21.2.data}/data/share/metaflow/devtools/Makefile +0 -0
  30. {ob_metaflow-2.15.18.1.data → ob_metaflow-2.15.21.2.data}/data/share/metaflow/devtools/pick_services.sh +0 -0
  31. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.15.21.2.dist-info}/WHEEL +0 -0
  32. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.15.21.2.dist-info}/entry_points.txt +0 -0
  33. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.15.21.2.dist-info}/licenses/LICENSE +0 -0
  34. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.15.21.2.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
@@ -2210,8 +2136,17 @@ class ArgoWorkflows(object):
2210
2136
  .node_selectors(resources.get("node_selector"))
2211
2137
  # Set tolerations
2212
2138
  .tolerations(resources.get("tolerations"))
2213
- # Set image pull secrets
2214
- .image_pull_secrets(resources.get("image_pull_secrets"))
2139
+ # Set image pull secrets if present. We need to use pod_spec_patch due to Argo not supporting this on a template level.
2140
+ .pod_spec_patch(
2141
+ {
2142
+ "imagePullSecrets": [
2143
+ {"name": secret}
2144
+ for secret in resources["image_pull_secrets"]
2145
+ ]
2146
+ }
2147
+ if resources["image_pull_secrets"]
2148
+ else None
2149
+ )
2215
2150
  # Set container
2216
2151
  .container(
2217
2152
  # TODO: Unify the logic with kubernetes.py
@@ -2359,40 +2294,190 @@ class ArgoWorkflows(object):
2359
2294
  templates.append(self._heartbeat_daemon_template())
2360
2295
  return templates
2361
2296
 
2362
- # Return exit hook templates for workflow execution notifications.
2363
- def _exit_hook_templates(self):
2364
- templates = []
2297
+ # Return lifecycle hooks for workflow execution notifications.
2298
+ def _lifecycle_hooks(self):
2299
+ hooks = []
2365
2300
  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())
2301
+ hooks.append(self._slack_error_template())
2302
+ hooks.append(self._pager_duty_alert_template())
2303
+ hooks.append(self._incident_io_alert_template())
2369
2304
  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())
2305
+ hooks.append(self._slack_success_template())
2306
+ hooks.append(self._pager_duty_change_template())
2307
+ hooks.append(self._incident_io_change_template())
2308
+
2309
+ exit_hook_decos = self.flow._flow_decorators.get("exit_hook", [])
2310
+
2311
+ for deco in exit_hook_decos:
2312
+ hooks.extend(self._lifecycle_hook_from_deco(deco))
2373
2313
 
2374
2314
  # 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(
2315
+ hooks = list(filter(None, hooks))
2316
+
2317
+ if hooks:
2318
+ hooks.append(
2319
+ ExitHookHack(
2320
+ url=(
2388
2321
  self.notify_slack_webhook_url
2389
2322
  or "https://events.pagerduty.com/v2/enqueue"
2390
2323
  )
2391
- .success_condition("true == true")
2392
2324
  )
2393
2325
  )
2326
+ return hooks
2327
+
2328
+ def _lifecycle_hook_from_deco(self, deco):
2329
+ from kubernetes import client as kubernetes_sdk
2330
+
2331
+ start_step = [step for step in self.graph if step.name == "start"][0]
2332
+ # We want to grab the base image used by the start step, as this is known to be pullable from within the cluster,
2333
+ # and it might contain the required libraries, allowing us to start up faster.
2334
+ start_kube_deco = [
2335
+ deco for deco in start_step.decorators if deco.name == "kubernetes"
2336
+ ][0]
2337
+ resources = dict(start_kube_deco.attributes)
2338
+ kube_defaults = dict(start_kube_deco.defaults)
2339
+
2340
+ # OBP Configs
2341
+ additional_obp_configs = {
2342
+ "OBP_PERIMETER": self.initial_configs["OBP_PERIMETER"],
2343
+ "OBP_INTEGRATIONS_URL": self.initial_configs["OBP_INTEGRATIONS_URL"],
2344
+ }
2345
+
2346
+ run_id_template = "argo-{{workflow.name}}"
2347
+ metaflow_version = self.environment.get_environment_info()
2348
+ metaflow_version["flow_name"] = self.graph.name
2349
+ metaflow_version["production_token"] = self.production_token
2350
+ env = {
2351
+ # These values are needed by Metaflow to set it's internal
2352
+ # state appropriately.
2353
+ "METAFLOW_CODE_URL": self.code_package_url,
2354
+ "METAFLOW_CODE_SHA": self.code_package_sha,
2355
+ "METAFLOW_CODE_DS": self.flow_datastore.TYPE,
2356
+ "METAFLOW_SERVICE_URL": SERVICE_INTERNAL_URL,
2357
+ "METAFLOW_SERVICE_HEADERS": json.dumps(SERVICE_HEADERS),
2358
+ "METAFLOW_USER": "argo-workflows",
2359
+ "METAFLOW_DEFAULT_DATASTORE": self.flow_datastore.TYPE,
2360
+ "METAFLOW_DEFAULT_METADATA": DEFAULT_METADATA,
2361
+ "METAFLOW_OWNER": self.username,
2362
+ }
2363
+ # pass on the Run pathspec for script
2364
+ env["RUN_PATHSPEC"] = f"{self.graph.name}/{run_id_template}"
2365
+
2366
+ # support Metaflow sandboxes
2367
+ env["METAFLOW_INIT_SCRIPT"] = KUBERNETES_SANDBOX_INIT_SCRIPT
2368
+
2369
+ # support fetching secrets
2370
+ env.update(additional_obp_configs)
2371
+
2372
+ env["METAFLOW_WORKFLOW_NAME"] = "{{workflow.name}}"
2373
+ env["METAFLOW_WORKFLOW_NAMESPACE"] = "{{workflow.namespace}}"
2374
+ env = {
2375
+ k: v
2376
+ for k, v in env.items()
2377
+ if v is not None
2378
+ and k not in set(ARGO_WORKFLOWS_ENV_VARS_TO_SKIP.split(","))
2379
+ }
2380
+
2381
+ def _cmd(fn_name):
2382
+ mflog_expr = export_mflog_env_vars(
2383
+ datastore_type=self.flow_datastore.TYPE,
2384
+ stdout_path="$PWD/.logs/mflog_stdout",
2385
+ stderr_path="$PWD/.logs/mflog_stderr",
2386
+ flow_name=self.flow.name,
2387
+ run_id=run_id_template,
2388
+ step_name=f"_hook_{fn_name}",
2389
+ task_id="1",
2390
+ retry_count="0",
2391
+ )
2392
+ cmds = " && ".join(
2393
+ [
2394
+ # For supporting sandboxes, ensure that a custom script is executed
2395
+ # before anything else is executed. The script is passed in as an
2396
+ # env var.
2397
+ '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"}',
2398
+ "mkdir -p $PWD/.logs",
2399
+ mflog_expr,
2400
+ ]
2401
+ + self.environment.get_package_commands(
2402
+ self.code_package_url, self.flow_datastore.TYPE
2403
+ )[:-1]
2404
+ # Replace the line 'Task in starting'
2405
+ + [f"mflog 'Lifecycle hook {fn_name} is starting.'"]
2406
+ + [
2407
+ f"python -m metaflow.plugins.exit_hook.exit_hook_script {metaflow_version['script']} {fn_name} $RUN_PATHSPEC"
2408
+ ]
2409
+ )
2410
+
2411
+ cmds = shlex.split('bash -c "%s"' % cmds)
2412
+ return cmds
2413
+
2414
+ def _container(cmds):
2415
+ return to_camelcase(
2416
+ kubernetes_sdk.V1Container(
2417
+ name="main",
2418
+ command=cmds,
2419
+ image=deco.attributes["options"].get("image", None)
2420
+ or resources["image"],
2421
+ env=[
2422
+ kubernetes_sdk.V1EnvVar(name=k, value=str(v))
2423
+ for k, v in env.items()
2424
+ ],
2425
+ env_from=[
2426
+ kubernetes_sdk.V1EnvFromSource(
2427
+ secret_ref=kubernetes_sdk.V1SecretEnvSource(
2428
+ name=str(k),
2429
+ # optional=True
2430
+ )
2431
+ )
2432
+ for k in list(
2433
+ []
2434
+ if not resources.get("secrets")
2435
+ else (
2436
+ [resources.get("secrets")]
2437
+ if isinstance(resources.get("secrets"), str)
2438
+ else resources.get("secrets")
2439
+ )
2440
+ )
2441
+ + KUBERNETES_SECRETS.split(",")
2442
+ + ARGO_WORKFLOWS_KUBERNETES_SECRETS.split(",")
2443
+ if k
2444
+ ],
2445
+ resources=kubernetes_sdk.V1ResourceRequirements(
2446
+ requests={
2447
+ "cpu": str(kube_defaults["cpu"]),
2448
+ "memory": "%sM" % str(kube_defaults["memory"]),
2449
+ }
2450
+ ),
2451
+ ).to_dict()
2452
+ )
2453
+
2454
+ # create lifecycle hooks from deco
2455
+ hooks = []
2456
+ for success_fn_name in deco.success_hooks:
2457
+ hook = ContainerHook(
2458
+ name=f"success-{success_fn_name.replace('_', '-')}",
2459
+ container=_container(cmds=_cmd(success_fn_name)),
2460
+ service_account_name=resources["service_account"],
2461
+ on_success=True,
2462
+ )
2463
+ hooks.append(hook)
2464
+
2465
+ for error_fn_name in deco.error_hooks:
2466
+ hook = ContainerHook(
2467
+ name=f"error-{error_fn_name.replace('_', '-')}",
2468
+ service_account_name=resources["service_account"],
2469
+ container=_container(cmds=_cmd(error_fn_name)),
2470
+ on_error=True,
2471
+ )
2472
+ hooks.append(hook)
2473
+
2474
+ return hooks
2475
+
2476
+ def _exit_hook_templates(self):
2477
+ templates = []
2394
2478
  if self.enable_error_msg_capture:
2395
2479
  templates.extend(self._error_msg_capture_hook_templates())
2480
+
2396
2481
  return templates
2397
2482
 
2398
2483
  def _error_msg_capture_hook_templates(self):
@@ -2541,30 +2626,30 @@ class ArgoWorkflows(object):
2541
2626
  # https://developer.pagerduty.com/docs/ZG9jOjExMDI5NTgx-send-an-alert-event
2542
2627
  if self.notify_pager_duty_integration_key is None:
2543
2628
  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
- },
2629
+ return HttpExitHook(
2630
+ name="notify-pager-duty-on-error",
2631
+ method="POST",
2632
+ url="https://events.pagerduty.com/v2/enqueue",
2633
+ headers={"Content-Type": "application/json"},
2634
+ body=json.dumps(
2635
+ {
2636
+ "event_action": "trigger",
2637
+ "routing_key": self.notify_pager_duty_integration_key,
2638
+ # "dedup_key": self.flow.name, # TODO: Do we need deduplication?
2639
+ "payload": {
2640
+ "source": "{{workflow.name}}",
2641
+ "severity": "info",
2642
+ "summary": "Metaflow run %s/argo-{{workflow.name}} failed!"
2643
+ % self.flow.name,
2644
+ "custom_details": {
2645
+ "Flow": self.flow.name,
2646
+ "Run ID": "argo-{{workflow.name}}",
2563
2647
  },
2564
- "links": self._pager_duty_notification_links(),
2565
- }
2566
- )
2567
- )
2648
+ },
2649
+ "links": self._pager_duty_notification_links(),
2650
+ }
2651
+ ),
2652
+ on_error=True,
2568
2653
  )
2569
2654
 
2570
2655
  def _incident_io_alert_template(self):
@@ -2575,50 +2660,52 @@ class ArgoWorkflows(object):
2575
2660
  "Creating alerts for errors requires a alert source config ID."
2576
2661
  )
2577
2662
  ui_links = self._incident_io_ui_urls_for_run()
2578
- return Template("notify-incident-io-on-error").http(
2579
- Http("POST")
2580
- .url(
2663
+ return HttpExitHook(
2664
+ name="notify-incident-io-on-error",
2665
+ method="POST",
2666
+ url=(
2581
2667
  "https://api.incident.io/v2/alert_events/http/%s"
2582
2668
  % 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
2669
+ ),
2670
+ headers={
2671
+ "Content-Type": "application/json",
2672
+ "Authorization": "Bearer %s" % self.notify_incident_io_api_key,
2673
+ },
2674
+ body=json.dumps(
2675
+ {
2676
+ "idempotency_key": "argo-{{workflow.name}}", # use run id to deduplicate alerts.
2677
+ "status": "firing",
2678
+ "title": "Flow %s has failed." % self.flow.name,
2679
+ "description": "Metaflow run {run_pathspec} failed!{urls}".format(
2680
+ run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
2681
+ urls=(
2682
+ "\n\nSee details for the run at:\n\n"
2683
+ + "\n\n".join(ui_links)
2684
+ if ui_links
2685
+ else ""
2610
2686
  ),
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
- },
2687
+ ),
2688
+ "source_url": (
2689
+ "%s/%s/%s"
2690
+ % (
2691
+ UI_URL.rstrip("/"),
2692
+ self.flow.name,
2693
+ "argo-{{workflow.name}}",
2694
+ )
2695
+ if UI_URL
2696
+ else None
2697
+ ),
2698
+ "metadata": {
2699
+ **(self.incident_io_metadata or {}),
2700
+ **{
2701
+ "run_status": "failed",
2702
+ "flow_name": self.flow.name,
2703
+ "run_id": "argo-{{workflow.name}}",
2618
2704
  },
2619
- }
2620
- )
2621
- )
2705
+ },
2706
+ }
2707
+ ),
2708
+ on_error=True,
2622
2709
  )
2623
2710
 
2624
2711
  def _incident_io_change_template(self):
@@ -2629,50 +2716,52 @@ class ArgoWorkflows(object):
2629
2716
  "Creating alerts for successes requires an alert source config ID."
2630
2717
  )
2631
2718
  ui_links = self._incident_io_ui_urls_for_run()
2632
- return Template("notify-incident-io-on-success").http(
2633
- Http("POST")
2634
- .url(
2719
+ return HttpExitHook(
2720
+ name="notify-incident-io-on-success",
2721
+ method="POST",
2722
+ url=(
2635
2723
  "https://api.incident.io/v2/alert_events/http/%s"
2636
2724
  % 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
2725
+ ),
2726
+ headers={
2727
+ "Content-Type": "application/json",
2728
+ "Authorization": "Bearer %s" % self.notify_incident_io_api_key,
2729
+ },
2730
+ body=json.dumps(
2731
+ {
2732
+ "idempotency_key": "argo-{{workflow.name}}", # use run id to deduplicate alerts.
2733
+ "status": "firing",
2734
+ "title": "Flow %s has succeeded." % self.flow.name,
2735
+ "description": "Metaflow run {run_pathspec} succeeded!{urls}".format(
2736
+ run_pathspec="%s/argo-{{workflow.name}}" % self.flow.name,
2737
+ urls=(
2738
+ "\n\nSee details for the run at:\n\n"
2739
+ + "\n\n".join(ui_links)
2740
+ if ui_links
2741
+ else ""
2664
2742
  ),
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
- },
2743
+ ),
2744
+ "source_url": (
2745
+ "%s/%s/%s"
2746
+ % (
2747
+ UI_URL.rstrip("/"),
2748
+ self.flow.name,
2749
+ "argo-{{workflow.name}}",
2750
+ )
2751
+ if UI_URL
2752
+ else None
2753
+ ),
2754
+ "metadata": {
2755
+ **(self.incident_io_metadata or {}),
2756
+ **{
2757
+ "run_status": "succeeded",
2758
+ "flow_name": self.flow.name,
2759
+ "run_id": "argo-{{workflow.name}}",
2672
2760
  },
2673
- }
2674
- )
2675
- )
2761
+ },
2762
+ }
2763
+ ),
2764
+ on_success=True,
2676
2765
  )
2677
2766
 
2678
2767
  def _incident_io_ui_urls_for_run(self):
@@ -2697,27 +2786,27 @@ class ArgoWorkflows(object):
2697
2786
  # https://developer.pagerduty.com/docs/ZG9jOjExMDI5NTgy-send-a-change-event
2698
2787
  if self.notify_pager_duty_integration_key is None:
2699
2788
  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
- },
2789
+ return HttpExitHook(
2790
+ name="notify-pager-duty-on-success",
2791
+ method="POST",
2792
+ url="https://events.pagerduty.com/v2/change/enqueue",
2793
+ headers={"Content-Type": "application/json"},
2794
+ body=json.dumps(
2795
+ {
2796
+ "routing_key": self.notify_pager_duty_integration_key,
2797
+ "payload": {
2798
+ "summary": "Metaflow run %s/argo-{{workflow.name}} Succeeded"
2799
+ % self.flow.name,
2800
+ "source": "{{workflow.name}}",
2801
+ "custom_details": {
2802
+ "Flow": self.flow.name,
2803
+ "Run ID": "argo-{{workflow.name}}",
2716
2804
  },
2717
- "links": self._pager_duty_notification_links(),
2718
- }
2719
- )
2720
- )
2805
+ },
2806
+ "links": self._pager_duty_notification_links(),
2807
+ }
2808
+ ),
2809
+ on_success=True,
2721
2810
  )
2722
2811
 
2723
2812
  def _pager_duty_notification_links(self):
@@ -2839,8 +2928,12 @@ class ArgoWorkflows(object):
2839
2928
  blocks = self._get_slack_blocks(message)
2840
2929
  payload = {"text": message, "blocks": blocks}
2841
2930
 
2842
- return Template("notify-slack-on-error").http(
2843
- Http("POST").url(self.notify_slack_webhook_url).body(json.dumps(payload))
2931
+ return HttpExitHook(
2932
+ name="notify-slack-on-error",
2933
+ method="POST",
2934
+ url=self.notify_slack_webhook_url,
2935
+ body=json.dumps(payload),
2936
+ on_error=True,
2844
2937
  )
2845
2938
 
2846
2939
  def _slack_success_template(self):
@@ -2855,8 +2948,12 @@ class ArgoWorkflows(object):
2855
2948
  blocks = self._get_slack_blocks(message)
2856
2949
  payload = {"text": message, "blocks": blocks}
2857
2950
 
2858
- return Template("notify-slack-on-success").http(
2859
- Http("POST").url(self.notify_slack_webhook_url).body(json.dumps(payload))
2951
+ return HttpExitHook(
2952
+ name="notify-slack-on-success",
2953
+ method="POST",
2954
+ url=self.notify_slack_webhook_url,
2955
+ body=json.dumps(payload),
2956
+ on_success=True,
2860
2957
  )
2861
2958
 
2862
2959
  def _heartbeat_daemon_template(self):
@@ -3777,6 +3874,14 @@ class Template(object):
3777
3874
  )
3778
3875
  return self
3779
3876
 
3877
+ def pod_spec_patch(self, pod_spec_patch=None):
3878
+ if pod_spec_patch is None:
3879
+ return self
3880
+
3881
+ self.payload["podSpecPatch"] = json.dumps(pod_spec_patch)
3882
+
3883
+ return self
3884
+
3780
3885
  def node_selectors(self, node_selectors):
3781
3886
  if "nodeSelector" not in self.payload:
3782
3887
  self.payload["nodeSelector"] = {}
@@ -3788,10 +3893,6 @@ class Template(object):
3788
3893
  self.payload["tolerations"] = tolerations
3789
3894
  return self
3790
3895
 
3791
- def image_pull_secrets(self, image_pull_secrets):
3792
- self.payload["image_pull_secrets"] = image_pull_secrets
3793
- return self
3794
-
3795
3896
  def to_json(self):
3796
3897
  return self.payload
3797
3898
 
@@ -4227,57 +4328,3 @@ class TriggerParameter(object):
4227
4328
 
4228
4329
  def __str__(self):
4229
4330
  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)