dag_template.py.jinja2 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. from airflow import DAG
  2. from datetime import datetime
  3. from airflow.operators.empty import EmptyOperator
  4. from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
  5. from airflow.configuration import conf
  6. from airflow.utils.task_group import TaskGroup
  7. import requests
  8. job_id = {{ job_id }}
  9. def task_finish_alert(context):
  10. print('############### task begin callback!###################')
  11. url = 'http://{{ af_backend_uri }}/af/af_run/notification'
  12. ti = context['ti']
  13. requests.post(url, json={"data": {
  14. "job_id": job_id,
  15. "dag_id": ti.dag_id,
  16. "task_id": ti.task.task_id,
  17. "run_ts": ti.execution_date.strftime('%Y%m%dT%H%M%S'),
  18. "af_run_id": context["run_id"],
  19. "start_time": ti.start_date.timestamp(),
  20. "end_time": ti.end_date.timestamp(),
  21. "status": ti.current_state()
  22. }})
  23. def dag_begin_alert(context):
  24. print('############### dag begin callback!###################')
  25. url = 'http://{{ af_backend_uri }}/af/af_run'
  26. ti = context['ti']
  27. requests.post(url, json={"data": {
  28. "job_id": job_id,
  29. "dag_id": ti.dag_id,
  30. "task_id": ti.task.task_id,
  31. "run_ts": ti.execution_date.strftime('%Y%m%dT%H%M%S'),
  32. "af_run_id": context["run_id"],
  33. "start_time": ti.start_date.timestamp(),
  34. "end_time": ti.end_date.timestamp(),
  35. "status": ti.current_state()
  36. }})
  37. print('enter dag run!')
  38. namespace = conf.get("kubernetes", "NAMESPACE")
  39. # set the name that will be printed
  40. name = "dag_user{{ user_name }}"
  41. # instantiate the DAG
  42. with DAG(start_date=datetime(2022,6,1),catchup=False,schedule_interval=None if '{{ interval }}' == 'None' else '{{ interval }}',dag_id="{{ dag_id }}") as dag:
  43. op_start = EmptyOperator(task_id='start', on_success_callback=dag_begin_alert)
  44. {% for spark_node in spark_nodes %}
  45. with TaskGroup("{{ spark_node['name'] }}", prefix_group_id=False, tooltip="{{ spark_node['desc'] }}") as op_{{ spark_node['id'] }}:
  46. {% for spark_sub_node in spark_node['sub_nodes'] %}
  47. op_{{ spark_sub_node['id'] }} = KubernetesPodOperator(
  48. task_id="{{ spark_sub_node["id"] }}",
  49. image="{{ spark_sub_node['image'] }}",
  50. in_cluster=True,
  51. namespace=namespace,
  52. name="{{ spark_sub_node['name'] }}",
  53. random_name_suffix=True,
  54. labels={'app':'backend', 'env':'dev', 'run_ts':{{ "'{{ ts_nodash }}'" }} },
  55. reattach_on_restart=True,
  56. is_delete_operator_pod=False,
  57. get_logs=True,
  58. {% if image_pull_key != None or image_pull_key != ""%}image_pull_secrets='{{ image_pull_key }}',{% endif %}
  59. log_events_on_failure=True,
  60. image_pull_policy='Always',
  61. cmds={{ spark_sub_node['cmds'] }},
  62. env_vars={{ spark_sub_node['env'] }},
  63. on_success_callback=task_finish_alert,
  64. on_failure_callback=task_finish_alert
  65. )
  66. {% endfor %}
  67. {% for edge in spark_node['edges'] %}
  68. op_{{edge[0] }} >> op_{{edge[1] }}
  69. {% endfor %}
  70. {% endfor %}
  71. {% for node in nodes %}
  72. {{ node['operator_name'] }} = KubernetesPodOperator(
  73. task_id="{{ node["id"] }}",
  74. image="{{ node['image'] }}",
  75. in_cluster=True,
  76. namespace=namespace,
  77. name="{{ node['name'] }}",
  78. random_name_suffix=True,
  79. labels={'app':'backend', 'env':'dev', 'run_ts':{{ "'{{ ts_nodash }}'" }} },
  80. reattach_on_restart=True,
  81. is_delete_operator_pod=False,
  82. get_logs=True,
  83. image_pull_policy='Always',
  84. log_events_on_failure=True,
  85. {% if image_pull_key not in ['None',"",None]%}image_pull_secrets='{{ image_pull_key }}',{% endif %}
  86. cmds={{ node['cmds'] }},
  87. env_vars={{ node['env'] }},
  88. on_success_callback=task_finish_alert,
  89. on_failure_callback=task_finish_alert
  90. )
  91. op_start >> {{ node['operator_name'] }}
  92. {% endfor %}
  93. {% for edge in edges%}
  94. {{ edge['source_operator_name'] }} >> {{ edge['target_operator_name'] }}
  95. {% endfor %}