From 94881b69913c48f95d3006bb3d173bbd481359a0 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Mon, 21 Jul 2025 06:20:50 +0000 Subject: [PATCH] add cuda gpu busy threshold --- .../single_gpu_controller/cuda_gpu_controller.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py b/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py index fda50142..1a0124cf 100644 --- a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py +++ b/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py @@ -171,7 +171,15 @@ def _keep_loop(self) -> None: raise RuntimeError("Failed to allocate matrix for GPU keeping") while not self._stop_evt.is_set(): try: - self._run_mat_batch(matrix) + gpu_utilization = self._monitor_utilization(self.rank) + if gpu_utilization > self.busy_threshold: + logger.debug( + "rank %s: GPU busy (%d%%), sleeping longer", + self.rank, + gpu_utilization, + ) + else: + self._run_mat_batch(matrix) time.sleep(self.interval) except RuntimeError as e: # Handle OOM by clearing cache; then sleep and continue