fix(pt): improve OOM detection

[njzjz]

See #4594 for more details.

Summary by CodeRabbit

• Bug Fixes
  • Strengthened error handling for memory-related issues to enhance reliability by better detecting and managing low-memory conditions.

PR Overview

This PR improves the out-of-memory (OOM) detection logic by extending error checks to cover additional CUDA error messages.

• Adds a new check for the error message "CUDA error: out of memory".
• Updates the conditional structure to include a direct check for torch.cuda.OutOfMemoryError.

Reviewed Changes

File	Description
deepmd/pt/utils/auto_batch_size.py	Enhanced error detection for OOM with an added error check

Copilot reviewed 1 out of 1 changed files in this pull request and generated no comments.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes update the is_oom_error method within the AutoBatchSize class in the deepmd/pt/utils/auto_batch_size.py file. The method now detects out-of-memory errors by additionally checking if the exception is an instance of torch.cuda.OutOfMemoryError and by looking for an extra error message ("CUDA error: out of memory"). The original checks using RuntimeError remain, and the method still calls torch.cuda.empty_cache() when an OOM error is detected.

Changes

File	Change Summary
deepmd/pt/utils/auto_batch_size.py	Enhanced the is_oom_error method by adding a condition for torch.cuda.OutOfMemoryError and an additional error message check.

Sequence Diagram(s)

sequenceDiagram
    participant Caller
    participant AutoBatchSize
    participant CUDA

    Caller->>AutoBatchSize: Call is_oom_error(exception)
    Note right of AutoBatchSize: Check if exception is a RuntimeError containing "CUDA error: out of memory"\nor an instance of torch.cuda.OutOfMemoryError
    alt OOM Error Detected
        AutoBatchSize->>CUDA: torch.cuda.empty_cache()
        AutoBatchSize-->>Caller: Return True
    else Not OOM Error
        AutoBatchSize-->>Caller: Return False
    end

Loading

Possibly related PRs

• fix(pt): improve out-of-memory capture #3857: Modifies the is_oom_error method similarly to broaden OOM error detection by adding extra conditions.
• fix(pt): improve out-of-memory handling #3836: Also enhances the is_oom_error function but targets different error scenarios for OOM handling.

Suggested labels

Python

Suggested reviewers

• iProzd

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c2843b7 and d1885dd.

📒 Files selected for processing (1)

• deepmd/pt/utils/auto_batch_size.py (1 hunks)

🔇 Additional comments (1)

deepmd/pt/utils/auto_batch_size.py (1)

52-61: Enhanced OOM detection looks good!

The changes effectively improve OOM detection by:

1. Adding support for the "CUDA error: out of memory" message pattern
2. Directly checking for torch.cuda.OutOfMemoryError instances

This implementation properly addresses the issue referenced in #4594 and makes the error detection more robust by covering additional OOM error formats.

✨ Finishing Touches

• 📝 Generate Docstrings

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 1 line in your changes missing coverage. Please review.

Project coverage is 84.77%. Comparing base (c2843b7) to head (d1885dd).
Report is 1 commits behind head on devel.

Files with missing lines	Patch %	Lines
deepmd/pt/utils/auto_batch_size.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##            devel    #4638      +/-   ##
==========================================
- Coverage   84.78%   84.77%   -0.01%     
==========================================
  Files         688      688              
  Lines       66091    66090       -1     
  Branches     3539     3538       -1     
==========================================
- Hits        56033    56031       -2     
+ Misses       8918     8917       -1     
- Partials     1140     1142       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.