Skip to content

Commit

Permalink
Cancel correctly tasks from sync context in Scheduler
Browse files Browse the repository at this point in the history
When stopping the executing from ee, which runs in another thread, we need to use the correct loop
when cancelling the job tasks. Further, we just signal to cancel therefore we don't need to await
for the tasks to finish. This is handled in the Scheduler.execute - asyncio.gather.
  • Loading branch information
xjules committed Jan 11, 2024
1 parent f27f7f8 commit ce5dd4f
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions src/ert/scheduler/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@ def __init__(
}

self._events: asyncio.Queue[Any] = asyncio.Queue()
self._loop = asyncio.get_running_loop()

self._average_job_runtime: float = 0
self._completed_jobs_num: int = 0
self.completed_jobs: asyncio.Queue[int] = asyncio.Queue()
Expand All @@ -77,6 +79,12 @@ def __init__(

def kill_all_jobs(self) -> None:
self._cancelled = True
assert self._loop.is_running()
asyncio.run_coroutine_threadsafe(
self._signal_kill_all_jobs(), self._loop
).result()

async def _signal_kill_all_jobs(self) -> None:
for task in self._tasks.values():
task.cancel()

Expand Down Expand Up @@ -166,8 +174,7 @@ async def execute(
)

start.set()
for task in self._tasks.values():
await task
await asyncio.gather(*self._tasks.values())

await self.driver.finish()

Expand Down

0 comments on commit ce5dd4f

Please sign in to comment.