dag1_calculate_hpc_worker.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255
  1. import os, json, logging, requests
  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 requests.auth import HTTPBasicAuth
  9. from igf_airflow.celery.check_celery_queue import fetch_queue_list_from_redis_server
  10. from igf_airflow.celery.check_celery_queue import calculate_new_workers
  11. CELERY_FLOWER_BASE_URL = Variable.get('celery_flower_base_url')
  12. args = {
  13. 'owner':'airflow',
  14. 'start_date':days_ago(2),
  15. 'provide_context': True,
  16. }
  17. hpc_hook = SSHHook(ssh_conn_id='hpc_conn')
  18. dag = DAG(
  19. dag_id='dag1_calculate_hpc_worker',
  20. catchup=False,
  21. max_active_runs=1,
  22. schedule_interval="*/15 * * * *",
  23. default_args=args,
  24. tags=['igf-lims',]
  25. )
  26. def airflow_utils_for_redis(**kwargs):
  27. """
  28. A function for dag1, TO DO
  29. """
  30. try:
  31. if 'redis_conf_file' not in kwargs:
  32. raise ValueError('redis_conf_file info is not present in the kwargs')
  33. redis_conf_file = kwargs.get('redis_conf_file')
  34. json_data = dict()
  35. with open(redis_conf_file,'r') as jp:
  36. json_data = json.load(jp)
  37. if 'redis_db' not in json_data:
  38. raise ValueError('redis_db key not present in the conf file')
  39. url = json_data.get('redis_db')
  40. queue_list = fetch_queue_list_from_redis_server(url=url)
  41. return queue_list
  42. except Exception as e:
  43. logging.error('Failed to run, error:{0}'.format(e))
  44. raise
  45. def get_new_workers(**kwargs):
  46. try:
  47. if 'ti' not in kwargs:
  48. raise ValueError('ti not present in kwargs')
  49. ti = kwargs.get('ti')
  50. active_tasks = ti.xcom_pull(task_ids='fetch_active_jobs_from_hpc')
  51. active_tasks = active_tasks.decode()
  52. active_tasks = json.loads(active_tasks)
  53. queued_tasks = ti.xcom_pull(task_ids='fetch_queue_list_from_redis')
  54. worker_to_submit,unique_queue_list = \
  55. calculate_new_workers(
  56. queue_list=queued_tasks,
  57. active_jobs_dict=active_tasks,
  58. max_workers_per_queue=Variable.get('hpc_max_workers_per_queue'),
  59. max_total_workers=Variable.get('hpc_max_total_workers'))
  60. for key,value in worker_to_submit.items():
  61. ti.xcom_push(key=key,value=value)
  62. unique_queue_list = \
  63. [q for q in unique_queue_list if q.startswith('hpc')]
  64. celery_worker_key = kwargs['params'].get('celery_worker_key')
  65. base_queue = kwargs['params'].get('base_queue')
  66. empty_celery_worker_key = kwargs['params'].get('empty_celery_worker_key')
  67. celery_workers = \
  68. ti.xcom_pull(task_ids='fetch_celery_workers',key=celery_worker_key)
  69. cleanup_list = list()
  70. for flower_entry in celery_workers:
  71. active_jobs_count = flower_entry.get('active_jobs')
  72. worker_id = flower_entry.get('worker_id')
  73. queue_list = flower_entry.get('queue_lists')
  74. if base_queue not in queue_list and \
  75. len(set(queue_list).intersection(set(unique_queue_list))) == 0 and \
  76. active_jobs_count == 0:
  77. cleanup_list.\
  78. append(worker_id)
  79. if len(cleanup_list) > 0:
  80. ti.xcom_push(key=empty_celery_worker_key,value=cleanup_list)
  81. unique_queue_list.\
  82. append('cleanup_celery_workers')
  83. return unique_queue_list
  84. except Exception as e:
  85. logging.error('Failed to get new workers, error: {0}'.format(e))
  86. raise
  87. def fetch_celery_worker_list(**context):
  88. """
  89. A function for fetching list of celery workers from flower server
  90. """
  91. try:
  92. ti = context.get('ti')
  93. celery_worker_key = context['params'].get('celery_worker_key')
  94. celery_basic_auth = os.environ.get('AIRFLOW__CELERY__FLOWER_BASIC_AUTH')
  95. if celery_basic_auth is None:
  96. raise ValueError('Missing env for flower basic auth')
  97. flower_user, flower_pass = celery_basic_auth.split(':')
  98. celery_url = '{0}/api/workers'.format(CELERY_FLOWER_BASE_URL)
  99. res = requests.get(celery_url, auth=HTTPBasicAuth(flower_user, flower_pass))
  100. if res.status_code != 200:
  101. raise ValueError('Failed to fetch celery workers')
  102. data = res.content.decode()
  103. data = json.loads(data)
  104. worker_list = list()
  105. for worker_id, val in data.items():
  106. worker_list.append({
  107. 'worker_id': worker_id,
  108. 'active_jobs': len(val.get('active')),
  109. 'queue_lists': [i.get('name') for i in val.get('active_queues')]})
  110. ti.xcom_push(key=celery_worker_key,value=worker_list)
  111. except Exception as e:
  112. logging.error('Failed to get celery workers, error: {0}'.format(e))
  113. raise
  114. def stop_celery_workers(**context):
  115. """
  116. A function for stopping celery workers
  117. """
  118. try:
  119. ti = context.get('ti')
  120. empty_celery_worker_key = context['params'].get('empty_celery_worker_key')
  121. celery_basic_auth = os.environ['AIRFLOW__CELERY__FLOWER_BASIC_AUTH']
  122. flower_user, flower_pass = celery_basic_auth.split(':')
  123. celery_workers = \
  124. ti.xcom_pull(
  125. task_ids='calculate_new_worker_size_and_branch',
  126. key=empty_celery_worker_key)
  127. for worker_id in celery_workers:
  128. flower_shutdown_url = \
  129. '{0}/api/worker/shutdown/{1}'.\
  130. format(CELERY_FLOWER_BASE_URL, worker_id)
  131. res = requests.post(
  132. flower_shutdown_url,
  133. auth=HTTPBasicAuth(flower_user, flower_pass))
  134. if res.status_code != 200:
  135. raise ValueError('Failed to delete worker {0}'.\
  136. format(worker_id))
  137. except Exception as e:
  138. logging.error('Failed to stop celery workers, error: {0}'.format(e))
  139. raise
  140. with dag:
  141. ## TASK
  142. fetch_queue_list_from_redis = \
  143. PythonOperator(
  144. task_id='fetch_queue_list_from_redis',
  145. dag=dag,
  146. python_callable=airflow_utils_for_redis,
  147. op_kwargs={"redis_conf_file":Variable.get('redis_conn_file')},
  148. queue='igf-lims')
  149. ## TASK
  150. check_hpc_queue = \
  151. SSHOperator(
  152. task_id='check_hpc_queue',
  153. ssh_hook=hpc_hook,
  154. dag=dag,
  155. command='source /etc/bashrc;qstat',
  156. queue='igf-lims')
  157. ## TASK
  158. fetch_active_jobs_from_hpc = \
  159. SSHOperator(
  160. task_id='fetch_active_jobs_from_hpc',
  161. ssh_hook=hpc_hook,
  162. dag=dag,
  163. command="""
  164. source /etc/bashrc;\
  165. source /project/tgu/data2/airflow_test/secrets/hpc_env.sh;\
  166. python /project/tgu/data2/airflow_test/github/data-management-python/scripts/hpc/count_active_jobs_in_hpc.py """,
  167. do_xcom_push=True,
  168. queue='igf-lims')
  169. ## TASK
  170. fetch_celery_workers = \
  171. PythonOperator(
  172. task_id='fetch_celery_workers',
  173. dag=dag,
  174. queue='igf-lims',
  175. python_callable=fetch_celery_worker_list,
  176. params={'celery_worker_key':'celery_workers'}
  177. )
  178. ## TASK
  179. calculate_new_worker_size_and_branch = \
  180. BranchPythonOperator(
  181. task_id='calculate_new_worker_size_and_branch',
  182. dag=dag,
  183. python_callable=get_new_workers,
  184. queue='igf-lims',
  185. params={'celery_worker_key':'celery_workers',
  186. 'empty_celery_worker_key':'empty_celery_worker',
  187. 'base_queue':'igf-lims'})
  188. ## TASK
  189. queue_tasks = list()
  190. hpc_queue_list = Variable.get('hpc_queue_list')
  191. for q,data in hpc_queue_list.items():
  192. pbs_resource = data.get('pbs_resource')
  193. airflow_queue = data.get('airflow_queue')
  194. t = SSHOperator(
  195. task_id=q,
  196. ssh_hook=hpc_hook,
  197. dag=dag,
  198. queue='igf-lims',
  199. command="""
  200. {% if ti.xcom_pull(key=params.job_name,task_ids="calculate_new_worker_size_and_branch" ) > 1 %}
  201. source /etc/bashrc; \
  202. qsub \
  203. -o /dev/null \
  204. -e /dev/null \
  205. -k n -m n \
  206. -N {{ params.job_name }} \
  207. -J 1-{{ ti.xcom_pull(key=params.job_name,task_ids="calculate_new_worker_size_and_branch" ) }} {{ params.pbs_resource }} -- \
  208. /project/tgu/data2/airflow_test/github/data-management-python/scripts/hpc/airflow_worker.sh {{ params.airflow_queue }} {{ params.job_name }}
  209. {% else %}
  210. source /etc/bashrc;\
  211. qsub \
  212. -o /dev/null \
  213. -e /dev/null \
  214. -k n -m n \
  215. -N {{ params.job_name }} {{ params.pbs_resource }} -- \
  216. /project/tgu/data2/airflow_test/github/data-management-python/scripts/hpc/airflow_worker.sh {{ params.airflow_queue }} {{ params.job_name }}
  217. {% endif %}
  218. """,
  219. params={'pbs_resource':pbs_resource,
  220. 'airflow_queue':airflow_queue,
  221. 'job_name':q})
  222. queue_tasks.\
  223. append(t)
  224. ## TASK
  225. cleanup_celery_workers = \
  226. PythonOperator(
  227. task_id='cleanup_celery_workers',
  228. dag=dag,
  229. queue='igf-lims',
  230. params={'empty_celery_worker_key':'empty_celery_worker'},
  231. python_callable=stop_celery_workers)
  232. ## PIPELINE
  233. check_hpc_queue >> fetch_active_jobs_from_hpc
  234. calculate_new_worker_size_and_branch << \
  235. [fetch_queue_list_from_redis,
  236. fetch_active_jobs_from_hpc,
  237. fetch_celery_workers]
  238. calculate_new_worker_size_and_branch >> queue_tasks
  239. calculate_new_worker_size_and_branch >> cleanup_celery_workers