NVIDIA / NVIDIA/NeMo-Agent-Toolkit

patch_with_retry's per-instance _in_retry_context flag races across concurrent requests sharing one client

Open
#2,177 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
2.6k
Forks
762
Avg merge
21h 28m
Merged PRs (30d)
27

Description

Describe the bug.

patch_with_retry wraps every LLM/embedder/memory client instance NAT builds, and always passes instance_context_aware=True (packages/nvidia_nat_core/src/nat/utils/exception_handlers/automatic_retries.py:477). The nested-call guard it relies on, _RetryContext, uses one plain instance attribute:

def __enter__(self):
    ...
    if getattr(obj, "_in_retry_context", False):
        return True   # already_in_context: call fn once, no retry logic
    object.__setattr__(obj, "_in_retry_context", True)
    ...

This attribute is per-instance, not per-call, so it can't tell "this call is nested inside another call on the same object" apart from "a completely unrelated call is concurrently in flight on the same object". NAT builds these clients once per workflow (e.g. packages/nvidia_nat_langchain/.../react_agent/register.py:133 calls builder.get_llm once, and the resulting closure is reused for every request), and a shared (non-per-user) workflow serves concurrent requests through one SessionManager semaphore with max_concurrency=8 by default (packages/nvidia_nat_core/src/nat/runtime/session.py:107-109,217-218). So up to 8 requests can call methods on the same patched instance at once.

When two such calls overlap, whichever __enter__ runs second sees the flag already True (set by the first, still in-flight call) and takes the already_in_context branch, which calls the wrapped function exactly once with none of the retry/backoff logic. A transient, genuinely retryable error on that second, unrelated concurrent request propagates immediately instead of being retried, for that request only, with nothing logged.

The mechanism was introduced by #803, which scoped it explicitly to nested same-task calls and cross-instance isolation; concurrent same-instance calls from different requests weren't part of that design.

Minimum reproducible example
import asyncio
from nat.utils.exception_handlers.automatic_retries import patch_with_retry

class FlakyClient:
    def __init__(self):
        self.calls_per_tag = {}
    async def ainvoke(self, tag):
        n = self.calls_per_tag[tag] = self.calls_per_tag.get(tag, 0) + 1
        await asyncio.sleep(0.05)  # simulated network latency
        if n == 1:
            raise RuntimeError(f"{tag} transient failure")
        return f"{tag} ok"

async def main():
    client = FlakyClient()
    patch_with_retry(client, retries=5, retry_on=(RuntimeError,))
    task_a = asyncio.create_task(client.ainvoke("A"))
    await asyncio.sleep(0.01)
    task_b = asyncio.create_task(client.ainvoke("B"))
    print(await asyncio.gather(task_a, task_b, return_exceptions=True))

asyncio.run(main())
Relevant log output
[A ok, RuntimeError('B transient failure')]

B is configured with retries=5 but its first-attempt transient error propagates unretried, because A's still-in-flight call on the same shared instance had already set _in_retry_context=True.

Other/Misc.

Verified against current develop (file still matches this description after #2146, which changed unrelated retry-budget-clamping behavior in the same file). Reproduced this session with a mock client exercising the real automatic_retries.py; I have not driven this through a live NAT workflow against a real or containerized LLM endpoint, since no credentials are available for that here, and I've confirmed the object-reuse setup by reading react_agent/register.py and runtime/session.py directly rather than exercising it end to end.

I did not check whether every downstream framework plugin (llama_index, crewai, autogen, agno, semantic_kernel, strands, the memory editors, the eval judges) shares this exact build-once-reuse pattern; the langchain path is confirmed directly, and the bug lives in the one shared helper all of them call through, so a fix there should cover all of them regardless.

Happy to send a PR if useful. I don't have a strong opinion yet on the right primitive (a per-call token/contextvar vs. some other approach), so flagging this first seemed better than guessing at the fix shape.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in packages/nvidia_nat_core/src/nat/utils/exception_handlers/automatic_retries.py at _RetryContext and patch_with_retry, then reproduce the provided asyncio example. Review the reuse paths in packages/nvidia_nat_langchain/.../react_agent/register.py and packages/nvidia_nat_core/src/nat/runtime/session.py. Done means concurrent calls on one patched instance retain independent retry behavior while nested same-task calls remain guarded.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
ai, backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.