Coverage for yaptide/celery/utils/manage_tasks.py: 42%

1import logging 

2 

3from celery import chord, group 

4from celery.result import AsyncResult 

5 

6from yaptide.celery.tasks import merge_results, run_single_simulation 

7from yaptide.celery.simulation_worker import celery_app 

8from yaptide.utils.enums import EntityState 

9 

10 

11def run_job(files_dict: dict, 

12 update_key: str, 

13 simulation_id: int, 

14 ntasks: int, 

15 celery_ids: list, 

16 sim_type: str = 'shieldhit') -> str: 

17 """Runs asynchronous simulation job""" 

18 logging.debug("Starting run_simulation task for %d tasks", ntasks) 

19 logging.debug("Simulation id: %d", simulation_id) 

20 logging.debug("Update key: %s", update_key) 

21 map_group = group([ 

22 run_single_simulation.s( 

23 files_dict=files_dict, # simulation input, keys: filenames, values: file contents 

24 task_id=i, 

25 update_key=update_key, 

26 simulation_id=simulation_id, 

27 sim_type=sim_type).set(task_id=celery_ids[i]) for i in range(ntasks) 

28 ]) 

29 

30 # By setup of simulation_worker all tasks from yaptide.celery.tasks are directed to simulations queue 

31 # For tests to work: putting signature as second task in chord requires specifying queue 

32 workflow = chord(map_group, merge_results.s().set(queue="simulations")) 

33 job: AsyncResult = workflow.delay() 

34 

35 return job.id 

36 

37 

38def get_task_status(job_id: str, state_key: str) -> dict: 

39 """Gets status of each task in the workflow""" 

40 job = AsyncResult(id=job_id, app=celery_app) 

41 job_state: str = translate_celery_state_naming(job.state) 

42 

43 # we still need to convert string to enum and operate later on Enum 

44 result = {state_key: job_state} 

45 if job_state == EntityState.FAILED.value: 

46 result["message"] = str(job.info) 

47 if "end_time" in job.info: 

48 result["end_time"] = job.info["end_time"] 

49 return result 

50 

51 

52def get_job_status(merge_id: str, celery_ids: list[str]) -> dict: 

53 """ 

54 Returns simulation state, results are not returned here 

55 Simulation may consist of multiple tasks, so we need to check all of them 

56 """ 

57 result = { 

58 "merge": get_task_status(merge_id, "job_state"), 

59 "tasks": [get_task_status(job_id, "task_state") for job_id in celery_ids] 

60 } 

61 

62 return result 

63 

64 

65def get_job_results(job_id: str) -> dict: 

66 """Returns simulation results""" 

67 job = AsyncResult(id=job_id, app=celery_app) 

68 if "result" not in job.info: 

69 return {} 

70 return job.info.get("result") 

71 

72 

73def translate_celery_state_naming(job_state: str) -> str: 

74 """Function translating celery states' names to ones used in YAPTIDE""" 

75 if job_state in ["RECEIVED", "RETRY"]: 

76 return EntityState.PENDING.value 

77 if job_state in ["PROGRESS", "STARTED"]: 

78 return EntityState.RUNNING.value 

79 if job_state in ["FAILURE"]: 

80 return EntityState.FAILED.value 

81 if job_state in ["REVOKED"]: 

82 return EntityState.CANCELED.value 

83 if job_state in ["SUCCESS"]: 

84 return EntityState.COMPLETED.value 

85 # Others are the same 

86 return job_state