Enable locking GPU in CI #250
Conversation
Summary of ChangesHello @fzyzcjy, 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 introduces a critical utility for managing GPU resources in CI environments. The new Highlights
Ignored Files
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 Python script for locking GPUs in a CI environment, which is a solid approach. My review focuses on improving the script's robustness, clarity, and adherence to Python best practices. I've suggested using a context manager to prevent potential resource leaks, refining error handling with more specific exception types, improving type hints for better static analysis, and removing some redundant code. These changes will make the script more reliable and maintainable.
| try: | ||
| fd_lock = FdLock(args.lock_path_pattern, i) | ||
| fd_lock.open() | ||
| try: | ||
| fd_lock.lock() | ||
| fcntl.flock(fd_lock.fd, fcntl.LOCK_UN) | ||
| free.append(i) | ||
| except BlockingIOError: | ||
| pass | ||
| fd_lock.close() | ||
| except Exception as e: | ||
| print(f"Warning: Error while probing lock: {e}", file=sys.stderr, flush=True) |
There was a problem hiding this comment.
The current implementation for probing locks in _execute_print_only can leak file descriptors if an exception occurs after opening the file but before closing it.
To make resource management safer and more Pythonic, you can turn FdLock into a context manager. This ensures that file descriptors are always closed, even if errors occur.
First, add __enter__ and __exit__ methods to the FdLock class:
class FdLock:
# ... existing methods ...
def __enter__(self):
self.open()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()Then, you can simplify _execute_print_only to use a with statement, which will automatically handle opening and closing the lock file.
| try: | |
| fd_lock = FdLock(args.lock_path_pattern, i) | |
| fd_lock.open() | |
| try: | |
| fd_lock.lock() | |
| fcntl.flock(fd_lock.fd, fcntl.LOCK_UN) | |
| free.append(i) | |
| except BlockingIOError: | |
| pass | |
| fd_lock.close() | |
| except Exception as e: | |
| print(f"Warning: Error while probing lock: {e}", file=sys.stderr, flush=True) | |
| try: | |
| with FdLock(args.lock_path_pattern, i) as fd_lock: | |
| try: | |
| fd_lock.lock() | |
| fcntl.flock(fd_lock.fd, fcntl.LOCK_UN) | |
| free.append(i) | |
| except BlockingIOError: | |
| pass | |
| except Exception as e: | |
| print(f"Warning: Error while probing lock: {e}", file=sys.stderr, flush=True) |
| if "{gpu_id}" not in args.lock_path_pattern: | ||
| raise Exception("ERROR: --lock-path-pattern must contain '{i}' placeholder.") | ||
|
|
||
| if not args.cmd and not args.print_only: | ||
| raise Exception("ERROR: missing command to run. Use -- before command.") |
There was a problem hiding this comment.
The argument validation in _parse_args can be improved for clarity and correctness:
- The error message on line 59 contains a typo. It refers to
'{i}'while the check is for'{gpu_id}'. - Using the generic
Exceptionis not ideal. Raising a more specific exception likeValueErroris better practice as it allows for more granular error handling if this script is ever imported as a module.
I've suggested changes to address both points.
| if "{gpu_id}" not in args.lock_path_pattern: | |
| raise Exception("ERROR: --lock-path-pattern must contain '{i}' placeholder.") | |
| if not args.cmd and not args.print_only: | |
| raise Exception("ERROR: missing command to run. Use -- before command.") | |
| if "{gpu_id}" not in args.lock_path_pattern: | |
| raise ValueError("ERROR: --lock-path-pattern must contain '{gpu_id}' placeholder.") | |
| if not args.cmd and not args.print_only: | |
| raise ValueError("ERROR: missing command to run. Use -- before command.") |
| start = time.time() | ||
| _ensure_lock_files(path_pattern, total_gpus) | ||
| while True: | ||
| fd_locks: List = [] |
There was a problem hiding this comment.
The type hint for fd_locks is List, which is generic. It should be List[FdLock] to improve readability and allow static analysis tools to catch potential bugs. Since FdLock is defined later in the file, you'll need to use a string forward reference.
| fd_locks: List = [] | |
| fd_locks: List["FdLock"] = [] |
|
Thanks! |
|
you are welcome! |
2 participants
Motivation
not tested yet
Modifications
Related Issues
Accuracy Test
Benchmark & Profiling
Checklist