notify optimizers scheduler of the budget change #6040

generall · 2025-02-22T18:49:38Z

Fixes #6015 for the case, where there are no running updates.

Summary by CodeRabbit

New Features
- Introduced a unified resource release mechanism that triggers notifications for dynamic resource adjustments.
Refactor
- Enhanced callback handling for improved responsiveness during optimization.
- Streamlined resource management by removing redundant GPU resource allocation logic, resulting in more consistent system behavior.

generall · 2025-02-22T18:51:21Z

lib/segment/src/segment_constructor/segment_builder.rs

-            // If GPU is enabled, release all CPU cores except one.
-            if let Some(_gpu_device) = &gpu_device {
-                if permit.num_cpus > 1 {
-                    permit.release_cpu_count(permit.num_cpus - 1);
-                }
-            }


This behavior was broken and would have created problems with changes of this PR. E.g. we don't want to schedule CPU optimization if GPU is available. More correct implementation should involve dedicated GPU permit.

lib/common/common/src/budget.rs

coderabbitai · 2025-02-24T10:14:00Z

Walkthrough

This pull request refines resource management and optimization processes across multiple modules. In the update handler, the callback signature is modified to support multiple invocations and thread sharing, and terminology is generalized to "Resources." The common budget code consolidates resource release methods into one unified function while adding an optional release callback. Meanwhile, the segment builder removes GPU-specific CPU allocation logic and simplifies permit handling. All changes focus on streamlining control flow and enhancing the responsiveness of resource optimization procedures.

Changes

File(s)	Change Summary
lib/collection/src/update_handler.rs	- Updated `launch_optimization` callback signature from `FnOnce(bool)` to `Fn(bool) + Send + Clone + Sync + 'static` - Changed resource permit acquisition to mutable - Updated comments to replace "CPUs" with "Resources"
lib/common/common/src/budget.rs	- Introduced a unified `release` method for releasing CPU and IO counts - Added an `on_release` field and the `set_on_release` method in `ResourcePermit` - Changed previous release methods (`release_cpu`, `release_io`, etc.) from public to private
lib/segment/src/segment_constructor/segment_builder.rs	- Modified `build` method signature to remove mutable permit parameter - Removed GPU-related CPU core adjustment logic, simplifying the segment construction process

Sequence Diagram(s)

sequenceDiagram
    participant UH as UpdateHandler
    participant RP as ResourcePermit
    participant S as Scheduler

    UH->>RP: Call launch_optimization(callback)
    RP-->>UH: Acquire mutable permit and set on_release callback
    RP->>S: Notify scheduler when resources are released

sequenceDiagram
    participant SB as SegmentBuilder
    participant RP as ResourcePermit
    participant Flag as AtomicBool

    SB->>RP: Receive immutable permit
    SB->>Flag: Check stopped flag
    SB-->>SB: Build segment without GPU-specific CPU management

Poem

I'm a little rabbit, coding in the night,
Hopping through updates with pure delight.
The resources now sing a tune so true,
With callbacks bouncing like morning dew.
In streamlined code, my heart takes flight –
A joyous hop in every byte.

📜 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 08cf86c and 0d6f363.

📒 Files selected for processing (3)

lib/collection/src/update_handler.rs (3 hunks)
lib/common/common/src/budget.rs (6 hunks)
lib/segment/src/segment_constructor/segment_builder.rs (1 hunks)

⏰ Context from checks skipped due to timeout of 90000ms (6)

GitHub Check: test-snapshot-operations-s3-minio
GitHub Check: test-shard-snapshot-api-s3-minio
GitHub Check: test (macos-latest)
GitHub Check: test-consensus
GitHub Check: test (windows-latest)
GitHub Check: test (ubuntu-latest)

🔇 Additional comments (9)

lib/segment/src/segment_constructor/segment_builder.rs (1)

443-446: LGTM! Mutability change is safe.

The change in mutability of the permit parameter is correct as the permit is only used for reading its properties and is later wrapped in an Arc for sharing.

lib/collection/src/update_handler.rs (3)

277-277: LGTM! Callback trait bounds are improved.

The change from FnOnce(bool) to Fn(bool) + Sync is beneficial as it:

Allows multiple invocations of the callback through Fn instead of FnOnce.

Ensures thread safety with the Sync trait.

298-299: LGTM! Comments are more accurate.

The terminology change from "CPUs" to "Resources" better reflects the broader context of resource management.

302-326: LGTM! Resource permit callback is properly set.

The permit is correctly acquired as mutable to set a callback that notifies the scheduler when the resource budget changes, improving the responsiveness of the optimization process.

lib/common/common/src/budget.rs (5)

262-265: LGTM! Well-documented callback field.

The on_release field is properly documented, clearly explaining its purpose for task manager notifications.

289-291: LGTM! Callback setter is properly implemented.

The set_on_release method correctly accepts a callback function with appropriate trait bounds for thread safety and lifetime.

338-347: LGTM! Method visibility is properly restricted.

The release_cpu and release_io methods are correctly made private as they are now internal implementation details.

379-386: LGTM! Resource release is properly consolidated.

The new release method effectively consolidates the resource release functionality and ensures the callback is invoked when resources are released.

389-405: LGTM! Drop implementation ensures proper cleanup.

The Drop implementation correctly:

Takes ownership of the permits to ensure they are dropped.

Invokes the callback one final time.

✨ Finishing Touches

📝 Generate Docstrings (Beta)

🪧 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. (Beta)
@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.

generall · 2025-02-24T10:19:27Z

@coderabbitai review

coderabbitai · 2025-02-24T10:19:35Z

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

* notify optimizers scheduler of the budget change * Zero permit count when releasing CPU or IO --------- Co-authored-by: Tim Visée <tim+github@visee.me>

notify optimizers scheduler of the budget change

020624c

generall commented Feb 22, 2025

View reviewed changes

github-actions bot mentioned this pull request Feb 22, 2025

Flaky test multivector_filtrable_hnsw_test::test_multi_filterable_hnsw::case_5_recommend_eq #5902

Closed

timvisee approved these changes Feb 24, 2025

View reviewed changes

lib/common/common/src/budget.rs Show resolved Hide resolved

lib/common/common/src/budget.rs Show resolved Hide resolved

Zero permit count when releasing CPU or IO

0d6f363

timvisee merged commit e5ee447 into dev Feb 24, 2025
18 checks passed

timvisee deleted the permit-notify-on-release branch February 24, 2025 10:57

timvisee added a commit that referenced this pull request Mar 21, 2025

notify optimizers scheduler of the budget change (#6040)

f6d58b4

* notify optimizers scheduler of the budget change * Zero permit count when releasing CPU or IO --------- Co-authored-by: Tim Visée <tim+github@visee.me>

timvisee mentioned this pull request Mar 21, 2025

Bump version to 1.13.5 #6223

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

notify optimizers scheduler of the budget change #6040

notify optimizers scheduler of the budget change #6040

Uh oh!

generall commented Feb 22, 2025 •

edited by coderabbitai bot

Loading

Uh oh!

generall Feb 22, 2025

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot commented Feb 24, 2025 •

edited

Loading

Chat

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

CodeRabbit Configuration File (`.coderabbit.yaml`)

Documentation and Community

Uh oh!

generall commented Feb 24, 2025

Uh oh!

coderabbitai bot commented Feb 24, 2025

Uh oh!

Uh oh!

Uh oh!

notify optimizers scheduler of the budget change #6040

notify optimizers scheduler of the budget change #6040

Uh oh!

Conversation

generall commented Feb 22, 2025 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

generall Feb 22, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot commented Feb 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Poem

Chat

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

CodeRabbit Configuration File (.coderabbit.yaml)

Documentation and Community

Uh oh!

generall commented Feb 24, 2025

Uh oh!

coderabbitai bot commented Feb 24, 2025

Uh oh!

Uh oh!

Uh oh!

generall commented Feb 22, 2025 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Feb 24, 2025 •

edited

Loading

CodeRabbit Configuration File (`.coderabbit.yaml`)