Try to bypass HPU hangs

benchmarks/accelerate_opt/main.py

@@ -126,16 +126,6 @@ class CustomInitProcessGroupKwargs(InitProcessGroupKwargs):
             world_size=int(os.environ["WORLD_SIZE"]),
         )
 
-        # Accelerator SUCK, it is impossible to make it use hccl
-        # We can bypass Accelerator logic by initializing the group ourselves
-        if acc.device_type == "hpu":
-            acc.init_process_group(
-                init_method=f"tcp://{MASTER_ADDR}:{MASTER_PORT}",
-                timeout=timedelta(seconds=60),
-                rank=int(os.environ["RANK"]),
-                world_size=int(os.environ["WORLD_SIZE"]),
-            )
-
         accelerator = Accelerator(kwargs_handlers=[init_process_group_kwargs])
     else:
         accelerator = Accelerator()

milabench/commands/executors.py

@@ -69,16 +69,18 @@ async def execute_command(
 
             fut = execute(pack, *argv, **{**_kwargs, **kwargs})
             coro.append(fut)
-            warden.extend(pack.processes)
-
-            if timeout:
-                delay = pack.config.get("max_duration", timeout_delay)
-                timeout_task = asyncio.create_task(force_terminate(pack, delay))
-                timeout_tasks.append(timeout_task)
-
-        results = await asyncio.gather(*coro)
+            warden.add_process(*pack.processes)
 
         if timeout:
-            for task in timeout_tasks:
-                task.cancel()
-        return results
+            delay = pack.config.get("max_duration", timeout_delay)
+
+            try:
+                async with asyncio.timeout(delay):
+                    return await asyncio.gather(*coro)
+
+            except TimeoutError:
+                await force_terminate(pack, delay)
+                return [-1 for _ in coro]
+
+        return await asyncio.gather(*coro)
+