Improve engine customization interface

[merrymercy]

We introduce a new interface to allow users to easily use their own ServerArgs, TokenizerManager, or Scheduler.
They can override the Engine class or pass these arguments to launch_server

class Engine:
    # Some fields to allow people to override the server args
    # and launch processes for their private forks.
    server_args_class: ServerArgs = ServerArgs
    init_tokenizer_manager_func: Callable = staticmethod(init_tokenizer_manager)
    run_scheduler_process_func: Callable = staticmethod(run_scheduler_process)
    run_detokenizer_process_func: Callable = staticmethod(run_detokenizer_process)

def launch_server(
    server_args: ServerArgs,
    init_tokenizer_manager_func: Callable = init_tokenizer_manager,
    run_scheduler_process_func: Callable = run_scheduler_process,
    run_detokenizer_process_func: Callable = run_detokenizer_process,
    execute_warmup_func: Callable = _execute_server_warmup,
    launch_callback: Optional[Callable[[], None]] = None,
):

[gemini-code-assist]

Summary of Changes

Hello @merrymercy, 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 significantly improves the extensibility of the engine's initialization and subprocess management. It achieves this by refactoring the core subprocess launching logic to allow for easier overriding of key component initialization functions, such as those for the tokenizer, scheduler, and detokenizer. Additionally, it cleans up and simplifies the server warmup procedures by removing an unnecessary inter-process communication pipe, and updates internal scheduler method naming for better readability. These changes collectively make the engine more adaptable to custom implementations and easier to maintain.

Highlights

• Enhanced Engine Customization: The Engine class now offers more granular customization points by exposing init_tokenizer_manager_func, run_scheduler_process_func, and run_detokenizer_process_func as overrideable static methods.
• Refactored Subprocess Launching: The internal _launch_subprocesses function has been refactored to accept these new customizable functions as arguments, increasing flexibility in how engine components are initialized.
• Simplified Server Warmup Logic: The inter-process communication mechanism (pipe_finish_writer) previously used during server warmup in both gRPC and HTTP servers has been removed, streamlining the warmup process.
• Scheduler Method Renaming: A method in the scheduler mixin, check_during_pp_idle, has been renamed to self_check_during_idle for improved clarity and consistency.
• Updated Customization Documentation: The guideline.md file has been updated to reflect the new and improved ways to customize the engine's behavior.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 significantly improves the engine's customization interface by refactoring the subprocess launching logic. The monolithic _launch_subprocesses function has been broken down into smaller, overridable functions for the tokenizer, scheduler, and detokenizer, making it easier to extend the engine with custom components. The changes also include beneficial cleanups, such as removing an unused pipe_finish_writer mechanism and eliminating a duplicated method. The code is well-structured, and the changes are consistent across all affected files. I have one minor suggestion to correct a type hint.

[gemini-code-assist]

The type hint for scheduler_infos in the return value is Tuple[Dict], but the function returns a list of dictionaries. Tuple[Dict] indicates a tuple with a single dictionary element, whereas the function actually returns a list of dictionaries. To accurately reflect the return type, it should be List[Dict].

Suggested change

	) -> Tuple[TokenizerManager, TemplateManager, Tuple[Dict], PortArgs]:
	) -> Tuple[TokenizerManager, TemplateManager, List[Dict], PortArgs]: