diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 4267f9d0952..78f412ea5ab 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1609,6 +1609,13 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int] metadata_types = req.get_metadata_types() for m, m_type in zip(metadata, metadata_types): label, _, _ = m_type + if not isinstance(m, torch.Tensor): + m = torch.as_tensor( + m, + device=self.request_metadata[label].device, + dtype=self.request_metadata[label].dtype, + ) + self.request_metadata[label][current_id] = m # Handle length and block assignments. diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 1bb4ac73f44..86501ab3063 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import asyncio +import concurrent.futures import logging import multiprocessing import os @@ -1358,11 +1359,29 @@ async def async_step( # Keep for compatibility with current test suite. return ret + def _run_coroutine_sync(self, coro): + """Run a coroutine synchronously, handling the case when already in an event loop. + + This method safely runs an async coroutine from synchronous code, even when + called from within an already running event loop (e.g., when used with async + frameworks like pytriton). + """ + try: + # Check if there's already a running event loop + asyncio.get_running_loop() + # We're inside a running loop - run in a separate thread + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(asyncio.run, coro) + return future.result() + except RuntimeError: + # No running loop - safe to use run_until_complete + return self._loop.run_until_complete(coro) + def step_modern( self, ) -> Tuple[List[DynamicInferenceRequest], List[DynamicInferenceRequest], float]: """Synchronous wrapper for `self.async_step`.""" - return self._loop.run_until_complete(self.async_step()) + return self._run_coroutine_sync(self.async_step()) def step_legacy( self, sampling_params: SamplingParams @@ -1373,7 +1392,7 @@ def step_legacy( "0.16. Please use `step_modern()` going forward, which will eventually " "be renamed to `step()`." ) - result = self._loop.run_until_complete(self.async_step()) + result = self._run_coroutine_sync(self.async_step()) active_requests = [self.get_request(i) for i in result["active_request_ids"]] finished_requests = [r.merge() for r in result["finished_request_records"]] return active_requests, finished_requests, result["step_time"]