From 504e04d4de1a4168aabb5436f1b236edef1dadd7 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 15:15:31 +0800 Subject: [PATCH 1/4] Cache process and thread identities during precondition inference --- tests/test_identity_cache.py | 194 +++++++++++++++++++++++++++++++ traincheck/trace/trace_pandas.py | 32 ++++- 2 files changed, 224 insertions(+), 2 deletions(-) create mode 100644 tests/test_identity_cache.py diff --git a/tests/test_identity_cache.py b/tests/test_identity_cache.py new file mode 100644 index 00000000..4093d74a --- /dev/null +++ b/tests/test_identity_cache.py @@ -0,0 +1,194 @@ +"""Identity lookup equivalence during precondition inference.""" + +import copy +import json +import logging +import pickle +import unittest +from unittest.mock import patch + +import numpy as np +import pandas as pd + +from traincheck.invariant.precondition import find_precondition_from_single_group +from traincheck.trace.trace_pandas import TracePandas + + +class LegacyIdentityTrace(TracePandas): + def get_process_ids(self): + if "process_id" not in self.events.columns: + return [] + return self.events["process_id"].dropna().unique().tolist() + + def get_thread_ids(self): + if "thread_id" not in self.events.columns: + return [] + return self.events["thread_id"].dropna().unique().tolist() + + +def make(frame, cls=TracePandas): + return cls(frame.copy(deep=True), truncate_incomplete_func_calls=False) + + +def frame(process_ids=(1, 2, 1), thread_ids=(4, 3, 4)): + return pd.DataFrame( + { + "time": range(len(process_ids)), + "process_id": process_ids, + "thread_id": thread_ids, + } + ) + + +class IdentityCacheTests(unittest.TestCase): + def setUp(self): + previous = logging.root.manager.disable + logging.disable(logging.CRITICAL) + self.addCleanup(logging.disable, previous) + + def assert_equivalent(self, data): + old, new = make(data, LegacyIdentityTrace), make(data) + for _ in range(2): + for method in ("get_process_ids", "get_thread_ids"): + expected, actual = getattr(old, method)(), getattr(new, method)() + self.assertEqual(expected, actual) + self.assertEqual(list(map(type, expected)), list(map(type, actual))) + return old, new + + def test_values_types_and_encounter_order(self): + for values in ( + pd.Series([8, 2, 8, -1], dtype="int64"), + pd.Series([2**63 + 7, 2**63 + 2, None, 2**63 + 7], dtype="UInt64"), + pd.Series([8, pd.NA, 2, 8], dtype="Int64"), + pd.Series([8.5, np.nan, 2.5, 8.5]), + pd.Series(["z", None, "a", "z"], dtype="string"), + pd.Series(["z", None, "a", "z"], dtype="category"), + pd.Series([True, 1, False, 0], dtype=object), + pd.Series([None, np.nan, pd.NA, None], dtype=object), + ): + with self.subTest(dtype=str(values.dtype), values=repr(values.tolist())): + self.assert_equivalent(frame(values, values)) + + def test_randomized_identity_columns(self): + rng = np.random.default_rng(42) + for _ in range(80): + values = rng.integers(-8, 8, size=50).astype(float) + values[rng.random(50) < 0.2] = np.nan + self.assert_equivalent(frame(values, values[::-1])) + + def test_empty_and_missing_columns(self): + for data in (frame([], []), pd.DataFrame({"time": [0]})): + self.assert_equivalent(data) + + def test_returned_lists_are_independent(self): + trace = make(frame()) + result = trace.get_process_ids() + result[:] = [999] + self.assertEqual(trace.get_process_ids(), [1, 2]) + result = trace.get_thread_ids() + result.clear() + self.assertEqual(trace.get_thread_ids(), [4, 3]) + + def test_frame_replacement_rebuilds_same_size_cache(self): + trace = make(frame()) + trace.get_process_ids() + trace.get_thread_ids() + trace.events = frame((9, 7, 9), (5, 6, 5)) + self.assertEqual(trace.get_process_ids(), [9, 7]) + self.assertEqual(trace.get_thread_ids(), [5, 6]) + + def test_same_frame_length_change_rebuilds_cache(self): + trace = make(frame()) + trace.get_process_ids() + trace.get_thread_ids() + trace.events.loc[3] = [3, 9, 7] + self.assertEqual(trace.get_process_ids(), [1, 2, 9]) + self.assertEqual(trace.get_thread_ids(), [4, 3, 7]) + trace.events.drop(index=[1, 3], inplace=True) + self.assertEqual(trace.get_process_ids(), [1]) + self.assertEqual(trace.get_thread_ids(), [4]) + + def test_stage_traces_and_instances_have_separate_caches(self): + data = frame() + data["meta_vars.stage"] = ["init", "train", "init"] + trace = make(data) + trace.get_process_ids() + stages = trace.get_traces_for_stage() + self.assertEqual(stages["init"].get_process_ids(), [1]) + self.assertEqual(stages["train"].get_process_ids(), [2]) + self.assertEqual(make(frame((9,), (7,))).get_process_ids(), [9]) + + def test_pickled_initialized_cache(self): + trace = make(frame()) + trace.get_process_ids() + restored = pickle.loads(pickle.dumps(trace)) + self.assertIs(restored._identity_values_source, restored.events) + self.assertEqual(restored.get_process_ids(), [1, 2]) + self.assertEqual(restored.get_thread_ids(), [4, 3]) + + def test_meta_vars_membership_preserves_independent_id_checks(self): + old, new = self.assert_equivalent(frame((1, 2), (3, 4))) + for pid, tid in ((1, 3), (1, 4), (2, 3), (9, 3), (1, 9)): + self.assertEqual( + old.get_meta_vars(0, pid, tid), new.get_meta_vars(0, pid, tid) + ) + self.assertEqual(new.get_meta_vars(0, 1, 4), {}) + self.assertIsNone(new.get_meta_vars(0, 9, 3)) + + def test_repeated_meta_vars_queries_only_scan_each_column_once(self): + trace = make(frame()) + original = pd.Series.unique + scanned = [] + + def unique(series): + scanned.append(series.name) + return original(series) + + with patch.object(pd.Series, "unique", unique): + for _ in range(100): + self.assertEqual(trace.get_meta_vars(0, 1, 4), {}) + self.assertEqual(scanned, ["process_id", "thread_id"]) + + def test_real_precondition_inference_matches_legacy(self): + old, new = self.assert_equivalent(frame()) + positive = [ + [ + { + "time": i, + "process_id": 1, + "thread_id": 4, + "function": "fixtures.step", + "meta_vars.phase": "train", + } + ] + for i in range(12) + ] + negative = [ + [ + { + "time": i, + "process_id": 2, + "thread_id": 3, + "function": "fixtures.step", + "meta_vars.phase": "eval", + } + ] + for i in range(8) + ] + results = [] + for trace in (old, new): + pos, neg = copy.deepcopy(positive), copy.deepcopy(negative) + clauses = find_precondition_from_single_group(pos, neg, [trace]) + results.append( + ( + sorted(json.dumps(c.to_dict(), sort_keys=True) for c in clauses), + pos, + neg, + ) + ) + self.assertTrue(results[0][0]) + self.assertEqual(*results) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/trace/trace_pandas.py b/traincheck/trace/trace_pandas.py index 2ab42924..e7e20cfe 100644 --- a/traincheck/trace/trace_pandas.py +++ b/traincheck/trace/trace_pandas.py @@ -111,6 +111,12 @@ def __init__(self, events, truncate_incomplete_func_calls=True): self._func_call_positions_length = -1 self._func_call_positions: dict[str, tuple[list[int], list[int]]] | None = None + # Inference treats event identities as immutable after initialization. + # Keep the frame alive so object identity cannot be reused across frames. + self._identity_values_source: pd.DataFrame | None = None + self._identity_values_length = -1 + self._identity_values: dict[str, tuple] = {} + self.column_dtypes_cached = {} self.stage_traces: dict[str, "TracePandas"] = {} @@ -573,6 +579,28 @@ def get_end_time(self, process_id=None, thread_id=None) -> float: return self.events["time"].max() + def _get_identity_values(self, column: str) -> list: + """Cache stable identity columns; return a fresh list for every caller. + + As with call-record indexes, frame replacement or length changes rebuild + the cache. Same-size in-place identity edits, row reordering, or dtype + changes require a new TracePandas. + Keep pandas' dropna/unique conversion to preserve dtype and encounter + order instead of replacing it with Python set semantics. + """ + if ( + self._identity_values_source is not self.events + or self._identity_values_length != len(self.events) + ): + self._identity_values = {} + self._identity_values_source = self.events + self._identity_values_length = len(self.events) + if column not in self._identity_values: + self._identity_values[column] = tuple( + self.events[column].dropna().unique().tolist() + ) + return list(self._identity_values[column]) + def get_process_ids(self) -> list[int]: """Find all process ids from the trace.""" if "process_id" not in self.events.columns: @@ -580,7 +608,7 @@ def get_process_ids(self) -> list[int]: "process_id column not found in the events, no process related invariants will be extracted." ) return [] - return self.events["process_id"].dropna().unique().tolist() + return self._get_identity_values("process_id") def get_thread_ids(self) -> list[int]: """Find all thread ids from the trace.""" @@ -589,7 +617,7 @@ def get_thread_ids(self) -> list[int]: "thread_id column not found in the events, no thread related invariants will be extracted." ) return [] - return self.events["thread_id"].dropna().unique().tolist() + return self._get_identity_values("thread_id") def get_func_call_ids(self, func_name: str = "") -> list[str]: """Find all function call ids from the trace.""" From fa277435dbff801f3b51ef3660f4eb003b1e5f2c Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 15:25:11 +0800 Subject: [PATCH 2/4] Document identity cache measurements and validation --- .../inference-identity-cache-20261025.json | 126 +++++++++ docs/measurements/inference-identity-cache.md | 65 +++++ tools/benchmarks/identity_cache.py | 246 ++++++++++++++++++ 3 files changed, 437 insertions(+) create mode 100644 docs/measurements/inference-identity-cache-20261025.json create mode 100644 docs/measurements/inference-identity-cache.md create mode 100644 tools/benchmarks/identity_cache.py diff --git a/docs/measurements/inference-identity-cache-20261025.json b/docs/measurements/inference-identity-cache-20261025.json new file mode 100644 index 00000000..fa3df296 --- /dev/null +++ b/docs/measurements/inference-identity-cache-20261025.json @@ -0,0 +1,126 @@ +{ + "measurement_date": "2026-10-25", + "scope": "construction plus repeated identity-checked get_meta_vars calls", + "settings": { + "rows": 573494, + "calls": 1000, + "repetitions": 3 + }, + "runtime": { + "python": "3.10.16", + "platform": "macOS-27.0.1-arm64-arm-64bit", + "host_class": "local Darwin arm64 workstation", + "logical_cpu_count": 12, + "pandas": "2.2.3", + "numpy": "1.26.4", + "native_thread_settings": { + "OMP_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1" + }, + "candidate_source_sha256": "41040a288284dc10041a4c06d8598219605732dd5dcd2ef080b9ed0799387f8a", + "candidate_commit": "504e04d4de1a4168aabb5436f1b236edef1dadd7", + "candidate_source_matches_commit": true, + "legacy_methods_source_sha256": "22a7976d0fc8558b5a3dcc77e9e527572ec1f1e36b15d5e244a4a0b7ec6f44e7" + }, + "input": { + "sha256_pandas_row_hashes": "ce7af99430be7c9494ef8ee397049b77847897dd965c1082a3cda9e88a159f8b", + "query_sha256": "cd7c7c854fec82747a9704c95dc1e897d6f2dca5664b1b9167a8bd845f227384", + "dtypes": { + "time": "int64", + "process_id": "float64", + "thread_id": "float64" + }, + "dataframe_deep_bytes": 13763984, + "context_manager_count": 0, + "process_count": 8, + "thread_count": 16 + }, + "medians": { + "legacy": { + "construction_seconds": 8.187499770428985e-05, + "queries_seconds": 9.489749166001275, + "total_seconds": 9.489831040998979 + }, + "candidate": { + "construction_seconds": 5.3582996770273894e-05, + "queries_seconds": 0.01196087500284193, + "total_seconds": 0.01200404100382002 + } + }, + "trials": [ + { + "repetition": 1, + "position": 1, + "variant": "legacy", + "total_seconds": 9.2562226250011 + }, + { + "repetition": 1, + "position": 2, + "variant": "candidate", + "total_seconds": 0.01413512499857461 + }, + { + "repetition": 2, + "position": 1, + "variant": "candidate", + "total_seconds": 0.010777290997793898 + }, + { + "repetition": 2, + "position": 2, + "variant": "legacy", + "total_seconds": 9.489831040998979 + }, + { + "repetition": 3, + "position": 1, + "variant": "legacy", + "total_seconds": 9.548994832999597 + }, + { + "repetition": 3, + "position": 2, + "variant": "candidate", + "total_seconds": 0.01200404100382002 + } + ], + "correctness": { + "all_results_equal_expected": true, + "result_sha256": "481e48248fde7cbd3011bd60ecb24d4b0ccaa470a6e5a1986acf27bcc3fa9570", + "input_unchanged": true + }, + "validation": { + "local_combined_tests_passed": 43, + "matching_runtime_linux_identity_tests_passed": 11, + "independent_differential_comparisons": 78, + "independent_dtype_cases": 13, + "pickle_protocols": 6, + "spawn_worker_tasks_passed": 2, + "matching_linux_runtime": { + "python": "3.11.4", + "pytorch": "2.16.0.dev20261012+cpu", + "numpy": "2.1.0", + "pandas": "2.3.3" + } + }, + "full_inference": { + "status": "paired_unprofiled_replay_pending", + "profile_attempts": "Both attempts terminated with SIGSEGV at scheduled diagnostic timers; no completed profile/output. Candidate stopped before the identity-cache phase.", + "performance_comparison_valid": false + }, + "limitations": [ + "Synthetic identity-lookup microbenchmark, not completed inference.", + "No context-manager work or trace-file loading; input generation excluded.", + "No RSS/PSS, peak-memory, or multiprocessing memory-scaling measurement.", + "Cache stores O(P+T) values and returns fresh O(P+T) lists.", + "Stable identity values, dtypes, and row order required within one frame." + ], + "source_evidence_sha256": { + "benchmark-full.json": "7c3b84db5645226b4d841b929185edad33c075c7868ad9951607a57b2852febe", + "review-results.json": "19e4c7bf403a8a3732e8dbf0fda3ab4f6902b99c0e997dc490d9383ffc406376", + "combined-tests.log": "6334a444075d2bd539f965998354922ea2d940d74536d61790e39e840bcc345c", + "matching-runtime-linux-tests.log": "56b695709855bb8a17cd0fbf58e1bf3e7d17ed228719d5cc2350d3718f8e3604" + } +} diff --git a/docs/measurements/inference-identity-cache.md b/docs/measurements/inference-identity-cache.md new file mode 100644 index 00000000..09c9919e --- /dev/null +++ b/docs/measurements/inference-identity-cache.md @@ -0,0 +1,65 @@ +# Cache process and thread identities during precondition inference + +`get_meta_vars` checks process and thread membership for each request. Previously, +each check repeated `dropna().unique().tolist()` over the full event column. +The change caches those exact pandas results lazily per trace and returns a fresh +list to each caller. Null filtering, scalar types, encounter order, and independent +process/thread membership remain unchanged. Frame replacement or length changes +invalidate both identity caches. + +## Isolated measurement + +On a deterministic 573,494-row synthetic trace, the median time for construction +plus 1,000 `get_meta_vars` calls fell from **9.489831 s to 0.012004 s** across three +alternating repetitions. Each trial starts with a fresh trace and includes the +first cache fill. The baseline restores only the two original identity methods; +both variants use the same constructor and other methods. + +| Repetition | Legacy | Cache | +|---|---:|---:| +| 1 (legacy first) | 9.256223 s | 0.014135 s | +| 2 (cache first) | 9.489831 s | 0.010777 s | +| 3 (legacy first) | 9.548995 s | 0.012004 s | + +The shared input has eight process IDs, sixteen thread IDs, deterministic missing +values, and no context managers. All six ordered outputs match: 750 empty +dictionaries and 250 `None` values; input hashes remain unchanged. Input generation, +garbage collection, and result hashing are excluded. The run used a macOS arm64 +workstation, Python 3.10.16, pandas 2.2.3, NumPy 1.26.4, and native thread settings +of one. No CPU affinity or exclusive machine access was imposed. + +This is an identity-lookup microbenchmark, not completed invariant inference or +training overhead. Paired full-inference results remain pending. Earlier profiling +attempts terminated with SIGSEGV at scheduled diagnostic timers and produced no +completed profile/output; the candidate stopped before the cache phase. Those +attempts provide no valid performance comparison. + +## Correctness and contract + +All 43 local context/call/identity regression tests and all 11 identity tests in the +matching Linux runtime pass. Independent review adds 78 differential comparisons +across 13 dtype/value cases, six pickle protocols, and two spawned-worker tasks. +The checks cover nulls, integer precision, scalar types, encounter order, missing +columns, independent process/thread membership, fresh return lists, stage isolation, +frame/length invalidation, serialization, and actual precondition inference. + +Identity values, dtypes, and row order must remain stable within a frame. Same-size +in-place edits, column replacement, and reordering require a new `TracePandas`. +The cache stores O(P + T) values; returned lists and membership checks still cost +O(P + T). Its source reference is not a copied DataFrame, but can retain the old +frame until invalidation. No RSS/PSS, peak-memory, or multiprocessing memory-scaling +improvement is claimed. + +From the checkout, reproduce the isolated measurement with: + +```sh +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \ +python -m tools.benchmarks.identity_cache --rows 573494 --calls 1000 --repetitions 3 + +python -m unittest tests.test_context_index tests.test_call_index tests.test_identity_cache +``` + +Use `--rows 1000 --calls 20` for a smoke test. The script needs no external trace; +it records the checkout commit and source hash and checks output equality. +[Compact evidence and provenance](inference-identity-cache-20261025.json) retain the +measured source commit, runtime, inputs, per-trial timings, and limitations. diff --git a/tools/benchmarks/identity_cache.py b/tools/benchmarks/identity_cache.py new file mode 100644 index 00000000..3a82e1d2 --- /dev/null +++ b/tools/benchmarks/identity_cache.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +"""Deterministic identity-lookup microbenchmark, with no external trace input. + +Run against the candidate checkout, for example: + python -m tools.benchmarks.identity_cache > results.json +Small smoke test: + python -m tools.benchmarks.identity_cache --rows 1000 --calls 20 + +The legacy subclass changes only the two identity methods back to their +pre-cache implementation. Each trial constructs a fresh trace over the same +read-only DataFrame, then makes the same ordered get_meta_vars calls. The first +lookup/cache build is included. Three repetitions alternate execution order. +Input generation, correctness hashing, and garbage collection are not timed. + +This measures identity scans, not full inference. The synthetic trace has no +context managers. No RSS/PSS, peak memory, or multiprocessing scaling is measured. +The frame-memory estimate is not whole-process memory. Cache tuples add O(P+T) +stored values; each return still allocates a fresh O(P+T) list. The extra source +reference retains a frame until invalidation and is not a DataFrame copy. +""" + +import argparse +import gc +import hashlib +import inspect +import json +import logging +import os +import platform +import statistics +import subprocess +import sys +import time +from pathlib import Path + +import numpy as np +import pandas as pd + +import traincheck.trace.trace_pandas as trace_module +from traincheck.trace.trace_pandas import TracePandas + + +class LegacyIdentityTrace(TracePandas): + """Original lookup expressions; all other behavior uses the same checkout.""" + + def get_process_ids(self): + if "process_id" not in self.events.columns: + return [] + return self.events["process_id"].dropna().unique().tolist() + + def get_thread_ids(self): + if "thread_id" not in self.events.columns: + return [] + return self.events["thread_id"].dropna().unique().tolist() + + +def synthetic_frame(rows): + sequence = np.arange(rows, dtype=np.int64) + process_ids = ((sequence // 16) % 8 + 100).astype(np.float64) + thread_ids = (sequence % 16 + 1000).astype(np.float64) + # NaNs exercise the same dropna work as mixed event traces. They never + # remove every occurrence of an identity at the minimum supported size. + process_ids[sequence % 97 == 0] = np.nan + thread_ids[sequence % 193 == 0] = np.nan + return pd.DataFrame( + {"time": sequence, "process_id": process_ids, "thread_id": thread_ids} + ) + + +def fixed_queries(calls): + # Six valid, one unknown-process, and one unknown-thread call per cycle. + # Process/thread presence remains independent, matching get_meta_vars. + return [ + ( + i, + -1 if i % 8 == 6 else 100 + i % 8, + -1 if i % 8 == 7 else 1000 + (i * 3) % 16, + ) + for i in range(calls) + ] + + +def digest(value): + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def trial(cls, data, queries): + gc.collect() + begin = time.perf_counter() + trace = cls(data, truncate_incomplete_func_calls=False) + constructed = time.perf_counter() + results = [trace.get_meta_vars(*query) for query in queries] + finished = time.perf_counter() + record = { + "construction_seconds": constructed - begin, + "queries_seconds": finished - constructed, + "total_seconds": finished - begin, + "result_sha256": digest(results), + "none_results": sum(value is None for value in results), + "empty_dict_results": sum(value == {} for value in results), + } + if cls is TracePandas: + assert trace._identity_values_source is trace.events + record["cached_identity_counts"] = { + key: len(values) for key, values in trace._identity_values.items() + } + record["cache_tuple_containers_shallow_bytes"] = sum( + sys.getsizeof(values) for values in trace._identity_values.values() + ) + return record + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, default=573494) + parser.add_argument("--calls", type=int, default=1000) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if args.rows < 256 or args.calls < 8 or args.repetitions < 1: + parser.error("require rows >= 256, calls >= 8, repetitions >= 1") + if not hasattr(TracePandas, "_get_identity_values"): + parser.error("import a candidate checkout containing the identity cache") + + logging.disable(logging.CRITICAL) + data = synthetic_frame(args.rows) + queries = fixed_queries(args.calls) + input_hash = hashlib.sha256( + pd.util.hash_pandas_object(data, index=True).to_numpy().tobytes() + ).hexdigest() + expected = [None if i % 8 in (6, 7) else {} for i in range(args.calls)] + expected_hash = digest(expected) + measurements = [] + variants = {"legacy": LegacyIdentityTrace, "candidate": TracePandas} + for repetition in range(args.repetitions): + order = ["legacy", "candidate"] + if repetition % 2: + order.reverse() + for position, variant in enumerate(order): + record = trial(variants[variant], data, queries) + assert record["result_sha256"] == expected_hash, record + record.update( + repetition=repetition + 1, position=position + 1, variant=variant + ) + measurements.append(record) + + # Catch accidental mutation of the shared input by either implementation. + assert ( + hashlib.sha256( + pd.util.hash_pandas_object(data, index=True).to_numpy().tobytes() + ).hexdigest() + == input_hash + ) + medians = { + variant: { + metric: statistics.median( + record[metric] + for record in measurements + if record["variant"] == variant + ) + for metric in ("construction_seconds", "queries_seconds", "total_seconds") + } + for variant in variants + } + report = { + "scope": "construction plus repeated identity-checked get_meta_vars calls", + "settings": vars(args), + "input": { + "sha256_pandas_row_hashes": input_hash, + "query_sha256": digest(queries), + "dtypes": {key: str(value) for key, value in data.dtypes.items()}, + "dataframe_deep_bytes": int(data.memory_usage(index=True, deep=True).sum()), + "context_manager_count": 0, + "process_count": 8, + "thread_count": 16, + }, + "runtime": { + "python": platform.python_version(), + "platform": platform.platform(), + "host_class": f"local {platform.system()} {platform.machine()} workstation", + "logical_cpu_count": os.cpu_count(), + "pandas": pd.__version__, + "numpy": np.__version__, + "native_thread_settings": { + name: os.environ.get(name) + for name in ( + "OMP_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "MKL_NUM_THREADS", + ) + }, + "candidate_source_sha256": hashlib.sha256( + Path(trace_module.__file__).read_bytes() + ).hexdigest(), + "candidate_commit": subprocess.check_output( + [ + "git", + "-C", + str(Path(trace_module.__file__).resolve().parents[2]), + "rev-parse", + "HEAD", + ], + text=True, + ).strip(), + "candidate_source_matches_commit": subprocess.run( + [ + "git", + "-C", + str(Path(trace_module.__file__).resolve().parents[2]), + "diff", + "--quiet", + "HEAD", + "--", + "traincheck/trace/trace_pandas.py", + ], + check=False, + ).returncode + == 0, + "legacy_methods_source_sha256": hashlib.sha256( + inspect.getsource(LegacyIdentityTrace).encode() + ).hexdigest(), + }, + "measurements": measurements, + "medians": medians, + "median_total_speedup": ( + medians["legacy"]["total_seconds"] / medians["candidate"]["total_seconds"] + ), + "correctness": { + "all_results_equal_expected": True, + "result_sha256": expected_hash, + "input_unchanged": True, + }, + "limitations": [ + "Synthetic identity-lookup microbenchmark, not completed inference.", + "No context-manager work or trace-file loading; input generation excluded.", + "No RSS/PSS, peak-memory, or multiprocessing memory-scaling measurement.", + "Cache stores O(P+T) values and returns fresh O(P+T) lists.", + "Stable identity values, dtypes, and row order required within one frame.", + ], + } + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() From 99d16905296d83495570fdcafa91370b242d2528 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 15:57:10 +0800 Subject: [PATCH 3/4] Document complete identity-cache inference comparison --- .../inference-identity-cache-20261025.json | 14 +- ...ce-identity-cache-full-trace-20261025.json | 247 ++++++++++++++++++ docs/measurements/inference-identity-cache.md | 68 ++++- 3 files changed, 319 insertions(+), 10 deletions(-) create mode 100644 docs/measurements/inference-identity-cache-full-trace-20261025.json diff --git a/docs/measurements/inference-identity-cache-20261025.json b/docs/measurements/inference-identity-cache-20261025.json index fa3df296..690fa7fa 100644 --- a/docs/measurements/inference-identity-cache-20261025.json +++ b/docs/measurements/inference-identity-cache-20261025.json @@ -106,9 +106,17 @@ } }, "full_inference": { - "status": "paired_unprofiled_replay_pending", - "profile_attempts": "Both attempts terminated with SIGSEGV at scheduled diagnostic timers; no completed profile/output. Candidate stopped before the identity-cache phase.", - "performance_comparison_valid": false + "status": "completed_unprofiled_pair", + "evidence": "inference-identity-cache-full-trace-20261025.json", + "baseline_wall_seconds": 1385.656381815672, + "candidate_wall_seconds": 98.55111291632056, + "wall_time_ratio": 14.060281419573906, + "invariant_records": 835, + "failed_hypothesis_records": 105, + "complete_json_multisets_equal": true, + "performance_comparison_valid": true, + "scope": "One sequential saved-trace CPU pair with six relations enabled; no repeated-run distribution or training-overhead claim.", + "profile_attempts": "Two earlier periodic-faulthandler plus cProfile attempts exited SIGSEGV without completed outputs and are excluded." }, "limitations": [ "Synthetic identity-lookup microbenchmark, not completed inference.", diff --git a/docs/measurements/inference-identity-cache-full-trace-20261025.json b/docs/measurements/inference-identity-cache-full-trace-20261025.json new file mode 100644 index 00000000..99c1b289 --- /dev/null +++ b/docs/measurements/inference-identity-cache-full-trace-20261025.json @@ -0,0 +1,247 @@ +{ + "measurement_date": "2026-10-25", + "scope": "One baseline-first sequential unprofiled pair on the same saved training trace; no repeats, affinity, or exclusive host reservation.", + "workload": { + "model_parameters": 6163712, + "training_steps": 10, + "sequence_length": 256, + "batch_size": 1, + "seed": 42, + "events": 573494, + "raw_trace_bytes": 286335289, + "training_run": "20261012T163111Z-88c4d357", + "trace_files": { + "trace_API_1088885_23438542501568.log": { + "bytes": 82668375, + "sha256": "883d288d6955b12c49dc741c166bcca7d191c890d19565802dad60708979278d" + }, + "trace_API_1088885_23456244348224.log": { + "bytes": 203666914, + "sha256": "416b8ca45b4df361fc73052447573d7cb23d433d16c0947d74bbc18f7cc6c9a8" + } + } + }, + "settings": { + "backend": "pandas", + "workers": 1, + "native_threads": 1, + "python_hash_seed": 0, + "gpu": false, + "disabled_relations": [ + "FunctionLeadRelation", + "FunctionCoverRelation" + ], + "timeout_seconds": 1800, + "sampled_child_session_pss_limit_gib": 12, + "memory_sampling_seconds": 1, + "profile_enabled": false, + "periodic_traceback_dumps": false, + "eligible_api_count": 41 + }, + "runtime": { + "host": "labmach", + "python": "3.11.4", + "torch": "2.16.0.dev20261012+cpu", + "torch_git_version": "8f9045d72f71e7aa9e16b6d7a055ec5bc1f720b5", + "pandas": "2.3.3", + "numpy": "2.1.0", + "original_trace_runtime_difference": "Original training used Python 3.12.3 and a CUDA 13.0 Torch build at the same Torch source revision. Original eligible API name set was not logged; only its count of 41 is available." + }, + "sources": { + "baseline": "Frozen call-index snapshot underlying PR141; before publication-only typing/comment formatting changes.", + "candidate_commit": "504e04d4de1a4168aabb5436f1b236edef1dadd7", + "trace_pandas_sha256": { + "baseline": "8e4c3589371449b78430ed5039d0159d3d0435664b218f3a96c81935be0d75d5", + "candidate": "41040a288284dc10041a4c06d8598219605732dd5dcd2ef080b9ed0799387f8a" + }, + "provenance_manifest_sha256": { + "baseline": "7e8e267085ece31844a820752004af6a1c5bb293a3a2733990b039d72fcd872a", + "candidate": "b48d705bd9dccbdf7bce17e8704e7ae5554a7d6239cfd545d7b5ffc66a888b09" + }, + "candidate_production_delta": "Identity caching plus publication-only import/type-annotation changes, comments/docstrings and formatting; all other production behavior unchanged. Independent pruning-state fix PR145 excluded from both variants." + }, + "comparison": { + "complete_pair": true, + "scope": "one sequential unprofiled run per source; no replicate/median claim", + "harness_sha256": "2dd8d63a8bae80a944defee70be6bf33e9c9c8934e22e4b909ff222b4f7d6777", + "baseline": { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "f3872b21c9ed7103e1020ccb83ef54b41009f2ec06348ce1b7dc5893e545f291", + "distinct_records": 105, + "records": 105 + }, + "invariants": { + "canonical_multiset_sha256": "f4f5db6772b857074d113e1c6ea091d5c4ed2452e091607fa425efab4edeed65", + "distinct_records": 835, + "records": 835 + }, + "sampled_peak_concurrent_memory": { + "pss_bytes": 2283671552, + "rss_bytes": 2291527680, + "uss_bytes": 2280935424 + }, + "source_unchanged": true, + "stage_times": { + "marker_durations": { + "cross_trace_collection_seconds": 0.001, + "generation_and_pruning_seconds": 619.58, + "precondition_and_output_sort_seconds": 751.403 + }, + "markers": { + "collection_start": "2026-10-25T03:37:57.781000", + "generation_start": "2026-10-25T03:27:38.201000", + "inference_end": "2026-10-25T03:50:29.185000", + "precondition_start": "2026-10-25T03:37:57.782000" + }, + "note": "Log marker intervals; missing endpoints mean incomplete phases, not zero duration.", + "reported": { + "inference_seconds": 1371.0323550701141, + "load_seconds": 10.503035306930542 + } + }, + "stop_reason": null, + "success": true, + "traces_unchanged": true, + "wall_seconds": 1385.656381815672 + }, + "candidate": { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "f3872b21c9ed7103e1020ccb83ef54b41009f2ec06348ce1b7dc5893e545f291", + "distinct_records": 105, + "records": 105 + }, + "invariants": { + "canonical_multiset_sha256": "f4f5db6772b857074d113e1c6ea091d5c4ed2452e091607fa425efab4edeed65", + "distinct_records": 835, + "records": 835 + }, + "sampled_peak_concurrent_memory": { + "pss_bytes": 2286718976, + "rss_bytes": 2294530048, + "uss_bytes": 2283991040 + }, + "source_unchanged": true, + "stage_times": { + "marker_durations": { + "cross_trace_collection_seconds": 0.0, + "generation_and_pruning_seconds": 72.759, + "precondition_and_output_sort_seconds": 10.913 + }, + "markers": { + "collection_start": "2026-10-25T03:51:57.205000", + "generation_start": "2026-10-25T03:50:44.446000", + "inference_end": "2026-10-25T03:52:08.118000", + "precondition_start": "2026-10-25T03:51:57.205000" + }, + "note": "Log marker intervals; missing endpoints mean incomplete phases, not zero duration.", + "reported": { + "inference_seconds": 83.71693992614746, + "load_seconds": 10.698827981948853 + } + }, + "stop_reason": null, + "success": true, + "traces_unchanged": true, + "wall_seconds": 98.55111291632056 + }, + "invariants": { + "comparison": "structural_json_multiset", + "equal": true, + "reference": { + "records": 835, + "distinct_records": 835, + "canonical_multiset_sha256": "f4f5db6772b857074d113e1c6ea091d5c4ed2452e091607fa425efab4edeed65" + }, + "candidate": { + "records": 835, + "distinct_records": 835, + "canonical_multiset_sha256": "f4f5db6772b857074d113e1c6ea091d5c4ed2452e091607fa425efab4edeed65" + }, + "missing_records": 0, + "extra_records": 0, + "missing_examples": [], + "extra_examples": [] + }, + "failed_hypotheses": { + "comparison": "structural_json_multiset", + "equal": true, + "reference": { + "records": 105, + "distinct_records": 105, + "canonical_multiset_sha256": "f3872b21c9ed7103e1020ccb83ef54b41009f2ec06348ce1b7dc5893e545f291" + }, + "candidate": { + "records": 105, + "distinct_records": 105, + "canonical_multiset_sha256": "f3872b21c9ed7103e1020ccb83ef54b41009f2ec06348ce1b7dc5893e545f291" + }, + "missing_records": 0, + "extra_records": 0, + "missing_examples": [], + "extra_examples": [] + }, + "wall_time_ratio_baseline_over_candidate": 14.060281419573906 + }, + "context_index": { + "baseline": [ + { + "call": 0, + "contexts": 711, + "rows": 573494, + "seconds": 0.46406279876828194, + "states_sha256": "02a9f08ad19e890f92d26367b4a8c433c4953770d5359554bf496fc6f17e3b07" + }, + { + "call": 1, + "contexts": 711, + "rows": 573494, + "seconds": 0.4610072821378708, + "states_sha256": "02a9f08ad19e890f92d26367b4a8c433c4953770d5359554bf496fc6f17e3b07" + } + ], + "candidate": [ + { + "call": 0, + "contexts": 711, + "rows": 573494, + "seconds": 0.4815337397158146, + "states_sha256": "02a9f08ad19e890f92d26367b4a8c433c4953770d5359554bf496fc6f17e3b07" + }, + { + "call": 1, + "contexts": 711, + "rows": 573494, + "seconds": 0.46630632504820824, + "states_sha256": "02a9f08ad19e890f92d26367b4a8c433c4953770d5359554bf496fc6f17e3b07" + } + ] + }, + "verification": { + "complete_structural_json_multisets_equal": true, + "duplicate_multiplicity_and_array_order_preserved": true, + "json_object_key_order_and_record_line_order_ignored": true, + "both_full_and_stage_context_files_byte_identical": true, + "all_source_and_input_hashes_unchanged": true + }, + "timer_scope": { + "wall": "Child startup through exit and monitor polling, including imports and JSON serialization; initial/final integrity hashing excluded.", + "load": "Read/concatenate/TracePandas construction, including harness context serialization.", + "inference": "Engine construction, hypothesis generation/pruning, collection, final preconditions and sort; ends before invariant JSON serialization.", + "marker_intervals": "Millisecond log timestamp resolution; generation/pruning and final preconditions/sort are parts of inference, not additional wall time." + }, + "memory_scope": "Peak concurrent child-session samples, excluding monitor; PSS depends on shared pages. No memory improvement or multiprocessing scaling claim.", + "evidence_archive": { + "name": "matching-runtime-unprofiled-pair-evidence.tar.gz", + "bytes": 314914, + "sha256": "d78a815e1516131b28342c7d7e9251bc195f9708bab247c9f2d184558ec26f9e" + }, + "limitations": [ + "One sequential pair, no distribution or confidence interval.", + "Offline saved-trace CPU inference, not training, distributed scaling, online checking overhead, or new trace collection.", + "Six enabled relations; FunctionLeadRelation and FunctionCoverRelation disabled identically.", + "This preserves existing inference behavior; it does not certify inferred invariants as ground truth.", + "Earlier periodic-faulthandler plus cProfile attempts crashed and are excluded; this successful pair alone does not isolate the crash cause." + ] +} diff --git a/docs/measurements/inference-identity-cache.md b/docs/measurements/inference-identity-cache.md index 09c9919e..7cf62e09 100644 --- a/docs/measurements/inference-identity-cache.md +++ b/docs/measurements/inference-identity-cache.md @@ -7,6 +7,64 @@ list to each caller. Null filtering, scalar types, encounter order, and independ process/thread membership remain unchanged. Frame replacement or length changes invalidate both identity caches. +## Complete saved-trace inference + +On labmach, one sequential unprofiled pair reduced total wall time from +**1,385.656 s to 98.551 s (14.06x)**. Both runs completed successfully. All **835 +invariant records and 105 saved failed-hypothesis records** match as complete +structural JSON multisets, preserving duplicate occurrences and array order. +Object key order and record line order are ignored. No records were missing or +extra. Both full/stage context files are byte-identical (711 states each), and +all source and input hashes remained unchanged. + +| Measurement | Baseline | Identity cache | +|---|---:|---:| +| Total child wall time | 1,385.656 s | 98.551 s | +| Trace loading/construction | 10.503 s | 10.699 s | +| Reported inference | 1,371.032 s | 83.717 s | +| Generation and pruning (part of inference) | 619.580 s | 72.759 s | +| Final preconditions and sort (part of inference) | 751.403 s | 10.913 s | +| Peak sampled child-session PSS | 2.127 GiB | 2.130 GiB | + +The workload is the two complete trace files from the saved 6,163,712-parameter, +10-step reference training run: 573,494 events and 286,335,289 raw bytes, sequence +length 256, batch size 1, seed 42. No training was rerun. Both variants use the +context/call indexes from PR #141; this comparison isolates the identity cache, +with only publication-time imports/type annotations, comments/docstrings, and +formatting also differing. +The independent pruning-state correctness fix in PR #145 is excluded from both +variants and evaluated separately. + +Both runs used pandas, one worker, native thread settings of one, no GPU, +`PYTHONHASHSEED=0`, and the same six enabled relations. `FunctionLeadRelation` +and `FunctionCoverRelation` were disabled. Both had a 1,800-second cutoff and a +12-GiB sampled child-session PSS limit. The runtime was Python 3.11.4, Torch +2.16.0.dev20261012+cpu, pandas 2.3.3, and NumPy 2.1.0. Torch's source revision +matches the original CUDA wheel, but the original training used Python 3.12.3 +and CUDA 13.0. The eligible API count is 41 in both logs; the original name set +was not recorded, so exact original-environment coverage is not established. + +This is one baseline-first pair without CPU affinity or exclusive host access; +it is not a repeated-run speedup estimate. Total wall time includes child startup, +imports, JSON serialization, and monitor polling (about one-second resolution), +but excludes initial/final integrity hashing. The load timer includes harness +context serialization. The inference timer ends before invariant JSON writes; +its two large phases above are log-marker intervals, not additional time. Memory +is sampled every second for the concurrent child session, excludes the monitor, +and depends on shared pages; no memory reduction is established. + +The result concerns offline inference over this saved trace and enabled relation +set. It does not establish training overhead, distributed scaling, or correctness +of the inferred invariants against ground truth. Earlier attempts combining +cProfile and periodic traceback dumps exited with SIGSEGV and are excluded. Both +were disabled for this timed pair; its success alone does not identify the crash +cause. + +[Full-trace measurements and provenance](inference-identity-cache-full-trace-20261025.json) +include phase times, output digests, trace/source hashes, runtime, and archive +identity. The retained evidence archive contains both complete output sets, +context files, stdout logs, memory samples, and invocation/provenance manifests. + ## Isolated measurement On a deterministic 573,494-row synthetic trace, the median time for construction @@ -28,11 +86,8 @@ garbage collection, and result hashing are excluded. The run used a macOS arm64 workstation, Python 3.10.16, pandas 2.2.3, NumPy 1.26.4, and native thread settings of one. No CPU affinity or exclusive machine access was imposed. -This is an identity-lookup microbenchmark, not completed invariant inference or -training overhead. Paired full-inference results remain pending. Earlier profiling -attempts terminated with SIGSEGV at scheduled diagnostic timers and produced no -completed profile/output; the candidate stopped before the cache phase. Those -attempts provide no valid performance comparison. +This isolates identity lookup. The separate complete inference comparison above +measures the saved training trace. ## Correctness and contract @@ -47,8 +102,7 @@ Identity values, dtypes, and row order must remain stable within a frame. Same-s in-place edits, column replacement, and reordering require a new `TracePandas`. The cache stores O(P + T) values; returned lists and membership checks still cost O(P + T). Its source reference is not a copied DataFrame, but can retain the old -frame until invalidation. No RSS/PSS, peak-memory, or multiprocessing memory-scaling -improvement is claimed. +frame until invalidation. No memory or multiprocessing memory-scaling improvement is claimed. From the checkout, reproduce the isolated measurement with: From c125d060e5f6461b89ffc8cb6542015ed9ff9f7c Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 21:29:34 +0800 Subject: [PATCH 4/4] Describe schemaless identity values and benchmark results honestly --- tools/benchmarks/identity_cache.py | 78 +++++++++++++++++++++++------- traincheck/trace/trace.py | 9 ++-- traincheck/trace/trace_pandas.py | 9 ++-- 3 files changed, 70 insertions(+), 26 deletions(-) diff --git a/tools/benchmarks/identity_cache.py b/tools/benchmarks/identity_cache.py index 3a82e1d2..e42582e9 100644 --- a/tools/benchmarks/identity_cache.py +++ b/tools/benchmarks/identity_cache.py @@ -32,6 +32,7 @@ import sys import time from pathlib import Path +from typing import Literal, TypedDict import numpy as np import pandas as pd @@ -39,22 +40,46 @@ import traincheck.trace.trace_pandas as trace_module from traincheck.trace.trace_pandas import TracePandas +TimingField = Literal["construction_seconds", "queries_seconds", "total_seconds"] +TIMING_FIELDS: tuple[TimingField, ...] = ( + "construction_seconds", + "queries_seconds", + "total_seconds", +) + + +class TrialResult(TypedDict): + repetition: int + position: int + variant: str + construction_seconds: float + queries_seconds: float + total_seconds: float + result_sha256: str + none_results: int + empty_dict_results: int + + +class Measurement(TrialResult, total=False): + cached_identity_counts: dict[str, int] + cache_tuple_containers_shallow_bytes: int + class LegacyIdentityTrace(TracePandas): """Original lookup expressions; all other behavior uses the same checkout.""" - def get_process_ids(self): + def get_process_ids(self) -> list[object]: if "process_id" not in self.events.columns: return [] return self.events["process_id"].dropna().unique().tolist() - def get_thread_ids(self): + def get_thread_ids(self) -> list[object]: if "thread_id" not in self.events.columns: return [] return self.events["thread_id"].dropna().unique().tolist() -def synthetic_frame(rows): +def synthetic_frame(rows: int) -> pd.DataFrame: sequence = np.arange(rows, dtype=np.int64) process_ids = ((sequence // 16) % 8 + 100).astype(np.float64) thread_ids = (sequence % 16 + 1000).astype(np.float64) @@ -67,7 +92,7 @@ def synthetic_frame(rows): ) -def fixed_queries(calls): +def fixed_queries(calls: int) -> list[tuple[int, int, int]]: # Six valid, one unknown-process, and one unknown-thread call per cycle. # Process/thread presence remains independent, matching get_meta_vars. return [ @@ -80,20 +105,31 @@ def fixed_queries(calls): ] -def digest(value): +def digest(value: object) -> str: return hashlib.sha256( json.dumps(value, sort_keys=True, separators=(",", ":")).encode() ).hexdigest() -def trial(cls, data, queries): +def trial( + cls: type[TracePandas], + data: pd.DataFrame, + queries: list[tuple[int, int, int]], + *, + repetition: int, + position: int, + variant: str, +) -> Measurement: gc.collect() begin = time.perf_counter() trace = cls(data, truncate_incomplete_func_calls=False) constructed = time.perf_counter() results = [trace.get_meta_vars(*query) for query in queries] finished = time.perf_counter() - record = { + record: Measurement = { + "repetition": repetition, + "position": position, + "variant": variant, "construction_seconds": constructed - begin, "queries_seconds": finished - constructed, "total_seconds": finished - begin, @@ -112,7 +148,7 @@ def trial(cls, data, queries): return record -def main(): +def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--rows", type=int, default=573494) parser.add_argument("--calls", type=int, default=1000) @@ -129,20 +165,26 @@ def main(): input_hash = hashlib.sha256( pd.util.hash_pandas_object(data, index=True).to_numpy().tobytes() ).hexdigest() - expected = [None if i % 8 in (6, 7) else {} for i in range(args.calls)] + expected: list[dict[str, object] | None] = [ + None if i % 8 in (6, 7) else {} for i in range(args.calls) + ] expected_hash = digest(expected) - measurements = [] + measurements: list[Measurement] = [] variants = {"legacy": LegacyIdentityTrace, "candidate": TracePandas} for repetition in range(args.repetitions): order = ["legacy", "candidate"] if repetition % 2: order.reverse() for position, variant in enumerate(order): - record = trial(variants[variant], data, queries) - assert record["result_sha256"] == expected_hash, record - record.update( - repetition=repetition + 1, position=position + 1, variant=variant + record = trial( + variants[variant], + data, + queries, + repetition=repetition + 1, + position=position + 1, + variant=variant, ) + assert record["result_sha256"] == expected_hash, record measurements.append(record) # Catch accidental mutation of the shared input by either implementation. @@ -159,7 +201,7 @@ def main(): for record in measurements if record["variant"] == variant ) - for metric in ("construction_seconds", "queries_seconds", "total_seconds") + for metric in TIMING_FIELDS } for variant in variants } @@ -191,13 +233,13 @@ def main(): ) }, "candidate_source_sha256": hashlib.sha256( - Path(trace_module.__file__).read_bytes() + Path(inspect.getfile(trace_module)).read_bytes() ).hexdigest(), "candidate_commit": subprocess.check_output( [ "git", "-C", - str(Path(trace_module.__file__).resolve().parents[2]), + str(Path(inspect.getfile(trace_module)).resolve().parents[2]), "rev-parse", "HEAD", ], @@ -207,7 +249,7 @@ def main(): [ "git", "-C", - str(Path(trace_module.__file__).resolve().parents[2]), + str(Path(inspect.getfile(trace_module)).resolve().parents[2]), "diff", "--quiet", "HEAD", diff --git a/traincheck/trace/trace.py b/traincheck/trace/trace.py index c04ebeda..d5636d29 100644 --- a/traincheck/trace/trace.py +++ b/traincheck/trace/trace.py @@ -1,4 +1,5 @@ import logging +from collections.abc import Sequence import polars as pl @@ -183,14 +184,14 @@ def get_end_time(self, process_id=None, thread_id=None) -> float: "This function should be implemented in the child class." ) - def get_process_ids(self) -> list[int]: - """Find all process ids from the trace.""" + def get_process_ids(self) -> Sequence[object]: + """Find process identities, preserving the backend scalar types.""" raise NotImplementedError( "This function should be implemented in the child class." ) - def get_thread_ids(self) -> list[int]: - """Find all thread ids from the trace.""" + def get_thread_ids(self) -> Sequence[object]: + """Find thread identities, preserving the backend scalar types.""" raise NotImplementedError( "This function should be implemented in the child class." ) diff --git a/traincheck/trace/trace_pandas.py b/traincheck/trace/trace_pandas.py index d5e93acc..bedcfd9d 100644 --- a/traincheck/trace/trace_pandas.py +++ b/traincheck/trace/trace_pandas.py @@ -115,7 +115,7 @@ def __init__(self, events, truncate_incomplete_func_calls=True): # Keep the frame alive so object identity cannot be reused across frames. self._identity_values_source: pd.DataFrame | None = None self._identity_values_length = -1 - self._identity_values: dict[str, tuple] = {} + self._identity_values: dict[str, tuple[object, ...]] = {} self.column_dtypes_cached = {} self.stage_traces: dict[str, "TracePandas"] = {} @@ -594,12 +594,13 @@ def get_end_time(self, process_id=None, thread_id=None) -> float: return self.events["time"].max() - def _get_identity_values(self, column: str) -> list: + def _get_identity_values(self, column: str) -> list[object]: """Cache stable identity columns; return a fresh list for every caller. As with call-record indexes, frame replacement or length changes rebuild the cache. Same-size in-place identity edits, row reordering, or dtype changes require a new TracePandas. + Identity columns are schemaless: their scalars are not necessarily ints. Keep pandas' dropna/unique conversion to preserve dtype and encounter order instead of replacing it with Python set semantics. """ @@ -616,7 +617,7 @@ def _get_identity_values(self, column: str) -> list: ) return list(self._identity_values[column]) - def get_process_ids(self) -> list[int]: + def get_process_ids(self) -> list[object]: """Find all process ids from the trace.""" if "process_id" not in self.events.columns: logger.warning( @@ -625,7 +626,7 @@ def get_process_ids(self) -> list[int]: return [] return self._get_identity_values("process_id") - def get_thread_ids(self) -> list[int]: + def get_thread_ids(self) -> list[object]: """Find all thread ids from the trace.""" if "thread_id" not in self.events.columns: logger.warning(