diff --git a/docs/measurements/trace-pandas-indexes-20261025.json b/docs/measurements/trace-pandas-indexes-20261025.json new file mode 100644 index 00000000..94cf3a94 --- /dev/null +++ b/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." + ] +} diff --git a/docs/measurements/trace-pandas-indexes.md b/docs/measurements/trace-pandas-indexes.md new file mode 100644 index 00000000..03e5d3b0 --- /dev/null +++ b/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. diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/legacy_trace_pandas.provenance.json b/tests/fixtures/legacy_trace_pandas.provenance.json new file mode 100644 index 00000000..00498a4d --- /dev/null +++ b/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." +} diff --git a/tests/fixtures/legacy_trace_pandas.py b/tests/fixtures/legacy_trace_pandas.py new file mode 100644 index 00000000..c7256089 --- /dev/null +++ b/tests/fixtures/legacy_trace_pandas.py @@ -0,0 +1,244 @@ +"""Frozen legacy query methods used as a differential oracle. + +Source: OrderLab/TrainCheck commit 9eb4f320d3db34d372145261704b805b88e39bda, +traincheck/trace/trace_pandas.py (SHA256 + e1bb62eae653f95c5498dc921d991af013e8b5717398f94cfbc718812bb600ad). +Only four exact original methods are retained. The production constructor and +supporting types are shared so tests isolate these query/index changes. Defining +this class as TracePandas intentionally makes the original stage factory create +legacy instances too. See legacy_trace_pandas.provenance.json for method hashes. +""" + +import logging + +import pandas as pd + +from traincheck.instrumentor.tracer import TraceLineType +from traincheck.instrumentor.types import PTID +from traincheck.trace.trace_pandas import STAGE_KEY +from traincheck.trace.trace_pandas import TracePandas as ProductionTracePandas +from traincheck.trace.types import ContextManagerState, Liveness +from traincheck.trace.utils import ( + BindedFuncInput, + bind_args_kwargs_to_signature, + load_signature_from_class_method_name, +) +from traincheck.utils import safe_isnan + +logger = logging.getLogger(__name__) + + +class TracePandas(ProductionTracePandas): + def _index_context_manager_meta_vars(self): + """Identify context manager entry and exit events, and add them to the meta_vars.""" + # 1. Find all trace records that are related to __enter__ and __exit__ functions + if "function" not in self.events.columns: + self.context_manager_states = None + return + + if ( + hasattr(self, "context_manager_states") + and self.context_manager_states is not None + ): + return self.context_manager_states + + all_context_managers: dict[PTID, list[ContextManagerState]] = {} + + context_manager_names = self.events[ + (self.events["function"].str.contains("__enter__", na=False)) + & ( + ~self.events["function"].str.contains( + "torch.autograd.grad_mode", na=False + ) + ) # HACK: ignore all autograd related context managers for now as it doesn't have an __init__ event + & ( + ~self.events["function"].str.contains( + "torch.autograd.profiler.record_function", na=False + ) + ) + ]["function"].unique() + + context_manager_init_pre_records = [] + for context_manager_name in context_manager_names: + context_manager_name = context_manager_name.removesuffix(".__enter__") + context_manager_init_pre_records.extend( + self.events[ + (self.events["function"] == f"{context_manager_name}.__init__") + & (self.events["type"] == TraceLineType.FUNC_CALL_PRE) + ].to_dict(orient="records") + ) + + for context_manager_init_pre_record in context_manager_init_pre_records: + init_pre_record = context_manager_init_pre_record + process_id = init_pre_record["process_id"] + thread_id = init_pre_record["thread_id"] + + context_manager_name = init_pre_record["function"].removesuffix(".__init__") + + # find nearest enter and exit events with the same obj_id + obj_id = init_pre_record["obj_id"] + try: + enter_post_record = ( + self.events[ + (self.events["type"] == TraceLineType.FUNC_CALL_POST) + & ( + self.events["function"] + == f"{context_manager_name}.__enter__" + ) + & (self.events["obj_id"] == obj_id) + & (self.events["time"] > init_pre_record["time"]) + & (self.events["process_id"] == process_id) + & (self.events["thread_id"] == thread_id) + ] + .iloc[0] + .to_dict() + ) + + exit_pre_record = ( + self.events[ + (self.events["type"] == TraceLineType.FUNC_CALL_PRE) + & ( + self.events["function"] + == f"{context_manager_name}.__exit__" + ) + & (self.events["obj_id"] == obj_id) + & (self.events["time"] > enter_post_record["time"]) + & (self.events["process_id"] == process_id) + & (self.events["thread_id"] == thread_id) + ] + .iloc[0] + .to_dict() + ) + except IndexError: + # sometimes this happens. an enter might not have a corresponding exit, vice versa. e.g. torch.serialization._opener.__enter__ and torch.serialization._open_zipfile_writer_file.__exit__ + logger.warning( + f"Context manager {context_manager_name} not used properly. Skipping." + ) + continue + + assert ( + enter_post_record["time"] < exit_pre_record["time"] + ), f"enter not before exit: {enter_post_record.to_dict()} {exit_pre_record.to_dict()}" + + start_time, end_time = ( + enter_post_record["time"], + exit_pre_record["time"], + ) + + args = init_pre_record["args"] + kwargs = init_pre_record["kwargs"] + if not safe_isnan(args): + signature = load_signature_from_class_method_name( + init_pre_record["function"] + ) + + binded_args_and_kwargs = bind_args_kwargs_to_signature( + args, kwargs, signature + ) + else: + # create an empty BindedFuncInput if args is NaN, as it indicates + # that we did not record the args and kwargs for this function call + binded_args_and_kwargs = BindedFuncInput({}) + + ptid = PTID(process_id, thread_id) + if ptid not in all_context_managers: + all_context_managers[ptid] = [] + + logger.debug(f"Adding context manager: {context_manager_name} {ptid}") + all_context_managers[ptid].append( + ContextManagerState( + name=init_pre_record["function"].removesuffix(".__init__"), + ptid=ptid, + liveness=Liveness(start_time, end_time), + input=binded_args_and_kwargs, + ) + ) + + logger.info(f"Found {len(all_context_managers)} context managers.") + + self.context_manager_states = all_context_managers + + def get_pre_func_call_record(self, func_call_id: str) -> dict: + """Get the pre call record of a function given its func_call_id.""" + # TODO: check if this is correct + pre_records = self.events[ + (self.events["func_call_id"] == func_call_id) + & (self.events["type"] == TraceLineType.FUNC_CALL_PRE) + ] + + assert ( + len(pre_records) == 1 + ), f"{len(pre_records)} pre call events found for {func_call_id}, expected 1" + + pre_record = pre_records.iloc[0].to_dict() + + return pre_record + + def get_post_func_call_record(self, func_call_id: str) -> dict | None: + """Get the post call record of a function given its func_call_id. + Returns None if the post call event is not found and the pre-call event is the outermost function call. (see the doc of `_rm_incomplete_trailing_func_calls`) + """ + # TODO: check if this is correct + post_records = self.events[ + (self.events["func_call_id"] == func_call_id) + & ( + self.events["type"].isin( + [ + TraceLineType.FUNC_CALL_POST, + TraceLineType.FUNC_CALL_POST_EXCEPTION, + ] + ) + ) + ] + + assert ( + len(post_records) <= 1 + ), f"{post_records.shape[0]} post call events found for {func_call_id}, expected 1" + + if len(post_records) == 0: + logger.warning(f"No post call event found for {func_call_id}") + # check if the pre-call event is the outermost function call + pre_record = self.get_pre_func_call_record(func_call_id) + + outermost_func_call_pre = ( + self.events[ + (self.events["type"] == TraceLineType.FUNC_CALL_PRE) + & (self.events["process_id"] == pre_record["process_id"]) + & (self.events["thread_id"] == pre_record["thread_id"]) + ] + .iloc[0] + .to_dict() + ) + + if pre_record["func_call_id"] == outermost_func_call_pre["func_call_id"]: + return None + else: + raise ValueError( + f"No post call event found for {func_call_id}, but it is not the outermost function call." + ) + + post_record = post_records.iloc[ + 0 + ].to_dict() # order of events doesn't matter as the result here is a single row + + return post_record + + def get_traces_for_stage(self) -> dict[str, "TracePandas"]: # type: ignore + """Get the traces split by stages.""" + + if not self.is_stage_annotated(): + raise ValueError("Trace is not annotated with stages.") + + if self.stage_traces: + return self.stage_traces + + traces = {} + for stage in self.events[STAGE_KEY].unique(): + traces[stage] = TracePandas( + self.events[self.events[STAGE_KEY] == stage], + truncate_incomplete_func_calls=False, + ) + + self.stage_traces = traces + + return traces diff --git a/tests/test_call_index.py b/tests/test_call_index.py new file mode 100644 index 00000000..acf8f030 --- /dev/null +++ b/tests/test_call_index.py @@ -0,0 +1,251 @@ +"""Call-record index differential tests against frozen legacy query methods.""" + +import copy +import logging +import unittest + +import numpy as np +import pandas as pd + +import traincheck.trace.trace_pandas as NEW +from tests.fixtures import legacy_trace_pandas as OLD + +PRE = OLD.TraceLineType.FUNC_CALL_PRE +POST = OLD.TraceLineType.FUNC_CALL_POST +EXCEPTION = OLD.TraceLineType.FUNC_CALL_POST_EXCEPTION + + +def record(call_id, when, kind, *, pid=1, tid=2, stage="train", **extra): + return { + "func_call_id": call_id, + "function": "fixtures.work", + "time": when, + "type": kind, + "process_id": pid, + "thread_id": tid, + "args": {"0": {"int": 7}}, + "kwargs": {}, + "return_values": {"int": 9}, + "exception": "fixture exception", + "meta_vars.stage": stage, + **extra, + } + + +def pair(call_id="call", start=0, **extra): + return [ + record(call_id, start, PRE, **extra), + record(call_id, start + 1, POST, **extra), + ] + + +def make(module, frame, truncate=False): + if not isinstance(frame, pd.DataFrame): + frame = pd.DataFrame(copy.deepcopy(frame)) + else: + frame = frame.copy(deep=True) + return module.TracePandas(frame, truncate_incomplete_func_calls=truncate) + + +def outcome(trace, method, call_id): + try: + return ("ok", getattr(trace, method)(call_id)) + except Exception as exc: + return (type(exc).__name__, str(exc)) + + +class StringAlias: + def __eq__(self, other): + return other == "alias" + + def __repr__(self): + return "StringAlias()" + + +class CallIndexTests(unittest.TestCase): + def setUp(self): + previous = logging.root.manager.disable + logging.disable(logging.WARNING) + self.addCleanup(logging.disable, previous) + + def equivalent(self, rows, call_ids, *, truncate=False): + old, new = [make(module, rows, truncate) for module in [OLD, NEW]] + for call_id in call_ids: + for method in ["get_pre_func_call_record", "get_post_func_call_record"]: + with self.subTest(call_id=repr(call_id), method=method): + expected = outcome(old, method, call_id) + actual = outcome(new, method, call_id) + self.assertEqual(expected, actual) + return old, new + + def test_exact_pre_post_records_and_row_order(self): + rows = pair("z", 10) + pair("a", 0, pid=9, tid=8) + _, new = self.equivalent(rows, ["a", "z", "missing"]) + self.assertEqual(new.get_pre_func_call_record("z"), rows[0]) + self.assertEqual(new.get_post_func_call_record("a"), rows[3]) + self.assertIs(type(new.get_pre_func_call_record("z")["time"]), int) + + def test_exception_post_record(self): + rows = [record("call", 0, PRE), record("call", 1, EXCEPTION)] + old, new = self.equivalent(rows, ["call"]) + self.assertEqual( + old.query_func_call_event("call").__dict__, + new.query_func_call_event("call").__dict__, + ) + + def test_duplicate_pre_count_assertions(self): + rows = [record("bad", 0, PRE), record("bad", 1, PRE), record("bad", 2, POST)] + old, new = self.equivalent(rows, ["bad"]) + self.assertEqual( + outcome(new, "get_pre_func_call_record", "bad"), + ("AssertionError", "2 pre call events found for bad, expected 1"), + ) + + def test_duplicate_post_and_exception_count_assertions(self): + rows = [ + record("bad", 0, PRE), + record("bad", 1, POST), + record("bad", 2, EXCEPTION), + ] + self.equivalent(rows, ["bad"]) + + def test_unqueried_duplicates_do_not_raise_early(self): + rows = pair("good") + [record("bad", 4, PRE), record("bad", 5, PRE)] + _, new = self.equivalent(rows, ["good"]) + self.assertIsNotNone(new._func_call_positions) + self.assertEqual(len(new._func_call_positions["bad"][0]), 2) + + def test_missing_outermost_and_nonoutermost_posts_per_thread(self): + rows = [ + record("outer", 100, PRE), + record("inner", 0, PRE), + record("other_thread", 2, PRE, tid=3), + ] + old, new = self.equivalent(rows, ["outer", "inner", "other_thread", "absent"]) + self.assertIsNone(new.get_post_func_call_record("outer")) + self.assertIsNone(new.get_post_func_call_record("other_thread")) + self.assertEqual( + outcome(new, "get_post_func_call_record", "inner"), + ( + "ValueError", + "No post call event found for inner, but it is not the outermost function call.", + ), + ) + + def test_duplicate_dataframe_labels_use_positions(self): + rows = pair("a") + pair("b", 4) + self.equivalent(pd.DataFrame(rows, index=[17, 17, 17, 17]), ["a", "b"]) + + def test_null_ids_never_use_dictionary_equality(self): + rows = pair("good") + for value in [None, float("nan"), pd.NA]: + rows += pair(value, 5) + _, new = self.equivalent(rows, ["good", None, float("nan"), pd.NA]) + self.assertIsNotNone(new._func_call_positions) + + def test_numeric_ids_and_custom_equality_fall_back(self): + for rows, ids in [ + (pair(7), [7, "7"]), + (pair(StringAlias()), ["alias"]), + (pair(np.str_("numpy")), ["numpy"]), + (pair("good") + pair(7, 5), ["good", 7]), + ]: + old, new = self.equivalent(rows, ids) + self.assertIsNone(new._func_call_positions) + + def test_pandas_string_and_unsupported_categorical_dtypes(self): + for dtype in ["string", "category"]: + frame = pd.DataFrame(pair("a") + pair("b", 4)) + for column in ["func_call_id", "type"]: + frame[column] = frame[column].astype(dtype) + _, new = self.equivalent(frame, ["a", "b"]) + if dtype == "category": + self.assertIsNone(new._func_call_positions) + else: + self.assertIsNotNone(new._func_call_positions) + + def test_null_event_types_are_ignored(self): + rows = pair("good") + rows += [record("extra", 3, value) for value in [None, float("nan"), pd.NA]] + self.equivalent(rows, ["good", "extra"]) + + def test_stage_split_indexes_are_independent(self): + rows = pair("init", stage="init") + pair("train", 5, stage="train") + rows += [ + record("cross", 10, PRE, stage="init"), + record("cross", 11, POST, stage="train"), + ] + old, new = self.equivalent(rows, ["init", "train", "cross"]) + old_stages = old.get_traces_for_stage() + new_stages = new.get_traces_for_stage() + for stage in old_stages: + self.assertIs(type(old_stages[stage]), OLD.TracePandas) + self.assertIs(type(new_stages[stage]), NEW.TracePandas) + for call_id in ["init", "train", "cross"]: + for method in ["get_pre_func_call_record", "get_post_func_call_record"]: + self.assertEqual( + outcome(old_stages[stage], method, call_id), + outcome(new_stages[stage], method, call_id), + ) + self.assertIsNot( + new._func_call_positions, new_stages["train"]._func_call_positions + ) + self.assertIsNot( + new_stages["init"]._func_call_positions, + new_stages["train"]._func_call_positions, + ) + + def test_top_level_return_dicts_are_fresh(self): + for module in [OLD, NEW]: + trace = make(module, pair()) + for method in ["get_pre_func_call_record", "get_post_func_call_record"]: + first = getattr(trace, method)("call") + second = getattr(trace, method)("call") + self.assertIsNot(first, second) + first["function"] = "mutated" + first["new_key"] = True + self.assertEqual(getattr(trace, method)("call"), second) + + def test_index_is_lazy_reused_and_rebuilt_on_frame_replacement(self): + old, new = [make(module, pair()) for module in [OLD, NEW]] + self.assertIsNone(new._func_call_positions) + new.get_pre_func_call_record("call") + cached = new._func_call_positions + new.get_post_func_call_record("call") + self.assertIs(new._func_call_positions, cached) + for trace in [old, new]: + trace.events = pd.DataFrame(pair("replacement", 9)) + for call_id in ["call", "replacement"]: + for method in ["get_pre_func_call_record", "get_post_func_call_record"]: + self.assertEqual( + outcome(old, method, call_id), outcome(new, method, call_id) + ) + self.assertIsNot(new._func_call_positions, cached) + + def test_in_place_row_count_change_rebuilds(self): + old, new = [make(module, pair()) for module in [OLD, NEW]] + new.get_pre_func_call_record("call") + cached = new._func_call_positions + for trace in [old, new]: + trace.events.loc[len(trace.events)] = record("call", 4, PRE) + self.assertEqual( + outcome(old, "get_pre_func_call_record", "call"), + outcome(new, "get_pre_func_call_record", "call"), + ) + self.assertIsNot(new._func_call_positions, cached) + + def test_index_built_after_incomplete_trailing_call_truncation(self): + rows = [ + record("outer", 0, PRE), + record("inner", 1, PRE), + record("inner", 2, POST), + record("outer", 10, POST), + record("trailing", 11, PRE), + ] + old, new = self.equivalent(rows, ["outer", "inner", "trailing"], truncate=True) + self.assertEqual(len(new.events), 4) + self.assertIs(new._func_call_positions_source, new.events) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_context_index.py b/tests/test_context_index.py new file mode 100644 index 00000000..a6390d8e --- /dev/null +++ b/tests/test_context_index.py @@ -0,0 +1,287 @@ +"""Differential context-index tests against frozen legacy query methods.""" + +import inspect +import logging +import random +import unittest +from unittest.mock import patch + +import numpy as np +import pandas as pd + +import traincheck.trace.trace_pandas as NEW +from tests.fixtures import legacy_trace_pandas as OLD + +PRE = OLD.TraceLineType.FUNC_CALL_PRE +POST = OLD.TraceLineType.FUNC_CALL_POST + + +def event( + time, + method, + *, + name="fixtures.Ctx", + obj=1, + pid=10, + tid=20, + kind=None, + args=float("nan"), + kwargs=None, + stage="train", +): + return { + "time": time, + "function": f"{name}.{method}", + "type": kind or (POST if method == "__enter__" else PRE), + "obj_id": obj, + "process_id": pid, + "thread_id": tid, + "args": args, + "kwargs": kwargs or {}, + "meta_vars.stage": stage, + } + + +def cycle(start, **kwargs): + return [ + event(start, "__init__", **kwargs), + event(start + 1, "__enter__", **kwargs), + event(start + 2, "__exit__", **kwargs), + ] + + +def run(module, frame): + trace = module.TracePandas( + frame.copy(deep=True), truncate_incomplete_func_calls=False + ) + return trace, encode(trace) + + +def encode(trace): + states = trace.context_manager_states + if states is None: + return None + return [ + (tuple(key), [state.to_dict() for state in value]) + for key, value in states.items() + ] + + +class ContextIndexTests(unittest.TestCase): + def setUp(self): + previous = logging.root.manager.disable + logging.disable(logging.WARNING) + self.addCleanup(logging.disable, previous) + + def equivalent(self, rows, **frame_options): + frame = ( + rows + if isinstance(rows, pd.DataFrame) + else pd.DataFrame(rows, **frame_options) + ) + original, expected = run(OLD, frame) + optimized, actual = run(NEW, frame) + self.assertEqual(expected, actual) + # Native scalar types remain suitable for the same JSON serialization. + for _, states in actual or []: + for state in states: + self.assertNotIsInstance(state["start_time"], np.generic) + self.assertNotIsInstance(state["end_time"], np.generic) + return original, optimized + + def test_nested_reused_ids_processes_threads_and_class_order(self): + rows = cycle(0, name="fixtures.B", obj=1) + rows += cycle(5, name="fixtures.A", obj=1) + rows += cycle(10, name="fixtures.B", obj=1) + rows += cycle(1, name="fixtures.B", obj=2) + rows += cycle(0, name="fixtures.B", obj=1, pid=11) + rows += cycle(0, name="fixtures.B", obj=1, tid=21) + old, new = self.equivalent(rows) + self.assertEqual( + [s.name for s in new.context_manager_states[OLD.PTID(10, 20)]], + ["fixtures.B", "fixtures.B", "fixtures.B", "fixtures.A"], + ) + for when in [0, 1, 2, 6, 11, 12, 100]: + self.assertEqual( + [s.to_dict() for s in old.query_active_context_managers(when, 10, 20)], + [s.to_dict() for s in new.query_active_context_managers(when, 10, 20)], + ) + + def test_first_row_not_nearest_timestamp(self): + rows = [ + event(0, "__init__"), + event(8, "__enter__"), + event(2, "__enter__"), + event(15, "__exit__"), + event(9, "__exit__"), + ] + _, new = self.equivalent(rows) + state = next(iter(new.context_manager_states.values()))[0] + self.assertEqual((state.liveness.start_time, state.liveness.end_time), (8, 15)) + + def test_equal_timestamps_strictly_later(self): + rows = [ + event(5, "__init__"), + event(5, "__enter__"), + event(6, "__enter__"), + event(6, "__exit__"), + event(7, "__exit__"), + ] + _, new = self.equivalent(rows) + self.assertEqual( + next(iter(new.context_manager_states.values()))[0].liveness.start_time, 6 + ) + + def test_null_keys_do_not_match_each_other(self): + for field in ["obj_id", "process_id", "thread_id"]: + for missing in [None, float("nan"), pd.NA]: + with self.subTest(field=field, missing=missing): + rows = cycle(0) + for row in rows: + row[field] = missing + rows += cycle(5, obj=2) + self.equivalent(rows) + + def test_md_none_key_matches_original_custom_equality(self): + from traincheck.trace.types import MD_NONE + + for field in ["obj_id", "process_id", "thread_id"]: + rows = cycle(0) + for row in rows: + row[field] = MD_NONE() + self.equivalent(rows) + + def test_mixed_md_none_and_null_keys(self): + from traincheck.trace.types import MD_NONE + + rows = [ + event(0, "__init__", obj=None), + event(1, "__enter__", obj=MD_NONE()), + event(2, "__exit__", obj=MD_NONE()), + ] + self.equivalent(rows) + + def test_nullable_numeric_keys(self): + frame = pd.DataFrame(cycle(0) + cycle(5, obj=None) + cycle(10, obj=2)) + for col in ["obj_id", "process_id", "thread_id"]: + frame[col] = frame[col].astype("Int64") + self.equivalent(frame) + + def test_missing_pairs_and_wrong_event_types(self): + rows = [ + event(0, "__init__"), + event(1, "__enter__", kind=PRE), + event(2, "__exit__"), + ] + rows += [event(4, "__init__", obj=2), event(5, "__enter__", obj=2)] + rows += [event(6, "__enter__", obj=3), event(7, "__exit__", obj=3)] + rows += cycle(10, obj=4) + self.equivalent(rows) + + def test_infinite_missing_and_nullable_timestamps(self): + rows = [ + event(float("nan"), "__init__", obj=2), + event(0, "__init__"), + event(float("nan"), "__enter__"), + event(1, "__enter__"), + event(float("inf"), "__exit__"), + ] + self.equivalent(rows) + frame = pd.DataFrame(rows) + frame["time"] = frame["time"].astype("Float64") + for module in [OLD, NEW]: + with self.assertRaises(TypeError): + run(module, frame) + + def test_large_integer_timestamps_preserve_precision(self): + for dtype in ["int64", "uint64", "Int64", "UInt64", "float64", "object"]: + for start in [2**53, 2**60, 2**63 - 4]: + with self.subTest(dtype=dtype, start=start): + frame = pd.DataFrame(cycle(start)) + frame["time"] = frame["time"].astype(dtype) + self.equivalent(frame) + for start in [2**53, 2**60, 2**63 - 8]: + # Object dtype retains a float init query against exact integers. + frame = pd.DataFrame(cycle(start)) + frame["time"] = pd.Series( + [float(start), start + 1, start + 2], dtype=object + ) + self.equivalent(frame) + + def test_duplicate_dataframe_indexes_and_mixed_key_types(self): + rows = cycle(0, obj="1") + cycle(5, obj=1) + cycle(10, obj=True) + self.equivalent(pd.DataFrame(rows, index=[4] * len(rows))) + + def test_no_context_or_no_function(self): + self.equivalent(pd.DataFrame({"time": [0]})) + self.equivalent([event(0, "unrelated")]) + self.equivalent([event(0, "__enter__")]) + + def test_ignore_autograd_contexts(self): + rows = cycle(0, name="torch.autograd.grad_mode.enable_grad") + rows += cycle(3, name="torch.autograd.profiler.record_function") + rows += cycle(6) + _, new = self.equivalent(rows) + self.assertEqual(len(next(iter(new.context_manager_states.values()))), 1) + + def test_stage_split_rebuild_and_cache(self): + rows = cycle(0, stage="init") + cycle(3, stage="train") + rows += [ + event(10, "__init__", obj=2, stage="init"), + event(11, "__enter__", obj=2, stage="train"), + event(12, "__exit__", obj=2, stage="train"), + ] + old, new = self.equivalent(rows) + expected = old.get_traces_for_stage() + actual = new.get_traces_for_stage() + self.assertEqual(list(expected), list(actual)) + for stage in expected: + self.assertIs(type(expected[stage]), OLD.TracePandas) + self.assertIs(type(actual[stage]), NEW.TracePandas) + self.assertEqual(encode(expected[stage]), encode(actual[stage])) + self.assertIs(new.get_traces_for_stage(), actual) + cached = new.context_manager_states + self.assertIs(new._index_context_manager_meta_vars(), cached) + + def test_bindings_and_signature_lookup_remain_per_init(self): + rows = cycle(0, args={"0": {"object": {}}, "1": {"int": 7}}) + rows += cycle(5, args={"0": {"object": {}}, "1": {"int": 9}}) + signatures = [ + inspect.Signature( + [ + inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD), + inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD), + ] + ) + for name in ["first", "second"] + ] + outputs = [] + for module in [OLD, NEW]: + with patch.object( + module, "load_signature_from_class_method_name", side_effect=signatures + ) as loader: + outputs.append(run(module, pd.DataFrame(rows))[1]) + self.assertEqual(loader.call_count, 2) + self.assertEqual(*outputs) + + def test_randomized_row_order_and_reused_identity_oracle(self): + for seed in range(80): + rng = random.Random(seed) + rows = [] + for _ in range(20): + rows.extend( + cycle( + rng.randrange(20), + obj=rng.randrange(3), + pid=rng.randrange(2), + tid=rng.randrange(2), + name=rng.choice(["fixtures.A", "fixtures.B"]), + ) + ) + rng.shuffle(rows) + with self.subTest(seed=seed): + self.equivalent(rows) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tools/benchmarks/trace_pandas_indexes.py b/tools/benchmarks/trace_pandas_indexes.py new file mode 100644 index 00000000..6f29b7de --- /dev/null +++ b/tools/benchmarks/trace_pandas_indexes.py @@ -0,0 +1,156 @@ +"""Reproduce bounded synthetic TracePandas indexing measurements from a checkout. + +Run: python -m tools.benchmarks.trace_pandas_indexes --output measurements.json +No training, cluster access, or saved trace files are required. Each lookup run +includes its first lazy-index construction. The legacy oracle preserves the four +original query methods; supporting runtime types are identical on both sides. +""" + +import argparse +import hashlib +import inspect +import json +import logging +import platform +import statistics +import time +from collections.abc import Callable +from pathlib import Path +from typing import TypedDict + +import numpy as np +import pandas as pd + +import traincheck.trace.trace_pandas as candidate +from tests.fixtures import legacy_trace_pandas as legacy +from tests.test_call_index import pair +from tests.test_context_index import cycle, encode, event +from traincheck.trace.trace_pandas import TracePandas + + +class ComparisonResult(TypedDict): + elapsed_seconds: dict[str, list[float]] + median_seconds: dict[str, float] + speedup: float + exact_results_equal: bool + + +def compare( + frame: pd.DataFrame, + operation: Callable[[type[TracePandas], pd.DataFrame], object], + repetitions: int, +) -> ComparisonResult: + seconds: dict[str, list[float]] = {"legacy": [], "candidate": []} + expected: object = None + for repetition in range(repetitions): + variants = [ + ("legacy", legacy.TracePandas), + ("candidate", candidate.TracePandas), + ] + if repetition % 2: + variants.reverse() + for name, trace_type in variants: + events = frame.copy(deep=True) + started = time.perf_counter() + value = operation(trace_type, events) + seconds[name].append(time.perf_counter() - started) + if expected is None: + expected = value + assert value == expected, f"{name} result differs from the legacy oracle" + medians = {name: statistics.median(values) for name, values in seconds.items()} + return { + "elapsed_seconds": seconds, + "median_seconds": medians, + "speedup": medians["legacy"] / medians["candidate"], + "exact_results_equal": True, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=["all", "context", "call"], default="all") + parser.add_argument("--repetitions", type=int, default=3) + parser.add_argument("--contexts", type=int, default=500) + parser.add_argument("--noise", type=int, default=50_000) + parser.add_argument("--calls", type=int, default=20_000) + parser.add_argument( + "--queries", type=int, default=500, help="Pre/post pairs queried" + ) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + for name in ["repetitions", "contexts", "calls", "queries"]: + if getattr(args, name) < 1: + parser.error(f"--{name} must be positive") + if args.noise < 0: + parser.error("--noise must be nonnegative") + logging.disable(logging.WARNING) + results: dict[str, object] = {} + if args.mode in ("all", "context"): + rows = [] + for index in range(args.contexts): + rows.extend( + cycle(index * 3, obj=index % 20, name=f"fixtures.Ctx{index % 4}") + ) + rows.extend( + event(args.contexts * 3 + index, "math_call") for index in range(args.noise) + ) + frame = pd.DataFrame(rows) + + def prepare(trace_type: type[TracePandas], events: pd.DataFrame) -> object: + return encode(trace_type(events, truncate_incomplete_func_calls=False)) + + results["context"] = { + "rows": len(frame), + "contexts": args.contexts, + "timed_work": "constructor plus context-state serialization", + **compare(frame, prepare, args.repetitions), + } + if args.mode in ("all", "call"): + frame = pd.DataFrame( + [ + record + for index in range(args.calls) + for record in pair(f"call-{index:08}", start=index * 2) + ] + ) + wanted = [ + f"call-{(index * 117) % args.calls:08}" for index in range(args.queries) + ] + + def lookup(trace_type: type[TracePandas], events: pd.DataFrame) -> object: + trace = trace_type(events, truncate_incomplete_func_calls=False) + return [ + ( + trace.get_pre_func_call_record(call_id), + trace.get_post_func_call_record(call_id), + ) + for call_id in wanted + ] + + results["call"] = { + "rows": len(frame), + "record_queries": args.queries * 2, + "timed_work": "constructor plus record lookups, including lazy-index construction", + **compare(frame, lookup, args.repetitions), + } + output = { + "description": "Synthetic fixtures, alternating variant order; no training or inference", + "repetitions": args.repetitions, + "runtime": { + "python": platform.python_version(), + "pandas": pd.__version__, + "numpy": np.__version__, + }, + "candidate_source_sha256": hashlib.sha256( + Path(inspect.getfile(candidate)).read_bytes() + ).hexdigest(), + "results": results, + } + rendered = json.dumps(output, indent=2) + "\n" + if args.output: + args.output.write_text(rendered) + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/traincheck/trace/trace_pandas.py b/traincheck/trace/trace_pandas.py index 8bb89ca7..61e933ee 100644 --- a/traincheck/trace/trace_pandas.py +++ b/traincheck/trace/trace_pandas.py @@ -1,5 +1,6 @@ import logging import re +from numbers import Real from typing import Any import pandas as pd @@ -15,6 +16,7 @@ ContextManagerState, FuncCallEvent, FuncCallExceptionEvent, + IncompleteFuncCallEvent, Liveness, VarChangeEvent, VarInstId, @@ -65,11 +67,11 @@ def __init__(self, events, truncate_incomplete_func_calls=True): self.function_pool = None self.cache_id_2_func_call_event: dict[ - str, FuncCallEvent | FuncCallExceptionEvent + str, FuncCallEvent | FuncCallExceptionEvent | IncompleteFuncCallEvent ] = {} self.all_func_call_ids: dict[str, list[str]] = {} - self.context_manager_states: dict[PTID, list[ContextManagerState]] + self.context_manager_states: dict[PTID, list[ContextManagerState]] | None self.cache_high_level_events_within_id: dict[ str, list[FuncCallEvent | FuncCallExceptionEvent | VarChangeEvent] ] = {} @@ -101,6 +103,14 @@ def __init__(self, events, truncate_incomplete_func_calls=True): "Done truncating incomplete trailing function calls from the trace." ) + # Build call-record indexes lazily, after incomplete-call truncation. + # Like the other trace caches, these assume stable event values after + # initialization. Replacing the frame or changing its length rebuilds + # this index; same-shape in-place cell edits require a new TracePandas. + self._func_call_positions_source: pd.DataFrame | None = None + self._func_call_positions_length = -1 + self._func_call_positions: dict[str, tuple[list[int], list[int]]] | None = None + self.column_dtypes_cached = {} self.stage_traces: dict[str, "TracePandas"] = {} @@ -283,12 +293,14 @@ def _rm_incomplete_trailing_func_calls(self): # test_dump(self.events) - def _index_context_manager_meta_vars(self): + def _index_context_manager_meta_vars( + self, + ) -> dict[PTID, list[ContextManagerState]] | None: """Identify context manager entry and exit events, and add them to the meta_vars.""" # 1. Find all trace records that are related to __enter__ and __exit__ functions if "function" not in self.events.columns: self.context_manager_states = None - return + return None if ( hasattr(self, "context_manager_states") @@ -312,15 +324,103 @@ def _index_context_manager_meta_vars(self): ) ]["function"].unique() - context_manager_init_pre_records = [] - for context_manager_name in context_manager_names: - context_manager_name = context_manager_name.removesuffix(".__enter__") - context_manager_init_pre_records.extend( - self.events[ - (self.events["function"] == f"{context_manager_name}.__init__") - & (self.events["type"] == TraceLineType.FUNC_CALL_PRE) - ].to_dict(orient="records") + # Group relevant events once, retaining their original row order. Context + # names and init records keep the old class-first encounter order. + init_functions = [ + f"{name.removesuffix('.__enter__')}.__init__" + for name in context_manager_names + ] + if not init_functions: + self.context_manager_states = all_context_managers + return None + + init_records = self.events[ + self.events["function"].isin(init_functions) + & (self.events["type"] == TraceLineType.FUNC_CALL_PRE) + ] + init_by_function = { + name: records.to_dict(orient="records") + for name, records in init_records.groupby("function", sort=False) + } + context_manager_init_pre_records = [ + record + for name in init_functions + for record in init_by_function.get(name, []) + ] + if not context_manager_init_pre_records: + self.context_manager_states = all_context_managers + return None + + context_names = [ + name.removesuffix(".__enter__") for name in context_manager_names + ] + group_columns = ["function", "type", "process_id", "thread_id", "obj_id"] + matching_events = self.events.loc[ + ( + self.events["function"].isin( + [f"{name}.__enter__" for name in context_names] + ) + & (self.events["type"] == TraceLineType.FUNC_CALL_POST) + ) + | ( + self.events["function"].isin( + [f"{name}.__exit__" for name in context_names] + ) + & (self.events["type"] == TraceLineType.FUNC_CALL_PRE) + ), + group_columns + ["time"], + ] + event_times: dict[ + tuple[object, object, object, object, object], tuple[pd.Series, bool] + ] = {} + for key, records in matching_events.groupby( + group_columns, sort=False, dropna=True, observed=True + ): + times = records["time"] + # Binary search is equivalent to the original first matching row + # only when this group's numeric times are sorted and non-null. + # Nullable extension arrays may promote UInt64 search values to + # floats and lose precision above 2**53, so keep their comparison. + ordered = ( + not isinstance(times.dtype, pd.api.extensions.ExtensionDtype) + and times.dtype.kind in "biuf" + and not times.isna().any() + and times.is_monotonic_increasing ) + event_times[key] = (times, ordered) + + # pandas object/extension columns may return arbitrary comparable scalar + # types. Keep Any only for this legacy comparison boundary: narrowing to + # float/int (or coercing) changes nullable, UInt64, and object semantics. + def first_time_after( + function: str, + event_type: str, + process_id: object, + thread_id: object, + obj_id: object, + time: Any, + ) -> Any: + if not isinstance(time, Real): + # Keep pandas validation for nonnumeric/null scalar thresholds, + # including malformed nullable timestamps, even with no group. + self.events["time"] > time + if any(pd.isna(key) for key in (process_id, thread_id, obj_id)): + # pandas equality never matches missing identity values, even + # when a custom non-null key compares equal to None. + raise IndexError("No matching context event") + group = event_times.get( + (function, event_type, process_id, thread_id, obj_id) + ) + if group is None: + raise IndexError("No matching context event") + times, ordered = group + if ordered and isinstance(time, Real) and not pd.isna(time): + position = times.searchsorted(time, side="right") + return times.iloc[position : position + 1].tolist()[0] + # Unsorted times must retain row-order selection, not nearest time. + # This also preserves pandas comparison behavior for nullable and + # nonnumeric timestamp dtypes. + return times[times > time].iloc[:1].tolist()[0] for context_manager_init_pre_record in context_manager_init_pre_records: init_pre_record = context_manager_init_pre_record @@ -329,55 +429,35 @@ def _index_context_manager_meta_vars(self): context_manager_name = init_pre_record["function"].removesuffix(".__init__") - # find nearest enter and exit events with the same obj_id + # Select the first later event in row order with the same identity. obj_id = init_pre_record["obj_id"] try: - enter_post_record = ( - self.events[ - (self.events["type"] == TraceLineType.FUNC_CALL_POST) - & ( - self.events["function"] - == f"{context_manager_name}.__enter__" - ) - & (self.events["obj_id"] == obj_id) - & (self.events["time"] > init_pre_record["time"]) - & (self.events["process_id"] == process_id) - & (self.events["thread_id"] == thread_id) - ] - .iloc[0] - .to_dict() + start_time = first_time_after( + f"{context_manager_name}.__enter__", + TraceLineType.FUNC_CALL_POST, + process_id, + thread_id, + obj_id, + init_pre_record["time"], ) - - exit_pre_record = ( - self.events[ - (self.events["type"] == TraceLineType.FUNC_CALL_PRE) - & ( - self.events["function"] - == f"{context_manager_name}.__exit__" - ) - & (self.events["obj_id"] == obj_id) - & (self.events["time"] > enter_post_record["time"]) - & (self.events["process_id"] == process_id) - & (self.events["thread_id"] == thread_id) - ] - .iloc[0] - .to_dict() + end_time = first_time_after( + f"{context_manager_name}.__exit__", + TraceLineType.FUNC_CALL_PRE, + process_id, + thread_id, + obj_id, + start_time, ) except IndexError: - # sometimes this happens. an enter might not have a corresponding exit, vice versa. e.g. torch.serialization._opener.__enter__ and torch.serialization._open_zipfile_writer_file.__exit__ + # Some traces have an unmatched entry or exit event. logger.warning( f"Context manager {context_manager_name} not used properly. Skipping." ) continue assert ( - enter_post_record["time"] < exit_pre_record["time"] - ), f"enter not before exit: {enter_post_record.to_dict()} {exit_pre_record.to_dict()}" - - start_time, end_time = ( - enter_post_record["time"], - exit_pre_record["time"], - ) + start_time < end_time + ), f"enter not before exit: {start_time} {end_time}" args = init_pre_record["args"] kwargs = init_pre_record["kwargs"] @@ -411,6 +491,7 @@ def _index_context_manager_meta_vars(self): logger.info(f"Found {len(all_context_managers)} context managers.") self.context_manager_states = all_context_managers + return None def is_stage_annotated(self): # ideally we want to have a static manifest for the trace produced by the instrumentor according to the args @@ -558,7 +639,7 @@ def get_column_dtype(self, column_name: str) -> type: return self.column_dtypes_cached[column_name] filtered_values = self.events[column_name].dropna() - #filtered_values = filtered_values[filtered_values != MD_NONE()] + # filtered_values = filtered_values[filtered_values != MD_NONE()] if filtered_values.dtype == object: filtered_values = filtered_values[filtered_values != MD_NONE()] @@ -1007,13 +1088,83 @@ def query_var_changes_within_func_call( (start_time, end_time), process_id=pre_record["process_id"] ) + def _get_func_call_record_positions( + self, func_call_id: str, *, post: bool + ) -> list[int] | None: + """Return indexed positions, or None to use the original pandas filter. + + Index only ordinary strings/nulls: pandas equality for arbitrary objects + and extension dtypes can differ from Python dictionary equality. + """ + if type(func_call_id) is not str: # noqa: E721 - exclude custom equality + return None + + if ( + self._func_call_positions_source is not self.events + or self._func_call_positions_length != len(self.events) + ): + ids = self.events["func_call_id"] + event_types = self.events["type"] + positions: dict[str, tuple[list[int], list[int]]] | None = None + if all( + isinstance(column, pd.Series) + and ( + pd.api.types.is_object_dtype(column.dtype) + or isinstance(column.dtype, pd.StringDtype) + ) + for column in (ids, event_types) + ): + positions = {} + for position, (call_id, event_type) in enumerate(zip(ids, event_types)): + # None, NaN, and pd.NA cannot equal an ordinary string in + # the original pandas filter. All other nonstrings use it. + if any( + type(value) is not str # noqa: E721 - exclude custom equality + and value is not None + and value is not pd.NA + and not (isinstance(value, float) and pd.isna(value)) + for value in (call_id, event_type) + ): + positions = None + break + if ( + type(call_id) is not str # noqa: E721 - exact scalar types + or type(event_type) is not str # noqa: E721 + ): + continue + if event_type == TraceLineType.FUNC_CALL_PRE: + side = 0 + elif event_type in ( + TraceLineType.FUNC_CALL_POST, + TraceLineType.FUNC_CALL_POST_EXCEPTION, + ): + side = 1 + else: + continue + if call_id not in positions: + positions[call_id] = ([], []) + positions[call_id][side].append(position) + + self._func_call_positions = positions + self._func_call_positions_source = self.events + self._func_call_positions_length = len(self.events) + + if self._func_call_positions is None: + return None + records = self._func_call_positions.get(func_call_id) + return records[int(post)] if records is not None else [] + def get_pre_func_call_record(self, func_call_id: str) -> dict: """Get the pre call record of a function given its func_call_id.""" # TODO: check if this is correct - pre_records = self.events[ - (self.events["func_call_id"] == func_call_id) - & (self.events["type"] == TraceLineType.FUNC_CALL_PRE) - ] + positions = self._get_func_call_record_positions(func_call_id, post=False) + if positions is None: + pre_records = self.events[ + (self.events["func_call_id"] == func_call_id) + & (self.events["type"] == TraceLineType.FUNC_CALL_PRE) + ] + else: + pre_records = self.events.iloc[positions] assert ( len(pre_records) == 1 @@ -1028,17 +1179,21 @@ def get_post_func_call_record(self, func_call_id: str) -> dict | None: Returns None if the post call event is not found and the pre-call event is the outermost function call. (see the doc of `_rm_incomplete_trailing_func_calls`) """ # TODO: check if this is correct - post_records = self.events[ - (self.events["func_call_id"] == func_call_id) - & ( - self.events["type"].isin( - [ - TraceLineType.FUNC_CALL_POST, - TraceLineType.FUNC_CALL_POST_EXCEPTION, - ] + positions = self._get_func_call_record_positions(func_call_id, post=True) + if positions is None: + post_records = self.events[ + (self.events["func_call_id"] == func_call_id) + & ( + self.events["type"].isin( + [ + TraceLineType.FUNC_CALL_POST, + TraceLineType.FUNC_CALL_POST_EXCEPTION, + ] + ) ) - ) - ] + ] + else: + post_records = self.events.iloc[positions] assert ( len(post_records) <= 1