From 1d29c01ec31f702959d843fca32ff897608d259f Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Tue, 26 Aug 2025 07:13:46 +0000 Subject: [PATCH 1/4] refactor keep loop to global controller --- src/keep_gpu/keep_gpu.py | 53 +++++++++++++++++++--------------------- 1 file changed, 25 insertions(+), 28 deletions(-) diff --git a/src/keep_gpu/keep_gpu.py b/src/keep_gpu/keep_gpu.py index 8d9a0443..08f4c80a 100644 --- a/src/keep_gpu/keep_gpu.py +++ b/src/keep_gpu/keep_gpu.py @@ -1,10 +1,10 @@ import argparse import os -import subprocess import time + import torch -from .benchmark import run_benchmark +from keep_gpu.global_gpu_controller.global_gpu_controller import GlobalGPUController from keep_gpu.utilities.logger import setup_logger logger = setup_logger(__name__) @@ -26,22 +26,19 @@ def parse_args(): default=None, help="Comma-separated list of GPU IDs to monitor and benchmark on (default: all)", ) - return parser.parse_args() - - -def check_gpu_usage(gpu_ids=None): - result = subprocess.run( - ["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"], - capture_output=True, - text=True, + parser.add_argument( + "--vram", + type=str, + default="1GiB", + help="Amount of VRAM to keep occupied (e.g., '500MB', '1GiB', or integer in bytes)", ) - usage_lines = result.stdout.strip().split("\n") - usages = [int(line.strip()) for line in usage_lines] - - if gpu_ids is not None: - usages = [usages[i] for i in gpu_ids if i < len(usages)] - - return any(usage > 0 for usage in usages) + parser.add_argument( + "--threshold", + type=int, + default=-1, + help="Max gpu utilization threshold to trigger keeping GPU awake", + ) + return parser.parse_args() def run(): @@ -57,16 +54,16 @@ def run(): gpu_count = torch.cuda.device_count() logger.info("Using all available GPUs") - idle_count = 0 logger.info(f"GPU count: {gpu_count}") + logger.info(f"VRAM to keep occupied: {args.vram}") + logger.info(f"Check interval: {args.interval} seconds") + logger.info(f"Busy threshold {args.threshold}%") + global_controller = GlobalGPUController( + gpu_ids=gpu_ids, + interval=args.interval, + vram_to_keep=args.vram, + busy_threshold=args.threshold, + ) + 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) From fc7e0bce07d992b9a51774b32e319c3d41fd1c20 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Tue, 26 Aug 2025 07:20:52 +0000 Subject: [PATCH 2/4] fix main logic --- src/keep_gpu/keep_gpu.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/keep_gpu/keep_gpu.py b/src/keep_gpu/keep_gpu.py index 08f4c80a..a8f7fcac 100644 --- a/src/keep_gpu/keep_gpu.py +++ b/src/keep_gpu/keep_gpu.py @@ -64,6 +64,10 @@ def run(): vram_to_keep=args.vram, busy_threshold=args.threshold, ) - global_controller.keep() - while True: - 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...") From 1f8d0ba8973b8cd4b10615ac0e1858bdad6da203 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Tue, 26 Aug 2025 07:29:21 +0000 Subject: [PATCH 3/4] use new entry in toml --- pyproject.toml | 2 +- src/keep_gpu/cli.py | 66 ++++++++++++++++++++++++++++++++---- src/keep_gpu/keep_gpu.py | 73 ---------------------------------------- 3 files changed, 61 insertions(+), 80 deletions(-) delete mode 100644 src/keep_gpu/keep_gpu.py diff --git a/pyproject.toml b/pyproject.toml index a750462b..744d7149 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ ] [project.scripts] -keep-gpu = "keep_gpu.cli:main" +keep-gpu = "keep_gpu.cli:app" [project.optional-dependencies] dev = [ diff --git a/src/keep_gpu/cli.py b/src/keep_gpu/cli.py index 18e1ee54..899c18ae 100644 --- a/src/keep_gpu/cli.py +++ b/src/keep_gpu/cli.py @@ -1,20 +1,74 @@ """Console script for keep_gpu.""" -from .keep_gpu import run +import os +import time +from typing import Optional +import torch import typer from rich.console import Console +from keep_gpu.global_gpu_controller.global_gpu_controller import GlobalGPUController +from keep_gpu.utilities.logger import setup_logger + app = typer.Typer() console = Console() +logger = setup_logger(__name__) @app.command() -def main(): - """Console script for keep_gpu.""" - console.print("Replace this message by putting your code into " "keep_gpu.cli.main") - console.print("See Typer documentation at https://typer.tiangolo.com/") - run() +def main( + interval: int = typer.Option( + 300, help="Interval in seconds between GPU usage checks" + ), + gpu_ids: Optional[str] = typer.Option( + None, + help="Comma-separated list of GPU IDs to monitor and benchmark on (default: all)", + ), + vram: str = typer.Option( + "1GiB", + help="Amount of VRAM to keep occupied (e.g., '500MB', '1GiB', or integer in bytes)", + ), + threshold: int = typer.Option( + -1, + help="Max GPU utilization threshold to trigger keeping GPU awake", + ), +): + """ + Keep specified GPUs awake by allocating VRAM and monitoring usage. + """ + # Process GPU IDs + if gpu_ids: + gpu_id_list = [int(i.strip()) for i in gpu_ids.split(",")] + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpu_id_list)) + logger.info(f"Using specified GPUs: {gpu_id_list}") + gpu_count = len(gpu_id_list) + else: + gpu_id_list = None + gpu_count = torch.cuda.device_count() + logger.info("Using all available GPUs") + + # Log settings + logger.info(f"GPU count: {gpu_count}") + logger.info(f"VRAM to keep occupied: {vram}") + logger.info(f"Check interval: {interval} seconds") + logger.info(f"Busy threshold: {threshold}%") + + # Create and start Global GPU Controller + global_controller = GlobalGPUController( + gpu_ids=gpu_id_list, + interval=interval, + vram_to_keep=vram, + busy_threshold=threshold, + ) + + 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...") if __name__ == "__main__": diff --git a/src/keep_gpu/keep_gpu.py b/src/keep_gpu/keep_gpu.py deleted file mode 100644 index a8f7fcac..00000000 --- a/src/keep_gpu/keep_gpu.py +++ /dev/null @@ -1,73 +0,0 @@ -import argparse -import os -import time - -import torch - -from keep_gpu.global_gpu_controller.global_gpu_controller import GlobalGPUController -from keep_gpu.utilities.logger import setup_logger - -logger = setup_logger(__name__) - - -def parse_args(): - parser = argparse.ArgumentParser( - description="GPU Idle Monitor and Benchmark Trigger" - ) - parser.add_argument( - "--interval", - type=int, - default=300, - help="Interval in seconds between GPU usage checks", - ) - parser.add_argument( - "--gpu-ids", - type=str, - default=None, - help="Comma-separated list of GPU IDs to monitor and benchmark on (default: all)", - ) - parser.add_argument( - "--vram", - type=str, - default="1GiB", - help="Amount of VRAM to keep occupied (e.g., '500MB', '1GiB', or integer in bytes)", - ) - parser.add_argument( - "--threshold", - type=int, - default=-1, - help="Max gpu utilization threshold to trigger keeping GPU awake", - ) - return parser.parse_args() - - -def run(): - args = parse_args() - - if args.gpu_ids: - gpu_ids = [int(i.strip()) for i in args.gpu_ids.split(",")] - os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpu_ids)) - logger.info(f"Using specified GPUs: {gpu_ids}") - gpu_count = len(gpu_ids) - else: - gpu_ids = None - gpu_count = torch.cuda.device_count() - logger.info("Using all available GPUs") - - logger.info(f"GPU count: {gpu_count}") - logger.info(f"VRAM to keep occupied: {args.vram}") - logger.info(f"Check interval: {args.interval} seconds") - logger.info(f"Busy threshold {args.threshold}%") - global_controller = GlobalGPUController( - gpu_ids=gpu_ids, - interval=args.interval, - vram_to_keep=args.vram, - busy_threshold=args.threshold, - ) - 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...") From 5f9c820131bbe579bfdc3a92e33ad13430558f53 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Tue, 26 Aug 2025 16:25:01 +0800 Subject: [PATCH 4/4] Update src/keep_gpu/cli.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/keep_gpu/cli.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/keep_gpu/cli.py b/src/keep_gpu/cli.py index 899c18ae..d64e1abf 100644 --- a/src/keep_gpu/cli.py +++ b/src/keep_gpu/cli.py @@ -39,7 +39,11 @@ def main( """ # Process GPU IDs if gpu_ids: - gpu_id_list = [int(i.strip()) for i in gpu_ids.split(",")] + try: + gpu_id_list = [int(i.strip()) for i in gpu_ids.split(",")] + except ValueError: + console.print(f"[bold red]Error: Invalid characters in --gpu-ids '{gpu_ids}'. Please use comma-separated integers.[/bold red]") + raise typer.Exit(code=1) os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpu_id_list)) logger.info(f"Using specified GPUs: {gpu_id_list}") gpu_count = len(gpu_id_list)