Index trace contexts and call records during invariant inference

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

@@ -0,0 +1,319 @@
+{
+  "measurement_date": "2026-10-25",
+  "source": {
+    "repository": "https://github.com/OrderLab/TrainCheck",
+    "base_commit": "9eb4f320d3db34d372145261704b805b88e39bda",
+    "unchanged_prerequisites": [
+      "deterministic inference sampling",
+      "example reservoir cap fixes"
+    ],
+    "trace_pandas_sha256": {
+      "legacy": "e1bb62eae653f95c5498dc921d991af013e8b5717398f94cfbc718812bb600ad",
+      "context_index": "13d5d1273595f8183d0511ae288a998cb57e3ea66dc232e4bfec503739de825c",
+      "context_and_call_index": "8e4c3589371449b78430ed5039d0159d3d0435664b218f3a96c81935be0d75d5"
+    },
+    "pr_copies": {
+      "variants": {
+        "context_index": {
+          "measured_source_sha256": "13d5d1273595f8183d0511ae288a998cb57e3ea66dc232e4bfec503739de825c",
+          "pr_source_sha256": "52c296917bba877b1e87aebb54cfee1e05813e2ae32f85744d5ddfe27383d5e0",
+          "equivalent_ast_after_normalizing_only_documented_typing_change": true
+        },
+        "context_and_call_index": {
+          "measured_source_sha256": "8e4c3589371449b78430ed5039d0159d3d0435664b218f3a96c81935be0d75d5",
+          "pr_source_sha256": "58aebdba9955834f4d6b152527e8ba2f9166f6e62f43a098c48b115b0bbc48f5",
+          "equivalent_ast_after_normalizing_only_documented_typing_change": true
+        }
+      },
+      "changes_from_measured_sources": [
+        "Both PR copies widen the inherited call-event cache annotation to include IncompleteFuncCallEvent, matching the base Trace contract, and import that existing type.",
+        "Final PR copy adds narrow E721 suppressions to preserve intentional exact-string checks; only comments/parentheses changed in lookup logic.",
+        "Existing explicit None-or-cache annotations are retained for the new lazy index."
+      ],
+      "no_frozen_source_edits": true,
+      "local_validation_tools": {
+        "ruff": "0.15.7",
+        "mypy": "1.19.1",
+        "black": "26.3.1",
+        "isort": "6.0.1"
+      },
+      "pinned_repository_hooks_need_final_worktree_verification": true
+    }
+  },
+  "saved_trace": {
+    "events": 573494,
+    "model_parameters": 6163712,
+    "training_steps": 10,
+    "archive_sha256": "62022da5a997c16ffe49aa44486a495b651c8758199f29732a216c7cf257559d",
+    "raw_inputs_unchanged": true,
+    "collection_pytorch": "2.16.0.dev20261012+cu130"
+  },
+  "runtime": {
+    "platform": "Linux",
+    "python": "3.10.16",
+    "pandas": "2.3.3",
+    "numpy": "1.26.4",
+    "torch": "2.2.2+cu121",
+    "workers": 1,
+    "native_threads": 1
+  },
+  "preparation": {
+    "runs": {
+      "legacy": {
+        "context_index_seconds": 159.43910144269466,
+        "load_seconds": 90.64373006299138,
+        "rows": 573494,
+        "stage_rows": {
+          "training": 573494
+        },
+        "stage_split_seconds": 80.21087664365768,
+        "total_seconds": 170.85460681840777,
+        "sampled_peak_pss_bytes": 2409772032,
+        "context_states": [
+          {
+            "call": 0,
+            "contexts": 711,
+            "rows": 573494,
+            "seconds": 79.72032228112221,
+            "states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
+          },
+          {
+            "call": 1,
+            "contexts": 711,
+            "rows": 573494,
+            "seconds": 79.71877916157246,
+            "states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
+          }
+        ]
+      },
+      "context_index": {
+        "context_index_seconds": 0.8803581297397614,
+        "load_seconds": 10.837436098605394,
+        "rows": 573494,
+        "stage_rows": {
+          "training": 573494
+        },
+        "stage_split_seconds": 0.8764883801341057,
+        "total_seconds": 11.713924750685692,
+        "sampled_peak_pss_bytes": 2411297792,
+        "context_states": [
+          {
+            "call": 0,
+            "contexts": 711,
+            "rows": 573494,
+            "seconds": 0.4527508243918419,
+            "states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
+          },
+          {
+            "call": 1,
+            "contexts": 711,
+            "rows": 573494,
+            "seconds": 0.42760730534791946,
+            "states_sha256": "a5ddb221a3679a50c476776b19aaa3cbc5277174775c8559eb0dcef0a28e2dd0"
+          }
+        ]
+      }
+    },
+    "comparison": {
+      "same_number_of_calls": true,
+      "calls": [
+        {
+          "call": 0,
+          "rows": 573494,
+          "contexts": 711,
+          "equal": true
+        },
+        {
+          "call": 1,
+          "rows": 573494,
+          "contexts": 711,
+          "equal": true
+        }
+      ],
+      "all_equal": true
+    },
+    "paired_repetitions": 1,
+    "timing_scope": "Load/init plus stage splitting, excluding process startup; includes context-state serialization. Context-index timers exclude serialization. PSS peaks are sampled."
+  },
+  "large_trace_diagnostic": {
+    "runs": {
+      "optimized-infer-profile": {
+        "wall_seconds": 546.4971170388162,
+        "exit_code": 1,
+        "stop_reason": null,
+        "source_unchanged": true,
+        "traces_unchanged": true,
+        "error": "ValueError: Too many positional arguments for function torch._prims_common.is_contiguous, expecting 1 (['a']) but got 2",
+        "lookups": [
+          {
+            "function": "get_pre_func_call_record",
+            "primitive_calls": 6891,
+            "calls": 6891,
+            "self_seconds": 0.087277701,
+            "cumulative_seconds": 187.91447304200003
+          },
+          {
+            "function": "get_post_func_call_record",
+            "primitive_calls": 13498,
+            "calls": 13498,
+            "self_seconds": 0.18074669200000001,
+            "cumulative_seconds": 289.47232359500003
+          }
+        ]
+      },
+      "call-index-infer-profile": {
+        "wall_seconds": 71.98188185319304,
+        "exit_code": 1,
+        "stop_reason": null,
+        "source_unchanged": true,
+        "traces_unchanged": true,
+        "error": "ValueError: Too many positional arguments for function torch._prims_common.is_contiguous, expecting 1 (['a']) but got 2",
+        "lookups": [
+          {
+            "function": "get_pre_func_call_record",
+            "primitive_calls": 6891,
+            "calls": 6891,
+            "self_seconds": 0.020590995,
+            "cumulative_seconds": 1.034124108
+          },
+          {
+            "function": "get_post_func_call_record",
+            "primitive_calls": 13498,
+            "calls": 13498,
+            "self_seconds": 0.038829666000000006,
+            "cumulative_seconds": 3.888496781
+          }
+        ]
+      }
+    },
+    "same_error": true,
+    "completed_inference": false
+  },
+  "diagnostic_scope": {
+    "completed_inference": false,
+    "profiled": true,
+    "disabled_relations": [
+      "FunctionLeadRelation",
+      "FunctionCoverRelation"
+    ],
+    "same_error_is_framework_version_mismatch": true,
+    "eligible_containment_apis_in_replay": 35,
+    "eligible_containment_apis_in_collection_environment_log": 41,
+    "time_to_error_is_not_completed_inference_throughput": true
+  },
+  "correctness": {
+    "differential_tests": 32,
+    "context_randomized_traces": 80,
+    "independent_extra_call_query_comparisons": 640,
+    "small_replay_run_count": 9,
+    "all_full_json_comparisons_equal": true,
+    "comparison": "Structural JSON multisets; object-key and file-line order ignored, multiplicity and array order retained.",
+    "runs": [
+      {
+        "case": "mnist-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "base"
+      },
+      {
+        "case": "mnist-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "optimized"
+      },
+      {
+        "case": "transformer-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "optimized"
+      },
+      {
+        "case": "transformer-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "base"
+      },
+      {
+        "case": "dcgan-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "base"
+      },
+      {
+        "case": "dcgan-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "optimized"
+      },
+      {
+        "case": "mnist-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "context_and_call_index"
+      },
+      {
+        "case": "transformer-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "context_and_call_index"
+      },
+      {
+        "case": "dcgan-fixed-seed42",
+        "exit_code": 0,
+        "workers": 1,
+        "valid_outputs": true,
+        "source": "context_and_call_index"
+      }
+    ],
+    "outputs_by_workload": {
+      "mnist-fixed-seed42": {
+        "invariants": {
+          "canonical_multiset_sha256": "5691e461a1fba5bf5314b18a62bf72f82fc95e56f71b26eca87218e53f321e84",
+          "distinct_records": 193,
+          "records": 193
+        },
+        "failed_hypotheses": {
+          "canonical_multiset_sha256": "8e39cf710d943898ca007478953584d363db135ce3cb458619e71836ae179820",
+          "distinct_records": 21,
+          "records": 21
+        }
+      },
+      "transformer-fixed-seed42": {
+        "invariants": {
+          "canonical_multiset_sha256": "21332736553d2fe20c1fdb1e2c852597950e3b799bc399d47007a2b1f7a71b09",
+          "distinct_records": 518,
+          "records": 518
+        },
+        "failed_hypotheses": {
+          "canonical_multiset_sha256": "b3424febf9358e01c3e079362e17bc68a939a933e14f1cfab954884482073bd7",
+          "distinct_records": 10,
+          "records": 10
+        }
+      },
+      "dcgan-fixed-seed42": {
+        "invariants": {
+          "canonical_multiset_sha256": "596f82fcfb946a7fe2fe138fa9e61302bafb93465fdc529317ccf93d2da1022a",
+          "distinct_records": 512,
+          "records": 512
+        },
+        "failed_hypotheses": {
+          "canonical_multiset_sha256": "dfba6a7ed565ce66163b9eb98921ad05857649e37f78075b48fdf9341b912204",
+          "distinct_records": 86,
+          "records": 86
+        }
+      }
+    }
+  },
+  "limitations": [
+    "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.",
+    "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."
+  ]
+}

docs/measurements/trace-pandas-indexes.md

@@ -0,0 +1,38 @@
+# TracePandas indexing measurements
+
+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.
+
+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.
+
+| Measurement | Before | After |
+|---|---:|---:|
+| Preparation, original → context index | 170.855 s | 11.714 s |
+| Context-index work within preparation | 159.439 s | 0.880 s |
+| Pre-record profiling, context → combined | 187.914 s / 6,891 calls | 1.034 s / 6,891 calls |
+| Post-record profiling, context → combined | 289.472 s / 13,498 calls | 3.888 s / 13,498 calls |
+| Diagnostic wall time to the same error | 546.50 s | 71.98 s |
+
+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.
+
+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.
+
+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.
+
+## Reproduce local correctness and synthetic measurements
+
+From a checkout with TrainCheck dependencies installed:
+
+```sh
+OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m unittest tests.test_context_index tests.test_call_index
+OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m tools.benchmarks.trace_pandas_indexes --output measurements.json
+```
+
+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.
+
+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.
+
+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.
+
+## Cache contract
+
+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/legacy_trace_pandas.provenance.json

@@ -0,0 +1,13 @@
+{
+  "repository": "https://github.com/OrderLab/TrainCheck",
+  "commit": "9eb4f320d3db34d372145261704b805b88e39bda",
+  "source_path": "traincheck/trace/trace_pandas.py",
+  "source_sha256": "e1bb62eae653f95c5498dc921d991af013e8b5717398f94cfbc718812bb600ad",
+  "method_sha256": {
+    "_index_context_manager_meta_vars": "6b7e5646260ee0ce7377c7cf19edb4774bdf69825e15cf10ff3c29cd4555c13c",
+    "get_pre_func_call_record": "bf89fcd11b98e08a3029ee6ff3b2a4b6df980aa49bfe6c13bcdbafcb10e8d74c",
+    "get_post_func_call_record": "a36cd5bc217f8e85f0f30775c24480f0060b6f3496e8fb2a2cb45a4214d46c07",
+    "get_traces_for_stage": "583c5adfd122d86027f427beea7c7823783e14b37f078d01740f55cdd72f04f3"
+  },
+  "hash_convention": "Original indented method source including its final newline. No source transformations."
+}