rltracer 0.1.0__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.
- rltracer-0.1.0/LICENSE +21 -0
- rltracer-0.1.0/PKG-INFO +331 -0
- rltracer-0.1.0/README.md +287 -0
- rltracer-0.1.0/pyproject.toml +58 -0
- rltracer-0.1.0/setup.cfg +4 -0
- rltracer-0.1.0/src/rltracer/__init__.py +42 -0
- rltracer-0.1.0/src/rltracer/__main__.py +81 -0
- rltracer-0.1.0/src/rltracer/analysis.py +283 -0
- rltracer-0.1.0/src/rltracer/criteria_suffix.py +231 -0
- rltracer-0.1.0/src/rltracer/export.py +1188 -0
- rltracer-0.1.0/src/rltracer/model.py +59 -0
- rltracer-0.1.0/src/rltracer/primerl.py +312 -0
- rltracer-0.1.0/src/rltracer/py.typed +0 -0
- rltracer-0.1.0/src/rltracer/retemplate.py +389 -0
- rltracer-0.1.0/src/rltracer/tracer.py +35 -0
- rltracer-0.1.0/src/rltracer/viz/__init__.py +1 -0
- rltracer-0.1.0/src/rltracer/viz/app.js +193 -0
- rltracer-0.1.0/src/rltracer/viz/browse.html +9 -0
- rltracer-0.1.0/src/rltracer/viz/browse.js +33 -0
- rltracer-0.1.0/src/rltracer/viz/index.html +1 -0
- rltracer-0.1.0/src/rltracer/viz/server.py +272 -0
- rltracer-0.1.0/src/rltracer/viz/style.css +15 -0
- rltracer-0.1.0/src/rltracer/workflow_jsonl.py +126 -0
- rltracer-0.1.0/src/rltracer.egg-info/PKG-INFO +331 -0
- rltracer-0.1.0/src/rltracer.egg-info/SOURCES.txt +33 -0
- rltracer-0.1.0/src/rltracer.egg-info/dependency_links.txt +1 -0
- rltracer-0.1.0/src/rltracer.egg-info/entry_points.txt +4 -0
- rltracer-0.1.0/src/rltracer.egg-info/requires.txt +30 -0
- rltracer-0.1.0/src/rltracer.egg-info/top_level.txt +1 -0
- rltracer-0.1.0/tests/test_criteria_suffix.py +115 -0
- rltracer-0.1.0/tests/test_export.py +543 -0
- rltracer-0.1.0/tests/test_packaging.py +24 -0
- rltracer-0.1.0/tests/test_primerl.py +19 -0
- rltracer-0.1.0/tests/test_retemplate.py +136 -0
- rltracer-0.1.0/tests/test_workflow_jsonl.py +22 -0
rltracer-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Goncalo Faria
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
rltracer-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,331 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: rltracer
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Lazy inspection, analysis, visualization, and export tooling for reinforcement-learning traces.
|
|
5
|
+
Author: Goncalo Faria
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/goncalorafaria/rltracer
|
|
8
|
+
Project-URL: Repository, https://github.com/goncalorafaria/rltracer
|
|
9
|
+
Project-URL: Issues, https://github.com/goncalorafaria/rltracer/issues
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Requires-Python: >=3.10
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
License-File: LICENSE
|
|
19
|
+
Requires-Dist: fire>=0.7.1
|
|
20
|
+
Provides-Extra: primerl
|
|
21
|
+
Requires-Dist: msgpack>=1.0; extra == "primerl"
|
|
22
|
+
Requires-Dist: transformers>=4.45; extra == "primerl"
|
|
23
|
+
Provides-Extra: retemplate
|
|
24
|
+
Requires-Dist: pyarrow>=16; extra == "retemplate"
|
|
25
|
+
Provides-Extra: plot
|
|
26
|
+
Requires-Dist: matplotlib>=3.8; extra == "plot"
|
|
27
|
+
Requires-Dist: numpy>=1.26; extra == "plot"
|
|
28
|
+
Requires-Dist: scipy>=1.12; extra == "plot"
|
|
29
|
+
Provides-Extra: speed
|
|
30
|
+
Requires-Dist: orjson>=3.10; extra == "speed"
|
|
31
|
+
Provides-Extra: test
|
|
32
|
+
Requires-Dist: pytest>=8; extra == "test"
|
|
33
|
+
Requires-Dist: pyarrow>=16; extra == "test"
|
|
34
|
+
Requires-Dist: msgpack>=1.0; extra == "test"
|
|
35
|
+
Provides-Extra: all
|
|
36
|
+
Requires-Dist: msgpack>=1.0; extra == "all"
|
|
37
|
+
Requires-Dist: transformers>=4.45; extra == "all"
|
|
38
|
+
Requires-Dist: pyarrow>=16; extra == "all"
|
|
39
|
+
Requires-Dist: matplotlib>=3.8; extra == "all"
|
|
40
|
+
Requires-Dist: numpy>=1.26; extra == "all"
|
|
41
|
+
Requires-Dist: scipy>=1.12; extra == "all"
|
|
42
|
+
Requires-Dist: orjson>=3.10; extra == "all"
|
|
43
|
+
Dynamic: license-file
|
|
44
|
+
|
|
45
|
+
# RLTracer
|
|
46
|
+
|
|
47
|
+
`rltracer` is a lazy inspection layer for trajectories produced by different
|
|
48
|
+
reinforcement-learning training systems. It keeps library-specific file formats
|
|
49
|
+
inside adapters while exposing one common navigation model.
|
|
50
|
+
|
|
51
|
+
It is a standalone package: it does not import JTC, Datadev, JTCEval, or
|
|
52
|
+
PrimeBeaker. Those projects can consume its public adapters and exporters.
|
|
53
|
+
Generated traces, indexes, datasets, and plots are deliberately excluded from
|
|
54
|
+
the distribution and Git repository.
|
|
55
|
+
|
|
56
|
+
## Install
|
|
57
|
+
|
|
58
|
+
```bash
|
|
59
|
+
pip install rltracer
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
PrimeRL token exports require the adapter extra:
|
|
63
|
+
|
|
64
|
+
```bash
|
|
65
|
+
pip install 'rltracer[primerl]'
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
Use `rltracer[retemplate]` for Parquet-backed SFT retemplating,
|
|
69
|
+
`rltracer[plot]` for pass-rate plots, or `rltracer[all]` for every feature.
|
|
70
|
+
All command-line interfaces use Fire:
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
rltracer inspect --run=/path/to/run --tokenizer=/path/to/tokenizer
|
|
74
|
+
rltracer-viz --browse-root=/path/to/runs --tokenizer=/path/to/tokenizer
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
## Development and release
|
|
78
|
+
|
|
79
|
+
```bash
|
|
80
|
+
uv lock
|
|
81
|
+
uv run --extra test pytest
|
|
82
|
+
uv build
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
The build produces a source distribution and universal wheel under `dist/`.
|
|
86
|
+
The `Publish to PyPI` GitHub Actions workflow builds with Python 3.12 and checks
|
|
87
|
+
both distributions with `twine check --strict`. Publishing a GitHub release
|
|
88
|
+
uploads them using PyPI Trusted Publishing; manual workflow runs only build
|
|
89
|
+
and validate artifacts.
|
|
90
|
+
|
|
91
|
+
For the first release:
|
|
92
|
+
|
|
93
|
+
1. Create the `pypi` environment in GitHub Settings → Environments.
|
|
94
|
+
2. Open [PyPI account publishing](https://pypi.org/manage/account/publishing/)
|
|
95
|
+
and add a pending GitHub publisher:
|
|
96
|
+
|
|
97
|
+
| Field | Value |
|
|
98
|
+
| --- | --- |
|
|
99
|
+
| PyPI project name | `rltracer` |
|
|
100
|
+
| Owner | `goncalorafaria` |
|
|
101
|
+
| Repository | `rltracer` |
|
|
102
|
+
| Workflow filename | `publish.yml` |
|
|
103
|
+
| Environment | `pypi` |
|
|
104
|
+
|
|
105
|
+
3. Commit and push the workflow and package source, then publish a GitHub
|
|
106
|
+
release tagged `v0.1.0` at that commit (matching the current package version).
|
|
107
|
+
|
|
108
|
+
The first successful upload creates the PyPI project automatically. No PyPI
|
|
109
|
+
API token or GitHub secret is required. For an existing project, add the same
|
|
110
|
+
publisher in its PyPI publishing settings.
|
|
111
|
+
|
|
112
|
+
For later releases, update `[project].version` in `pyproject.toml` and publish a
|
|
113
|
+
GitHub release tagged `v<version>` or `<version>` at that commit. The workflow
|
|
114
|
+
rejects mismatched tags. Choose a version not already uploaded to PyPI; existing
|
|
115
|
+
releases cannot be overwritten. Draft releases and tag pushes alone do not
|
|
116
|
+
publish a package.
|
|
117
|
+
|
|
118
|
+
## Object hierarchy
|
|
119
|
+
|
|
120
|
+
```text
|
|
121
|
+
RLTracer
|
|
122
|
+
└── RolloutStep
|
|
123
|
+
└── PromptGroup
|
|
124
|
+
└── TrajectoryRef
|
|
125
|
+
└── Trajectory
|
|
126
|
+
└── messages
|
|
127
|
+
```
|
|
128
|
+
|
|
129
|
+
- **RLTracer** is the entry point for one training run. It receives a configured
|
|
130
|
+
adapter in its constructor.
|
|
131
|
+
- **RolloutStep** represents one RL training step, such as `step_21`.
|
|
132
|
+
- **PromptGroup** represents all sampled trajectories sharing the same initial
|
|
133
|
+
prompt within a rollout step.
|
|
134
|
+
- **TrajectoryRef** is a lightweight SQLite-backed reference. Calling
|
|
135
|
+
`load()` decodes its source trajectory.
|
|
136
|
+
- **Trajectory** is the decoded multi-turn record. Its `messages` field is a
|
|
137
|
+
normalized list of `{"role": ..., "content": ...}` dictionaries.
|
|
138
|
+
|
|
139
|
+
The canonical internal word is **trajectory**. A rollout is the sampled
|
|
140
|
+
trajectory produced by the RL system.
|
|
141
|
+
|
|
142
|
+
## Prime-RL adapter
|
|
143
|
+
|
|
144
|
+
`PrimeRLAdapter` currently supports Prime-RL rollout artifacts:
|
|
145
|
+
|
|
146
|
+
```text
|
|
147
|
+
<run>/rollouts/
|
|
148
|
+
step_1/
|
|
149
|
+
rank_0.bin
|
|
150
|
+
rank_1.bin
|
|
151
|
+
step_2/
|
|
152
|
+
rank_0.bin
|
|
153
|
+
...
|
|
154
|
+
```
|
|
155
|
+
|
|
156
|
+
Each `rank_*.bin` file is a MessagePack payload. It contains groups of packed
|
|
157
|
+
Qwen token sequences. One packed sequence can contain multiple independent
|
|
158
|
+
conversations, so the adapter decodes it with the configured tokenizer and
|
|
159
|
+
splits it at each new Qwen system turn.
|
|
160
|
+
|
|
161
|
+
## Lazy behavior
|
|
162
|
+
|
|
163
|
+
Opening a tracer and listing its steps is cheap:
|
|
164
|
+
|
|
165
|
+
```python
|
|
166
|
+
adapter = PrimeRLAdapter(
|
|
167
|
+
rollout_dir="/path/to/run/rollouts",
|
|
168
|
+
tokenizer_name="/path/to/tokenizer",
|
|
169
|
+
index_path="/tmp/prime-rl.sqlite", # optional
|
|
170
|
+
)
|
|
171
|
+
tracer = RLTracer(adapter)
|
|
172
|
+
|
|
173
|
+
tracer.steps() # discovers step directories only
|
|
174
|
+
```
|
|
175
|
+
|
|
176
|
+
No shard is read and no tokenizer is loaded until a prompt query is made:
|
|
177
|
+
|
|
178
|
+
```python
|
|
179
|
+
step = tracer.step(21)
|
|
180
|
+
prompt_groups = step.prompts()
|
|
181
|
+
```
|
|
182
|
+
|
|
183
|
+
The first `prompts()` call for a step:
|
|
184
|
+
|
|
185
|
+
1. Reads only that step's `rank_*.bin` files.
|
|
186
|
+
2. Decodes packed sequences with the adapter tokenizer.
|
|
187
|
+
3. Splits them into conversations.
|
|
188
|
+
4. Computes a stable hash of each conversation's pre-assistant prompt context.
|
|
189
|
+
5. Stores only lookup metadata in SQLite.
|
|
190
|
+
|
|
191
|
+
The SQLite index stores shard location, packed group/row location, conversation
|
|
192
|
+
position, prompt key, preview, and token count. It does not copy full token
|
|
193
|
+
sequences or decoded conversation content.
|
|
194
|
+
|
|
195
|
+
Loading a selected trajectory re-reads only the shard named in its SQLite row,
|
|
196
|
+
decodes the selected packed sequence, and returns the requested conversation:
|
|
197
|
+
|
|
198
|
+
```python
|
|
199
|
+
group = step.prompts()[0]
|
|
200
|
+
trajectory = group.trajectory(0).load()
|
|
201
|
+
|
|
202
|
+
for message in trajectory.messages:
|
|
203
|
+
print(message["role"], message["content"])
|
|
204
|
+
```
|
|
205
|
+
|
|
206
|
+
The index is invalidated and rebuilt for a step if its shard paths, sizes, or
|
|
207
|
+
modification times change.
|
|
208
|
+
|
|
209
|
+
## Adapter boundary
|
|
210
|
+
|
|
211
|
+
New RL systems should implement the methods consumed by `RLTracer`:
|
|
212
|
+
|
|
213
|
+
```python
|
|
214
|
+
class Adapter:
|
|
215
|
+
rollout_dir: Path
|
|
216
|
+
|
|
217
|
+
def list_steps(self) -> list[int]: ...
|
|
218
|
+
def ensure_step_indexed(self, step: int) -> None: ...
|
|
219
|
+
def list_prompts(self, step: int) -> list[tuple[str, int, str]]: ...
|
|
220
|
+
def list_trajectory_ids(self, step: int, prompt_key: str) -> list[int]: ...
|
|
221
|
+
def load_trajectory(self, trajectory_id: int) -> Trajectory: ...
|
|
222
|
+
```
|
|
223
|
+
|
|
224
|
+
The adapter owns all source-specific choices: artifact format, tokenizer,
|
|
225
|
+
decoding, prompt grouping, optional rewards/logprobs, and index location.
|
|
226
|
+
`RLTracer` stays provider-neutral.
|
|
227
|
+
|
|
228
|
+
## Training-data export
|
|
229
|
+
|
|
230
|
+
TraceDataExporter builds JSONL directly from one or more structured PrimeRL
|
|
231
|
+
trace runs. A run can be an experiment directory, its run_default directory,
|
|
232
|
+
or its rollouts directory.
|
|
233
|
+
|
|
234
|
+
SFT mode selects complete trajectories:
|
|
235
|
+
|
|
236
|
+
```python
|
|
237
|
+
from rltracer import (
|
|
238
|
+
Acceptance,
|
|
239
|
+
DifficultyFilter,
|
|
240
|
+
ExportPreferences,
|
|
241
|
+
FieldFilter,
|
|
242
|
+
TraceDataExporter,
|
|
243
|
+
TraceRunSpec,
|
|
244
|
+
)
|
|
245
|
+
|
|
246
|
+
exporter = TraceDataExporter(
|
|
247
|
+
runs=[
|
|
248
|
+
TraceRunSpec("/path/to/first/run", step_end=800),
|
|
249
|
+
TraceRunSpec("/path/to/second/run", step_end=400),
|
|
250
|
+
],
|
|
251
|
+
preferences=ExportPreferences(
|
|
252
|
+
mode="sft",
|
|
253
|
+
acceptance=Acceptance.CORRECT,
|
|
254
|
+
# Empirical prompt pass rate of 0.25-0.75.
|
|
255
|
+
difficulty=DifficultyFilter(
|
|
256
|
+
min_pass_rate=0.25,
|
|
257
|
+
max_pass_rate=0.75,
|
|
258
|
+
),
|
|
259
|
+
field_filters=(
|
|
260
|
+
FieldFilter("metrics.num_turns", maximum=8),
|
|
261
|
+
FieldFilter("rewards.reward.score", minimum=0),
|
|
262
|
+
),
|
|
263
|
+
window_size=100,
|
|
264
|
+
per_prompt=2,
|
|
265
|
+
target_count=16_000,
|
|
266
|
+
seed=17,
|
|
267
|
+
),
|
|
268
|
+
)
|
|
269
|
+
result = exporter.write_jsonl("correct_sft_16k.jsonl")
|
|
270
|
+
```
|
|
271
|
+
|
|
272
|
+
RL mode writes datadev RL rows whose prompt is the context preceding the first
|
|
273
|
+
assistant turn:
|
|
274
|
+
|
|
275
|
+
```python
|
|
276
|
+
rl_exporter = TraceDataExporter(
|
|
277
|
+
[
|
|
278
|
+
TraceRunSpec(
|
|
279
|
+
"/path/to/run",
|
|
280
|
+
step_end=800,
|
|
281
|
+
rl_data_path="/path/to/datadev_rl_train.jsonl",
|
|
282
|
+
)
|
|
283
|
+
],
|
|
284
|
+
ExportPreferences(
|
|
285
|
+
mode="rl",
|
|
286
|
+
acceptance="any",
|
|
287
|
+
difficulty=DifficultyFilter(min_difficulty=0.25),
|
|
288
|
+
# Use one pass-rate estimate over the complete step range.
|
|
289
|
+
window_size=None,
|
|
290
|
+
unique_prompts=True,
|
|
291
|
+
target_count=5_000,
|
|
292
|
+
),
|
|
293
|
+
)
|
|
294
|
+
rl_exporter.write_jsonl("rl_prompts_5k.jsonl")
|
|
295
|
+
```
|
|
296
|
+
|
|
297
|
+
Difficulty is identical to prompt_pass_rate. Pass rate is calculated from all scored
|
|
298
|
+
rollouts for a prompt in each configured step window. RL mode emits at most one
|
|
299
|
+
copy of a prompt per run/window; SFT mode honors per_prompt. Global targets are
|
|
300
|
+
deterministic and fair: every eligible prompt group receives its first row
|
|
301
|
+
before any group receives a second.
|
|
302
|
+
|
|
303
|
+
Set `unique_prompts=True` to keep at most one row for a prompt across all runs
|
|
304
|
+
and step windows. Deduplication is deterministic under `seed` and happens
|
|
305
|
+
before applying `target_count`.
|
|
306
|
+
|
|
307
|
+
Set `tool_call_sampling_power` to softly reduce the prominence of traces with
|
|
308
|
+
few command/tool calls. Selection weight is
|
|
309
|
+
`tool_call_count ** tool_call_sampling_power`, applied both within prompt
|
|
310
|
+
groups and when filling the global target. `0` preserves uniform selection,
|
|
311
|
+
`0.5` is a gentle square-root bias, `1` is proportional to command count, and
|
|
312
|
+
`2` is a strong quadratic bias. Selection remains deterministic under `seed`,
|
|
313
|
+
and export summaries include available and selected tool-call histograms.
|
|
314
|
+
|
|
315
|
+
By default, output rows match `datadev data` exactly. SFT rows contain
|
|
316
|
+
`messages` and optional JSON-string `tools`. RL rows are losslessly joined
|
|
317
|
+
through `task.data.idx` to `rl_data_path` and contain `prompt`, `answer`,
|
|
318
|
+
`output`, `feedback`, `record_id`, `rubric_index`, `source`, and list-valued
|
|
319
|
+
`tools`. Set `include_provenance=True` only when an additional
|
|
320
|
+
`trace_provenance` field is acceptable.
|
|
321
|
+
|
|
322
|
+
SFT tool calls are canonicalized to OpenAI's nested representation:
|
|
323
|
+
`{"id": ..., "type": "function", "function": {"name": ..., "arguments":
|
|
324
|
+
...}}`. Legacy flattened and singular calls are converted during export, and
|
|
325
|
+
calls without a function name are rejected rather than emitted as corrupted
|
|
326
|
+
training targets.
|
|
327
|
+
|
|
328
|
+
## Next extensions
|
|
329
|
+
|
|
330
|
+
- Add reward, advantage, log-probability, tool-call, and completion metadata to
|
|
331
|
+
`Trajectory.metadata`.
|
rltracer-0.1.0/README.md
ADDED
|
@@ -0,0 +1,287 @@
|
|
|
1
|
+
# RLTracer
|
|
2
|
+
|
|
3
|
+
`rltracer` is a lazy inspection layer for trajectories produced by different
|
|
4
|
+
reinforcement-learning training systems. It keeps library-specific file formats
|
|
5
|
+
inside adapters while exposing one common navigation model.
|
|
6
|
+
|
|
7
|
+
It is a standalone package: it does not import JTC, Datadev, JTCEval, or
|
|
8
|
+
PrimeBeaker. Those projects can consume its public adapters and exporters.
|
|
9
|
+
Generated traces, indexes, datasets, and plots are deliberately excluded from
|
|
10
|
+
the distribution and Git repository.
|
|
11
|
+
|
|
12
|
+
## Install
|
|
13
|
+
|
|
14
|
+
```bash
|
|
15
|
+
pip install rltracer
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
PrimeRL token exports require the adapter extra:
|
|
19
|
+
|
|
20
|
+
```bash
|
|
21
|
+
pip install 'rltracer[primerl]'
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
Use `rltracer[retemplate]` for Parquet-backed SFT retemplating,
|
|
25
|
+
`rltracer[plot]` for pass-rate plots, or `rltracer[all]` for every feature.
|
|
26
|
+
All command-line interfaces use Fire:
|
|
27
|
+
|
|
28
|
+
```bash
|
|
29
|
+
rltracer inspect --run=/path/to/run --tokenizer=/path/to/tokenizer
|
|
30
|
+
rltracer-viz --browse-root=/path/to/runs --tokenizer=/path/to/tokenizer
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Development and release
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
uv lock
|
|
37
|
+
uv run --extra test pytest
|
|
38
|
+
uv build
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
The build produces a source distribution and universal wheel under `dist/`.
|
|
42
|
+
The `Publish to PyPI` GitHub Actions workflow builds with Python 3.12 and checks
|
|
43
|
+
both distributions with `twine check --strict`. Publishing a GitHub release
|
|
44
|
+
uploads them using PyPI Trusted Publishing; manual workflow runs only build
|
|
45
|
+
and validate artifacts.
|
|
46
|
+
|
|
47
|
+
For the first release:
|
|
48
|
+
|
|
49
|
+
1. Create the `pypi` environment in GitHub Settings → Environments.
|
|
50
|
+
2. Open [PyPI account publishing](https://pypi.org/manage/account/publishing/)
|
|
51
|
+
and add a pending GitHub publisher:
|
|
52
|
+
|
|
53
|
+
| Field | Value |
|
|
54
|
+
| --- | --- |
|
|
55
|
+
| PyPI project name | `rltracer` |
|
|
56
|
+
| Owner | `goncalorafaria` |
|
|
57
|
+
| Repository | `rltracer` |
|
|
58
|
+
| Workflow filename | `publish.yml` |
|
|
59
|
+
| Environment | `pypi` |
|
|
60
|
+
|
|
61
|
+
3. Commit and push the workflow and package source, then publish a GitHub
|
|
62
|
+
release tagged `v0.1.0` at that commit (matching the current package version).
|
|
63
|
+
|
|
64
|
+
The first successful upload creates the PyPI project automatically. No PyPI
|
|
65
|
+
API token or GitHub secret is required. For an existing project, add the same
|
|
66
|
+
publisher in its PyPI publishing settings.
|
|
67
|
+
|
|
68
|
+
For later releases, update `[project].version` in `pyproject.toml` and publish a
|
|
69
|
+
GitHub release tagged `v<version>` or `<version>` at that commit. The workflow
|
|
70
|
+
rejects mismatched tags. Choose a version not already uploaded to PyPI; existing
|
|
71
|
+
releases cannot be overwritten. Draft releases and tag pushes alone do not
|
|
72
|
+
publish a package.
|
|
73
|
+
|
|
74
|
+
## Object hierarchy
|
|
75
|
+
|
|
76
|
+
```text
|
|
77
|
+
RLTracer
|
|
78
|
+
└── RolloutStep
|
|
79
|
+
└── PromptGroup
|
|
80
|
+
└── TrajectoryRef
|
|
81
|
+
└── Trajectory
|
|
82
|
+
└── messages
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
- **RLTracer** is the entry point for one training run. It receives a configured
|
|
86
|
+
adapter in its constructor.
|
|
87
|
+
- **RolloutStep** represents one RL training step, such as `step_21`.
|
|
88
|
+
- **PromptGroup** represents all sampled trajectories sharing the same initial
|
|
89
|
+
prompt within a rollout step.
|
|
90
|
+
- **TrajectoryRef** is a lightweight SQLite-backed reference. Calling
|
|
91
|
+
`load()` decodes its source trajectory.
|
|
92
|
+
- **Trajectory** is the decoded multi-turn record. Its `messages` field is a
|
|
93
|
+
normalized list of `{"role": ..., "content": ...}` dictionaries.
|
|
94
|
+
|
|
95
|
+
The canonical internal word is **trajectory**. A rollout is the sampled
|
|
96
|
+
trajectory produced by the RL system.
|
|
97
|
+
|
|
98
|
+
## Prime-RL adapter
|
|
99
|
+
|
|
100
|
+
`PrimeRLAdapter` currently supports Prime-RL rollout artifacts:
|
|
101
|
+
|
|
102
|
+
```text
|
|
103
|
+
<run>/rollouts/
|
|
104
|
+
step_1/
|
|
105
|
+
rank_0.bin
|
|
106
|
+
rank_1.bin
|
|
107
|
+
step_2/
|
|
108
|
+
rank_0.bin
|
|
109
|
+
...
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
Each `rank_*.bin` file is a MessagePack payload. It contains groups of packed
|
|
113
|
+
Qwen token sequences. One packed sequence can contain multiple independent
|
|
114
|
+
conversations, so the adapter decodes it with the configured tokenizer and
|
|
115
|
+
splits it at each new Qwen system turn.
|
|
116
|
+
|
|
117
|
+
## Lazy behavior
|
|
118
|
+
|
|
119
|
+
Opening a tracer and listing its steps is cheap:
|
|
120
|
+
|
|
121
|
+
```python
|
|
122
|
+
adapter = PrimeRLAdapter(
|
|
123
|
+
rollout_dir="/path/to/run/rollouts",
|
|
124
|
+
tokenizer_name="/path/to/tokenizer",
|
|
125
|
+
index_path="/tmp/prime-rl.sqlite", # optional
|
|
126
|
+
)
|
|
127
|
+
tracer = RLTracer(adapter)
|
|
128
|
+
|
|
129
|
+
tracer.steps() # discovers step directories only
|
|
130
|
+
```
|
|
131
|
+
|
|
132
|
+
No shard is read and no tokenizer is loaded until a prompt query is made:
|
|
133
|
+
|
|
134
|
+
```python
|
|
135
|
+
step = tracer.step(21)
|
|
136
|
+
prompt_groups = step.prompts()
|
|
137
|
+
```
|
|
138
|
+
|
|
139
|
+
The first `prompts()` call for a step:
|
|
140
|
+
|
|
141
|
+
1. Reads only that step's `rank_*.bin` files.
|
|
142
|
+
2. Decodes packed sequences with the adapter tokenizer.
|
|
143
|
+
3. Splits them into conversations.
|
|
144
|
+
4. Computes a stable hash of each conversation's pre-assistant prompt context.
|
|
145
|
+
5. Stores only lookup metadata in SQLite.
|
|
146
|
+
|
|
147
|
+
The SQLite index stores shard location, packed group/row location, conversation
|
|
148
|
+
position, prompt key, preview, and token count. It does not copy full token
|
|
149
|
+
sequences or decoded conversation content.
|
|
150
|
+
|
|
151
|
+
Loading a selected trajectory re-reads only the shard named in its SQLite row,
|
|
152
|
+
decodes the selected packed sequence, and returns the requested conversation:
|
|
153
|
+
|
|
154
|
+
```python
|
|
155
|
+
group = step.prompts()[0]
|
|
156
|
+
trajectory = group.trajectory(0).load()
|
|
157
|
+
|
|
158
|
+
for message in trajectory.messages:
|
|
159
|
+
print(message["role"], message["content"])
|
|
160
|
+
```
|
|
161
|
+
|
|
162
|
+
The index is invalidated and rebuilt for a step if its shard paths, sizes, or
|
|
163
|
+
modification times change.
|
|
164
|
+
|
|
165
|
+
## Adapter boundary
|
|
166
|
+
|
|
167
|
+
New RL systems should implement the methods consumed by `RLTracer`:
|
|
168
|
+
|
|
169
|
+
```python
|
|
170
|
+
class Adapter:
|
|
171
|
+
rollout_dir: Path
|
|
172
|
+
|
|
173
|
+
def list_steps(self) -> list[int]: ...
|
|
174
|
+
def ensure_step_indexed(self, step: int) -> None: ...
|
|
175
|
+
def list_prompts(self, step: int) -> list[tuple[str, int, str]]: ...
|
|
176
|
+
def list_trajectory_ids(self, step: int, prompt_key: str) -> list[int]: ...
|
|
177
|
+
def load_trajectory(self, trajectory_id: int) -> Trajectory: ...
|
|
178
|
+
```
|
|
179
|
+
|
|
180
|
+
The adapter owns all source-specific choices: artifact format, tokenizer,
|
|
181
|
+
decoding, prompt grouping, optional rewards/logprobs, and index location.
|
|
182
|
+
`RLTracer` stays provider-neutral.
|
|
183
|
+
|
|
184
|
+
## Training-data export
|
|
185
|
+
|
|
186
|
+
TraceDataExporter builds JSONL directly from one or more structured PrimeRL
|
|
187
|
+
trace runs. A run can be an experiment directory, its run_default directory,
|
|
188
|
+
or its rollouts directory.
|
|
189
|
+
|
|
190
|
+
SFT mode selects complete trajectories:
|
|
191
|
+
|
|
192
|
+
```python
|
|
193
|
+
from rltracer import (
|
|
194
|
+
Acceptance,
|
|
195
|
+
DifficultyFilter,
|
|
196
|
+
ExportPreferences,
|
|
197
|
+
FieldFilter,
|
|
198
|
+
TraceDataExporter,
|
|
199
|
+
TraceRunSpec,
|
|
200
|
+
)
|
|
201
|
+
|
|
202
|
+
exporter = TraceDataExporter(
|
|
203
|
+
runs=[
|
|
204
|
+
TraceRunSpec("/path/to/first/run", step_end=800),
|
|
205
|
+
TraceRunSpec("/path/to/second/run", step_end=400),
|
|
206
|
+
],
|
|
207
|
+
preferences=ExportPreferences(
|
|
208
|
+
mode="sft",
|
|
209
|
+
acceptance=Acceptance.CORRECT,
|
|
210
|
+
# Empirical prompt pass rate of 0.25-0.75.
|
|
211
|
+
difficulty=DifficultyFilter(
|
|
212
|
+
min_pass_rate=0.25,
|
|
213
|
+
max_pass_rate=0.75,
|
|
214
|
+
),
|
|
215
|
+
field_filters=(
|
|
216
|
+
FieldFilter("metrics.num_turns", maximum=8),
|
|
217
|
+
FieldFilter("rewards.reward.score", minimum=0),
|
|
218
|
+
),
|
|
219
|
+
window_size=100,
|
|
220
|
+
per_prompt=2,
|
|
221
|
+
target_count=16_000,
|
|
222
|
+
seed=17,
|
|
223
|
+
),
|
|
224
|
+
)
|
|
225
|
+
result = exporter.write_jsonl("correct_sft_16k.jsonl")
|
|
226
|
+
```
|
|
227
|
+
|
|
228
|
+
RL mode writes datadev RL rows whose prompt is the context preceding the first
|
|
229
|
+
assistant turn:
|
|
230
|
+
|
|
231
|
+
```python
|
|
232
|
+
rl_exporter = TraceDataExporter(
|
|
233
|
+
[
|
|
234
|
+
TraceRunSpec(
|
|
235
|
+
"/path/to/run",
|
|
236
|
+
step_end=800,
|
|
237
|
+
rl_data_path="/path/to/datadev_rl_train.jsonl",
|
|
238
|
+
)
|
|
239
|
+
],
|
|
240
|
+
ExportPreferences(
|
|
241
|
+
mode="rl",
|
|
242
|
+
acceptance="any",
|
|
243
|
+
difficulty=DifficultyFilter(min_difficulty=0.25),
|
|
244
|
+
# Use one pass-rate estimate over the complete step range.
|
|
245
|
+
window_size=None,
|
|
246
|
+
unique_prompts=True,
|
|
247
|
+
target_count=5_000,
|
|
248
|
+
),
|
|
249
|
+
)
|
|
250
|
+
rl_exporter.write_jsonl("rl_prompts_5k.jsonl")
|
|
251
|
+
```
|
|
252
|
+
|
|
253
|
+
Difficulty is identical to prompt_pass_rate. Pass rate is calculated from all scored
|
|
254
|
+
rollouts for a prompt in each configured step window. RL mode emits at most one
|
|
255
|
+
copy of a prompt per run/window; SFT mode honors per_prompt. Global targets are
|
|
256
|
+
deterministic and fair: every eligible prompt group receives its first row
|
|
257
|
+
before any group receives a second.
|
|
258
|
+
|
|
259
|
+
Set `unique_prompts=True` to keep at most one row for a prompt across all runs
|
|
260
|
+
and step windows. Deduplication is deterministic under `seed` and happens
|
|
261
|
+
before applying `target_count`.
|
|
262
|
+
|
|
263
|
+
Set `tool_call_sampling_power` to softly reduce the prominence of traces with
|
|
264
|
+
few command/tool calls. Selection weight is
|
|
265
|
+
`tool_call_count ** tool_call_sampling_power`, applied both within prompt
|
|
266
|
+
groups and when filling the global target. `0` preserves uniform selection,
|
|
267
|
+
`0.5` is a gentle square-root bias, `1` is proportional to command count, and
|
|
268
|
+
`2` is a strong quadratic bias. Selection remains deterministic under `seed`,
|
|
269
|
+
and export summaries include available and selected tool-call histograms.
|
|
270
|
+
|
|
271
|
+
By default, output rows match `datadev data` exactly. SFT rows contain
|
|
272
|
+
`messages` and optional JSON-string `tools`. RL rows are losslessly joined
|
|
273
|
+
through `task.data.idx` to `rl_data_path` and contain `prompt`, `answer`,
|
|
274
|
+
`output`, `feedback`, `record_id`, `rubric_index`, `source`, and list-valued
|
|
275
|
+
`tools`. Set `include_provenance=True` only when an additional
|
|
276
|
+
`trace_provenance` field is acceptable.
|
|
277
|
+
|
|
278
|
+
SFT tool calls are canonicalized to OpenAI's nested representation:
|
|
279
|
+
`{"id": ..., "type": "function", "function": {"name": ..., "arguments":
|
|
280
|
+
...}}`. Legacy flattened and singular calls are converted during export, and
|
|
281
|
+
calls without a function name are rejected rather than emitted as corrupted
|
|
282
|
+
training targets.
|
|
283
|
+
|
|
284
|
+
## Next extensions
|
|
285
|
+
|
|
286
|
+
- Add reward, advantage, log-probability, tool-call, and completion metadata to
|
|
287
|
+
`Trajectory.metadata`.
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "rltracer"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Lazy inspection, analysis, visualization, and export tooling for reinforcement-learning traces."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
authors = [{name = "Goncalo Faria"}]
|
|
13
|
+
dependencies = ["fire>=0.7.1"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 3 - Alpha",
|
|
16
|
+
"Programming Language :: Python :: 3",
|
|
17
|
+
"Programming Language :: Python :: 3.10",
|
|
18
|
+
"Programming Language :: Python :: 3.11",
|
|
19
|
+
"Programming Language :: Python :: 3.12",
|
|
20
|
+
"Programming Language :: Python :: 3.13",
|
|
21
|
+
]
|
|
22
|
+
|
|
23
|
+
[project.optional-dependencies]
|
|
24
|
+
primerl = ["msgpack>=1.0", "transformers>=4.45"]
|
|
25
|
+
retemplate = ["pyarrow>=16"]
|
|
26
|
+
plot = ["matplotlib>=3.8", "numpy>=1.26", "scipy>=1.12"]
|
|
27
|
+
speed = ["orjson>=3.10"]
|
|
28
|
+
test = ["pytest>=8", "pyarrow>=16", "msgpack>=1.0"]
|
|
29
|
+
all = [
|
|
30
|
+
"msgpack>=1.0",
|
|
31
|
+
"transformers>=4.45",
|
|
32
|
+
"pyarrow>=16",
|
|
33
|
+
"matplotlib>=3.8",
|
|
34
|
+
"numpy>=1.26",
|
|
35
|
+
"scipy>=1.12",
|
|
36
|
+
"orjson>=3.10",
|
|
37
|
+
]
|
|
38
|
+
|
|
39
|
+
[project.urls]
|
|
40
|
+
Homepage = "https://github.com/goncalorafaria/rltracer"
|
|
41
|
+
Repository = "https://github.com/goncalorafaria/rltracer"
|
|
42
|
+
Issues = "https://github.com/goncalorafaria/rltracer/issues"
|
|
43
|
+
|
|
44
|
+
[project.scripts]
|
|
45
|
+
rltracer = "rltracer.__main__:main"
|
|
46
|
+
rltracer-viz = "rltracer.viz.server:main"
|
|
47
|
+
rltracer-retemplate = "rltracer.retemplate:cli"
|
|
48
|
+
|
|
49
|
+
[tool.setuptools.packages.find]
|
|
50
|
+
where = ["src"]
|
|
51
|
+
include = ["rltracer*"]
|
|
52
|
+
|
|
53
|
+
[tool.setuptools.package-data]
|
|
54
|
+
rltracer = ["py.typed"]
|
|
55
|
+
"rltracer.viz" = ["*.html", "*.js", "*.css"]
|
|
56
|
+
|
|
57
|
+
[tool.pytest.ini_options]
|
|
58
|
+
testpaths = ["tests"]
|
rltracer-0.1.0/setup.cfg
ADDED