fix: fix timeout error in vllmwrapper

[ChenZiHong-Gavin]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @ChenZiHong-Gavin, 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 enhances the reliability of the VLLMWrapper by integrating explicit timeout mechanisms for vLLM generation requests. The changes ensure that long-running or unresponsive LLM calls are gracefully terminated, preventing resource exhaustion and improving the overall stability and responsiveness of the system when interacting with the vLLM engine. This addresses potential issues where vLLM might hang, providing a more robust interaction layer.

Highlights

• Configurable Timeout: A new timeout parameter has been added to the VLLMWrapper constructor, allowing users to specify a maximum duration for vLLM generation requests, defaulting to 300 seconds.
• Asynchronous Timeout Handling: The generate_answer and generate_topk_per_token methods now utilize asyncio.wait_for to enforce the configured timeout during vLLM generation, preventing indefinite hangs.
• Robust Error Management: Comprehensive exception handling for asyncio.TimeoutError, asyncio.CancelledError, and general Exception has been implemented, ensuring that the vLLM engine's request is properly aborted (self.engine.abort) if an issue occurs during generation.
• Helper for Generator Consumption: A new private asynchronous method, _consume_generator, was introduced to efficiently iterate through and retrieve the final output from an asynchronous generator.

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 a timeout mechanism for vLLM generation calls using asyncio.wait_for, which is a good improvement to prevent requests from hanging. The implementation correctly handles timeouts and other exceptions by aborting the vLLM request. My review focuses on improving code maintainability by reducing duplication. I've identified repeated exception handling logic and a larger block of duplicated code for running the generator, and I've provided suggestions for refactoring. I also suggest adding type hints to the new helper method for better code clarity.

[gemini-code-assist]

Similar to the generate_answer method, this exception handling logic is duplicated. This can be simplified by catching (Exception, asyncio.CancelledError) in a single block.

More importantly, the entire try...except block for running the generator with a timeout is a duplicate of the one in generate_answer. This violates the DRY (Don't Repeat Yourself) principle. To improve maintainability, consider extracting this logic into a private helper method.

Suggested change

	except asyncio.TimeoutError:
	await self.engine.abort(request_id)
	raise
	except asyncio.CancelledError:
	await self.engine.abort(request_id)
	raise
	except Exception as e:
	await self.engine.abort(request_id)
	raise
	except (Exception, asyncio.CancelledError):
	await self.engine.abort(request_id)
	raise

[gemini-code-assist]

For better code clarity and maintainability, it's good practice to add type hints to method signatures. Since the specific vLLM types are not imported at the module level, using typing.Any is a reasonable approach here.

Suggested change

	async def _consume_generator(self, generator):
	async def _consume_generator(self, generator: Any) -> Any:

[gemini-code-assist]

The exception handling for TimeoutError, CancelledError, and Exception contains duplicated code (await self.engine.abort(request_id)). You can combine these into a single except block to make the code more concise and maintainable. asyncio.TimeoutError inherits from Exception, while asyncio.CancelledError inherits from BaseException, so catching (Exception, asyncio.CancelledError) covers all intended cases.

Suggested change

	except asyncio.TimeoutError:
	await self.engine.abort(request_id)
	raise
	except asyncio.CancelledError:
	await self.engine.abort(request_id)
	raise
	except Exception as e:
	await self.engine.abort(request_id)
	raise
	except (Exception, asyncio.CancelledError):
	await self.engine.abort(request_id)
	raise