[Essoz]

TrainCheck repeatedly scans the full trace for every context-manager instance and pre/post call-record lookup. On a complete 573,494-event tiny-Llama trace, context indexing reduced preparation from 170.855 to 11.714 seconds while preserving the same 711 ordered context states in both the full trace and its stage trace.

This change groups context events once by function, event type, process, thread, and object ID, retaining first-row selection and strict timestamp comparisons. A second commit adds a lazy per-trace call-ID index of pre/post row positions. Duplicate assertions, exception records, missing-post behavior, and fresh returned dictionaries remain intact; unsupported IDs/dtypes keep the original filters. Nullable and unsorted context timestamps retain pandas comparison semantics.

Measured work	Before	After	Comparison
Full-trace preparation	170.855 s	11.714 s	Original → context index
Pre-record cumulative profile time	187.914 s	1.034 s	Context index → context + call index; 6,891 calls each
Post-record cumulative profile time	289.472 s	3.888 s	Context index → context + call index; 13,498 calls each

Preparation is one paired measurement on the same Linux host with one native thread. The profiled diagnostic reached the same PyTorch-signature error in 546.50 versus 71.98 seconds. Neither large-trace inference completed: the trace was collected with PyTorch 2.16.0.dev20261012, but replay used 2.2.2. Time to that error is not a completed-inference throughput or online-checker overhead result.

Validation:

• 32 differential tests against four frozen legacy methods, including nullable UInt64 precision, unsorted/equal times, reused IDs, stage isolation, duplicate rows, exception/missing posts, and fresh dictionaries.
• Nine completed small-workload replays: original, context-only, and combined indexes on MNIST, small Transformer, and DCGAN. Full invariant/failed-hypothesis JSON matched, preserving multiplicity and array order (counts: 193/21, 518/10, 512/86).
• Source and input hashes remained unchanged. A portable synthetic benchmark and sanitized measurement JSON accompany the tests.

The call index assumes event row order, call IDs, and event types remain stable after initialization. Replacing the dataframe or changing its length rebuilds the index; same-size in-place edits or reordering are unsupported.

Run the portable tests with python -m unittest tests.test_context_index tests.test_call_index; run the synthetic benchmark with python -m tools.benchmarks.trace_pandas_indexes --output measurements.json.

Engineering hardening

Typed context-index keys, state, timestamps at the supported boundary, and the reusable benchmark's inputs/results (336fc1b), with signed parent-update commits. All 32 context/call differential tests and benchmark smoke checks pass; pinned hooks and stronger Mypy 1.9 checking of the benchmark pass. One narrowly documented Any comparison boundary remains for legacy pandas scalar semantics; values are not coerced or errors suppressed.

Separately composed updated production sources pass 140 affected Linux regression tests, whole-source Mypy (77 files), Ruff, and isort. Whole-source Black still reports three unchanged legacy files, down from six in the original base; changed-file hooks pass. Strict checks are scoped and do not establish repository-wide strict typing. Historical performance numbers above retain their original measured revisions; this hardening pass does not re-estimate them.