dag8_copy_ongoing_seqrun.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325
  1. from datetime import timedelta
  2. import os,json,logging
  3. from airflow.models import DAG,Variable
  4. from airflow.utils.dates import days_ago
  5. from airflow.operators.bash_operator import BashOperator
  6. from airflow.contrib.operators.ssh_operator import SSHOperator
  7. from airflow.operators.python_operator import PythonOperator,BranchPythonOperator
  8. from airflow.contrib.hooks.ssh_hook import SSHHook
  9. from airflow.operators.dummy_operator import DummyOperator
  10. from igf_airflow.seqrun.ongoing_seqrun_processing import fetch_ongoing_seqruns
  11. from igf_airflow.logging.upload_log_msg import send_log_to_channels,log_success,log_failure,log_sleep
  12. from igf_data.utils.fileutils import get_temp_dir,copy_remote_file,check_file_path,read_json_data
  13. ## DEFAULT ARGS
  14. default_args = {
  15. 'owner': 'airflow',
  16. 'depends_on_past': False,
  17. 'start_date': days_ago(2),
  18. 'email_on_failure': False,
  19. 'email_on_retry': False,
  20. 'retries': 1,
  21. 'retry_delay': timedelta(minutes=5),
  22. 'provide_context': True,
  23. }
  24. ## SSH HOOKS
  25. orwell_ssh_hook = \
  26. SSHHook(
  27. key_file=Variable.get('hpc_ssh_key_file'),
  28. username=Variable.get('hpc_user'),
  29. remote_host=Variable.get('seqrun_server'))
  30. ## DAG
  31. dag = \
  32. DAG(
  33. dag_id='dag8_copy_ongoing_seqrun',
  34. catchup=False,
  35. schedule_interval="0 */2 * * *",
  36. max_active_runs=1,
  37. tags=['hpc'],
  38. default_args=default_args,
  39. orientation='LR')
  40. ## FUNCTIONS
  41. def get_ongoing_seqrun_list(**context):
  42. """
  43. A function for fetching ongoing sequencing run ids
  44. """
  45. try:
  46. ti = context.get('ti')
  47. seqrun_server = Variable.get('seqrun_server')
  48. seqrun_base_path = Variable.get('seqrun_base_path')
  49. seqrun_server_user = Variable.get('seqrun_server_user')
  50. database_config_file = Variable.get('database_config_file')
  51. ongoing_seqruns = \
  52. fetch_ongoing_seqruns(
  53. seqrun_server=seqrun_server,
  54. seqrun_base_path=seqrun_base_path,
  55. user_name=seqrun_server_user,
  56. database_config_file=database_config_file)
  57. ti.xcom_push(key='ongoing_seqruns',value=ongoing_seqruns)
  58. branch_list = ['generate_seqrun_file_list_{0}'.format(i[0])
  59. for i in enumerate(ongoing_seqruns)]
  60. if len(branch_list) == 0:
  61. branch_list = ['no_ongoing_seqrun']
  62. else:
  63. send_log_to_channels(
  64. slack_conf=Variable.get('slack_conf'),
  65. ms_teams_conf=Variable.get('ms_teams_conf'),
  66. task_id=context['task'].task_id,
  67. dag_id=context['task'].dag_id,
  68. comment='Ongoing seqruns found: {0}'.format(ongoing_seqruns),
  69. reaction='pass')
  70. return branch_list
  71. except Exception as e:
  72. logging.error(e)
  73. send_log_to_channels(
  74. slack_conf=Variable.get('slack_conf'),
  75. ms_teams_conf=Variable.get('ms_teams_conf'),
  76. task_id=context['task'].task_id,
  77. dag_id=context['task'].dag_id,
  78. comment=e,
  79. reaction='fail')
  80. raise
  81. def copy_seqrun_manifest_file(**context):
  82. """
  83. A function for copying filesize manifest for ongoing sequencing runs to hpc
  84. """
  85. try:
  86. remote_file_path = context['params'].get('file_path')
  87. seqrun_server = Variable.get('seqrun_server')
  88. seqrun_server_user = Variable.get('seqrun_server_user')
  89. xcom_pull_task_ids = context['params'].get('xcom_pull_task_ids')
  90. ti = context.get('ti')
  91. remote_file_path = ti.xcom_pull(task_ids=xcom_pull_task_ids)
  92. if remote_file_path is not None and \
  93. not isinstance(remote_file_path,str):
  94. remote_file_path = remote_file_path.decode()
  95. tmp_work_dir = get_temp_dir(use_ephemeral_space=True)
  96. local_file_path = \
  97. os.path.join(
  98. tmp_work_dir,
  99. os.path.basename(remote_file_path))
  100. remote_address = \
  101. '{0}@{1}'.format(seqrun_server_user,seqrun_server)
  102. copy_remote_file(
  103. remote_file_path,
  104. local_file_path,
  105. source_address=remote_address)
  106. return local_file_path
  107. except Exception as e:
  108. logging.error(e)
  109. send_log_to_channels(
  110. slack_conf=Variable.get('slack_conf'),
  111. ms_teams_conf=Variable.get('ms_teams_conf'),
  112. task_id=context['task'].task_id,
  113. dag_id=context['task'].dag_id,
  114. comment=e,
  115. reaction='fail')
  116. raise
  117. def get_seqrun_chunks(**context):
  118. """
  119. A function for setting file chunk size for seqrun files copy
  120. """
  121. try:
  122. ti = context.get('ti')
  123. worker_size = context['params'].get('worker_size')
  124. child_task_prefix = context['params'].get('child_task_prefix')
  125. seqrun_chunk_size_key = context['params'].get('seqrun_chunk_size_key')
  126. xcom_pull_task_ids = context['params'].get('xcom_pull_task_ids')
  127. file_path = ti.xcom_pull(task_ids=xcom_pull_task_ids)
  128. if file_path is not None and \
  129. not isinstance(file_path,str):
  130. file_path = file_path.decode()
  131. check_file_path(file_path)
  132. file_data = read_json_data(file_path)
  133. chunk_size = None
  134. if worker_size is None or \
  135. worker_size == 0:
  136. raise ValueError(
  137. 'Incorrect worker size: {0}'.\
  138. format(worker_size))
  139. if len(file_data) == 0:
  140. raise ValueError(
  141. 'No data present in seqrun list file {0}'.\
  142. format(file_path))
  143. if len(file_data) < int(5 * worker_size):
  144. worker_size = 1 # setting worker size to 1 for low input
  145. if len(file_data) % worker_size == 0:
  146. chunk_size = int(len(file_data) / worker_size)
  147. else:
  148. chunk_size = int(len(file_data) / worker_size)+1
  149. ti.xcom_push(key=seqrun_chunk_size_key,value=chunk_size)
  150. worker_branchs = \
  151. ['{0}_{1}'.format(child_task_prefix,i)
  152. for i in range(worker_size)]
  153. return worker_branchs
  154. except Exception as e:
  155. logging.error(e)
  156. send_log_to_channels(
  157. slack_conf=Variable.get('slack_conf'),
  158. ms_teams_conf=Variable.get('ms_teams_conf'),
  159. task_id=context['task'].task_id,
  160. dag_id=context['task'].dag_id,
  161. comment=e,
  162. reaction='fail')
  163. raise
  164. def copy_seqrun_chunk(**context):
  165. """
  166. A function for copying seqrun chunks
  167. """
  168. try:
  169. ti = context.get('ti')
  170. file_path_task_ids = context['params'].get('file_path_task_ids')
  171. seqrun_chunk_size_key = context['params'].get('seqrun_chunk_size_key')
  172. seqrun_chunk_size_task_ids = context['params'].get('seqrun_chunk_size_task_ids')
  173. chunk_index_number = context['params'].get('chunk_index_number')
  174. run_index_number = context['params'].get('run_index_number')
  175. local_seqrun_path = context['params'].get('local_seqrun_path')
  176. seqrun_id_pull_key = context['params'].get('seqrun_id_pull_key')
  177. seqrun_id_pull_task_ids = context['params'].get('seqrun_id_pull_task_ids')
  178. seqrun_server = Variable.get('seqrun_server')
  179. seqrun_server_user = Variable.get('seqrun_server_user')
  180. seqrun_base_path = Variable.get('seqrun_base_path')
  181. seqrun_id = \
  182. ti.xcom_pull(key=seqrun_id_pull_key,task_ids=seqrun_id_pull_task_ids)[run_index_number]
  183. file_path = \
  184. ti.xcom_pull(task_ids=file_path_task_ids)
  185. chunk_size = \
  186. ti.xcom_pull(key=seqrun_chunk_size_key,task_ids=seqrun_chunk_size_task_ids)
  187. check_file_path(file_path)
  188. file_data = read_json_data(file_path)
  189. start_index = chunk_index_number*chunk_size
  190. finish_index = ((chunk_index_number+1)*chunk_size) - 1
  191. if finish_index > len(file_data) - 1:
  192. finish_index = len(file_data) - 1
  193. local_seqrun_path = \
  194. os.path.join(local_seqrun_path,seqrun_id)
  195. remote_address = \
  196. '{0}@{1}'.format(seqrun_server_user,seqrun_server)
  197. for entry in file_data[start_index:finish_index]:
  198. file_path = entry.get('file_path')
  199. file_size = entry.get('file_size')
  200. remote_path = \
  201. os.path.join(
  202. seqrun_base_path,
  203. file_path)
  204. local_path = \
  205. os.path.join(
  206. local_seqrun_path,
  207. file_path)
  208. if os.path.exists(local_path) and \
  209. os.path.getsize(local_path) == file_size:
  210. pass
  211. else:
  212. copy_remote_file(
  213. remote_path,
  214. local_path,
  215. source_address=remote_address,
  216. check_file=False)
  217. except Exception as e:
  218. logging.error(e)
  219. send_log_to_channels(
  220. slack_conf=Variable.get('slack_conf'),
  221. ms_teams_conf=Variable.get('ms_teams_conf'),
  222. task_id=context['task'].task_id,
  223. dag_id=context['task'].dag_id,
  224. comment=e,
  225. reaction='fail')
  226. raise
  227. with dag:
  228. ## TASK
  229. generate_seqrun_list = \
  230. BranchPythonOperator(
  231. task_id='generate_seqrun_list',
  232. dag=dag,
  233. queue='hpc_4G',
  234. python_callable=get_ongoing_seqrun_list)
  235. ## TASK
  236. no_ongoing_seqrun = \
  237. DummyOperator(
  238. task_id='no_ongoing_seqrun',
  239. dag=dag,
  240. queue='hpc_4G',
  241. on_success_callback=log_sleep)
  242. ## TASK
  243. tasks = list()
  244. for i in range(5):
  245. t1 = \
  246. SSHOperator(
  247. task_id='generate_seqrun_file_list_{0}'.format(i),
  248. dag=dag,
  249. pool='orwell_exe_pool',
  250. ssh_hook=orwell_ssh_hook,
  251. do_xcom_push=True,
  252. queue='hpc_4G',
  253. params={'source_task_id':'generate_seqrun_list',
  254. 'pull_key':'ongoing_seqruns',
  255. 'index_number':i},
  256. command="""
  257. source /home/igf/igf_code/airflow/env.sh; \
  258. python /home/igf/igf_code/airflow/data-management-python/scripts/seqrun_processing/create_file_list_for_ongoing_seqrun.py \
  259. --seqrun_base_dir /home/igf/seqrun/illumina \
  260. --output_path /home/igf/ongoing_run_tracking \
  261. --seqrun_id {{ ti.xcom_pull(key=params.pull_key,task_ids=params.source_task_id)[ params.index_number ] }}
  262. """)
  263. ## TASK
  264. t2 = \
  265. PythonOperator(
  266. task_id='copy_seqrun_file_list_{0}'.format(i),
  267. dag=dag,
  268. pool='orwell_scp_pool',
  269. queue='hpc_4G',
  270. params={'xcom_pull_task_ids':'generate_seqrun_file_list_{0}'.format(i)},
  271. python_callable=copy_seqrun_manifest_file)
  272. ## TASK
  273. t3 = \
  274. BranchPythonOperator(
  275. task_id='decide_copy_branch_{0}'.format(i),
  276. dag=dag,
  277. queue='hpc_4G',
  278. params={'xcom_pull_task_ids':'copy_seqrun_file_list_{0}'.format(i),
  279. 'worker_size':10,
  280. 'seqrun_chunk_size_key':'seqrun_chunk_size',
  281. 'child_task_prefix':'copy_file_run_{0}_chunk'.format(i)},
  282. python_callable=get_seqrun_chunks)
  283. ## TASK
  284. t4 = list()
  285. for j in range(10):
  286. t4j = \
  287. PythonOperator(
  288. task_id='copy_file_run_{0}_chunk_{1}'.format(i,j),
  289. dag=dag,
  290. queue='hpc_4G',
  291. pool='orwell_scp_pool',
  292. params={'file_path_task_ids':'copy_seqrun_file_list_{0}'.format(i),
  293. 'seqrun_chunk_size_key':'seqrun_chunk_size',
  294. 'seqrun_chunk_size_task_ids':'decide_copy_branch_{0}'.format(i),
  295. 'run_index_number':i,
  296. 'chunk_index_number':j,
  297. 'seqrun_id_pull_key':'ongoing_seqruns',
  298. 'seqrun_id_pull_task_ids':'generate_seqrun_list',
  299. 'local_seqrun_path':Variable.get('hpc_seqrun_path')},
  300. python_callable=copy_seqrun_chunk)
  301. t4.append(t4j)
  302. #tasks.append([ t1 >> t2 >> t3 >> t4 ])
  303. generate_seqrun_list >> t1 >> t2 >> t3 >> t4
  304. ## PIPELINE
  305. generate_seqrun_list >> no_ongoing_seqrun
  306. #generate_seqrun_list >> tasks