Introduce driver.finish that makes sure all the tasks started have be…

…en awaited.

src/ert/scheduler/driver.py

@@ -47,9 +47,5 @@ async def poll(self) -> None:
         """Poll for new job events"""
 
     @abstractmethod
-    async def finish(self, iens: int) -> None:
-        """wait until the job / realization execution is complete.
-
-        Args:
-          iens: Realization number.
-        """
+    async def finish(self) -> None:
+        """make sure that all the jobs / realizations are complete."""

src/ert/scheduler/job.py

@@ -80,8 +80,6 @@ async def _submit_and_run_once(self, sem: asyncio.BoundedSemaphore) -> None:
             while not self.returncode.done():
                 await asyncio.sleep(0.01)
             returncode = await self.returncode
-            # we need to make sure that the task has finished too
-            await self.driver.finish(self.real.iens)
 
             if (
                 returncode == 0

src/ert/scheduler/local_driver.py

@@ -13,18 +13,21 @@ def __init__(self) -> None:
         self._tasks: MutableMapping[int, asyncio.Task[None]] = {}
 
     async def submit(self, iens: int, executable: str, /, *args: str, cwd: str) -> None:
+        await self.kill(iens)
         self._tasks[iens] = asyncio.create_task(
             self._wait_until_finish(iens, executable, *args, cwd=cwd)
         )
 
     async def kill(self, iens: int) -> None:
         try:
             self._tasks[iens].cancel()
+            await self._tasks[iens]
+            del self._tasks[iens]
         except KeyError:
             return
 
-    async def finish(self, iens: int) -> None:
-        await self._tasks[iens]
+    async def finish(self) -> None:
+        await asyncio.gather(*self._tasks.values())
 
     async def _wait_until_finish(
         self, iens: int, executable: str, /, *args: str, cwd: str

src/ert/scheduler/scheduler.py

@@ -138,6 +138,8 @@ async def execute(
             for task in self._tasks.values():
                 await task
 
+            await self.driver.finish()
+
         if self._cancelled:
             return EVTYPE_ENSEMBLE_CANCELLED