UN-2801 [FIX] Fix file batch distribution algorithm to respect configured batch limits#1530
Merged
jaseemjaskp merged 1 commit intoSep 14, 2025
Conversation
…ured batch limits Fixed an issue where the file batching algorithm was creating fewer batches than configured (16 instead of 30), causing uneven distribution of files across parallel workers. The new implementation ensures files are distributed as evenly as possible across the target number of batches. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
Summary by CodeRabbit
WalkthroughReplaced fixed-size batching with an even-distribution algorithm in get_file_batches, handling zero files, ensuring non-empty batches, and updating the docstring. Removed an unused math import. No public signatures changed. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. 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 |
Contributor
|
|
muhammad-ali-e
requested review from
chandrasekharan-zipstack,
hari-kuriakose,
pk-zipstack and
vishnuszipstack
vishnuszipstack
approved these changes
Sep 14, 2025
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 (1)
backend/workflow_manager/workflow_v2/workflow_helper.py (1)
88-99: Return type annotation is incorrect (JSON dict vs FileHash).You convert FileHash to JSON before batching, so the function returns list[list[tuple[str, dict]]], not FileHash. Fix the annotation and clarify in the docstring.
- def get_file_batches( - cls, input_files: dict[str, FileHash] - ) -> list[list[tuple[str, FileHash]]]: - """Split input files into batches for parallel processing. - Distributes files as evenly as possible across the target number of batches. + def get_file_batches( + cls, input_files: dict[str, FileHash] + ) -> list[list[tuple[str, dict[str, Any]]]]: + """Split input files into batches for parallel processing. + Distributes files as evenly as possible across the target number of batches. + Returns batches of (file_name, serialized FileHash as dict).
🧹 Nitpick comments (1)
backend/workflow_manager/workflow_v2/workflow_helper.py (1)
128-133: Remove redundant guards to simplify loop.start_index < len(file_items) and the non-empty check are redundant once batch_size is computed as above and num_files > 0. Simplify for readability.
- if start_index < len(file_items): - end_index = min(start_index + batch_size, len(file_items)) - batch = file_items[start_index:end_index] - if batch: # Only add non-empty batches - batches.append(batch) - start_index = end_index + end_index = start_index + batch_size + batches.append(file_items[start_index:end_index]) + start_index = end_index
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (1)
backend/workflow_manager/workflow_v2/workflow_helper.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (2)
backend/workflow_manager/workflow_v2/workflow_helper.py (2)
124-127: Even distribution logic looks correct.Remainder-first allocation is the right approach to keep size deltas ≤ 1. Nice.
91-93: Summary/docstring mismatch: zero-files case not actually handled here.AI summary mentions explicit zero-input handling returning [], but the method would currently divide by zero without the guard suggested above.
Comment on lines
106
to
+119
| BATCH_SIZE = Configuration.get_value_by_organization( | ||
| config_key=ConfigKey.MAX_PARALLEL_FILE_BATCHES, organization=organization | ||
| ) # Max number of batches | ||
|
|
||
| file_items = list(json_serializable_files.items()) | ||
|
|
||
| # Calculate how many items per batch | ||
| # Calculate distribution | ||
| num_files = len(file_items) | ||
| # Target number of batches (can't exceed number of files) | ||
| num_batches = min(BATCH_SIZE, num_files) | ||
| items_per_batch = math.ceil(num_files / num_batches) | ||
|
|
||
| # Split into batches | ||
| # Distribute files as evenly as possible | ||
| base_items_per_batch = num_files // num_batches | ||
| remainder = num_files % num_batches | ||
|
|
Contributor
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Prevent ZeroDivisionError; cast and clamp MAX_PARALLEL_FILE_BATCHES; handle empty input.
If input_files is empty or the config resolves to 0/None/non‑int, num_batches becomes 0 and base_items_per_batch = num_files // num_batches will raise ZeroDivisionError. Also, Configuration.get_value_by_organization may return strings; min(str, int) will TypeError. Guard and normalize.
Apply:
- BATCH_SIZE = Configuration.get_value_by_organization(
- config_key=ConfigKey.MAX_PARALLEL_FILE_BATCHES, organization=organization
- ) # Max number of batches
+ raw_batches = Configuration.get_value_by_organization(
+ config_key=ConfigKey.MAX_PARALLEL_FILE_BATCHES, organization=organization
+ ) # Max number of batches (may be str/None)
+ try:
+ BATCH_SIZE = max(1, int(raw_batches))
+ except (TypeError, ValueError):
+ BATCH_SIZE = 1
file_items = list(json_serializable_files.items())
- # Calculate distribution
- num_files = len(file_items)
+ # Calculate distribution
+ num_files = len(file_items)
+ if num_files == 0:
+ return []
# Target number of batches (can't exceed number of files)
num_batches = min(BATCH_SIZE, num_files)Quick checks:
🏁 Script executed:
#!/bin/bash
# 1) Find any other callers that might pass empty input
rg -n -C2 '\bget_file_batches\s*\('
# 2) Sanity-test the distribution property on representative values (no repo import)
python - <<'PY'
def sizes(n, b):
from math import ceil
if n == 0: return []
b = max(1, int(b))
nb = min(b, n)
base, r = divmod(n, nb)
return [base + (1 if i < r else 0) for i in range(nb)]
for n in [0,1,2,3,10,31,60]:
for b in [0,1,2,3,30,64,"30","not-int",None]:
try:
s = sizes(n,b)
ok = sum(s)==n and (max(s, default=0)-min(s, default=0) <= 1)
print(f"n={n}, b={b!r} -> {s} | ok={ok}")
except Exception as e:
print(f"n={n}, b={b!r} -> EXC {e}")
PYLength of output: 1041
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== get_file_batches definition (first 300 lines) ==="
nl -ba backend/workflow_manager/workflow_v2/workflow_helper.py | sed -n '1,300p'
echo
echo "=== callers of get_file_batches ==="
rg -n -C2 '\bget_file_batches\s*\(' || true
echo
echo "=== occurrences of MAX_PARALLEL_FILE_BATCHES ==="
rg -n 'MAX_PARALLEL_FILE_BATCHES' -S || true
echo
echo "=== sanity python test ==="
python - <<'PY'
def sizes(n,b):
if n==0:
return []
try:
b_int = max(1, int(b))
except Exception:
b_int = 1
nb = min(b_int, n)
base, r = divmod(n, nb)
return [base + (1 if i < r else 0) for i in range(nb)]
tests_n = [0,1,2,3,10,31,60]
tests_b = [0,1,2,3,30,64,"30","not-int",None]
for n in tests_n:
for b in tests_b:
try:
s = sizes(n,b)
ok = sum(s)==n and (max(s, default=0)-min(s, default=0) <= 1)
print(f"n={n}, b={b!r} -> {s} | ok={ok}")
except Exception as e:
print(f"n={n}, b={b!r} -> EXC {e}")
PYLength of output: 94
🏁 Script executed:
#!/bin/bash
# Print relevant portions of the file and search for callers / config keys
echo "=== backend/workflow_manager/workflow_v2/workflow_helper.py (lines 1-300) ==="
sed -n '1,300p' backend/workflow_manager/workflow_v2/workflow_helper.py || true
echo
echo "=== callers of get_file_batches ==="
rg -n -C2 '\bget_file_batches\s*\(' || true
echo
echo "=== occurrences of MAX_PARALLEL_FILE_BATCHES ==="
rg -n 'MAX_PARALLEL_FILE_BATCHES' -S || true
echo
echo "=== quick sanity python test for distribution logic ==="
python - <<'PY'
def sizes(n,b):
if n==0:
return []
try:
b_int = max(1, int(b))
except Exception:
b_int = 1
nb = min(b_int, n)
base, r = divmod(n, nb)
return [base + (1 if i < r else 0) for i in range(nb)]
tests_n = [0,1,2,3,10,31,60]
tests_b = [0,1,2,3,30,64,"30","not-int",None]
for n in tests_n:
for b in tests_b:
try:
s = sizes(n,b)
ok = sum(s)==n and (max(s, default=0)-min(s, default=0) <= 1)
print(f"n={n}, b={b!r} -> {s} | ok={ok}")
except Exception as e:
print(f"n={n}, b={b!r} -> EXC {e}")
PYLength of output: 13726
Cast & clamp MAX_PARALLEL_FILE_BATCHES; prevent TypeError/ZeroDivisionError
- process_input_files already returns early when input_files is empty, so num_files == 0 is not reachable from current callers.
- Mandatory: normalize the configuration value before use — replace the direct read with a safe cast/fallback to ensure an int >= 1 (avoid TypeError in min() and possible division by zero). Example change:
raw_batches = Configuration.get_value_by_organization(
config_key=ConfigKey.MAX_PARALLEL_FILE_BATCHES, organization=organization
)
try:
BATCH_SIZE = max(1, int(raw_batches))
except (TypeError, ValueError):
BATCH_SIZE = 1- Optional hardening: keep
if num_files == 0: return []inside get_file_batches for defense-in-depth.
Location: backend/workflow_manager/workflow_v2/workflow_helper.py (around lines 106–119).
🤖 Prompt for AI Agents
In backend/workflow_manager/workflow_v2/workflow_helper.py around lines 106 to
119, the code reads ConfigKey.MAX_PARALLEL_FILE_BATCHES directly which can be
None or non-int and lead to TypeError/ZeroDivisionError; replace the direct read
with a safe normalization: fetch the raw config, attempt to cast to int inside a
try/except, clamp to at least 1 (e.g., BATCH_SIZE = max(1, int(raw))) and
fallback to 1 on TypeError/ValueError, then proceed using that BATCH_SIZE;
optionally also add a defensive early return if num_files == 0 to guard against
future callers.
jaseemjaskp
approved these changes
Sep 14, 2025
jaseemjaskp
deleted the
fix/UN-2801-FIX_improve_file_batch_distribution_for_parallel_processing
branch
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
WorkflowHelper.get_file_batches()methodWhy
math.ceil(num_files / num_batches)was causing uneven distributionMAX_PARALLEL_FILE_BATCHESsetting properlyHow
Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
Database Migrations
Env Config
Relevant Docs
Related Issues or PRs
Dependencies Versions
Notes on Testing
Screenshots
N/A
Checklist
I have read and understood the Contribution Guidelines.
🤖 Generated with Claude Code