dag1_calculate_hpc_worker.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. import json
  2. from airflow.models import DAG,Variable
  3. from airflow.operators.bash_operator import BashOperator
  4. from airflow.operators.python_operator import PythonOperator,BranchPythonOperator
  5. from airflow.contrib.operators.ssh_operator import SSHOperator
  6. from airflow.contrib.hooks.ssh_hook import SSHHook
  7. from airflow.utils.dates import days_ago
  8. from igf_airflow.check_celery_queue import fetch_queue_list_from_redis_server,airflow_utils_for_redis
  9. from igf_airflow.check_celery_queue import calculate_new_workers
  10. args = {
  11. 'owner':'airflow',
  12. 'start_date':days_ago(2),
  13. 'provide_context': True,
  14. }
  15. hpc_hook = SSHHook(ssh_conn_id='hpc_conn')
  16. dag = DAG(
  17. dag_id='dag1_calculate_hpc_worker',
  18. schedule_interval="*/15 * * * *",
  19. default_args=args,
  20. tags=['igf-lims',]
  21. )
  22. def get_new_workers(**kwargs):
  23. try:
  24. if 'ti' not in kwargs:
  25. raise ValueError('ti not present in kwargs')
  26. ti = kwargs.get('ti')
  27. active_tasks = ti.xcom_pull(task_ids='fetch_active_jobs_from_hpc')
  28. active_tasks = active_tasks.decode()
  29. active_tasks = json.loads(active_tasks)
  30. queued_tasks = ti.xcom_pull(task_ids='fetch_queue_list_from_redis')
  31. worker_to_submit,unique_queue_list = \
  32. calculate_new_workers(
  33. queue_list=queued_tasks,
  34. active_jobs_dict=active_tasks,
  35. max_workers_per_queue=3,
  36. max_total_workers=10)
  37. for key,value in worker_to_submit.items():
  38. ti.xcom_push(key=key,value=value)
  39. unique_queue_list = \
  40. [q for q in unique_queue_list if q.startswith('hpc')]
  41. return unique_queue_list
  42. except Exception as e:
  43. raise ValueError('Failed to get new workers, error: {0}'.format(e))
  44. with dag:
  45. fetch_queue_list_from_redis = \
  46. PythonOperator(
  47. task_id='fetch_queue_list_from_redis',
  48. dag=dag,
  49. python_callable=airflow_utils_for_redis,
  50. op_kwargs={"redis_conf_file":Variable.get('redis_conn_file')},
  51. queue='igf-lims'
  52. )
  53. check_hpc_queue = \
  54. SSHOperator(
  55. task_id='check_hpc_queue',
  56. ssh_hook=hpc_hook,
  57. dag=dag,
  58. command='source /etc/bashrc;qstat',
  59. queue='igf-lims'
  60. )
  61. fetch_active_jobs_from_hpc = \
  62. SSHOperator(
  63. task_id='fetch_active_jobs_from_hpc',
  64. ssh_hook=hpc_hook,
  65. dag=dag,
  66. command='source /etc/bashrc;bash /project/tgu/data2/airflow_test/github/igf-airflow-hpc/scripts/hpc/hpc_job_count_runner.sh ',
  67. do_xcom_push=True,
  68. queue='igf-lims'
  69. )
  70. calculate_new_worker_size_and_branch = \
  71. BranchPythonOperator(
  72. task_id='calculate_new_worker_size_and_branch',
  73. dag=dag,
  74. python_callable=get_new_workers,
  75. queue='igf-lims',
  76. )
  77. check_hpc_queue >> fetch_active_jobs_from_hpc
  78. calculate_new_worker_size_and_branch << [fetch_queue_list_from_redis,fetch_active_jobs_from_hpc]
  79. hpc_queue_list = Variable.get('hpc_queue_list')
  80. for q,data in hpc_queue_list.items():
  81. pbs_resource = data.get('pbs_resource')
  82. airflow_queue = data.get('airflow_queue')
  83. t = SSHOperator(
  84. task_id=q,
  85. ssh_hook=hpc_hook,
  86. dag=dag,
  87. queue='igf-lims',
  88. command="""
  89. {% if ti.xcom_pull(key=params.job_name,task_ids="calculate_new_worker_size_and_branch" ) > 1 %}
  90. source /etc/bashrc;qsub -o /dev/null -e /dev/null -k n -m n -N {{ params.job_name }} -J 1-{{ ti.xcom_pull(key=params.job_name,task_ids="calculate_new_worker_size_and_branch" ) }} {{ params.pbs_resource }} -- /project/tgu/data2/airflow_test/github/igf-airflow-hpc/scripts/hpc/airflow_worker.sh {{ params.airflow_queue }} {{ params.job_name }}
  91. {% else %}
  92. source /etc/bashrc;qsub -o /dev/null -e /dev/null -k n -m n -N {{ params.job_name }} {{ params.pbs_resource }} -- /project/tgu/data2/airflow_test/github/igf-airflow-hpc/scripts/hpc/airflow_worker.sh {{ params.airflow_queue }} {{ params.job_name }}
  93. {% endif %}
  94. """,
  95. params={'pbs_resource':pbs_resource,'airflow_queue':airflow_queue,'job_name':q}
  96. )
  97. calculate_new_worker_size_and_branch >> t