fix: fix mem leak

[blueswhen]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @blueswhen, 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 primarily addresses memory leaks and improves overall memory management within the LightLLM framework. It achieves this by centralizing pinned memory allocation through a dedicated manager, implementing explicit shared memory detachment and closing procedures for request-related data, and ensuring proper tensor handling (cloning) within the radix cache to avoid unintended shared references. Additionally, the changes include robustness improvements to the benchmark suite, such as retry logic for API calls and task timeouts, which will help in more reliably assessing system performance and stability.

Highlights

• Centralized Pinned Memory Management: Introduced g_pin_mem_manager to standardize and centralize the allocation of pinned memory for various tensors, replacing direct torch.tensor(...).pin_memory() calls across multiple files.
• Explicit Shared Memory Release: Implemented new mechanisms (release_shm_arrays, detach_shm) and integrated calls across different managers (detokenization, infer batch, HTTP server) to ensure shared memory objects are properly detached or closed upon request completion, preventing leaks.
• Radix Cache Tensor Cloning: Modified the radix cache implementation to clone tensors (token_id_key, token_mem_index_value) when splitting nodes or adding children, addressing potential shared state issues that could lead to incorrect behavior or memory problems.
• Benchmark Robustness Improvements: Enhanced benchmark scripts with retry mechanisms for network requests and added task timeouts in the response collector to prevent indefinite hangs and improve testing reliability.

🧠 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.

Changelog

• lightllm/common/req_manager.py
  • Integrated g_pin_mem_manager for managing pinned memory allocations for prompt IDs and sampling parameters.
  • Modified init_req_sampling_params to use g_pin_mem_manager.gen_from_list for prompt_ids.
  • Modified gen_cpu_out_token_counter_sampling_params to use g_pin_mem_manager.gen_from_list for various token-related tensors.
• lightllm/server/core/objs/req.py
  • Added a release_shm_arrays method to detach shared memory objects (shm_prompt_ids, shm_logprobs) associated with a request.
• lightllm/server/core/objs/shm_array.py
  • Introduced a detach_shm method for shared memory arrays, allowing detachment without unlinking.
• lightllm/server/detokenization/manager.py
  • Ensured release_shm_arrays is called on finished requests to detach shared memory.
• lightllm/server/httpserver/manager.py
  • Explicitly closed shared memory for prompt IDs and log probabilities using close_shm during resource recycling.
• lightllm/server/router/dynamic_prompt/radix_cache.py
  • Modified node splitting and child addition to use .clone() for token_id_key and token_mem_index_value to prevent shared tensor references.
  • Added a conditional check for update_refs before decrementing reference counters in match_prefix.
• lightllm/server/router/model_infer/infer_batch.py
  • Invoked release_shm_arrays on request objects when they are filtered out as finished.
• lightllm/server/router/model_infer/mode_backend/generic_post_process.py
  • Utilized g_pin_mem_manager for eos_ids and various sampling parameter tensors to centralize pinned memory handling.
• test/benchmark/service/benchmark_qps.py
  • Added a _DEFAULT_RETRY constant and implemented retry logic for async_post_stream_openai and async_post_stream_lightllm functions.
  • Enhanced response_collector with a task_timeout_s (600 seconds) and improved error handling to ensure robust task completion and queue management.
  • Removed a print(len(results)) statement.
  • Minor formatting changes.

Activity

• The author blueswhen created this pull request to address memory leaks within the system.
• The changes involve refactoring memory management logic across several core components to ensure proper resource deallocation.
• No specific review comments or approvals have been provided in the context.

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 effectively addresses a memory leak by introducing more robust lifecycle management for shared memory and utilizing a centralized manager for pinned memory. Key changes include differentiating between detaching from and unlinking shared memory segments, assigning clear ownership to httpserver for final cleanup, and replacing on-the-fly pinned memory allocations with a managed pool. Additionally, correctness issues in the RadixCache are fixed, and the robustness of the benchmark script is improved. The changes are well-implemented and significantly enhance the application's stability. I have one suggestion to reduce code duplication.

[gemini-code-assist]

The new detach_shm method has logic that is very similar to the existing close_shm method. To improve maintainability and reduce code duplication, you could extract the common logic into a private helper method. The suggestion below introduces _cleanup and refactors detach_shm to use it. You could then also update close_shm to call self._cleanup(unlink=True).

Suggested change

	def detach_shm(self):
	if self.shm is not None:
	self.shm.close()
	self.shm = None
	self.arr = None
	def detach_shm(self):
	self._cleanup(unlink=False)
	def _cleanup(self, unlink: bool):
	if self.shm is not None:
	self.shm.close()
	if unlink:
	self.shm.unlink()
	self.shm = None
	self.arr = None