awslabs / awslabs/agentcore-rl-toolkit

RolloutFuture.result_async can discard a fetched result when session cleanup times out

Open
#137 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.