| 81 | |
| 82 | |
| 83 | def wait_for_task_state(env, job_id: int, task_ids: Union[int, List[int]], states: Union[str, List[str]], **kwargs): |
| 84 | if isinstance(task_ids, int): |
| 85 | task_ids = [task_ids] |
| 86 | |
| 87 | if isinstance(states, str): |
| 88 | states = [states] |
| 89 | assert len(task_ids) == len(states) |
| 90 | |
| 91 | ids = ",".join(str(t) for t in task_ids) |
| 92 | states = [s.lower() for s in states] |
| 93 | result = None |
| 94 | |
| 95 | def check(): |
| 96 | nonlocal result |
| 97 | result = env.command(["--output-mode=json", "task", "info", str(job_id), ids], as_json=True) |
| 98 | return [r["state"] for r in result] == states |
| 99 | |
| 100 | def on_timeout(): |
| 101 | return f"most recent output:\n{result}" |
| 102 | |
| 103 | wait_until(check, on_timeout=on_timeout, **kwargs) |
| 104 | |
| 105 | |
| 106 | def wait_for_pid_exit(pid: int): |