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.
Files changed (20) hide show
  1. marin_rolloutengine-0.2.160.dev38062449553/.gitignore +253 -0
  2. marin_rolloutengine-0.2.160.dev38062449553/PKG-INFO +8 -0
  3. marin_rolloutengine-0.2.160.dev38062449553/README.md +27 -0
  4. marin_rolloutengine-0.2.160.dev38062449553/pyproject.toml +51 -0
  5. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/__init__.py +2 -0
  6. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/cleanup.py +78 -0
  7. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/contracts.py +132 -0
  8. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/engine.py +291 -0
  9. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/grading.py +60 -0
  10. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/lowering.py +106 -0
  11. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/machines.py +199 -0
  12. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/spec.py +59 -0
  13. marin_rolloutengine-0.2.160.dev38062449553/src/rolloutengine/task_session.py +168 -0
  14. marin_rolloutengine-0.2.160.dev38062449553/tests/__init__.py +2 -0
  15. marin_rolloutengine-0.2.160.dev38062449553/tests/fixtures/local-grader.in +3 -0
  16. marin_rolloutengine-0.2.160.dev38062449553/tests/fixtures/local-grader.lock +14 -0
  17. marin_rolloutengine-0.2.160.dev38062449553/tests/test_local_grading.py +129 -0
  18. marin_rolloutengine-0.2.160.dev38062449553/tests/test_optional_imports.py +28 -0
  19. marin_rolloutengine-0.2.160.dev38062449553/tests/test_rollout.py +1383 -0
  20. 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,2 @@
1
+ # Copyright The Marin Authors
2
+ # SPDX-License-Identifier: Apache-2.0
@@ -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: ...