diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py index 7b5de5c21f2..9e2b6bfa983 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py +++ b/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], diff --git a/megatron/core/inference/async_stream.py b/megatron/core/inference/async_stream.py index 1bf8775e368..6c3242a13db 100644 --- a/megatron/core/inference/async_stream.py +++ b/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""" diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index ea0560183d8..0045d5947a1 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/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: """ diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 2c43a7e2611..4bff4f85fa8 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/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() diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py index d084528b8f2..dc86eb775f9 100644 --- a/megatron/core/inference/engines/static_engine.py +++ b/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()) diff --git a/megatron/core/inference/inference_client.py b/megatron/core/inference/inference_client.py index 59b9144a207..53daac091b0 100644 --- a/megatron/core/inference/inference_client.py +++ b/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. diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index c2717767fed..65b133aa018 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/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) ) diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py index 985042f31e3..d58f3c3a652 100644 --- a/megatron/core/inference/utils.py +++ b/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." + ) diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 93b2e593d84..f8954feede8 100644 --- a/megatron/core/utils.py +++ b/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) diff --git a/megatron/rl/__init__.py b/megatron/rl/__init__.py index 035da465519..d3ae2fefd16 100644 --- a/megatron/rl/__init__.py +++ b/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 diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index 3e16f74599f..fce7c3073ee 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -8,7 +8,7 @@ import numpy as np from pydantic import BaseModel -from ..__init__ import Request, TypeLookupable, trace_async_exceptions +from ..__init__ import Request, TypeLookupable from ..inference import ( ChatInferenceInterface, ChatInferenceRequest, @@ -18,6 +18,8 @@ ReturnsRaw, ) +from megatron.core.utils import trace_async_exceptions + class AgentBaseModel(BaseModel, extra='allow'): pass @@ -192,7 +194,7 @@ async def get_grouped_rollouts(self, request: GroupedRolloutRequest): ) submitted_groups = 0 - @trace_async_exceptions + @trace_async_exceptions(verbose=True) async def group_task(): nonlocal submitted_groups while request.num_groups == -1 or submitted_groups < request.num_groups: