Repository navigation
feat(episode): record episodes in the run index - #457
polakamtejas wants to merge 1 commit into
Conversation
An episode called Task.run directly, so the run index never saw it and the explorer's run lists, which read only the index, never showed it: training episodes were reachable only by instance id. An episode now records itself as a runner records the runs it accepts (runner `episode`): RUNNING before Task.run, then COMPLETED, FAILED with the error, or CANCELED with the abandon reason once its teardown is done. `EpisodeSpec.start(run_group_id=...)` files it in a run group, so a trainer can show a group of rollouts together. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| if failure is not None: | ||
| return RunStatus.FAILED, f"{type(failure).__name__}: {failure}" | ||
| return RunStatus.COMPLETED, None |
There was a problem hiding this comment.
_outcome records COMPLETED when a step fails with fail_task_on_error=False. In that case, Task.run() returns without raising and leaves the failure in failed_steps. The explorer therefore counts a failed episode as completed.
After the cancellation checks, use run_outcome(self._context, failure) so episodes follow the same rule as other runners. Add a test for a failed step that allows the task to continue.
Knowledge Base Used: Task orchestration and execution lifecycle
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/task/episode/run.py
Line: 287-289
Comment:
**Failed episodes look successful**
`_outcome` records `COMPLETED` when a step fails with `fail_task_on_error=False`. In that case, `Task.run()` returns without raising and leaves the failure in `failed_steps`. The explorer therefore counts a failed episode as completed.
After the cancellation checks, use `run_outcome(self._context, failure)` so episodes follow the same rule as other runners. Add a test for a failed step that allows the task to continue.
**Knowledge Base Used:** [Task orchestration and execution lifecycle](https://app.greptile.com/-/switch-tenant?tenant=scale-ai&next=%2Fcustom-context%2Fknowledge-base%2Fscaleapi%2Fagentenv-framework%2F-%2Fdocs%2Ftask-orchestration.md)
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
| run_store.insert_run(RunRecord( | ||
| run_id=_run_id(instance_id), runner="episode", task_id=task.id, task_version=task.version, | ||
| instance_id=instance_id, status=RunStatus.RUNNING, started_at_utc=now, |
There was a problem hiding this comment.
The newly listed episodes get a Cancel button, but the explorer sends their run_id to get_runner().cancel(). With LocalRunner, cancellation always returns false because episodes never enter its _inflight map. The episode keeps running, and the page eventually asks the user to retry.
Either route cancellation to the episode's owner or hide cancellation for these records.
Knowledge Base Used: Runtime explorer APIs and user interface
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/task/episode/run.py
Line: 398-400
Comment:
**Cancel cannot stop episodes**
The newly listed episodes get a Cancel button, but the explorer sends their `run_id` to `get_runner().cancel()`. With `LocalRunner`, cancellation always returns false because episodes never enter its `_inflight` map. The episode keeps running, and the page eventually asks the user to retry.
Either route cancellation to the episode's owner or hide cancellation for these records.
**Knowledge Base Used:** [Runtime explorer APIs and user interface](https://app.greptile.com/-/switch-tenant?tenant=scale-ai&next=%2Fcustom-context%2Fknowledge-base%2Fscaleapi%2Fagentenv-framework%2F-%2Fdocs%2Fruntime-explorer.md)
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.| resolves.""" | ||
| running.set() | ||
| with namespace_routing(): | ||
| _record_run(self._task, self.id, self._context.metadata) # before any await: start() returns it listed |
There was a problem hiding this comment.
_record_run runs before the cleanup try. If the store raises during insertion, _drive exits without removing the episode from _live. Since running was already set, start() still returns an episode even though its task never started. Repeated failed starts retain more unused episodes in memory.
Protect registration with cleanup and pass the insertion error back through start().
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/task/episode/run.py
Line: 257
Comment:
**Failed starts retain episodes**
`_record_run` runs before the cleanup `try`. If the store raises during insertion, `_drive` exits without removing the episode from `_live`. Since `running` was already set, `start()` still returns an episode even though its task never started. Repeated failed starts retain more unused episodes in memory.
Protect registration with cleanup and pass the insertion error back through `start()`.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
What
An episode now records itself in the run index, as a runner records the runs it accepts:
Task.run: aRunRecordwith runnerepisode, statusRUNNING, the task, its version and the instance id. It's written before the first await, sostart()still returns with the run listed.mark_terminalwithCOMPLETED,FAILEDand the error, orCANCELEDand the abandon reason.EpisodeSpec.start(run_group_id=...): files the run in that run group (overrides.metadata.run_group_id), the key the explorer groups a task's runs by. Without it, each run is its own group.Why
The explorer's run lists read only the run index (
agent_env.runner.store). An episode calledTask.rundirectly, so its runs never appeared there: training episodes were reachable only by instance id. That holds for every trainer that plays through the Episode API, the OpenEnv-based helpers included. With a run group, a trainer can show a group of rollouts together, with their mean score.The terminal status is written after teardown, as the cancelled record already is. A scored episode still answers
result(teardown=False)at once (#393), and its run turnsCOMPLETEDonce its teardown is done.Known limit
An episode in a trainer process that dies stays
RUNNINGin the index.LocalRunnermarks its own interrupted runs failed on restart, and nothing does that for episodes yet.Tests
run_test.py::test_an_episode_is_in_the_run_index_running_then_ended_in_the_group_it_was_started_in. It checks running with its group, then completed, cancelled with the reason, and failed with the error.tst/unit/task,runner,explorer,openenvandtrainingpass, 532 of 532, on five runs. One run right afteruv synchad a single failure that didn't recur, likely a timing flake; the log didn't keep which test.run_group_id.🤖 Generated with Claude Code
Fix the incorrect completed status and unsupported explorer cancellation before merging.
Fix with agent prompt
Summary
Episodes now appear in the run index and accept an optional
run_group_id.Acknowledged limit:
polakamtejasstates that episodes left by a dead trainer process remain running because episode recovery is deferred.Diagram
sequenceDiagram participant Caller participant Episode participant Index as Run index participant Task participant Explorer Caller->>Episode: start(run_group_id) Episode->>Index: Insert RUNNING record Episode->>Task: run() Episode-->>Caller: Episode handle Task-->>Episode: Finish or raise Episode->>Task: Tear down resources Episode->>Index: Write terminal status Explorer->>Index: Read runs and groupsReviews (1) · Last reviewed commit: "feat(episode): record episodes in the ru..." · Reviewed by Greptile