Skip to content

Commit dd33de9

Browse files
EssozEssoz
authored
Merge pull request #141 from OrderLab/codex/inference-indexes
Index trace contexts and call records during invariant inference
2 parents 6759bd6 + 0866977 commit dd33de9

9 files changed

Lines changed: 1528 additions & 65 deletions

File tree

‎docs/measurements/trace-pandas-indexes-20261025.json‎

Lines changed: 319 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,319 @@
1+
{
2+
"measurement_date": "2026-10-25",
3+
"source": {
4+
"repository": "https://github.com/OrderLab/TrainCheck",
5+
"base_commit": "9eb4f320d3db34d372145261704b805b88e39bda",
6+
"unchanged_prerequisites": [
7+
"deterministic inference sampling",
8+
"example reservoir cap fixes"
9+
],
10+
"trace_pandas_sha256": {
11+
"legacy": "e1bb62eae653f95c5498dc921d991af013e8b5717398f94cfbc718812bb600ad",
12+
"context_index": "13d5d1273595f8183d0511ae288a998cb57e3ea66dc232e4bfec503739de825c",
13+
"context_and_call_index": "8e4c3589371449b78430ed5039d0159d3d0435664b218f3a96c81935be0d75d5"
14+
},
15+
"pr_copies": {
16+
"variants": {
17+
"context_index": {
18+
"measured_source_sha256": "13d5d1273595f8183d0511ae288a998cb57e3ea66dc232e4bfec503739de825c",
19+
"pr_source_sha256": "52c296917bba877b1e87aebb54cfee1e05813e2ae32f85744d5ddfe27383d5e0",
20+
"equivalent_ast_after_normalizing_only_documented_typing_change": true
21+
},
22+
"context_and_call_index": {
23+
"measured_source_sha256": "8e4c3589371449b78430ed5039d0159d3d0435664b218f3a96c81935be0d75d5",
24+
"pr_source_sha256": "58aebdba9955834f4d6b152527e8ba2f9166f6e62f43a098c48b115b0bbc48f5",
25+
"equivalent_ast_after_normalizing_only_documented_typing_change": true
26+
}
27+
},
28+
"changes_from_measured_sources": [
29+
"Both PR copies widen the inherited call-event cache annotation to include IncompleteFuncCallEvent, matching the base Trace contract, and import that existing type.",
30+
"Final PR copy adds narrow E721 suppressions to preserve intentional exact-string checks; only comments/parentheses changed in lookup logic.",
31+
"Existing explicit None-or-cache annotations are retained for the new lazy index."
32+
],
33+
"no_frozen_source_edits": true,
34+
"local_validation_tools": {
35+
"ruff": "0.15.7",
36+
"mypy": "1.19.1",
37+
"black": "26.3.1",
38+
"isort": "6.0.1"
39+
},
40+
"pinned_repository_hooks_need_final_worktree_verification": true
41+
}
42+
},
43+
"saved_trace": {
44+
"events": 573494,
45+
"model_parameters": 6163712,
46+
"training_steps": 10,
47+
"archive_sha256": "62022da5a997c16ffe49aa44486a495b651c8758199f29732a216c7cf257559d",
48+
"raw_inputs_unchanged": true,
49+
"collection_pytorch": "2.16.0.dev20261012+cu130"
50+
},
51+
"runtime": {
52+
"platform": "Linux",
53+
"python": "3.10.16",
54+
"pandas": "2.3.3",
55+
"numpy": "1.26.4",
56+
"torch": "2.2.2+cu121",
57+
"workers": 1,
58+
"native_threads": 1
59+
},
60+
"preparation": {
61+
"runs": {
62+
"legacy": {
63+
"context_index_seconds": 159.43910144269466,
64+
"load_seconds": 90.64373006299138,
65+
"rows": 573494,
66+
"stage_rows": {
67+
"training": 573494
68+
},
69+
"stage_split_seconds": 80.21087664365768,
70+
"total_seconds": 170.85460681840777,
71+
"sampled_peak_pss_bytes": 2409772032,
72+
"context_states": [
73+
{
74+
"call": 0,
75+
"contexts": 711,
76+
"rows": 573494,
77+
"seconds": 79.72032228112221,
78+
"states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
79+
},
80+
{
81+
"call": 1,
82+
"contexts": 711,
83+
"rows": 573494,
84+
"seconds": 79.71877916157246,
85+
"states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
86+
}
87+
]
88+
},
89+
"context_index": {
90+
"context_index_seconds": 0.8803581297397614,
91+
"load_seconds": 10.837436098605394,
92+
"rows": 573494,
93+
"stage_rows": {
94+
"training": 573494
95+
},
96+
"stage_split_seconds": 0.8764883801341057,
97+
"total_seconds": 11.713924750685692,
98+
"sampled_peak_pss_bytes": 2411297792,
99+
"context_states": [
100+
{
101+
"call": 0,
102+
"contexts": 711,
103+
"rows": 573494,
104+
"seconds": 0.4527508243918419,
105+
"states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
106+
},
107+
{
108+
"call": 1,
109+
"contexts": 711,
110+
"rows": 573494,
111+
"seconds": 0.42760730534791946,
112+
"states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
113+
}
114+
]
115+
}
116+
},
117+
"comparison": {
118+
"same_number_of_calls": true,
119+
"calls": [
120+
{
121+
"call": 0,
122+
"rows": 573494,
123+
"contexts": 711,
124+
"equal": true
125+
},
126+
{
127+
"call": 1,
128+
"rows": 573494,
129+
"contexts": 711,
130+
"equal": true
131+
}
132+
],
133+
"all_equal": true
134+
},
135+
"paired_repetitions": 1,
136+
"timing_scope": "Load/init plus stage splitting, excluding process startup; includes context-state serialization. Context-index timers exclude serialization. PSS peaks are sampled."
137+
},
138+
"large_trace_diagnostic": {
139+
"runs": {
140+
"optimized-infer-profile": {
141+
"wall_seconds": 546.4971170388162,
142+
"exit_code": 1,
143+
"stop_reason": null,
144+
"source_unchanged": true,
145+
"traces_unchanged": true,
146+
"error": "ValueError: Too many positional arguments for function torch._prims_common.is_contiguous, expecting 1 (['a']) but got 2",
147+
"lookups": [
148+
{
149+
"function": "get_pre_func_call_record",
150+
"primitive_calls": 6891,
151+
"calls": 6891,
152+
"self_seconds": 0.087277701,
153+
"cumulative_seconds": 187.91447304200003
154+
},
155+
{
156+
"function": "get_post_func_call_record",
157+
"primitive_calls": 13498,
158+
"calls": 13498,
159+
"self_seconds": 0.18074669200000001,
160+
"cumulative_seconds": 289.47232359500003
161+
}
162+
]
163+
},
164+
"call-index-infer-profile": {
165+
"wall_seconds": 71.98188185319304,
166+
"exit_code": 1,
167+
"stop_reason": null,
168+
"source_unchanged": true,
169+
"traces_unchanged": true,
170+
"error": "ValueError: Too many positional arguments for function torch._prims_common.is_contiguous, expecting 1 (['a']) but got 2",
171+
"lookups": [
172+
{
173+
"function": "get_pre_func_call_record",
174+
"primitive_calls": 6891,
175+
"calls": 6891,
176+
"self_seconds": 0.020590995,
177+
"cumulative_seconds": 1.034124108
178+
},
179+
{
180+
"function": "get_post_func_call_record",
181+
"primitive_calls": 13498,
182+
"calls": 13498,
183+
"self_seconds": 0.038829666000000006,
184+
"cumulative_seconds": 3.888496781
185+
}
186+
]
187+
}
188+
},
189+
"same_error": true,
190+
"completed_inference": false
191+
},
192+
"diagnostic_scope": {
193+
"completed_inference": false,
194+
"profiled": true,
195+
"disabled_relations": [
196+
"FunctionLeadRelation",
197+
"FunctionCoverRelation"
198+
],
199+
"same_error_is_framework_version_mismatch": true,
200+
"eligible_containment_apis_in_replay": 35,
201+
"eligible_containment_apis_in_collection_environment_log": 41,
202+
"time_to_error_is_not_completed_inference_throughput": true
203+
},
204+
"correctness": {
205+
"differential_tests": 32,
206+
"context_randomized_traces": 80,
207+
"independent_extra_call_query_comparisons": 640,
208+
"small_replay_run_count": 9,
209+
"all_full_json_comparisons_equal": true,
210+
"comparison": "Structural JSON multisets; object-key and file-line order ignored, multiplicity and array order retained.",
211+
"runs": [
212+
{
213+
"case": "mnist-fixed-seed42",
214+
"exit_code": 0,
215+
"workers": 1,
216+
"valid_outputs": true,
217+
"source": "base"
218+
},
219+
{
220+
"case": "mnist-fixed-seed42",
221+
"exit_code": 0,
222+
"workers": 1,
223+
"valid_outputs": true,
224+
"source": "optimized"
225+
},
226+
{
227+
"case": "transformer-fixed-seed42",
228+
"exit_code": 0,
229+
"workers": 1,
230+
"valid_outputs": true,
231+
"source": "optimized"
232+
},
233+
{
234+
"case": "transformer-fixed-seed42",
235+
"exit_code": 0,
236+
"workers": 1,
237+
"valid_outputs": true,
238+
"source": "base"
239+
},
240+
{
241+
"case": "dcgan-fixed-seed42",
242+
"exit_code": 0,
243+
"workers": 1,
244+
"valid_outputs": true,
245+
"source": "base"
246+
},
247+
{
248+
"case": "dcgan-fixed-seed42",
249+
"exit_code": 0,
250+
"workers": 1,
251+
"valid_outputs": true,
252+
"source": "optimized"
253+
},
254+
{
255+
"case": "mnist-fixed-seed42",
256+
"exit_code": 0,
257+
"workers": 1,
258+
"valid_outputs": true,
259+
"source": "context_and_call_index"
260+
},
261+
{
262+
"case": "transformer-fixed-seed42",
263+
"exit_code": 0,
264+
"workers": 1,
265+
"valid_outputs": true,
266+
"source": "context_and_call_index"
267+
},
268+
{
269+
"case": "dcgan-fixed-seed42",
270+
"exit_code": 0,
271+
"workers": 1,
272+
"valid_outputs": true,
273+
"source": "context_and_call_index"
274+
}
275+
],
276+
"outputs_by_workload": {
277+
"mnist-fixed-seed42": {
278+
"invariants": {
279+
"canonical_multiset_sha256": "5691e461a1fba5bf5314b18a62bf72f82fc95e56f71b26eca87218e53f321e84",
280+
"distinct_records": 193,
281+
"records": 193
282+
},
283+
"failed_hypotheses": {
284+
"canonical_multiset_sha256": "8e39cf710d943898ca007478953584d363db135ce3cb458619e71836ae179820",
285+
"distinct_records": 21,
286+
"records": 21
287+
}
288+
},
289+
"transformer-fixed-seed42": {
290+
"invariants": {
291+
"canonical_multiset_sha256": "21332736553d2fe20c1fdb1e2c852597950e3b799bc399d47007a2b1f7a71b09",
292+
"distinct_records": 518,
293+
"records": 518
294+
},
295+
"failed_hypotheses": {
296+
"canonical_multiset_sha256": "b3424febf9358e01c3e079362e17bc68a939a933e14f1cfab954884482073bd7",
297+
"distinct_records": 10,
298+
"records": 10
299+
}
300+
},
301+
"dcgan-fixed-seed42": {
302+
"invariants": {
303+
"canonical_multiset_sha256": "596f82fcfb946a7fe2fe138fa9e61302bafb93465fdc529317ccf93d2da1022a",
304+
"distinct_records": 512,
305+
"records": 512
306+
},
307+
"failed_hypotheses": {
308+
"canonical_multiset_sha256": "dfba6a7ed565ce66163b9eb98921ad05857649e37f78075b48fdf9341b912204",
309+
"distinct_records": 86,
310+
"records": 86
311+
}
312+
}
313+
}
314+
},
315+
"limitations": [
316+
"Event row order, call IDs, and event types must remain stable after initialization. Frame replacement or a length change rebuilds the call index; same-size in-place changes or reordering are unsupported.",
317+
"The large trace uses a different PyTorch version; both diagnostic variants end at the same signature error. No complete large-trace inference or online-checker overhead result is claimed."
318+
]
319+
}

‎docs/measurements/trace-pandas-indexes.md‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# TracePandas indexing measurements
2+
3+
The context index removes two full-frame scans per context instance. The call-ID index removes repeated full-frame scans for pre/post lookups, including lookups made before containment sampling. The optimizations preserve selection coverage and record-validation behavior.
4+
5+
The adjacent `trace-pandas-indexes-20261025.json` records sanitized measurements, source/input hashes, environment versions, state equality, profiler call counts, and nine successful small-workload replay results. The complete saved Llama trace contains 573,494 events from ten training steps of a 6,163,712-parameter model.
6+
7+
| Measurement | Before | After |
8+
|---|---:|---:|
9+
| Preparation, original → context index | 170.855 s | 11.714 s |
10+
| Context-index work within preparation | 159.439 s | 0.880 s |
11+
| Pre-record profiling, context → combined | 187.914 s / 6,891 calls | 1.034 s / 6,891 calls |
12+
| Post-record profiling, context → combined | 289.472 s / 13,498 calls | 3.888 s / 13,498 calls |
13+
| Diagnostic wall time to the same error | 546.50 s | 71.98 s |
14+
15+
Preparation is one paired measurement on the same Linux host using one native thread. Total preparation includes state serialization and excludes process startup; context-index timers exclude serialization. Both preparation variants produce the same 711 ordered states in the full trace and stage trace. The profiler's cumulative times overlap across callers and must not be added indiscriminately.
16+
17+
The large-trace diagnostics both fail at `torch._prims_common.is_contiguous` argument binding: replay PyTorch 2.2.2 expects one argument while the PyTorch 2.16.0.dev20261012 trace contains two. These results demonstrate reduced search cost, not completed large-trace inference throughput or online-checker overhead. The replay environment also resolves fewer eligible containment APIs than the collection environment (35 versus 41), so a final reference run needs compatible framework dependencies.
18+
19+
Nine small-workload runs complete successfully: original/context/combined variants for MNIST, small Transformer, and DCGAN. Full invariant and failed-hypothesis JSON match under structural multiset comparison, retaining multiplicity and array order. Per-workload counts are 193/21, 518/10, and 512/86.
20+
21+
## Reproduce local correctness and synthetic measurements
22+
23+
From a checkout with TrainCheck dependencies installed:
24+
25+
```sh
26+
OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m unittest tests.test_context_index tests.test_call_index
27+
OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m tools.benchmarks.trace_pandas_indexes --output measurements.json
28+
```
29+
30+
The synthetic benchmark constructs its inputs and needs no external trace or cluster. It alternates variant order, checks exact results, and reports every repetition and the median. Its context case times construction plus state serialization; its call case times construction and all lookups, including the first lazy-index build. Its sizes/repetition count are configurable; these synthetic timings are separate from the saved-trace measurements above.
31+
32+
The oracle fixture contains only the exact original context/pre/post/stage methods from commit `9eb4f320d3db34d372145261704b805b88e39bda`; supporting types and the constructor come from the package under test. Provenance includes the original file and method hashes. The fixture's `TracePandas` binding ensures stage splitting creates legacy instances, and the tests assert this explicitly.
33+
34+
The published source also corrects the inherited call-event cache annotation to include `IncompleteFuncCallEvent`; narrow lint comments preserve exact-string guards. These typing/comment changes do not alter lookup behavior.
35+
36+
## Cache contract
37+
38+
Event rows, their order, call IDs, and event types must remain stable after initialization. Frame replacement or a length change rebuilds the call index. Same-size in-place edits or row reordering can leave stale positions and require a new `TracePandas` instance. The tested inference path does not perform those mutations. This change does not introduce general cache invalidation for arbitrary dataframe mutations.

‎tests/fixtures/__init__.py‎

Whitespace-only changes.

‎tests/fixtures/legacy_trace_pandas.provenance.json‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
{
2+
"repository": "https://github.com/OrderLab/TrainCheck",
3+
"commit": "9eb4f320d3db34d372145261704b805b88e39bda",
4+
"source_path": "traincheck/trace/trace_pandas.py",
5+
"source_sha256": "e1bb62eae653f95c5498dc921d991af013e8b5717398f94cfbc718812bb600ad",
6+
"method_sha256": {
7+
"_index_context_manager_meta_vars": "6b7e5646260ee0ce7377c7cf19edb4774bdf69825e15cf10ff3c29cd4555c13c",
8+
"get_pre_func_call_record": "bf89fcd11b98e08a3029ee6ff3b2a4b6df980aa49bfe6c13bcdbafcb10e8d74c",
9+
"get_post_func_call_record": "a36cd5bc217f8e85f0f30775c24480f0060b6f3496e8fb2a2cb45a4214d46c07",
10+
"get_traces_for_stage": "583c5adfd122d86027f427beea7c7823783e14b37f078d01740f55cdd72f04f3"
11+
},
12+
"hash_convention": "Original indented method source including its final newline. No source transformations."
13+
}

0 commit comments

Comments
 (0)