awslabs / awslabs/agentcore-rl-toolkit
RolloutFuture.result_async can discard a fetched result when session cleanup times out
- Dominant language
- Python
- Stars
- 57
- Forks
- 12
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 21
Description
## Problem
`RolloutFuture.result_async(timeout=...)` can raise `TimeoutError` even after it has successfully fetched and cached the result from S3.
In `src/agentcore_rl_toolkit/client.py`, `_async_poll()` performs these operations in order:
```python
self._result = await asyncio.to_thread(self._fetch_result)
await self.cancel_async() # Calls StopRuntimeSession.
return self._result
```
`result_async()` wraps the entire sequence in `asyncio.wait_for()`. If session cleanup exceeds the remaining deadline, the caller receives `TimeoutError` instead of the already-fetched result. The result remains in `future._result`, but the current call does not return it.
The verl integration treats this exception as a rollout failure and assigns reward zero, so a successful agent result can be counted as a failure.
## Minimal reproduction
This uses the existing `RolloutFuture` with stub AWS clients; no AWS requests are made. S3 returns a successful result immediately, while session cleanup takes longer than the deadline.
```python
import asyncio
import io
import json
import time
from agentcore_rl_toolkit.client import RolloutFuture
class ReadyS3:
def head_object(self, **kwargs):
return {}
def get_object(self, **kwargs):
result = {"status_code": 200, "rewards": 1.0}
return {"Body": io.BytesIO(json.dumps(result).encode())}
class SlowCleanup:
def stop_runtime_session(self, **kwargs):
time.sleep(0.5)
return {"statusCode": 200}
async def main():
future = RolloutFuture(
s3_client=ReadyS3(),
s3_bucket="unused",
result_key="unused",
session_id="example-session",
agent_runtime_arn="unused",
agentcore_client=SlowCleanup(),
)
try:
await future.result_async(timeout=0.2)
except TimeoutError:
print("Timed out despite already fetching:", future._result)
assert future._result == {"status_code": 200, "rewards": 1.0}
else:
raise AssertionError("Expected the current implementation to time out")
asyncio.run(main())
```
Observed output:
```text
Timed out despite already fetching: {'status_code': 200, 'rewards': 1.0}
```
## Expected behavior
Once the result has been fetched successfully, session cleanup timing should not turn that result into a retrieval failure. Keep automatic session cleanup, but separate its timeout handling from the result deadline.
Regression coverage should verify both that slow cleanup preserves a fetched result and that waiting for an unavailable result still times out.
## Investigation context
This bug was reproduced while investigating a one-point difference between S3 rewards and logged MigrationBench validation rewards. The client behavior above is confirmed; whether it caused that particular discrepancy remains unconfirmed because the necessary per-session client records were not retained.
Contributor guide
Research direction
Start in src/agentcore_rl_toolkit/client.py with RolloutFuture.result_async() and _async_poll(), then use the provided ReadyS3 and SlowCleanup reproduction as the behavioral guide. Add regression coverage for preserving a fetched result during slow cleanup and for timing out when the result remains unavailable; done means both cases behave as expected.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- backend, cloud
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100