[cuda, rocm, ci, tests] fix: resolve CodeRabbit review findings for controller safety and test reliability#63
Conversation
📝 WalkthroughWalkthroughUpdates CI workflows (Python 3.13, pre-commit, pip install ordering), switches CUDA controller background workload from matmul to in-place ReLU with relu_iterations, adds allocation-retry and status tracking to ROCm controller, tightens VRAM parsing and logging, and adjusts tests for new APIs and safer GPU checks. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @Wangmerlyn, I'm Gemini Code Assist1! 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 addresses several findings from a code analysis tool, focusing on enhancing the stability, clarity, and maintainability of the GPU memory management system. The changes primarily involve refining the logic for keeping GPU VRAM occupied, improving error handling, standardizing logging practices, and making the test suite more robust and accurate. Highlights
🧠 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
Ignored Files
Activity
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a series of fixes and improvements across the codebase, addressing findings likely from a static analysis tool. The changes include correcting docstrings, improving error handling with retries in rocm_gpu_controller, enhancing logger robustness, and making tests more reliable and meaningful. The use of deferred formatting for logging and adding parents=True to mkdir are also good improvements.
I've added a couple of review comments with suggestions for further improvements:
- In
CudaGPUController, there's a naming inconsistency where a method and parameters refer tomatmulbut the operation isReLU. - In one of the tests, there's a redundant thread join that should be removed to avoid confusion.
| @@ -166,7 +169,7 @@ def _keep_loop(self) -> None: | |||
| # ------------------------------------------------------------------ | |||
| @torch.no_grad() | |||
| def _run_mat_batch(self, matrix: torch.Tensor) -> None: | |||
There was a problem hiding this comment.
The method name _run_mat_batch is misleading as the implementation uses torch.relu_ instead of a matrix multiplication. The updated docstring correctly identifies the operation as ReLU, which makes the method name even more confusing.
To improve clarity and maintainability, consider renaming this method to something like _run_relu_batch.
Additionally, the parameter matmul_iterations and the debug log message "mat ops batch done" are also inconsistent with the ReLU operation and should be updated accordingly throughout the class.
|
|
||
| assert ctrl._thread is not None | ||
| ctrl.release() | ||
| ctrl._thread.join(timeout=2) |
There was a problem hiding this comment.
This call to _thread.join(timeout=2) is redundant. The ctrl.release() method, called on the previous line, already blocks and waits for the thread to terminate by calling _thread.join() internally. Adding another join here is unnecessary and can be removed to simplify the test. This redundant pattern is repeated multiple times in this test file.
If there's a concern about the thread hanging, the timeout logic should be part of the release() method itself.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9491f5fc99
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if num_elements <= 0: | ||
| raise ValueError("vram_to_keep must be positive") |
There was a problem hiding this comment.
The new num_elements <= 0 guard raises ValueError inside _keep_loop, but this loop only catches RuntimeError, so keep() can report success and then the background thread immediately dies when vram_to_keep resolves to 0 (for example --vram 0 from CLI/API). This leaves the controller inactive without a synchronous failure path; validation should happen before starting the thread or ValueError should be handled and surfaced.
Useful? React with 👍 / 👎.
| tensor = None | ||
| while not self._stop_evt.is_set(): | ||
| attempts = 0 | ||
| while not self._stop_evt.is_set() and attempts < self.max_allocation_retries: |
There was a problem hiding this comment.
The startup allocation loop now stops after max_allocation_retries attempts and returns, which means the keeper thread permanently exits if VRAM is temporarily unavailable at startup; if memory frees up moments later, this controller never resumes unless the caller manually restarts it. This is a functional regression from the previous indefinite retry behavior and can cause keep-alive to silently stop in transiently busy environments.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py (2)
170-186:⚠️ Potential issue | 🟡 Minor
torch.relu_on[0, 1)data is a computational no-op; potentialZeroDivisionErroron line 185.
torch.randproduces values in[0, 1), sorelu_(matrix)(which clamps negatives to 0) leaves every element unchanged. The GPU still processes kernel launches, but there's essentially zero ALU work — a scheduler could optimize this away. If the goal is to generate real utilization, considertorch.neg_(matrix)alternated withtorch.relu_or retain the original matmul approach.Separately, if
self.matmul_iterationsis ever 0, line 185 divides by zero. The ROCm controller guards this withmax(1, self.iterations).Suggested guard for division
- (toc - tic) * 1000 / self.matmul_iterations, + (toc - tic) * 1000 / max(1, self.matmul_iterations),
127-144:⚠️ Potential issue | 🟠 MajorAllocation retry loop differs from ROCm controller in three important ways.
ValueErrorescapes the loop – Line 139 catches onlyRuntimeError, but the validation on lines 130-131 raisesValueError. Ifvram_to_keepresolves to a non-positive integer, theValueErrorkills the background thread silently. The ROCm controller catches(RuntimeError, ValueError).Unbounded retries – This loop retries forever (
while not self._stop_evt.is_set()) on allocationRuntimeError. The ROCm controller bounds retries withmax_allocation_retries. A persistent OOM condition here means an infinite retry loop.
raisein a daemon thread is silently lost – Line 144 raisesRuntimeErrorinside the background thread. Since it's a daemon thread, the exception is swallowed and the caller has no way to detect the failure. The ROCm controller stores it inself._failure_excand exposesallocation_status().Consider aligning with the ROCm controller's pattern for consistency.
Suggested diff
+ attempts = 0 - while not self._stop_evt.is_set(): + while not self._stop_evt.is_set() and attempts < self.max_allocation_retries: try: num_elements = int(self.vram_to_keep) if num_elements <= 0: raise ValueError("vram_to_keep must be positive") matrix = torch.rand( num_elements, device=self.device, dtype=torch.float32, requires_grad=False, ) break - except RuntimeError as e: - logger.error("rank %s: failed to allocate matrix: %s", self.rank, e) + except (RuntimeError, ValueError) as e: + attempts += 1 + logger.error("rank %s: failed to allocate matrix (attempt %d/%d): %s", self.rank, attempts, self.max_allocation_retries, e) time.sleep(self.interval) if matrix is None: - logger.error("rank %s: failed to allocate matrix, exiting loop", self.rank) - raise RuntimeError("Failed to allocate matrix for GPU keeping") + self._failure_exc = RuntimeError( + f"rank {self.rank}: failed to allocate matrix after {attempts} attempts" + ) + logger.error("%s", self._failure_exc) + returnsrc/keep_gpu/single_gpu_controller/rocm_gpu_controller.py (1)
68-80:⚠️ Potential issue | 🟡 Minor
release()logs "stopped & cache cleared" even when the thread was never running.Line 80's
logger.info("rank %s: keep thread stopped & cache cleared")executes unconditionally after theif/else. When the thread isn't alive (else branch), the user sees both the "not running" warning and the "stopped & cache cleared" info — which is contradictory.Suggested fix
def release(self) -> None: if self._thread and self._thread.is_alive(): self._stop_evt.set() self._thread.join() torch.cuda.empty_cache() + logger.info("rank %s: keep thread stopped & cache cleared", self.rank) else: logger.warning("rank %s: keep thread not running", self.rank) if self._rocm_smi: try: self._rocm_smi.rsmi_shut_down() except Exception as exc: # pragma: no cover - best effort logger.debug("rsmi_shut_down failed: %s", exc) - logger.info("rank %s: keep thread stopped & cache cleared", self.rank)
🤖 Fix all issues with AI agents
In `@tests/cuda_controller/context_manager.py`:
- Line 12: The test file tests/cuda_controller/context_manager.py containing the
test function test_cuda_controller_context_manager will not be collected by
pytest because its filename doesn't match the default discovery patterns; fix by
either renaming the file to a pytest-recognized name like
test_context_manager.py (so the existing function
test_cuda_controller_context_manager is discovered) or add a python_files
pattern to test configuration (pyproject.toml or conftest.py) to include
context_manager.py.
🧹 Nitpick comments (2)
tests/cuda_controller/test_keep_and_release.py (1)
22-25: Redundantjoin()afterrelease()– masks real bugs in production code.
CudaGPUController.release()already callsself._thread.join()(unbounded) internally. The extractrl._thread.join(timeout=2)on line 24 re-joins an already-joined thread, which is a no-op. More importantly, ifrelease()ever fails to join (a real bug), this test-side join would paper over it instead of exposing the defect.Similarly, the
assert ctrl._thread is not Noneon line 22 is trivially true given line 20 already assertedctrl._thread and ctrl._thread.is_alive().Consider removing the redundant join and None check here and at lines 30-32, keeping only the post-
release()liveness assertion which is the meaningful check.Suggested diff
assert ctrl._thread and ctrl._thread.is_alive() - assert ctrl._thread is not None ctrl.release() - ctrl._thread.join(timeout=2) assert not (ctrl._thread and ctrl._thread.is_alive())src/keep_gpu/single_gpu_controller/rocm_gpu_controller.py (1)
148-150:allocation_status()exposes thread-internal failure — consider documenting thread-safety semantics.Since
_failure_excis written by the background thread and read from the main thread, a brief note that the reference assignment is atomic in CPython (GIL) but not guaranteed in other Python implementations would help future maintainers. This is not a bug today but a minor documentation opportunity.
| @@ -10,15 +11,20 @@ | |||
| ) | |||
| def test_cuda_controller_context_manager(): | |||
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -e
echo "=== Checking pytest configuration for python_files setting ==="
for config_file in pyproject.toml setup.cfg pytest.ini tox.ini; do
if [ -f "$config_file" ]; then
echo "--- Found $config_file ---"
rg -n "python_files" "$config_file" 2>/dev/null || echo "No python_files setting in $config_file"
fi
done
echo -e "\n=== Checking conftest.py for custom collection ==="
find . -name "conftest.py" -type f | while read conftest; do
echo "--- Checking $conftest ---"
rg -n "collect_ignore\|pytest_collect_file\|python_files" "$conftest" 2>/dev/null || echo "No custom collection config in $conftest"
done
echo -e "\n=== Checking if context_manager.py file exists and its git status ==="
if [ -f "tests/cuda_controller/context_manager.py" ]; then
echo "File exists at tests/cuda_controller/context_manager.py"
git log --oneline --all -- "tests/cuda_controller/context_manager.py" 2>/dev/null | head -5 || echo "File not in git history yet (new file)"
else
echo "File does not exist at tests/cuda_controller/context_manager.py"
fi
echo -e "\n=== Listing test files in tests/cuda_controller directory ==="
ls -la tests/cuda_controller/ | grep -E "\.py$" || echo "No Python files found"Repository: Wangmerlyn/KeepGPU
Length of output: 958
Test file context_manager.py will not be discovered by pytest.
Pytest uses default collection patterns test_*.py or *_test.py. There is no python_files configuration in pyproject.toml or custom collection logic in conftest.py. Rename the file to test_context_manager.py or add a python_files setting to enable discovery of non-standard test filenames.
🤖 Prompt for AI Agents
In `@tests/cuda_controller/context_manager.py` at line 12, The test file
tests/cuda_controller/context_manager.py containing the test function
test_cuda_controller_context_manager will not be collected by pytest because its
filename doesn't match the default discovery patterns; fix by either renaming
the file to a pytest-recognized name like test_context_manager.py (so the
existing function test_cuda_controller_context_manager is discovered) or add a
python_files pattern to test configuration (pyproject.toml or conftest.py) to
include context_manager.py.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/keep_gpu/single_gpu_controller/rocm_gpu_controller.py (1)
68-80:⚠️ Potential issue | 🟡 Minor"stopped & cache cleared" logged even when the thread was not running.
When
release()takes theelsebranch (line 73-74), execution falls through to line 80, which logs the misleading message. Move the success log inside theifblock or make it conditional.Proposed fix
def release(self) -> None: if self._thread and self._thread.is_alive(): self._stop_evt.set() self._thread.join() torch.cuda.empty_cache() + logger.info("rank %s: keep thread stopped & cache cleared", self.rank) else: logger.warning("rank %s: keep thread not running", self.rank) if self._rocm_smi: try: self._rocm_smi.rsmi_shut_down() except Exception as exc: # pragma: no cover - best effort logger.debug("rsmi_shut_down failed: %s", exc) - logger.info("rank %s: keep thread stopped & cache cleared", self.rank)tests/cuda_controller/test_throttle.py (1)
37-38:⚠️ Potential issue | 🟡 MinorStale comment: still references "matmul" after the ReLU migration.
Line 37 says "no matmul batches should run" but the workload is now ReLU-based.
Proposed fix
- # When utilization is always high, no matmul batches should run + # When utilization is always high, no ReLU batches should run
🤖 Fix all issues with AI agents
In `@src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py`:
- Around line 196-200: The debug log divides by self.relu_iterations and can
raise ZeroDivisionError when relu_iterations is 0; update the calculation used
in logger.debug (the expression using (toc - tic) * 1000 / self.relu_iterations)
to guard against zero—e.g., use max(1, self.relu_iterations) (similar to the
ROCm controller's pattern) or ensure relu_iterations is validated in __init__ to
be >0; apply the same guard for matmul_iterations where applicable so the
division never uses zero.
In `@src/keep_gpu/single_gpu_controller/rocm_gpu_controller.py`:
- Around line 105-137: The ValueError raised for invalid vram_to_keep should not
be treated as transient: in the allocation loop inside rocm_gpu_controller (the
block that creates tensor with torch.rand using self.vram_to_keep and the
surrounding retry logic referencing self.max_allocation_retries, attempts, and
self._failure_exc), stop catching ValueError so it propagates immediately;
change the except clause to only catch RuntimeError (allocation failures), keep
the existing logging and retry/failure behavior for RuntimeError, and ensure
configuration errors (ValueError) are not swallowed or retried.
🧹 Nitpick comments (1)
src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py (1)
134-158: Allocation-failure path raises in a daemon thread on clean shutdown.When
release()is called while the allocation-retry loop (lines 144-151) is still spinning,_stop_evtgets set, thewhileexits withmatrix is None, and line 158 raisesRuntimeErrorinside the daemon thread. The exception is silently swallowed (daemon thread), but it generates a noisy traceback on stderr.Consider returning early (or logging) when
_stop_evtis set, similar to the ROCm controller's graceful-exit pattern.Proposed fix
if matrix is None: - logger.error("rank %s: failed to allocate matrix, exiting loop", self.rank) - raise RuntimeError("Failed to allocate matrix for GPU keeping") + if self._stop_evt.is_set(): + logger.info("rank %s: allocation cancelled (stop requested)", self.rank) + return + logger.error("rank %s: failed to allocate matrix, exiting loop", self.rank) + raise RuntimeError("Failed to allocate matrix for GPU keeping")
Wangmerlyn
changed the title
1 participant
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
Chores