Merge remote-tracking branch 'origin/main' into homa

experiments/simbricks/orchestration/exectools.py

@@ -94,7 +94,7 @@ class Component(object):
             self._read_stream(self._proc.stderr, self._consume_err)
         )
         rc = await self._proc.wait()
-        await asyncio.wait([stdout_handler, stderr_handler])
+        await asyncio.gather(stdout_handler, stderr_handler)
         await self.terminated(rc)
 
     async def send_input(self, bs: bytes, eof=False) -> None:
@@ -148,7 +148,8 @@ class Component(object):
         try:
             await asyncio.wait_for(self._proc.wait(), delay)
             return
-        except TimeoutError:
+        # before Python 3.11, asyncio.wait_for() throws asyncio.TimeoutError -_-
+        except (TimeoutError, asyncio.TimeoutError):
             print(
                 f'terminating component {self.cmd_parts[0]} '
                 f'pid {self._proc.pid}',
@@ -159,14 +160,13 @@ class Component(object):
         try:
             await asyncio.wait_for(self._proc.wait(), delay)
             return
-        except TimeoutError:
+        except (TimeoutError, asyncio.TimeoutError):
             print(
                 f'killing component {self.cmd_parts[0]} '
                 f'pid {self._proc.pid}',
                 flush=True
             )
             await self.kill()
-
         await self._proc.wait()
 
     async def started(self) -> None:
@@ -352,7 +352,7 @@ class Executor(abc.ABC):
             waiter = asyncio.create_task(self.await_file(p, *args, **kwargs))
             xs.append(waiter)
 
-        await asyncio.wait(xs)
+        await asyncio.gather(*xs)
 
 
 class LocalExecutor(Executor):

experiments/simbricks/orchestration/runners.py

@@ -128,7 +128,7 @@ class ExperimentBaseRunner(ABC):
             executor = self.sim_executor(host)
             task = asyncio.create_task(executor.send_file(path, self.verbose))
             copies.append(task)
-        await asyncio.wait(copies)
+        await asyncio.gather(*copies)
 
         # prepare all simulators in parallel
         sims = []
@@ -141,7 +141,7 @@ class ExperimentBaseRunner(ABC):
                 )
             )
             sims.append(task)
-        await asyncio.wait(sims)
+        await asyncio.gather(*sims)
 
     async def wait_for_sims(self) -> None:
         """Wait for simulators to terminate (the ones marked to wait on)."""
@@ -162,24 +162,24 @@ class ExperimentBaseRunner(ABC):
         scs = []
         for _, sc in self.running:
             scs.append(asyncio.create_task(sc.int_term_kill()))
-        await asyncio.shield(asyncio.wait(scs))
+        await asyncio.gather(*scs)
 
         # wait for all processes to terminate
         for _, sc in self.running:
-            await asyncio.shield(sc.wait())
+            await sc.wait()
 
         # remove all sockets
         scs = []
         for (executor, sock) in self.sockets:
             scs.append(asyncio.create_task(executor.rmtree(sock)))
         if scs:
-            await asyncio.shield(asyncio.wait(scs))
+            await asyncio.gather(*scs)
 
         # add all simulator components to the output
         for sim, sc in self.running:
             self.out.add_sim(sim, sc)
 
-        await asyncio.shield(self.after_cleanup())
+        await self.after_cleanup()
         return self.out
 
     async def run(self) -> ExpOutput:
@@ -197,7 +197,7 @@ class ExperimentBaseRunner(ABC):
                     sims.append(sim)
 
                 # wait for starts to complete
-                await asyncio.wait(starting)
+                await asyncio.gather(*starting)
 
                 for sim in sims:
                     ts.done(sim)
@@ -222,7 +222,8 @@ class ExperimentBaseRunner(ABC):
         while True:
             try:
                 return await asyncio.shield(terminate_collect_task)
-            except asyncio.CancelledError:
+            except asyncio.CancelledError as e:
+                print(e)
                 pass
 
 

experiments/simbricks/orchestration/runtime/local.py

@@ -207,7 +207,7 @@ class LocalParallelRuntime(Runtime):
             self._pending_jobs.add(job)
 
         # wait for all runs to finish
-        await asyncio.wait(self._pending_jobs)
+        await asyncio.gather(*self._pending_jobs)
 
     async def start(self) -> None:
         """Execute all defined runs."""
@@ -218,7 +218,7 @@ class LocalParallelRuntime(Runtime):
             for job in self._pending_jobs:
                 job.cancel()
             # wait for all runs to finish
-            await asyncio.wait(self._pending_jobs)
+            await asyncio.gather(*self._pending_jobs)
 
     def interrupt_handler(self) -> None:
         self._starter_task.cancel()