WIP: Add asynchronous runtime support

[oeqqwq]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @oeqqwq, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces comprehensive asynchronous runtime support, fundamentally enhancing the system's ability to handle distributed computations more efficiently. By integrating asynchronous communication primitives, an asynchronous AST evaluator, and an asynchronous HTTP communicator, it enables non-blocking operations and concurrent execution across different parties. This architectural shift is crucial for improving the performance and responsiveness of distributed machine learning tasks, particularly those involving significant I/O or waiting periods.

Highlights

• Asynchronous Communication Primitives: Introduced new interfaces (IAsyncCommunicator, IAsyncCollective) and base implementations (AsyncCollectiveMixin, AsyncCommunicatorBase, AsyncThreadCommunicator) for asynchronous point-to-point and collective communication operations.
• Asynchronous AST Evaluation: Extended the Abstract Syntax Tree (AST) Expr classes with an accept_async method and implemented AsyncIterativeEvaluator to enable non-blocking, asynchronous traversal and evaluation of computation graphs.
• Asynchronous HTTP Communicator: Added AsyncHttpCommunicator leveraging httpx.AsyncClient to provide efficient, non-blocking network communication for distributed environments, replacing the synchronous HTTP communicator in async mode.
• Integrated Asynchronous Execution: The Session and Simulator classes have been updated to support asynchronous execution of computations, utilizing the new async evaluators and communicators, and the runtime server now uses this async execution path.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces asynchronous communication and evaluation capabilities across the system. Key changes include the addition of async_comm.py which defines asynchronous communication interfaces (IAsyncCommunicator, IAsyncCollective), a mixin for collective operations (AsyncCollectiveMixin), and base implementations for message passing (AsyncCommunicatorBase, AsyncThreadCommunicator). The ast.py file is updated to include an accept_async method for asynchronous visitor pattern traversal. A new async_evaluator.py introduces AsyncEvalSemantic and AsyncIterativeEvaluator for asynchronous expression evaluation, leveraging an Executor for CPU-bound tasks and implementing asynchronous versions of shuffle and other operations. The runtime/communicator.py file gains AsyncHttpCommunicator for asynchronous HTTP-based communication. The runtime/server.py and runtime/session.py files are modified to support asynchronous execution of computations, including a new async_execute method in Session and updating FastAPI endpoints to be asynchronous. The runtime/simulation.py is refactored to use asynchronous evaluators and communicators, replacing synchronous thread pool execution with asyncio.gather. Finally, new unit tests for asynchronous communication are added in test_async_comm.py, and tutorial paths are updated. A review comment highlighted a typo (sefl instead of self) in an AsyncCommunicatorBase method. Another significant review comment pointed out that the new asynchronous evaluation in simulation.py lost the detailed exception handling and timeout mechanisms present in the previous synchronous implementation, suggesting the use of asyncio.wait_for and return_exceptions=True with asyncio.gather for improved robustness.