Supporting inference when called within an asyncio loop

[shanmugamr1992]

Export-Deploy faced this issue

WARNING:pytriton.proxy.inference:Exception while performing inference on requests=00000002: Traceback (most recent call last):
  File "/opt/venv/lib/python3.12/site-packages/pytriton/proxy/inference.py", line 393, in _handle_requests
    async for responses in self._model_callable(requests):
  File "/opt/venv/lib/python3.12/site-packages/pytriton/proxy/inference.py", line 85, in _callable
    yield inference_callable(requests)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.12/site-packages/pytriton/decorators.py", line 213, in batch
    outputs = wrapped(*args, **new_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.12/site-packages/pytriton/decorators.py", line 672, in wrapper
    return wrapped(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/Export-Deploy/nemo_deploy/llm/megatronllm_deployable.py", line 390, in triton_infer_fn
    output_infer = self._infer_fn(
                   ^^^^^^^^^^^^^^^
  File "/opt/Export-Deploy/nemo_deploy/llm/megatronllm_deployable.py", line 472, in _infer_fn
    results = self.generate(prompts, inference_params)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/Export-Deploy/nemo_deploy/llm/megatronllm_deployable.py", line 243, in generate
    results = self.mcore_engine.generate(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/megatron-lm/megatron/core/inference/engines/static_engine.py", line 323, in generate
    return self.generate_using_dynamic_engine(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/megatron-lm/megatron/core/inference/engines/static_engine.py", line 232, in generate_using_dynamic_engine
    return self.dynamic_engine.generate(prompts=prompts, sampling_params=sampling_params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/megatron-lm/megatron/core/inference/engines/dynamic_engine.py", line 1286, in generate
    result = self.step_modern()
             ^^^^^^^^^^^^^^^^^^
  File "/opt/megatron-lm/megatron/core/inference/engines/dynamic_engine.py", line 1255, in step_modern
    return self._loop.run_until_complete(self.async_step())
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/asyncio/base_events.py", line 663, in run_until_complete
    self._check_running()
  File "/usr/lib/python3.12/asyncio/base_events.py", line 624, in _check_running
    raise RuntimeError(
RuntimeError: Cannot run the event loop while another loop is running

This is a known issue with Python asyncio - you cannot call run_until_complete() from within an already running event loop. The error occurs because:
PyTriton runs its own async event loop internally
When it calls Megatron's generate() → step_modern(), the code tries to run self._loop.run_until_complete(self.async_step())
This fails because Python's asyncio doesn't allow nested run_until_complete() calls

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tdene]

@shanmugamr1992 @santhnm2 I am continuing this discussion in a separate thread.

This is the Pythonic way to implement this fix: #2831. Here are the differences between #2816 and #2831.

• This Supporting inference when called within an asyncio loop #2816 has generate work no matter how the user calls it. Supporting inference when called within an asyncio loop (Pythonic) #2831 instead throws an error (Please call the async_generate method instead.) if the user calls generate in an async context, asking the user to modify their code.
• This Supporting inference when called within an asyncio loop #2816 is blocking, preventing the user's other work from running in parallel in the same event loop, potentially slowing down the user. Supporting inference when called within an asyncio loop (Pythonic) #2831 instead is non-blocking, allowing the user's other work to run while the CPU idles during the forward pass.

After thinking about this yesterday, I decided on this current #2816 for its simplicity.
Perhaps it's worth having a brief conversation with the users though: why are they running inference from inside an asynchronous context, and do they actually expect to run other workloads in parallel in the same event loop?

[Phlip79]

/ok to test ef4b760

[oyilmaz-nvidia]

Can you please change the https://github.com/shanmugamr1992/Megatron-LM/blob/fix_asyncio_issue/megatron/core/inference/contexts/dynamic_context.py to something like

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
?

I just tested this in the nvcr.io/nvidian/nemo:26.02.rc1 and it works after this update.

[shanmugamr1992]

/ok to test 12abb38

[oyilmaz-nvidia]

Tested with triton based deployment and it's working. Thanks!

[shanmugamr1992]

/ok to test 3c7a941