Async task and event loop safety in Megatron Core

examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py

@@ -20,6 +20,10 @@
 from megatron.training.arguments import parse_args
 from megatron.core import parallel_state
 
+import logging
+
+logging.basicConfig(level=logging.INFO, force=True)
+
 async def main(
     engine: DynamicInferenceEngine,
     requests: List[Request],

megatron/core/inference/async_stream.py

@@ -9,6 +9,7 @@
 from typing import Any, AsyncGenerator, Callable, Optional, Type, Union
 
 from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.utils import get_asyncio_loop
 
 STOP_ITERATION = Exception()
 
@@ -20,12 +21,17 @@ class AsyncStream:
     Adopted from https://github.com/vllm-project/vllm/blob/eb881ed006ca458b052905e33f0d16dbb428063a/vllm/v1/engine/async_stream.py # pylint: disable=line-too-long
     """
 
-    def __init__(self, request_id: int, cancel: Callable[[str], None]) -> None:
+    def __init__(
+        self,
+        request_id: int,
+        cancel: Callable[[str], None],
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ) -> None:
         self._request_id = request_id
         self._cancel = cancel
         self._queue: asyncio.Queue = asyncio.Queue()
         self._finished = False
-        self._loop = asyncio.get_running_loop()
+        self._loop = get_asyncio_loop(loop)
 
     def put(self, item: Union[InferenceRequest, Exception]) -> None:
         """Adds a new value to the stream"""

megatron/core/inference/data_parallel_inference_coordinator.py

@@ -1,6 +1,8 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import faulthandler
 import logging
+import signal
 from collections import deque
 from itertools import cycle
 from multiprocessing import Event
@@ -23,6 +25,11 @@
 except:
     HAVE_MSGPACK = False
 
+# Register faulthandler to emit stack traces upon process kill.
+faulthandler.enable()
+faulthandler.register(signal.SIGTERM, all_threads=False, chain=True)
+faulthandler.register(signal.SIGINT, all_threads=False, chain=True)
+
 
 class DataParallelInferenceCoordinator:
     """

megatron/core/inference/engines/dynamic_engine.py

@@ -33,8 +33,8 @@
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
 )
-from megatron.core.inference.utils import Counter
-from megatron.core.utils import get_asyncio_loop
+from megatron.core.inference.utils import Counter, await_process_event
+from megatron.core.utils import get_asyncio_loop, trace_async_exceptions
 
 try:
     from tqdm import tqdm
@@ -293,7 +293,11 @@ def create_cuda_graphs(self, reset_context: bool = True):
         self.capture_stats = capture_stats
 
     async def start_listening_to_data_parallel_coordinator(
-        self, inference_coordinator_port: int, launch_inference_coordinator: bool = True
+        self,
+        inference_coordinator_port: int,
+        launch_inference_coordinator: bool = True,
+        *,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
     ):
         """Initializes ZMQ communication to connect the engine with an inference coordinator.
 
@@ -407,12 +411,14 @@ async def start_listening_to_data_parallel_coordinator(
         torch.distributed.barrier(parallel_state.get_tensor_model_parallel_group())
 
         if launch_inference_coordinator and torch.distributed.get_rank() == 0:
-            coordinator_ready_event.wait()
+            await await_process_event(coordinator_ready_event, self.inference_coordinator_process)
             logging.info("Inference co-ordinator is ready to receive requests!")
 
         # Finally run the engine infinite loop
-        self.engine_loop_task = asyncio.create_task(self.run_engine_with_coordinator())
+        loop = get_asyncio_loop(loop)
+        self.engine_loop_task = loop.create_task(self.run_engine_with_coordinator(loop=loop))
 
+    @trace_async_exceptions
     async def _notify_cond_for_new_request(self):
         """Helper function to notify condition variable when a new request is added."""
         async with self._cond:
@@ -466,7 +472,7 @@ def _add_request(
             self.waiting_request_ids.append(request_id)
 
         # Create a new asyncio Future to notify the user when the request has completed.
-        self.request_completion_futures[request_id] = asyncio.Future()
+        self.request_completion_futures[request_id] = self._loop.create_future()
         return self.request_completion_futures[request_id]
 
     def add_request(
@@ -641,7 +647,7 @@ def schedule_non_chunked_prefill(self):
             if request_can_be_added and request_tokens_can_be_added and kv_cache_available:
                 self.context.add_request(req)
                 self._loop.call_soon_threadsafe(
-                    asyncio.create_task, self._notify_cond_for_new_request()
+                    self._loop.create_task, self._notify_cond_for_new_request()
                 )
                 req.remaining_prompt_tokens = req.remaining_prompt_tokens.new_empty(0)
                 req.add_event_add()
@@ -720,7 +726,7 @@ def schedule_chunked_prefill(self):
                     self.context.chunked_prefill_request_id = -1
                     self.context.add_request(req)
                     self._loop.call_soon_threadsafe(
-                        asyncio.create_task, self._notify_cond_for_new_request()
+                        self._loop.create_task, self._notify_cond_for_new_request()
                     )
                     req.remaining_prompt_tokens = req.remaining_prompt_tokens.new_empty(0)
                     req.add_event_add()
@@ -732,7 +738,7 @@ def schedule_chunked_prefill(self):
                     chunk_length = self.context.max_tokens - self.context.active_token_count
                     self.context.add_request(req, chunk_length=chunk_length)
                     self._loop.call_soon_threadsafe(
-                        asyncio.create_task, self._notify_cond_for_new_request()
+                        self._loop.create_task, self._notify_cond_for_new_request()
                     )
                     self.context.chunked_prefill_request_id = req.request_id
                     req.remaining_prompt_tokens = req.remaining_prompt_tokens[chunk_length:]
@@ -1039,8 +1045,12 @@ def stop(self):
         self.zmq_context.term()
         parallel_state.destroy_model_parallel()
 
-    async def run_engine(self, *, verbose: Optional[bool] = False):
+    @trace_async_exceptions
+    async def run_engine(
+        self, *, loop: Optional[asyncio.AbstractEventLoop] = None, verbose: Optional[bool] = False
+    ):
         """Continually steps the engine asynchronously."""
+        self._loop = get_asyncio_loop(loop)
         try:
             while True:
                 # Wait until there are active requests before proceeding.
@@ -1054,8 +1064,12 @@ async def run_engine(self, *, verbose: Optional[bool] = False):
         except asyncio.CancelledError:
             pass
 
-    async def run_engine_with_coordinator(self, *, verbose: Optional[bool] = False):
+    @trace_async_exceptions
+    async def run_engine_with_coordinator(
+        self, *, loop: Optional[asyncio.AbstractEventLoop] = None, verbose: Optional[bool] = False
+    ):
         """Continually steps the engine asynchronously."""
+        self._loop = get_asyncio_loop(loop)
         try:
             while True:
                 self.schedule_requests()

megatron/core/inference/engines/static_engine.py

@@ -17,6 +17,7 @@
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
 )
+from megatron.core.utils import get_asyncio_loop
 
 try:
     from tqdm import tqdm
@@ -217,11 +218,6 @@ def generate_using_dynamic_engine(
             generated tokens, texts and log probs if required
         """
         assert hasattr(self, 'dynamic_engine'), "Dynamic engine not initialized"
-        try:
-            loop = asyncio.get_running_loop()
-        except RuntimeError:  # 'RuntimeError: There is no current event loop...'
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
 
         if common_inference_params:
             sampling_params = common_inference_params
@@ -385,8 +381,8 @@ def _wrapped_run_engine(self, cuda_device):
         torch.cuda.set_device(cuda_device)
         self.run_engine()
 
-    async def run_engine_async(self):
+    async def run_engine_async(self, loop: Optional[asyncio.AbstractEventLoop] = None):
         """Runs the engine asynchronously using asyncio"""
-        loop = asyncio.get_running_loop()
+        loop = get_asyncio_loop(loop)
 
         await loop.run_in_executor(None, self._wrapped_run_engine, torch.cuda.current_device())

megatron/core/inference/inference_client.py

@@ -8,6 +8,7 @@
 
 from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.utils import get_asyncio_loop, trace_async_exceptions
 
 from .headers import Headers
 
@@ -103,10 +104,11 @@ def add_request(
         payload_serialized = msgpack.packb(payload, use_bin_type=True)
         self.socket.send(payload_serialized)
         assert request_id not in self.completion_futures
-        self.completion_futures[request_id] = asyncio.get_event_loop().create_future()
+        self.completion_futures[request_id] = get_asyncio_loop().create_future()
         self.request_submission_times[request_id] = time.perf_counter()
         return self.completion_futures[request_id]
 
+    @trace_async_exceptions
     async def _listen_for_completed_requests(self):
         """
         Listens for completed inference requests from the coordinator.

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -762,10 +762,12 @@ async def async_generate_output_tokens_dynamic_batch(
 
     @torch.inference_mode()
     def generate_output_tokens_dynamic_batch(
-        self, active_sampling_map: List[Tuple[SamplingParams, List[int]]]
+        self,
+        active_sampling_map: List[Tuple[SamplingParams, List[int]]],
+        loop: Optional[asyncio.AbstractEventLoop] = None,
     ) -> Optional[Dict]:
         """Synchronous wrapper for `self.async_generate_output_tokens_dynamic_batch."""
-        loop = get_asyncio_loop()
+        loop = get_asyncio_loop(loop)
         return loop.run_until_complete(
             self.async_generate_output_tokens_dynamic_batch(active_sampling_map)
         )

megatron/core/inference/utils.py

@@ -1,5 +1,8 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
+import asyncio
+import multiprocessing
+
 import torch
 
 from megatron.core.transformer.moe.moe_layer import MoELayer
@@ -133,3 +136,28 @@ def tensor_swap(x, src_idxs, dst_idxs):
     Swap x[src_idxs] and x[dst_idxs]
     """
     x[dst_idxs], x[src_idxs] = x[src_idxs], x[dst_idxs]
+
+
+async def await_process_event(
+    event: multiprocessing.Event, process: multiprocessing.Process, timeout: float = 1.0
+) -> None:
+    """Repeatedly wait for a multiprocessing event to be set, aborting upon process failure.
+
+    Note that the timeout in this function is only for checking process liveness.
+    Its value should be set to a relatively high number. The only problem a high timeout
+    introduces is that an error is raised slighly later.
+    The timeout does not have any effect on the event-waiting, only on process failure detection.
+
+    Args:
+        event: The multiprocessing event to wait on.
+        process: The process to monitor for failure.
+        timeout: The timeout for each wait iteration in seconds.
+    """
+    while True:
+        signal = await asyncio.to_thread(event.wait, timeout)
+        if signal:
+            return
+        if not process.is_alive():
+            raise RuntimeError(
+                f"Process {process.name} (pid {process.pid}) has exited unexpectedly."
+            )

megatron/core/utils.py

@@ -17,14 +17,16 @@
 import time
 import traceback
 import warnings
+from collections import defaultdict
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from datetime import datetime
 from functools import lru_cache, reduce, wraps
 from importlib.metadata import version
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Coroutine, Dict, List, Optional, Tuple, Type, Union
 
+import numpy
 import torch
 
 from megatron.core import config
@@ -2095,3 +2097,58 @@ def get_asyncio_loop(loop: asyncio.AbstractEventLoop | None = None) -> asyncio.A
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
     return loop
+
+
+_ASYNC_TASK_STATS = defaultdict(lambda: [0, 0.0])  # cnt, total_time
+
+
+def trace_async_exceptions(
+    func: Optional[Callable[..., Coroutine]], *, verbose: bool = False
+) -> Callable[..., Coroutine]:
+    """Decorator to be applied to every coroutine that runs in a separate task.
+
+    This is needed because asyncio tasks do not propagate exceptions.
+    Coroutines running inside separate tasks will fail silently if not decorated.
+
+    Passing in `verbose=True` will print additional lifetime logging information about the task.
+    Such functionality is relied on by some users, and can be enabled as shown below:
+    ```
+        @trace_async_exceptions(verbose=True)
+        async def my_coroutine(...):
+            ...
+    ```
+    """
+
+    def _decorate(fn):
+        if not asyncio.iscoroutinefunction(fn):
+            raise TypeError("trace_async_exceptions can only be used with async functions")
+
+        @functools.wraps(fn)
+        async def wrapper(*args, **kwargs):
+            if verbose:
+                start = time.perf_counter()
+            try:
+                return await fn(*args, **kwargs)
+            except Exception as e:
+                logger.error(f"Exception in async function {fn.__name__}: {e}")
+                traceback.print_exc()
+                sys.exit(1)
+            finally:
+                if verbose:
+                    elapsed = (time.perf_counter() - start) * 1000.0
+                    name = fn.__qualname__
+                    cnt, tot = _ASYNC_TASK_STATS[name]
+                    _ASYNC_TASK_STATS[name] = [cnt + 1, tot + elapsed]
+                    avg = _ASYNC_TASK_STATS[name][1] / _ASYNC_TASK_STATS[name][0]
+
+                    log10 = numpy.log10(max(cnt, 1))
+                    if numpy.isclose(log10, round(log10)):
+                        logger.info(
+                            f"{name} completed in {elapsed:.3f} ms, "
+                            f"lifetime avg: {avg:.3f} ms, "
+                            f"lifetime cnt: {cnt + 1}"
+                        )
+
+        return wrapper
+
+    return _decorate if func is None else _decorate(func)

megatron/rl/__init__.py

@@ -12,7 +12,6 @@
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Self, Type
 
-
 def import_class(class_path: str) -> Type:
     """Import a class from a string path.
 
@@ -76,43 +75,3 @@ class Request(BaseModel):
     """Generation Request."""
 
     generation_args: GenericGenerationArgs = GenericGenerationArgs()
-
-
-from collections import defaultdict
-
-_STATS = defaultdict(lambda: [0, 0.0])  # cnt, total_time
-
-
-def trace_async_exceptions(fn: Callable[..., Coroutine]) -> Callable[..., Coroutine]:
-    """Decorator to be applied to every coroutine that runs in a separate task.
-
-    This is needed because asyncio tasks do not propagate exceptions.
-    Coroutines running inside separate tasks will fail silently if not decorated.
-    """
-    if not asyncio.iscoroutinefunction(fn):
-        raise TypeError("trace_async_exceptions can only be used with async functions")
-
-    @functools.wraps(fn)
-    async def wrapper(*args, **kwargs):
-        start = time.perf_counter()
-        try:
-            return await fn(*args, **kwargs)
-        except Exception as e:
-            print(f"Exception in async function {fn.__name__}: {e}")
-            traceback.print_exc()
-            sys.exit(1)
-        finally:
-            elapsed = (time.perf_counter() - start) * 1000.0
-            name = fn.__qualname__
-            cnt, tot = _STATS[name]
-            _STATS[name] = [cnt + 1, tot + elapsed]
-            avg = _STATS[name][1] / _STATS[name][0]
-            import numpy as np
-
-            log10 = np.log10(max(cnt, 1))
-            if np.isclose(log10, round(log10)):
-                print(
-                    f"{name} completed in {elapsed:.3f} ms, lifetime avg: {avg:.3f} ms, lifetime cnt: {cnt + 1}"
-                )
-
-    return wrapper