Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

core: Add test for AsyncCallbackManager context loss (relates to #23909) #26857

Closed
wants to merge 8 commits into from
Closed

core: Add test for AsyncCallbackManager context loss (relates to #23909) #26857

wants to merge 8 commits into from

Conversation

parambharat
Copy link
Contributor

  • Description:
    This PR adds a failing test case to demonstrate the context loss issue in AsyncCallbackManager. The current implementation doesn't honor the run_inline attribute of callback handlers, leading to potential race conditions and ordering issues, especially when handlers depend on shared state or execution sequence.

    Key points:

    1. Issue with contextvars:

      • Since contextvars maintain context across asynchronous tasks within the same event loop.
      • While contextvars ensure state consistency, they don't reflect the order of execution or concurrency issues.
      • A test using only contextvars would pass even if tasks execute concurrently, masking the actual problem.

    2. Test Case Approach:

      • Introduces shared_stack to capture the sequence of operations.
      • Uses three handlers: StateModifier and StateReader (both run_inline=True), and NonInlineHandler (run_inline=False).
      • Handlers manipulate a shared counter and append states to shared_stack.
      • Expected behavior: Sequential execution for inline handlers, resulting in a predictable state sequence.
      • Actual behavior: Concurrent execution leads to a non-deterministic order in shared_stack.

    3. Effectiveness of the Approach:

      • shared_stack reveals concurrent execution issues that contextvars alone cannot detect.
      • Demonstrates the need to honor run_inline=True in AsyncCallbackManager.
      • Highlights potential problems in scenarios where handlers rely on sequential execution.

    4. Outcome and Significance:

      • The test fails, proving that AsyncCallbackManager executes tasks concurrently even when they should be inline.
      • This failure effectively demonstrates the issue addressed by the proposed fix.
      • Emphasizes the importance of maintaining execution order for handlers dependent on shared state or sequential logic.

    This test case clearly demonstrates the problem and justifies the proposed changes to AsyncCallbackManager and ensures it respects handlers' run_inline attribute.

  • Issue: This test case relates to the issue being addressed in PR core: ensure run_inline is honored in AsyncCallback Handler #23909 (core: ensure run_inline is honored in AsyncCallback Handler)

  • Dependencies: No new dependencies are required. This PR only adds a test case.

  • Twitter handle: @parambharat

@eyurtsev

@parambharat
fix: Add test for AsyncCallbackManager context loss (relates to langc…
eb939a3 
…hain-ai#23909)

- Add failing test to demonstrate issue fixed in PR langchain-ai#23909
- Show `AsyncCallbackManager` doesn't honor run_inline attribute
- Use `shared_stack` to capture non-deterministic execution order
- Illustrate why `contextvars` alone can't detect this issue
- Test expects to fail until fix for langchain-ai#23909 is implemented
@vercel Vercel
Copy link

vercel bot commented Sep 25, 2024 •
edited
Loading

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Skipped Deployment
Name Status Preview Comments Updated (UTC)
langchain ⬜️ Ignored (Inspect) Visit Preview Oct 3, 2024 7:33am

@dosubot dosubot bot added size:M This PR changes 30-99 lines, ignoring generated files. Ɑ: core Related to langchain-core labels Sep 25, 2024
@eyurtsev eyurtsev self-assigned this Sep 25, 2024
eyurtsev
eyurtsev reviewed Sep 25, 2024
View reviewed changes
libs/core/tests/unit_tests/callbacks/test_async_callback_manager.py Outdated
) -> None:
event_name = "on_llm_start"
prompt = prompts[0]
await asyncio.sleep(self.delay)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is delay necessary?

Copy link
Collaborator

@eyurtsev eyurtsev Sep 25, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is most of the test case necessary? Can we capture what we need in ~10 lines of test code only verifying that context vars is being set and can be read?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@eyurtsev , Thank you for your questions. After careful consideration:

  1. The delay is unnecessary and can be removed without affecting the test's ability to demonstrate the issue. I've done this now.

  2. Regarding simplifying the test case, it's important to note that a simpler test using only contextvars would actually pass, even if the run_inline attribute is not being honored. This is because contextvars maintain context across async tasks but don't reflect execution order or concurrency issues.

The test case using shared_stack is necessary because:

  • It reveals concurrent execution issues that contextvars alone cannot detect.
  • It demonstrates the need to honor run_inline=True in AsyncCallbackManager.
  • It highlights potential problems in scenarios where handlers rely on sequential execution.

While we could simplify some aspects of the test, maintaining the core structure with shared_stack is crucial for effectively demonstrating the issue. I highlighted this in the PR description, but now I'm unsure how to simplify it further while still demonstrating the issue. I'm open to suggestions on how to make it more concise while still effectively capturing the problem.

@eyurtsev
Copy link
Collaborator

@parambharat this should be fine -- i can take over, would you make the 2nd PR that updates the callback dispatcher?

parambharat added a commit to parambharat/langchain that referenced this pull request Sep 26, 2024
@parambharat
fix(callbacks): Ensure correct execution order in AsyncCallbackManager
ab635e6 
- Separate inline and non-inline handlers in `on_llm_start` and `on_chat_model_start`
- Execute inline handlers sequentially
- Run non-inline handlers concurrently
- Maintain context integrity for stateful handlers
- Address issue langchain-ai#23909 and fix test in PR langchain-ai#26857
@parambharat parambharat mentioned this pull request Sep 26, 2024
Merged
@parambharat
Copy link
Contributor Author

hi @eyurtsev . Just checking in. Any update on merging this here ?

eyurtsev added a commit that referenced this pull request Oct 7, 2024
@parambharat @eyurtsev
core[patch]: Update AsyncCallbackManager to honor run_inline attr…
931ce8d 
…ibute and prevent context loss (#26885)

## Description

This PR fixes the context loss issue in `AsyncCallbackManager`,
specifically in `on_llm_start` and `on_chat_model_start` methods. It
properly honors the `run_inline` attribute of callback handlers,
preventing race conditions and ordering issues.

Key changes:
1. Separate handlers into inline and non-inline groups.
2. Execute inline handlers sequentially for each prompt.
3. Execute non-inline handlers concurrently across all prompts.
4. Preserve context for stateful handlers.
5. Maintain performance benefits for non-inline handlers.

**These changes are implemented in `AsyncCallbackManager` rather than
`ahandle_event` because the issue occurs at the prompt and message_list
levels, not within individual events.**

## Testing

- Test case implemented in #26857 now passes, verifying execution order
for inline handlers.

## Related Issues

- Fixes issue discussed in #23909 

## Dependencies

No new dependencies are required.

---

@eyurtsev: This PR implements the discussed changes to respect
`run_inline` in `AsyncCallbackManager`. Please review and advise on any
needed changes.

Twitter handle: @parambharat

---------

Co-authored-by: Eugene Yurtsev <[email protected]>
@eyurtsev
Copy link
Collaborator

eyurtsev commented Oct 7, 2024 •
edited
Loading

Related PR was merged. This PR no longer needed.

@eyurtsev eyurtsev closed this Oct 7, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@eyurtsev eyurtsev eyurtsev left review comments

Assignees

@eyurtsev eyurtsev

Labels
Ɑ: core Related to langchain-core size:M This PR changes 30-99 lines, ignoring generated files.
Projects
PR Reviews
Archived in project
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@parambharat @eyurtsev