[Scheduler] improve psweep LB #1492
Conversation
Summary of ChangesHello @tdavidcl, 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 significantly enhances the scheduler's load balancing capabilities by integrating a new parallel sweep algorithm. It also refines the existing load balancing mechanism to intelligently compare and select the most effective strategy among several options, aiming to improve overall system performance and resource utilization through dynamic workload allocation. Highlights
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 an improved load balancing strategy, parallel_sweep2, which uses rounding for better load distribution. The main load_balance function has been refactored to execute multiple strategies and select the best-performing one. This is a solid improvement. My review focuses on some critical correctness issues in the new strategy, such as handling of empty inputs and calculation errors, which could lead to crashes or incorrect behavior. I have also provided suggestions to improve code maintainability and readability by reducing code duplication and using standard C++ library features.
|
|
||
| template<class Torder, class Tweight> | ||
| inline std::vector<i32> lb_startegy_parallel_sweep2( | ||
| const std::vector<TileWithLoad<Torder, Tweight>> &lb_vector, i32 wsize) { |
There was a problem hiding this comment.
This function can crash if lb_vector is empty. An empty lb_vector will result in res being empty, and the access res[res.size() - 1] on line 192 will be out-of-bounds. Please add a check at the beginning of the function to handle this case.
For example:
if (lb_vector.empty()) {
return {};
}
| accum += cur_val; | ||
| } | ||
|
|
||
| double target_datacnt = double(res[res.size() - 1].accumulated_load_value) / wsize; |
There was a problem hiding this comment.
The calculation of target_datacnt is incorrect. It uses res[res.size() - 1].accumulated_load_value, which is the result of an exclusive scan, meaning it's the sum of all loads except for the last one. To get the total load, you should use the accum variable, which holds the total sum after the loop on lines 186-190 completes. I've also added a check for wsize > 0 to prevent a potential division by zero.
| double target_datacnt = double(res[res.size() - 1].accumulated_load_value) / wsize; | |
| double target_datacnt = (wsize > 0) ? (double(accum) / wsize) : 0.0; |
| u64 cur_val = tile.load_value; | ||
| tile.accumulated_load_value = accum; | ||
| accum += cur_val; |
There was a problem hiding this comment.
The local variable cur_val is redundant. You can use tile.load_value directly in the accumulation to make the code slightly simpler and more direct.
| u64 cur_val = tile.load_value; | |
| tile.accumulated_load_value = accum; | |
| accum += cur_val; | |
| tile.accumulated_load_value = accum; | |
| accum += tile.load_value; |
| for (LBTileResult t : res) { | ||
| shamlog_debug_ln( | ||
| "HilbertLoadBalance", | ||
| t.ordering_val, | ||
| t.load_value, | ||
| t.accumulated_load_value, | ||
| t.index, | ||
| (target_datacnt == 0) | ||
| ? 0 | ||
| : sycl::clamp( | ||
| i32((t.accumulated_load_value / target_datacnt) + 0.5), | ||
| 0, | ||
| i32(wsize) - 1), | ||
| (target_datacnt == 0) ? 0 | ||
| : ((t.accumulated_load_value / target_datacnt) + 0.5)); | ||
| } |
There was a problem hiding this comment.
The debug logging loop re-calculates the new_owner value, which has already been computed and stored in tile.new_owner. This is code duplication and can lead to inconsistencies if the calculation logic changes in one place but not the other. Additionally, the loop for (LBTileResult t : res) creates a copy of each element in res, which is inefficient. You can simplify the logging by using t.new_owner and iterate by const reference to avoid copies.
for (const LBTileResult &t : res) {
shamlog_debug_ln(
"HilbertLoadBalance",
t.ordering_val,
t.load_value,
t.accumulated_load_value,
t.index,
t.new_owner,
(target_datacnt == 0)
? 0.0
: ((t.accumulated_load_value / target_datacnt) + 0.5));
}| { | ||
| auto tmpres = details::lb_startegy_parallel_sweep(lb_vector, world_size); | ||
| auto metric = details::compute_LB_metric(lb_vector, tmpres, world_size); | ||
| results.push_back(LBResult{tmpres, metric, "parallel sweep"}); | ||
| } | ||
|
|
||
| auto tmpres_2 = details::lb_startegy_roundrobin(lb_vector, world_size); | ||
| auto metric_rrobin = details::compute_LB_metric(lb_vector, tmpres_2, world_size); | ||
| { | ||
| auto tmpres = details::lb_startegy_roundrobin(lb_vector, world_size); | ||
| auto metric = details::compute_LB_metric(lb_vector, tmpres, world_size); | ||
| results.push_back(LBResult{tmpres, metric, "round robin"}); | ||
| } | ||
|
|
||
| { | ||
| auto tmpres = details::lb_startegy_parallel_sweep2(lb_vector, world_size); | ||
| auto metric = details::compute_LB_metric(lb_vector, tmpres, world_size); | ||
| results.push_back(LBResult{tmpres, metric, "parallel sweep 2"}); | ||
| } |
There was a problem hiding this comment.
There is repeated code for running each load balancing strategy and storing the results. This can be refactored into a helper lambda to improve readability and maintainability, and to avoid potential vector copies when storing results.
using strategy_fn = std::vector<i32>(
const std::vector<TileWithLoad<Torder, Tweight>> &, i32);
auto run_and_store_strategy = [&](strategy_fn *strategy, const std::string &name) {
auto ranks = strategy(lb_vector, world_size);
auto metric = details::compute_LB_metric(lb_vector, ranks, world_size);
results.emplace_back(LBResult{std::move(ranks), metric, name});
};
run_and_store_strategy(&details::lb_startegy_parallel_sweep, "parallel sweep");
run_and_store_strategy(&details::lb_startegy_roundrobin, "round robin");
run_and_store_strategy(&details::lb_startegy_parallel_sweep2, "parallel sweep 2");| for (const auto &result : results) { | ||
| if (shamcomm::world_rank() == 0) { | ||
| logger::info_ln( | ||
| "LoadBalance", | ||
| " - strategy \"", | ||
| result.strategy, | ||
| "\" : max =", | ||
| result.metric.max, | ||
| "min =", | ||
| result.metric.min); | ||
| } | ||
| if (result.metric.max < metric_best.max) { | ||
| metric_best = result.metric; | ||
| res_best = result.ranks; | ||
| strategy_best = result.strategy; | ||
| } | ||
| } |
There was a problem hiding this comment.
This loop for finding the best strategy can be made more concise and idiomatic by using std::min_element. This also separates the concern of logging from the logic of finding the best result, improving readability. The current implementation also copies the ranks vector every time a better strategy is found, which can be inefficient. Using std::min_element will find the best result, and then you can copy the data just once.
for (const auto &result : results) {
if (shamcomm::world_rank() == 0) {
logger::info_ln(
"LoadBalance",
" - strategy \"",
result.strategy,
"\" : max =",
result.metric.max,
"min =",
result.metric.min);
}
}
auto best_it = std::min_element(
results.begin(),
results.end(),
[](const LBResult &a, const LBResult &b) { return a.metric.max < b.metric.max; });
if (best_it != results.end()) {
res_best = best_it->ranks;
strategy_best = best_it->strategy;
}
|
Thanks @tdavidcl for opening this PR! You can do multiple things directly here: Once the workflow completes a message will appear displaying informations related to the run. Also the PR gets automatically reviewed by gemini, you can: |
Workflow reportworkflow report corresponding to commit e18a67c Light CI is enabled. This will only run the basic tests and not the full tests. Pre-commit check reportPre-commit check: ✅ Test pipeline can run. |
1 participant
No description provided.