Fix the loop exit condition in the agenerator by yulangz · Pull Request #858 · areal-project/AReaL

yulangz · 2026-01-27T03:33:25Z

Description

This test verifies the fix for the bug where the generation loop would
prematurely exit when stop_reason='abort' because:

BUG (before fix):

Loop condition: len(accumulated_output_tokens) < gconfig.max_new_tokens
Problem: gconfig.max_new_tokens gets decremented in the loop
Example: original_max=16384, after 12000 tokens -> gconfig.max_new_tokens=4384
Condition: 12000 < 4384 = False -> loop exits prematurely with abort

FIX:

Save original max_new_tokens before the loop
Use original_max_new_tokens in the condition instead of gconfig.max_new_tokens
This ensures the loop continues until we actually reach the generation limit

Related Issue

Fixes #(issue)

Type of Change

Bug fix (non-breaking change that fixes an issue)
New feature (non-breaking change that adds functionality)
Breaking change (fix or feature that would cause existing functionality to not
work as expected)
Documentation update
Code refactoring (no functional changes)
Performance improvement
Test coverage improvement

Checklist

I have read the Contributing Guide
I have run formatting tools (pre-commit or manual)
I have run relevant unit tests and they pass
I have added tests for new functionality
I have updated documentation if needed
My branch is up to date with main
This PR introduces breaking changes (if yes, fill out details below)
If this PR changes documentation, I have built and previewed it locally with
jb build docs
No critical issues raised by AI reviewers (/gemini review)

Breaking Change Details (if applicable):

Additional Context

Need help? Check the Contributing Guide or ask in
GitHub Discussions!

Add comprehensive tests to verify the fix for the bug where the generation loop would prematurely exit when stop_reason='abort'. The bug was caused by using gconfig.max_new_tokens (which gets decremented in the loop) instead of the original value in the loop condition. Tests include: - test_abort_should_continue_loop_with_fix: verify fix works correctly - test_abort_prematurely_exits_without_fix: reproduce the original bug - test_abort_multiple_times_continues_with_fix: multiple abort scenarios - test_abort_reaches_length_limit_with_fix: length limit handling - Additional tests for tool_calls, length stop reasons, and gconfig behavior Reference: commit 445dca7216e66c24d7a5b47a272761ea81749f8d Co-Authored-By: Claude (claude-opus-4-5) <noreply@anthropic.com>

Rewrite tests to directly use RemoteInfEngine instance instead of simulating the loop logic. This provides more realistic testing by: - Creating actual RemoteInfEngine with MockBackend - Mocking arequest_with_retry to avoid needing a real SGLang server - Testing the actual agenerate() method behavior The tests now verify that the fix works correctly in the real codebase, not just in isolated simulation. Co-Authored-By: Claude (claude-opus-4-5) <noreply@anthropic.com>

gemini-code-assist · 2026-01-27T03:33:39Z

Summary of Changes

Hello @yulangz, I'm Gemini Code Assist¹! 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 addresses a critical issue in the remote inference engine's token generation process. The previous implementation could lead to incomplete token sequences when an 'abort' signal was encountered, as the loop's termination condition was inadvertently modified during execution. The changes ensure that the generation process reliably reaches its intended token limit, significantly improving the robustness of the inference engine.

Highlights

Bug Fix: Premature Loop Exit: Resolved a bug in the agenerate function where the token generation loop would prematurely exit when stop_reason='abort'. This occurred because the loop condition incorrectly used a decremented gconfig.max_new_tokens value.
Corrected Loop Condition: The fix involves storing the original max_new_tokens value before the generation loop begins and using this immutable value in the loop's exit condition. This ensures the loop continues until the true generation limit is met, even if intermediate 'abort' signals are received.
New Test Coverage: A new comprehensive test file, test_remote_inf_engine_abort_loop.py, has been added. It includes tests to verify the corrected behavior in various scenarios, as well as simulations to clearly document the original buggy behavior and the impact of the fix.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request correctly fixes a bug in the agenerate method where the generation loop would exit prematurely due to a mutable loop condition variable. The fix involves caching the original max_new_tokens value and using it for the loop termination check. The change is accompanied by a comprehensive new test suite that validates the fix and documents the original bug, which is excellent. My only suggestion is to improve the variable naming for better code clarity.

Copilot

Pull request overview

This PR fixes a bug in RemoteInfEngine.agenerate where the generation loop could exit prematurely when stop_reason="abort" due to using a decremented max_new_tokens in the loop condition, and adds tests to lock in the correct behavior.

Changes:

Adjusts the agenerate loop condition to use the original max_new_tokens value captured before the loop, instead of the mutable gconfig.max_new_tokens.
Adds a dedicated test suite exercising RemoteInfEngine.agenerate behavior for various stop_reason values and verifying the abort/length-limit interaction.
Adds simulation tests documenting the original buggy behavior and confirming that ModelRequest.copy() and gconfig.max_new_tokens updates behave as expected.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
`areal/core/remote_inf_engine.py`	Fixes the generation loop exit condition by capturing `gconfig.max_new_tokens` into `ori_max_new_tokens` before the loop and using that constant in the `while` condition, preventing premature exit when `stop_reason="abort"`.
`areal/tests/test_remote_inf_engine_abort_loop.py`	Introduces a comprehensive test suite for the abort-loop behavior of `RemoteInfEngine.agenerate`, including real-engine tests against a mocked backend, simulations of the pre-fix bug, and tests for `gconfig` mutation and `ModelRequest.copy()` semantics.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-01-27T03:36:57Z

+        - Each call returns 80 tokens with abort (staying within remaining limit)
+        - After 3 calls: 240 tokens, remaining = 10
+        - 4th call: returns 80 tokens with abort, but 240 < 250, so we enter loop
+        - After 4th call: 320 tokens >= 250, exit loop


In test_abort_reaches_length_limit_with_fix, the scenario described in the docstring (4th call returning 80 tokens and total reaching 320) no longer matches the actual test setup, where the 4th response returns 10 tokens and total tokens are 250 as noted in the inline comment. Please update the docstring steps and totals to reflect the actual response sequence so the test documentation is accurate and not confusing.

Suggested change

- Each call returns 80 tokens with abort (staying within remaining limit)

- After 3 calls: 240 tokens, remaining = 10

- 4th call: returns 80 tokens with abort, but 240 < 250, so we enter loop

- After 4th call: 320 tokens >= 250, exit loop

- First 3 calls each return 80 tokens with abort (staying within remaining limit)

- After 3 calls: 240 tokens, remaining = 10

- 4th call: returns 10 tokens with abort, still respecting remaining max_new_tokens

- After 4th call: 250 tokens == max_new_tokens, so we exit loop due to length limit

garrett4wade

I'm surprised that we had this kind of stupid bug. Nice catch and thanks a lot!!

…t#858) * test: add unit tests for abort loop condition fix in RemoteInfEngine Add comprehensive tests to verify the fix for the bug where the generation loop would prematurely exit when stop_reason='abort'. The bug was caused by using gconfig.max_new_tokens (which gets decremented in the loop) instead of the original value in the loop condition. Tests include: - test_abort_should_continue_loop_with_fix: verify fix works correctly - test_abort_prematurely_exits_without_fix: reproduce the original bug - test_abort_multiple_times_continues_with_fix: multiple abort scenarios - test_abort_reaches_length_limit_with_fix: length limit handling - Additional tests for tool_calls, length stop reasons, and gconfig behavior Reference: commit 445dca7216e66c24d7a5b47a272761ea81749f8d Co-Authored-By: Claude (claude-opus-4-5) <noreply@anthropic.com> * refactor: use real RemoteInfEngine with mocked arequest_with_retry Rewrite tests to directly use RemoteInfEngine instance instead of simulating the loop logic. This provides more realistic testing by: - Creating actual RemoteInfEngine with MockBackend - Mocking arequest_with_retry to avoid needing a real SGLang server - Testing the actual agenerate() method behavior The tests now verify that the fix works correctly in the real codebase, not just in isolated simulation. Co-Authored-By: Claude (claude-opus-4-5) <noreply@anthropic.com> --------- Co-authored-by: Claude (claude-opus-4-5) <noreply@anthropic.com>

yulangz and others added 4 commits January 27, 2026 11:25

fix code and ut

e4bf4e0

format

5e8dfe3

Copilot AI review requested due to automatic review settings January 27, 2026 03:33

update doc

8052d22

Copilot started reviewing on behalf of yulangz January 27, 2026 03:33 View session

yulangz added the safe-to-test Ready to run unit-tests in a PR. label Jan 27, 2026

gemini-code-assist Bot reviewed Jan 27, 2026

View reviewed changes

Comment thread areal/core/remote_inf_engine.py

Comment thread areal/core/remote_inf_engine.py

Copilot AI reviewed Jan 27, 2026

View reviewed changes

yulangz temporarily deployed to AReaL-unittests January 27, 2026 04:09 — with GitHub Actions Inactive

garrett4wade reviewed Jan 27, 2026

View reviewed changes

Comment thread areal/tests/test_remote_inf_engine_abort_loop.py Outdated

remove ut

2502b06

garrett4wade merged commit a4d01ec into main Jan 27, 2026
1 check passed

garrett4wade deleted the wht/fix_abort_loop branch January 27, 2026 07:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix the loop exit condition in the agenerator#858

Fix the loop exit condition in the agenerator#858
garrett4wade merged 6 commits into
mainfrom
wht/fix_abort_loop

yulangz commented Jan 27, 2026

Uh oh!

gemini-code-assist Bot commented Jan 27, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jan 27, 2026

Uh oh!

garrett4wade left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

yulangz commented Jan 27, 2026

Description

Related Issue

Type of Change

Checklist

Additional Context

Uh oh!

gemini-code-assist Bot commented Jan 27, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Jan 27, 2026

Choose a reason for hiding this comment

Uh oh!

garrett4wade left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

garrett4wade left a comment •

edited

Loading