[CLI] fix: refactor keep loop to global controller

[Wangmerlyn]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @Wangmerlyn, 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 refactors the GPU management logic within the CLI tool. The primary goal is to centralize the functionality responsible for keeping GPUs active by moving it into a dedicated GlobalGPUController class. This change enhances the maintainability and organization of the codebase. Additionally, it introduces new configuration options via command-line arguments, allowing users to specify the amount of VRAM to occupy and a utilization threshold for triggering the GPU-keeping mechanism. The overall effect is a cleaner, more robust, and more configurable GPU management system.

Highlights

• Refactoring to GlobalGPUController: The core logic for maintaining GPU activity, previously handled directly within keep_gpu.py, has been extracted and encapsulated within a new GlobalGPUController class. This improves modularity and separation of concerns.
• New Command-Line Arguments: New command-line arguments, --vram and --threshold, have been introduced. These allow users to specify the desired amount of VRAM to keep occupied and a maximum GPU utilization threshold, respectively, providing more granular control over the GPU keeping mechanism.
• Delegation of GPU Management Logic: The direct calls to nvidia-smi via subprocess and the run_benchmark function have been removed from keep_gpu.py. Their functionality is now managed internally by the GlobalGPUController, streamlining the main script.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 refactors the GPU keeping logic into a GlobalGPUController, which is a good improvement for modularity. I've identified a couple of areas for improvement. First, the default value for the new --threshold argument is potentially problematic. Second, the script doesn't handle graceful shutdown, which could lead to leaked GPU resources. Using the GlobalGPUController as a context manager would solve this.

[gemini-code-assist]

The current implementation starts the GPU keeping process but lacks a graceful shutdown mechanism. If the script is terminated with Ctrl+C, the release() method on the controller is never called, which could leave GPU resources allocated.

The GlobalGPUController is designed as a context manager (with __enter__ and __exit__), which is the perfect pattern for this scenario. Using a with statement will ensure that release() is always called on exit, making the script more robust.

Suggested change

	global_controller.keep()
	while True:
	if not check_gpu_usage(gpu_ids):
	idle_count += 1
	else:
	idle_count = 0
	if idle_count >= 1:
	run_benchmark(gpu_count)
	idle_count = 0
	time.sleep(args.interval)
	time.sleep(1)
	with global_controller:
	logger.info("Keeping GPUs awake. Press Ctrl+C to exit.")
	try:
	while True:
	time.sleep(3600)
	except KeyboardInterrupt:
	logger.info("Interruption received. Releasing GPUs...")

[gemini-code-assist]

The default value for --threshold is set to -1, which is not a valid GPU utilization percentage. This value is passed to GlobalGPUController, overriding its own default of 10. This could lead to unexpected behavior. It's better to use a sensible default here, like 10, to match the controller's default.

Suggested change

	parser.add_argument(
	"--threshold",
	type=int,
	default=-1,
	help="Max gpu utilization threshold to trigger keeping GPU awake",
	)
	parser.add_argument(
	"--threshold",
	type=int,
	default=10,
	help="Max gpu utilization threshold to trigger keeping GPU awake",
	)

[Wangmerlyn]

/gemini review

[gemini-code-assist]

Code Review

This pull request is a significant improvement, refactoring the command-line interface to use typer and encapsulating the GPU-keeping logic within a GlobalGPUController. This enhances modularity and maintainability. My review focuses on improving the robustness and user experience of the new CLI by adding better error handling for user inputs and making the application more responsive to termination signals.

[gemini-code-assist]

The GlobalGPUController may raise a ValueError if the vram string is in an invalid format. This exception is currently unhandled in the main function, which will cause the application to crash with a stack trace. It's better to catch this exception and display a user-friendly error message.

    try:
        global_controller = GlobalGPUController(
            gpu_ids=gpu_id_list,
            interval=interval,
            vram_to_keep=vram,
            busy_threshold=threshold,
        )
    except ValueError as e:
        console.print(f"[bold red]Error: {e}[/bold red]")
        raise typer.Exit(code=1)

[gemini-code-assist]

The main thread sleeps for 3600 seconds (1 hour). This makes the application unresponsive to termination signals like Ctrl+C. Since the actual work is done in background threads managed by the GlobalGPUController, the main thread only needs to wait. Using a much shorter sleep interval in the loop will allow for a prompt shutdown when requested.

                time.sleep(1)

[Wangmerlyn]

mark this for later.