marin-rolloutengine 0.2.160.dev38062449553__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- marin_rolloutengine-0.2.160.dev38062449553/.gitignore +253 -0
- marin_rolloutengine-0.2.160.dev38062449553/PKG-INFO +8 -0
- marin_rolloutengine-0.2.160.dev38062449553/README.md +27 -0
- marin_rolloutengine-0.2.160.dev38062449553/pyproject.toml +51 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/__init__.py +2 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/cleanup.py +78 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/contracts.py +132 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/engine.py +291 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/grading.py +60 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/lowering.py +106 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/machines.py +199 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/spec.py +59 -0
- marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/task_session.py +168 -0
- marin_rolloutengine-0.2.160.dev38062449553/tests/__init__.py +2 -0
- marin_rolloutengine-0.2.160.dev38062449553/tests/fixtures/local-grader.in +3 -0
- marin_rolloutengine-0.2.160.dev38062449553/tests/fixtures/local-grader.lock +14 -0
- marin_rolloutengine-0.2.160.dev38062449553/tests/test_local_grading.py +129 -0
- marin_rolloutengine-0.2.160.dev38062449553/tests/test_optional_imports.py +28 -0
- marin_rolloutengine-0.2.160.dev38062449553/tests/test_rollout.py +1383 -0
- marin_rolloutengine-0.2.160.dev38062449553/tests/test_swe_rollout.py +139 -0
|
@@ -0,0 +1,253 @@
|
|
|
1
|
+
# redundant, but Ray looks for this otherwise.
|
|
2
|
+
.git
|
|
3
|
+
|
|
4
|
+
logs/
|
|
5
|
+
|
|
6
|
+
# CPU profiles
|
|
7
|
+
prof/
|
|
8
|
+
|
|
9
|
+
# Downloaded build tools (zig, etc.)
|
|
10
|
+
.tools/
|
|
11
|
+
|
|
12
|
+
tests/snapshots/outputs
|
|
13
|
+
tests/snapshots/diffs
|
|
14
|
+
|
|
15
|
+
# don't log data/MD outputs to git
|
|
16
|
+
data/*
|
|
17
|
+
output/*
|
|
18
|
+
outputs/*
|
|
19
|
+
|
|
20
|
+
# Snapshot diffs and outputs
|
|
21
|
+
tests/snapshots/*/outputs/*
|
|
22
|
+
tests/snapshots/*/diffs/*
|
|
23
|
+
|
|
24
|
+
# This is mainly for Ray and using submodule
|
|
25
|
+
*/**/.git
|
|
26
|
+
|
|
27
|
+
### Python template
|
|
28
|
+
# Byte-compiled / optimized / DLL files
|
|
29
|
+
__pycache__/
|
|
30
|
+
*.py[cod]
|
|
31
|
+
*$py.class
|
|
32
|
+
|
|
33
|
+
# C extensions
|
|
34
|
+
*.so
|
|
35
|
+
|
|
36
|
+
# pypa/gh-action-pypi-publish caches its Docker action manifest here.
|
|
37
|
+
.github/.tmp/
|
|
38
|
+
|
|
39
|
+
# Distribution / packaging
|
|
40
|
+
.Python
|
|
41
|
+
build/
|
|
42
|
+
develop-eggs/
|
|
43
|
+
dist/
|
|
44
|
+
downloads/
|
|
45
|
+
eggs/
|
|
46
|
+
.eggs/
|
|
47
|
+
lib64/
|
|
48
|
+
parts/
|
|
49
|
+
sdist/
|
|
50
|
+
local_store/
|
|
51
|
+
wheels/
|
|
52
|
+
share/python-wheels/
|
|
53
|
+
*.egg-info/
|
|
54
|
+
.installed.cfg
|
|
55
|
+
*.egg
|
|
56
|
+
MANIFEST
|
|
57
|
+
|
|
58
|
+
# PyInstaller
|
|
59
|
+
# Usually these files are written by a python script from a template
|
|
60
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
61
|
+
*.manifest
|
|
62
|
+
*.spec
|
|
63
|
+
|
|
64
|
+
# Installer logs
|
|
65
|
+
pip-log.txt
|
|
66
|
+
pip-delete-this-directory.txt
|
|
67
|
+
|
|
68
|
+
# Unit test / coverage reports
|
|
69
|
+
htmlcov/
|
|
70
|
+
.tox/
|
|
71
|
+
.nox/
|
|
72
|
+
.coverage
|
|
73
|
+
.coverage.*
|
|
74
|
+
.cache
|
|
75
|
+
nosetests.xml
|
|
76
|
+
coverage.xml
|
|
77
|
+
*.cover
|
|
78
|
+
*.py,cover
|
|
79
|
+
.hypothesis/
|
|
80
|
+
.pytest_cache/
|
|
81
|
+
cover/
|
|
82
|
+
|
|
83
|
+
# Translations
|
|
84
|
+
*.mo
|
|
85
|
+
*.pot
|
|
86
|
+
|
|
87
|
+
# Django stuff:
|
|
88
|
+
*.log
|
|
89
|
+
local_settings.py
|
|
90
|
+
db.sqlite3
|
|
91
|
+
db.sqlite3-journal
|
|
92
|
+
|
|
93
|
+
# Flask stuff:
|
|
94
|
+
instance/
|
|
95
|
+
.webassets-cache
|
|
96
|
+
|
|
97
|
+
# Scrapy stuff:
|
|
98
|
+
.scrapy
|
|
99
|
+
|
|
100
|
+
# Sphinx documentation
|
|
101
|
+
docs/_build/
|
|
102
|
+
|
|
103
|
+
# PyBuilder
|
|
104
|
+
.pybuilder/
|
|
105
|
+
target/
|
|
106
|
+
|
|
107
|
+
# Jupyter Notebook
|
|
108
|
+
.ipynb_checkpoints
|
|
109
|
+
|
|
110
|
+
# IPython
|
|
111
|
+
profile_default/
|
|
112
|
+
ipython_config.py
|
|
113
|
+
|
|
114
|
+
# pyenv
|
|
115
|
+
# For a library or package, you might want to ignore these files since the code is
|
|
116
|
+
# intended to run in multiple environments; otherwise, check them in:
|
|
117
|
+
# .python-version
|
|
118
|
+
|
|
119
|
+
# pipenv
|
|
120
|
+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
|
121
|
+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
|
122
|
+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
|
123
|
+
# install all needed dependencies.
|
|
124
|
+
#Pipfile.lock
|
|
125
|
+
|
|
126
|
+
# poetry
|
|
127
|
+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
|
128
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
129
|
+
# commonly ignored for libraries.
|
|
130
|
+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
|
131
|
+
#poetry.lock
|
|
132
|
+
|
|
133
|
+
# pdm
|
|
134
|
+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
|
135
|
+
#pdm.lock
|
|
136
|
+
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
|
137
|
+
# in version control.
|
|
138
|
+
# https://pdm.fming.dev/#use-with-ide
|
|
139
|
+
.pdm.toml
|
|
140
|
+
|
|
141
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
142
|
+
__pypackages__/
|
|
143
|
+
|
|
144
|
+
# Celery stuff
|
|
145
|
+
celerybeat-schedule
|
|
146
|
+
celerybeat.pid
|
|
147
|
+
|
|
148
|
+
# SageMath parsed files
|
|
149
|
+
*.sage.py
|
|
150
|
+
|
|
151
|
+
# Environments
|
|
152
|
+
.env
|
|
153
|
+
.venv
|
|
154
|
+
env/
|
|
155
|
+
venv/
|
|
156
|
+
ENV/
|
|
157
|
+
env.bak/
|
|
158
|
+
venv.bak/
|
|
159
|
+
|
|
160
|
+
# Spyder project settings
|
|
161
|
+
.spyderproject
|
|
162
|
+
.spyproject
|
|
163
|
+
|
|
164
|
+
# Rope project settings
|
|
165
|
+
.ropeproject
|
|
166
|
+
|
|
167
|
+
# mkdocs documentation
|
|
168
|
+
/site
|
|
169
|
+
|
|
170
|
+
# mypy
|
|
171
|
+
.mypy_cache/
|
|
172
|
+
.dmypy.json
|
|
173
|
+
dmypy.json
|
|
174
|
+
|
|
175
|
+
# Pyre type checker
|
|
176
|
+
.pyre/
|
|
177
|
+
|
|
178
|
+
# Ruff
|
|
179
|
+
.ruff_cache/
|
|
180
|
+
|
|
181
|
+
# pytype static type analyzer
|
|
182
|
+
.pytype/
|
|
183
|
+
|
|
184
|
+
# Cython debug symbols
|
|
185
|
+
cython_debug/
|
|
186
|
+
|
|
187
|
+
# PyCharm
|
|
188
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
189
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
190
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
191
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
192
|
+
.idea/
|
|
193
|
+
*.iml
|
|
194
|
+
|
|
195
|
+
# IDE Config
|
|
196
|
+
.vscode/
|
|
197
|
+
|
|
198
|
+
# Mac OS
|
|
199
|
+
.DS_Store
|
|
200
|
+
|
|
201
|
+
# Secrets
|
|
202
|
+
credentials.json
|
|
203
|
+
marin/crawl/bigquery-gcs-key.json
|
|
204
|
+
|
|
205
|
+
# Archive
|
|
206
|
+
archive/
|
|
207
|
+
|
|
208
|
+
# Caches and Outputs
|
|
209
|
+
!/scripts/web/output/
|
|
210
|
+
!/output/
|
|
211
|
+
|
|
212
|
+
# csv
|
|
213
|
+
*.csv
|
|
214
|
+
|
|
215
|
+
# wandb logs
|
|
216
|
+
wandb
|
|
217
|
+
artifacts
|
|
218
|
+
|
|
219
|
+
# Ignore generated credentials from google-github-actions/auth
|
|
220
|
+
gha-creds-*.json
|
|
221
|
+
|
|
222
|
+
.aider*
|
|
223
|
+
.git/*
|
|
224
|
+
|
|
225
|
+
*.jsonl
|
|
226
|
+
**/*.jsonl
|
|
227
|
+
# Echo's judged search corpus is source, not a run artifact.
|
|
228
|
+
!infra/marina/apps/echo/benchmarks/*.jsonl
|
|
229
|
+
scr/*
|
|
230
|
+
.weaver/
|
|
231
|
+
|
|
232
|
+
# Local host Marin config
|
|
233
|
+
.marin.yaml
|
|
234
|
+
|
|
235
|
+
/scratch
|
|
236
|
+
|
|
237
|
+
.forge
|
|
238
|
+
.claude/*
|
|
239
|
+
!.claude/skills
|
|
240
|
+
!.claude/agents
|
|
241
|
+
!.claude/settings.json
|
|
242
|
+
.agents/tmp/
|
|
243
|
+
.codex
|
|
244
|
+
.entire
|
|
245
|
+
.beads
|
|
246
|
+
|
|
247
|
+
.worktrees
|
|
248
|
+
.obsidian
|
|
249
|
+
.cw_env
|
|
250
|
+
.localbin/
|
|
251
|
+
|
|
252
|
+
# Stamped into built artifacts by lib/iris/hatch_build.py; never tracked.
|
|
253
|
+
lib/iris/src/iris/_build_info.py
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: marin-rolloutengine
|
|
3
|
+
Version: 0.2.160.dev38062449553
|
|
4
|
+
Summary: Task execution and exact-token rollout records
|
|
5
|
+
Requires-Python: <3.14,>=3.12
|
|
6
|
+
Requires-Dist: marin-shellbox==0.2.160.dev38062449553
|
|
7
|
+
Requires-Dist: marin-taskcompendium==0.2.160.dev38062449553
|
|
8
|
+
Requires-Dist: pydantic>=2
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
# Rollout engine
|
|
2
|
+
|
|
3
|
+
`marin-rolloutengine` executes one TaskCompendium task through Shellbox and returns
|
|
4
|
+
`RolloutData`: the conversation, exact tokens, loss mask, optional log probabilities, grade, and per-turn records.
|
|
5
|
+
|
|
6
|
+
Callers supply a model callable, machine factories, and optional `TaskSession`
|
|
7
|
+
factories for task actions and grading.
|
|
8
|
+
The default session handles shell tools.
|
|
9
|
+
For text, number, JSON, and native-action answers, it adds the task's answer-format instruction and tools to the model request.
|
|
10
|
+
For workspace-state tasks, this session adds shell and final-response instructions to the model request.
|
|
11
|
+
Importers retain the task problem without these interface instructions.
|
|
12
|
+
The engine owns model calls, conversation and token accumulation, deadlines, and resource cleanup.
|
|
13
|
+
|
|
14
|
+
`ShellboxRolloutEngine.run(lowered)` accepts a `LoweredTaskSpec` and asynchronously returns one rollout.
|
|
15
|
+
The lowered record preserves its `TaskSpec` and adds machine selections and session limits.
|
|
16
|
+
The supported scope is single-stage tasks with prebuilt, digest-pinned images, local package-lock environments, or the built-in ShellSim filesystem.
|
|
17
|
+
After the turn loop, the default session grades an in-process `VerifyitGrader` on the host.
|
|
18
|
+
A `ScriptGrader`, or a `VerifyitGrader` with an environment, grades in a separate verifier machine using its digest-pinned image or a local environment built from its package lock.
|
|
19
|
+
A `SessionGrader` requires a registered task session. A `NoGrader` task receives an `unavailable` grade.
|
|
20
|
+
See the [task rollout reference](../../docs/references/task-rollouts.md)
|
|
21
|
+
for the session lifecycle, exact-token contract, failure handling, and backend configuration.
|
|
22
|
+
|
|
23
|
+
From the Marin repository root:
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
uv run --package marin-rolloutengine --frozen --group test pytest lib/rolloutengine/tests -q
|
|
27
|
+
```
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "marin-rolloutengine"
|
|
7
|
+
version = "0.2.160.dev38062449553"
|
|
8
|
+
description = "Task execution and exact-token rollout records"
|
|
9
|
+
requires-python = ">=3.12,<3.14"
|
|
10
|
+
dependencies = [
|
|
11
|
+
"pydantic>=2",
|
|
12
|
+
"marin-taskcompendium==0.2.160.dev38062449553",
|
|
13
|
+
"marin-shellbox==0.2.160.dev38062449553",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
[dependency-groups]
|
|
17
|
+
test = [
|
|
18
|
+
"pytest>=8.3",
|
|
19
|
+
"pytest-asyncio",
|
|
20
|
+
"pytest-timeout",
|
|
21
|
+
"marin-shellbox[daytona,shellsim]",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[tool.hatch.build.targets.wheel]
|
|
25
|
+
packages = ["src/rolloutengine"]
|
|
26
|
+
|
|
27
|
+
[tool.pytest.ini_options]
|
|
28
|
+
asyncio_mode = "auto"
|
|
29
|
+
timeout = 60
|
|
30
|
+
addopts = "--import-mode=importlib -m 'not slow and not integration and not data_integration and not cluster and not requires_cluster and not docker and not manual'"
|
|
31
|
+
markers = [
|
|
32
|
+
"slow: long-running test",
|
|
33
|
+
"integration: external service",
|
|
34
|
+
"data_integration: external data",
|
|
35
|
+
"cluster: live cluster",
|
|
36
|
+
"requires_cluster: live cluster",
|
|
37
|
+
"docker: Docker daemon",
|
|
38
|
+
"manual: explicit validation",
|
|
39
|
+
]
|
|
40
|
+
|
|
41
|
+
[tool.ruff]
|
|
42
|
+
extend = "../../pyproject.toml"
|
|
43
|
+
|
|
44
|
+
[tool.ruff.lint.isort]
|
|
45
|
+
known-first-party = ["rolloutengine"]
|
|
46
|
+
|
|
47
|
+
[tool.pyrefly]
|
|
48
|
+
disable-project-excludes-heuristics = true
|
|
49
|
+
project-includes = ["src"]
|
|
50
|
+
search-path = ["src"]
|
|
51
|
+
python-interpreter = "../../.venv/bin/python"
|
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
# Copyright The Marin Authors
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
|
|
4
|
+
"""Bounded cleanup and ownership of operations that outlive cancellation."""
|
|
5
|
+
|
|
6
|
+
import asyncio
|
|
7
|
+
import logging
|
|
8
|
+
from collections.abc import Callable, Coroutine
|
|
9
|
+
from dataclasses import dataclass, field
|
|
10
|
+
from typing import Any
|
|
11
|
+
|
|
12
|
+
logger = logging.getLogger(__name__)
|
|
13
|
+
_background_tasks: set[asyncio.Task[None]] = set()
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def _retain_task(task: asyncio.Task[None]) -> None:
|
|
17
|
+
_background_tasks.add(task)
|
|
18
|
+
task.add_done_callback(_background_done)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _background_done(task: asyncio.Task[None]) -> None:
|
|
22
|
+
_background_tasks.discard(task)
|
|
23
|
+
if not task.cancelled() and (error := task.exception()) is not None:
|
|
24
|
+
logger.error("Background resource operation failed", exc_info=error)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
@dataclass(frozen=True)
|
|
28
|
+
class _CleanupError:
|
|
29
|
+
operation: str
|
|
30
|
+
exception_type: str
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@dataclass
|
|
34
|
+
class _Cleanup:
|
|
35
|
+
timeout: float
|
|
36
|
+
errors: list[_CleanupError] = field(default_factory=list)
|
|
37
|
+
|
|
38
|
+
async def run(self, operation: str, action: Callable[[], Coroutine[Any, Any, None]]) -> Exception | None:
|
|
39
|
+
pending = asyncio.create_task(action())
|
|
40
|
+
loop = asyncio.get_running_loop()
|
|
41
|
+
end = loop.time() + self.timeout
|
|
42
|
+
cancellation = None
|
|
43
|
+
failure = None
|
|
44
|
+
while not pending.done():
|
|
45
|
+
remaining = end - loop.time()
|
|
46
|
+
if remaining <= 0:
|
|
47
|
+
pending.cancel()
|
|
48
|
+
_retain_task(pending)
|
|
49
|
+
failure = TimeoutError("Cleanup deadline expired")
|
|
50
|
+
break
|
|
51
|
+
try:
|
|
52
|
+
await asyncio.wait((pending,), timeout=remaining)
|
|
53
|
+
except asyncio.CancelledError as error:
|
|
54
|
+
# Repeated cancellation cannot extend the cleanup deadline.
|
|
55
|
+
cancellation = error
|
|
56
|
+
if failure is None:
|
|
57
|
+
try:
|
|
58
|
+
pending.result()
|
|
59
|
+
except asyncio.CancelledError as error:
|
|
60
|
+
failure = RuntimeError("Cleanup operation cancelled")
|
|
61
|
+
failure.__cause__ = error
|
|
62
|
+
except Exception as error:
|
|
63
|
+
failure = error
|
|
64
|
+
if failure is not None:
|
|
65
|
+
self.errors.append(_CleanupError(operation, type(failure).__name__))
|
|
66
|
+
logger.warning("Cleanup failed during %s", operation, exc_info=failure)
|
|
67
|
+
if cancellation is not None:
|
|
68
|
+
if failure is not None:
|
|
69
|
+
raise cancellation from failure
|
|
70
|
+
raise cancellation
|
|
71
|
+
return failure
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
async def finish_cleanup(action: Callable[[], Coroutine[Any, Any, None]], *, timeout: float) -> None:
|
|
75
|
+
"""Finish resource cleanup within its deadline and propagate failures or repeated cancellation."""
|
|
76
|
+
failure = await _Cleanup(timeout).run("resource_close", action)
|
|
77
|
+
if failure is not None:
|
|
78
|
+
raise failure
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
# Copyright The Marin Authors
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
|
|
4
|
+
"""Public rollout protocols, records, and exceptions."""
|
|
5
|
+
|
|
6
|
+
from dataclasses import dataclass, field
|
|
7
|
+
from enum import StrEnum
|
|
8
|
+
from typing import Any, Protocol
|
|
9
|
+
|
|
10
|
+
from taskcompendium.grading_result import GradeResult
|
|
11
|
+
|
|
12
|
+
LENGTH_STOP_REASON = "length"
|
|
13
|
+
MAX_TURNS_STOP_REASON = "max_turns"
|
|
14
|
+
TOTAL_TURN_TIMEOUT_STOP_REASON = "total_turn_timeout"
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
@dataclass(frozen=True)
|
|
18
|
+
class ModelRequest:
|
|
19
|
+
messages: tuple[dict[str, Any], ...]
|
|
20
|
+
options: dict[str, Any]
|
|
21
|
+
prefix_token_ids: tuple[int, ...]
|
|
22
|
+
assistant_message_index: int | None
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
@dataclass(frozen=True)
|
|
26
|
+
class SessionStart:
|
|
27
|
+
"""The initial conversation and model options for one task session."""
|
|
28
|
+
|
|
29
|
+
messages: tuple[dict[str, Any], ...]
|
|
30
|
+
options: dict[str, Any]
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@dataclass(frozen=True)
|
|
34
|
+
class ModelTurn:
|
|
35
|
+
"""Exact tokens and the parsed message from one inference request."""
|
|
36
|
+
|
|
37
|
+
message: dict[str, Any]
|
|
38
|
+
prompt_token_ids: tuple[int, ...]
|
|
39
|
+
response_token_ids: tuple[int, ...]
|
|
40
|
+
logprobs: tuple[float, ...] | None
|
|
41
|
+
stop_reason: str
|
|
42
|
+
text: str = ""
|
|
43
|
+
metadata: dict[str, Any] = field(default_factory=dict)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class GenerationLimitReached(Exception):
|
|
47
|
+
"""The rendered prompt leaves no permitted generation budget."""
|
|
48
|
+
|
|
49
|
+
def __init__(self, prompt_token_ids: tuple[int, ...]):
|
|
50
|
+
self.prompt_token_ids = prompt_token_ids
|
|
51
|
+
super().__init__("The prompt reached the configured generation limit")
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@dataclass(frozen=True)
|
|
55
|
+
class RolloutFailure:
|
|
56
|
+
"""Safe failure fields for consumers of a rollout record."""
|
|
57
|
+
|
|
58
|
+
exception_type: str
|
|
59
|
+
diagnostics: dict[str, Any] = field(default_factory=dict)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
@dataclass(frozen=True)
|
|
63
|
+
class RolloutData:
|
|
64
|
+
task_id: str
|
|
65
|
+
messages: tuple[dict[str, Any], ...]
|
|
66
|
+
prompt_token_ids: tuple[int, ...]
|
|
67
|
+
response_token_ids: tuple[int, ...]
|
|
68
|
+
loss_mask: tuple[int, ...]
|
|
69
|
+
logprobs: tuple[float, ...] | None
|
|
70
|
+
grade: GradeResult
|
|
71
|
+
stop_reason: str
|
|
72
|
+
steps: tuple["RolloutStep", ...] = ()
|
|
73
|
+
metrics: dict[str, Any] = field(default_factory=dict)
|
|
74
|
+
failure: RolloutFailure | None = None
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
class RolloutOperation(StrEnum):
|
|
78
|
+
ATTEMPT = "attempt"
|
|
79
|
+
START = "start"
|
|
80
|
+
PREPARE = "prepare"
|
|
81
|
+
MODEL = "model"
|
|
82
|
+
ADVANCE = "advance"
|
|
83
|
+
GRADE = "grade"
|
|
84
|
+
CLEANUP = "cleanup"
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class RolloutInterrupted(RuntimeError):
|
|
88
|
+
"""An execution failure with the last completed rollout and original exception cause."""
|
|
89
|
+
|
|
90
|
+
def __init__(self, rollout: RolloutData, operation: RolloutOperation):
|
|
91
|
+
self.rollout = rollout
|
|
92
|
+
self.operation = operation
|
|
93
|
+
super().__init__(f"Rollout interrupted during {operation}")
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class RolloutContractError(ValueError):
|
|
97
|
+
"""Rollout evidence violates the exact-token contract."""
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
@dataclass(frozen=True)
|
|
101
|
+
class Transition:
|
|
102
|
+
"""One task transition after a model response."""
|
|
103
|
+
|
|
104
|
+
done: bool
|
|
105
|
+
observations: tuple[dict[str, Any], ...] = ()
|
|
106
|
+
reset_conversation: tuple[dict[str, Any], ...] | None = None
|
|
107
|
+
reward: float | None = None
|
|
108
|
+
token_rewards: tuple[float, ...] | None = None
|
|
109
|
+
token_credit: tuple[float, ...] | None = None
|
|
110
|
+
reward_components: dict[str, float] = field(default_factory=dict)
|
|
111
|
+
grade: GradeResult | None = None
|
|
112
|
+
metrics: dict[str, Any] = field(default_factory=dict)
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
@dataclass(frozen=True)
|
|
116
|
+
class RolloutStep:
|
|
117
|
+
turn: ModelTurn
|
|
118
|
+
transition: Transition
|
|
119
|
+
response_end: int
|
|
120
|
+
messages: tuple[dict[str, Any], ...]
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
class TaskSession(Protocol):
|
|
124
|
+
"""Task operations without model inference or token processing."""
|
|
125
|
+
|
|
126
|
+
async def prepare(self) -> SessionStart: ...
|
|
127
|
+
|
|
128
|
+
async def advance(self, turn: ModelTurn) -> Transition: ...
|
|
129
|
+
|
|
130
|
+
async def grade(self, messages: tuple[dict[str, Any], ...]) -> GradeResult: ...
|
|
131
|
+
|
|
132
|
+
async def close(self) -> None: ...
|