From b6b0e5b6b8af88c08d1a7ab269c03be8e15d312f Mon Sep 17 00:00:00 2001 From: danielryckman Date: Thu, 4 Jun 2026 14:05:21 -0400 Subject: [PATCH 01/17] fixed module parsing if the module and current directory are named the same --- traincheck/instrumentor/tracer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 8a6ac1e1..86621e1f 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -523,7 +523,7 @@ def get_module_path_from_file_path(file_path: str, root_module: str) -> str | No ): return None # get the path of the module from the file path - path_after_root_module = file_path.split(f"/{root_module}/")[1].split(".py")[0] + path_after_root_module = file_path.rsplit(f"/{root_module}/", 1)[-1].split(".py")[0] module_path = f"{root_module}.{path_after_root_module}".replace("/", ".") return module_path From fec69179f6fbfd906789ffa6278c61a627229179 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Mon, 8 Jun 2026 18:21:25 -0400 Subject: [PATCH 02/17] temp fixes --- traincheck/instrumentor/dumper.py | 2 +- traincheck/instrumentor/proxy_wrapper/hash.py | 2 ++ traincheck/instrumentor/proxy_wrapper/subclass.py | 3 ++- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/traincheck/instrumentor/dumper.py b/traincheck/instrumentor/dumper.py index 27c458f8..b808e986 100644 --- a/traincheck/instrumentor/dumper.py +++ b/traincheck/instrumentor/dumper.py @@ -296,7 +296,7 @@ def dump_tensor(value): param_list = tensor_stats(value) elif tensor_dump_format["dump_tensor_hash"]: if value.is_cuda: - param_list = tensor_hash(value, with_parallel=True, with_cuda=True) + param_list = tensor_hash(value, with_parallel=True, with_cuda=False) else: # TODO: support quick hashing methods for MPS tensors param_list = tensor_hash(value, with_parallel=True, with_cuda=False) diff --git a/traincheck/instrumentor/proxy_wrapper/hash.py b/traincheck/instrumentor/proxy_wrapper/hash.py index 1c1b3088..d0dae64f 100644 --- a/traincheck/instrumentor/proxy_wrapper/hash.py +++ b/traincheck/instrumentor/proxy_wrapper/hash.py @@ -95,6 +95,8 @@ def _reduce_last_axis(x: Tensor) -> Tensor: def tensor_hash(x: Tensor, with_parallel: bool = True, with_cuda: bool = True) -> int: if hasattr(x, "_traincheck_tensor_hash"): return x._traincheck_tensor_hash + if hasattr(x, "to_local"): + x = x.to_local() if with_parallel: if x.dtype in [ torch.float32, diff --git a/traincheck/instrumentor/proxy_wrapper/subclass.py b/traincheck/instrumentor/proxy_wrapper/subclass.py index ac7ae393..1e059fb7 100644 --- a/traincheck/instrumentor/proxy_wrapper/subclass.py +++ b/traincheck/instrumentor/proxy_wrapper/subclass.py @@ -279,7 +279,8 @@ def subclass_setattr_hook(self, name, value): @functools.wraps(orig_setattr) def wrapped_setattr(self, name, value): - hook = getattr(self, SUBCLASS_HOOK_KEY, None) + #hook = getattr(self, SUBCLASS_HOOK_KEY, None) + hook = self.__dict__.get(SUBCLASS_HOOK_KEY) if hook is not None: # If hook returns True, skip the original setattr; otherwise continue. hook(self, name, value) From 173297ef5bc955d6e50873183df12721ff2fbe1b Mon Sep 17 00:00:00 2001 From: danielryckman Date: Tue, 9 Jun 2026 01:51:47 -0400 Subject: [PATCH 03/17] fixed epoch looping and kernel on cpu errors --- traincheck/config/config.py | 1 + traincheck/instrumentor/tracer.py | 8 ++++++++ traincheck/trace/trace_pandas.py | 4 +++- 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/traincheck/config/config.py b/traincheck/config/config.py index ddbc4a2a..4aee9942 100644 --- a/traincheck/config/config.py +++ b/traincheck/config/config.py @@ -21,6 +21,7 @@ "torch.fx", "torch._dynamo", "torch._sources", # FIXME: cannot handle this module, instrumenting it will lead to exceptions: TypeError: module, class, method, function, traceback, frame, or code object was expected, got builtin_function_or_method + "torchtitan.trainer.Trainer.train_step", # "torch.autocast", # "torch.amp", # "torch.matmul", diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 86621e1f..154e51f8 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -634,6 +634,14 @@ def instrument(self) -> int: global IS_INSTRUMENTING IS_INSTRUMENTING = True + + try: + import torch.utils._device as _tc_torch_device + + _tc_torch_device._device_constructors() + except Exception: + pass + visited_file_paths: set[str] = set() first_pass_instrumented_count = 0 diff --git a/traincheck/trace/trace_pandas.py b/traincheck/trace/trace_pandas.py index 6263bea6..8bb89ca7 100644 --- a/traincheck/trace/trace_pandas.py +++ b/traincheck/trace/trace_pandas.py @@ -558,7 +558,9 @@ 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()] if filtered_values.empty: self.column_dtypes_cached[column_name] = MD_NONE From f91cc5a8f10f2f073cd4895fbf8f1e590bad17db Mon Sep 17 00:00:00 2001 From: danielryckman Date: Mon, 15 Jun 2026 14:33:01 -0400 Subject: [PATCH 04/17] new module --- traincheck/instrumentor/tracer.py | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 154e51f8..de1e55b1 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -2,6 +2,7 @@ import importlib import inspect import os +import pathlib import threading import time import traceback @@ -522,11 +523,31 @@ def get_module_path_from_file_path(file_path: str, root_module: str) -> str | No or f"/{root_module}/" not in file_path ): return None - # get the path of the module from the file path - path_after_root_module = file_path.rsplit(f"/{root_module}/", 1)[-1].split(".py")[0] - module_path = f"{root_module}.{path_after_root_module}".replace("/", ".") - return module_path + # get the location of the root module + spec = importlib.find_spec(root_module) + if spec is None or spec.origin is None: + # raise error + raise ImportError(f"Cannot locate root module {root_module!r}") + + # get the path to the root module + root_module_path = pathlib.PurePath(spec.origin).parent + + # parse file_path into a PurePath object + file_path_obj = pathlib.PurePath(file_path).resolve() + + try: + # remove the root module from file_path_obj + relative_path = file_path_obj.relative_to(root_module_path) + except ValueError: + # file_path is not under root module + raise ImportError(f"File path is not under root module {root_module!r}") + + # strip off .py + module_name = str(relative_path.with_suffix('')) + + # replace / with . + return f"{root_module}.{module_name.replace('/', '.')}" class Instrumentor: def __init__( From b5518c687c5edd6c29ef8dd2b004f3b7cd4f6cb1 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Mon, 15 Jun 2026 14:35:54 -0400 Subject: [PATCH 05/17] import error --- traincheck/instrumentor/tracer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index de1e55b1..e736a80a 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -1,5 +1,5 @@ import functools -import importlib +from importlib.util import find_spec import inspect import os import pathlib From 8f486d5e91d948bd2dddb1da94c3ead1b9657555 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Mon, 15 Jun 2026 14:42:06 -0400 Subject: [PATCH 06/17] importlib find_spec --- traincheck/instrumentor/tracer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index e736a80a..743aee7d 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -525,7 +525,7 @@ def get_module_path_from_file_path(file_path: str, root_module: str) -> str | No return None # get the location of the root module - spec = importlib.find_spec(root_module) + spec = find_spec(root_module) if spec is None or spec.origin is None: # raise error raise ImportError(f"Cannot locate root module {root_module!r}") From 4ce010b81ad63257596a87b154439554ef162d73 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Mon, 15 Jun 2026 14:44:00 -0400 Subject: [PATCH 07/17] pathlib issue --- traincheck/instrumentor/tracer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 743aee7d..dfc05e41 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -533,8 +533,8 @@ def get_module_path_from_file_path(file_path: str, root_module: str) -> str | No # get the path to the root module root_module_path = pathlib.PurePath(spec.origin).parent - # parse file_path into a PurePath object - file_path_obj = pathlib.PurePath(file_path).resolve() + # parse file_path into a Path object (requires file I/O) + file_path_obj = pathlib.Path(file_path).resolve() try: # remove the root module from file_path_obj From 554ee3330bbfc827774282bfdafec57b11477840 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Mon, 15 Jun 2026 14:45:07 -0400 Subject: [PATCH 08/17] pathlib issue --- traincheck/instrumentor/tracer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index dfc05e41..5ba9d984 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -1,4 +1,5 @@ import functools +import importlib from importlib.util import find_spec import inspect import os From a9e0acd28adb1588466a261957841f39c253c4a2 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Mon, 15 Jun 2026 14:59:01 -0400 Subject: [PATCH 09/17] sys.exit issue --- traincheck/instrumentor/tracer.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 5ba9d984..1f9d4f1e 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -4,6 +4,7 @@ import inspect import os import pathlib +import sys import threading import time import traceback @@ -170,6 +171,14 @@ def function_wrapper( # TODO: all meta vars update should be done outside the function_wrapper (e.g. step increment) by applying a separate wrapper return original_function(*args, **kwargs) + if sys.meta_path is None: + # Interpreter is shutting down: the import system is already torn down, + # so any serialization path that imports a module would raise + # "ImportError: sys.meta_path is None". Instrumented functions can still + # be invoked from finalizers (e.g. CheckpointManager.__del__), so skip + # tracing here and just delegate to the original function. + return original_function(*args, **kwargs) + if COLLECT_OVERHEAD_METRICS: ENTER_PERF_TIME = time.perf_counter() From 56180d9febe526c1a9d9946dd4449d571507cac4 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Tue, 30 Jun 2026 10:34:40 -0400 Subject: [PATCH 10/17] initial parallelized implementation --- traincheck/infer_engine.py | 311 +++++++++++++++--- .../instrumentor/proxy_wrapper/subclass.py | 4 +- 2 files changed, 259 insertions(+), 56 deletions(-) diff --git a/traincheck/infer_engine.py b/traincheck/infer_engine.py index 917cb530..73b095be 100644 --- a/traincheck/infer_engine.py +++ b/traincheck/infer_engine.py @@ -5,6 +5,7 @@ import os import random import time +from concurrent.futures import ProcessPoolExecutor, as_completed from tqdm import tqdm @@ -29,9 +30,71 @@ random.seed(0) +# === parallel inference workers === +# Used by the parallel paths in both generate_hypothesis and infer_precondition. +# Each worker process keeps its own copy of the traces in this module global so that +# tasks reuse them instead of re-pickling the (potentially large) traces per task. +_WORKER_TRACES: list | None = None + + +def _config_snapshot() -> dict: + """CLI-overridable config values that must be propagated to workers. + + These are set in main() at runtime; workers started with the 'spawn' method + begin from a fresh interpreter and would otherwise see only the module + defaults (harmless under 'fork'). Add any future CLI-overridable config here. + """ + return { + "ENABLE_PRECOND_SAMPLING": config.ENABLE_PRECOND_SAMPLING, + "PRECOND_SAMPLING_THRESHOLD": config.PRECOND_SAMPLING_THRESHOLD, + } + + +def _worker_init(traces: list, config_snapshot: dict) -> None: + """Initializer run once per worker process. + + Stores the traces in a module global (one copy per worker, not per task) and + restores CLI-overridable config values (see _config_snapshot). Also suppresses + the per-example inner progress bars, otherwise every worker would fight over + the terminal. + """ + global _WORKER_TRACES + _WORKER_TRACES = traces + for name, value in config_snapshot.items(): + setattr(config, name, value) + _tc_utils._suppress_inner_progress = True + + +def _genhypo_worker_task(task: tuple): + """Generate hypotheses for one (trace_idx, relation) pair in a worker process. + + Returns (trace_idx, relation, inferred_hypos); the parent merges them in a + deterministic order so the result matches the serial path. + """ + trace_idx, relation = task + assert _WORKER_TRACES is not None, "Worker traces were not initialized" + return trace_idx, relation, relation.generate_hypothesis(_WORKER_TRACES[trace_idx]) + + +def _precond_worker_task(hypothesis: Hypothesis): + """Run precondition inference for a single hypothesis in a worker process. + + Returns just the precondition (picklable); the parent applies it to its own + copy of the hypothesis. + """ + assert _WORKER_TRACES is not None, "Worker traces were not initialized" + return find_precondition(hypothesis, _WORKER_TRACES) + + class InferEngine: - def __init__(self, traces: list, disabled_relations: list[Relation] = []): + def __init__( + self, + traces: list, + disabled_relations: list[Relation] = [], + num_workers: int = 1, + ): self.traces = traces + self.num_workers = max(1, num_workers) self.all_stages = set() for trace in traces: # FIXME: we don't fully support multi-stage traces with inconsistent stage annotations yet, not sure about the impact to invariant correctness. @@ -54,12 +117,72 @@ def generate_hypothesis(self) -> dict[Hypothesis, list[int]]: """ logger.info("============= GENERATING HYPOTHESIS =============") hypotheses_and_trace_idxs: dict[Hypothesis, list[int]] = {} - hypo_lookup = {} - n_traces = len(self.traces) + hypo_lookup: dict[Hypothesis, Hypothesis] = {} active_relations = [ r for r in relation_pool if r not in self.disabled_relations ] + if self.num_workers <= 1: + self._generate_hypothesis_serial( + active_relations, hypotheses_and_trace_idxs, hypo_lookup + ) + else: + self._generate_hypothesis_parallel( + active_relations, hypotheses_and_trace_idxs, hypo_lookup + ) + + total = len(hypotheses_and_trace_idxs) + print(f"\n {total} hypotheses generated across all relations") + logger.info(f"Finished generating hypotheses, found {total} hypotheses") + return hypotheses_and_trace_idxs + + def _merge_hypotheses( + self, + inferred_hypos: list[Hypothesis], + trace_idx: int, + hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], + hypo_lookup: dict[Hypothesis, Hypothesis], + ) -> None: + """Merge one trace's inferred hypotheses into the running collection. + + Mutates shared state (the dicts and the stored hypotheses' example lists), + so this is always run in the parent process, never in a worker. + """ + for hypo in inferred_hypos: + if hypo not in hypotheses_and_trace_idxs: + hypotheses_and_trace_idxs[hypo] = [trace_idx] + hypo_lookup[hypo] = hypo + else: + hypotheses_and_trace_idxs[hypo].append(trace_idx) + original_hypo = hypo_lookup[hypo] + orig_num_pos_exps = len(original_hypo.positive_examples) + orig_num_neg_exps = len(original_hypo.negative_examples) + original_hypo.positive_examples.examples.extend( + hypo.positive_examples.examples + ) + original_hypo.negative_examples.examples.extend( + hypo.negative_examples.examples + ) + + assert len( + hypo_lookup[hypo].positive_examples + ) == orig_num_pos_exps + len( + hypo.positive_examples + ), f"Expected {orig_num_pos_exps} + {len(hypo.positive_examples)} positive examples, got {len(hypo_lookup[hypo].positive_examples)}" + assert len( + hypo_lookup[hypo].negative_examples + ) == orig_num_neg_exps + len( + hypo.negative_examples + ), f"Expected {orig_num_neg_exps} + {len(hypo.negative_examples)} negative examples, got {len(hypo_lookup[hypo].negative_examples)}" + + def _generate_hypothesis_serial( + self, + active_relations: list, + hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], + hypo_lookup: dict[Hypothesis, Hypothesis], + ) -> None: + """Serial hypothesis generation (preserves the original behavior exactly).""" + n_traces = len(self.traces) rel_bar_fmt = "{desc} {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]" for trace_idx, trace in enumerate(self.traces): tqdm.write(f"\n[Trace {trace_idx + 1}/{n_traces}] Generating hypotheses") @@ -84,37 +207,70 @@ def generate_hypothesis(self) -> dict[Hypothesis, list[int]]: f"Found {len(inferred_hypos)} hypotheses for {relation.__name__} " f"on trace {trace_idx + 1}/{n_traces}" ) - for hypo in inferred_hypos: - if hypo not in hypotheses_and_trace_idxs: - hypotheses_and_trace_idxs[hypo] = [trace_idx] - hypo_lookup[hypo] = hypo - else: - hypotheses_and_trace_idxs[hypo].append(trace_idx) - original_hypo = hypo_lookup[hypo] - orig_num_pos_exps = len(original_hypo.positive_examples) - orig_num_neg_exps = len(original_hypo.negative_examples) - original_hypo.positive_examples.examples.extend( - hypo.positive_examples.examples - ) - original_hypo.negative_examples.examples.extend( - hypo.negative_examples.examples - ) - - assert len( - hypo_lookup[hypo].positive_examples - ) == orig_num_pos_exps + len( - hypo.positive_examples - ), f"Expected {orig_num_pos_exps} + {len(hypo.positive_examples)} positive examples, got {len(hypo_lookup[hypo].positive_examples)}" - assert len( - hypo_lookup[hypo].negative_examples - ) == orig_num_neg_exps + len( - hypo.negative_examples - ), f"Expected {orig_num_neg_exps} + {len(hypo.negative_examples)} negative examples, got {len(hypo_lookup[hypo].negative_examples)}" + self._merge_hypotheses( + inferred_hypos, + trace_idx, + hypotheses_and_trace_idxs, + hypo_lookup, + ) + + def _generate_hypothesis_parallel( + self, + active_relations: list, + hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], + hypo_lookup: dict[Hypothesis, Hypothesis], + ) -> None: + """Parallel hypothesis generation. + + Each (trace, relation) pair is an independent task: relation.generate_hypothesis + is CPU-bound and pure w.r.t. its (per-worker) trace, so we fan them out across + processes. Workers only PRODUCE hypothesis lists; the merge into shared state + stays in the parent. Results are merged in trace-then-relation order (identical + to the serial path) regardless of completion order, so the output is unchanged. + """ + n_traces = len(self.traces) + tasks = [ + (trace_idx, relation) + for trace_idx in range(n_traces) + for relation in active_relations + ] + print( + f"\nGenerating hypotheses for {n_traces} trace(s) x {len(active_relations)} " + f"relations using {self.num_workers} worker processes" + ) + + # collect results keyed by (trace_idx, relation name) so we can merge them in a + # deterministic order afterwards (completion order from as_completed is arbitrary). + results: dict[tuple[int, str], list[Hypothesis]] = {} + bar_fmt = "{desc} {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]" + with tqdm(total=len(tasks), bar_format=bar_fmt, unit="task") as pbar: + pbar.set_description("generating") + with ProcessPoolExecutor( + max_workers=self.num_workers, + initializer=_worker_init, + initargs=(self.traces, _config_snapshot()), + ) as executor: + future_to_task = { + executor.submit(_genhypo_worker_task, task): task for task in tasks + } + for future in as_completed(future_to_task): + trace_idx, relation, inferred_hypos = future.result() + results[(trace_idx, relation.__name__)] = inferred_hypos + logger.info( + f"Found {len(inferred_hypos)} hypotheses for {relation.__name__} " + f"on trace {trace_idx + 1}/{n_traces}" + ) + pbar.update(1) - total = len(hypotheses_and_trace_idxs) - print(f"\n {total} hypotheses generated across all relations") - logger.info(f"Finished generating hypotheses, found {total} hypotheses") - return hypotheses_and_trace_idxs + # deterministic merge: trace order, then relation order (matching the serial path) + for trace_idx in range(n_traces): + for relation in active_relations: + self._merge_hypotheses( + results[(trace_idx, relation.__name__)], + trace_idx, + hypotheses_and_trace_idxs, + hypo_lookup, + ) def collect_examples(self, hypotheses: dict[Hypothesis, list[int]]): logger.info("============= COLLECTING EXAMPLES =============") @@ -162,29 +318,62 @@ def infer_precondition(self, hypotheses: dict[Hypothesis, list[int]]): invariants = [] failed_hypos = [] bar_fmt = "{desc} {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]" - _tc_utils._suppress_inner_progress = True - try: + + def _record_result(hypothesis: Hypothesis, precondition): + """Apply a worker/serial precondition result to the parent's hypothesis.""" + if precondition is None: + failed_hypos.append( + FailedHypothesis(hypothesis, "Precondition not found") + ) + else: + hypothesis.invariant.precondition = precondition + invariants.append(hypothesis.get_invariant(self.all_stages)) + + if self.num_workers <= 1: + # Serial path (preserves the original behavior exactly). + _tc_utils._suppress_inner_progress = True + try: + with tqdm(total=total, bar_format=bar_fmt, unit="hypo") as pbar: + pbar.set_description("0 done · 0 failed") + for hypo_idx, hypothesis in enumerate(all_hypotheses): + logger.info( + f"Inferring precondition for hypothesis {hypo_idx + 1}/{total}: " + f"{hypothesis.invariant.text_description}" + ) + precondition = find_precondition(hypothesis, self.traces) + _record_result(hypothesis, precondition) + pbar.set_description( + f"{len(invariants)} done · {len(failed_hypos)} failed" + ) + pbar.update(1) + finally: + _tc_utils._suppress_inner_progress = False + else: + # Parallel path: each hypothesis's precondition inference is independent and + # CPU-bound, so we fan it out across processes (threads would serialize on the + # GIL). Workers load the traces once via the initializer (one copy per worker, + # not per task); the parent applies each returned precondition to its own + # hypothesis copy. + print(f" Using {self.num_workers} worker processes") with tqdm(total=total, bar_format=bar_fmt, unit="hypo") as pbar: pbar.set_description("0 done · 0 failed") - for hypo_idx, hypothesis in enumerate(all_hypotheses): - logger.info( - f"Inferring precondition for hypothesis {hypo_idx + 1}/{total}: " - f"{hypothesis.invariant.text_description}" - ) - precondition = find_precondition(hypothesis, self.traces) - if precondition is None: - failed_hypos.append( - FailedHypothesis(hypothesis, "Precondition not found") + with ProcessPoolExecutor( + max_workers=self.num_workers, + initializer=_worker_init, + initargs=(self.traces, _config_snapshot()), + ) as executor: + future_to_hypo = { + executor.submit(_precond_worker_task, hypothesis): hypothesis + for hypothesis in all_hypotheses + } + for future in as_completed(future_to_hypo): + hypothesis = future_to_hypo[future] + precondition = future.result() + _record_result(hypothesis, precondition) + pbar.set_description( + f"{len(invariants)} done · {len(failed_hypos)} failed" ) - else: - hypothesis.invariant.precondition = precondition - invariants.append(hypothesis.get_invariant(self.all_stages)) - pbar.set_description( - f"{len(invariants)} done · {len(failed_hypos)} failed" - ) - pbar.update(1) - finally: - _tc_utils._suppress_inner_progress = False + pbar.update(1) print(f" {len(invariants)} invariants · {len(failed_hypos)} failed") return invariants, failed_hypos @@ -262,6 +451,16 @@ def main(): default="pandas", help="Specify the backend to use for Trace", ) + parser.add_argument( + "-j", + "--num-workers", + type=int, + default=1, + help="Number of worker processes for hypothesis generation and precondition " + "inference. 1 (default) runs serially. Use a value > 1 to parallelize across " + "cores (0 = use all available cores). Note: each worker keeps its own copy of " + "the traces, so memory scales with the number of workers.", + ) args = parser.parse_args() # check if either traces or trace folders are provided @@ -334,8 +533,12 @@ def main(): time_end = time.time() logger.info(f"Traces read successfully in {time_end - time_start} seconds.") + num_workers = args.num_workers + if num_workers == 0: + num_workers = os.cpu_count() or 1 + time_start = time.time() - engine = InferEngine(traces, disabled_relations) + engine = InferEngine(traces, disabled_relations, num_workers=num_workers) invs, failed_hypos = engine.infer_multi_trace() # sort the invariants by the text description diff --git a/traincheck/instrumentor/proxy_wrapper/subclass.py b/traincheck/instrumentor/proxy_wrapper/subclass.py index 1e059fb7..4c79207d 100644 --- a/traincheck/instrumentor/proxy_wrapper/subclass.py +++ b/traincheck/instrumentor/proxy_wrapper/subclass.py @@ -279,8 +279,8 @@ def subclass_setattr_hook(self, name, value): @functools.wraps(orig_setattr) def wrapped_setattr(self, name, value): - #hook = getattr(self, SUBCLASS_HOOK_KEY, None) - hook = self.__dict__.get(SUBCLASS_HOOK_KEY) + hook = getattr(self, SUBCLASS_HOOK_KEY, None) + #hook = self.__dict__.get(SUBCLASS_HOOK_KEY) if hook is not None: # If hook returns True, skip the original setattr; otherwise continue. hook(self, name, value) From c6ebfeccb0cbee7a5ddd32e73a6d27fd35651a24 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Tue, 7 Jul 2026 08:46:09 -0400 Subject: [PATCH 11/17] implemented checker parallelization --- traincheck/checker.py | 108 ++++++++++++++++++++++++++++++++++++- traincheck/infer_engine.py | 51 ++++-------------- traincheck/parallel.py | 54 +++++++++++++++++++ 3 files changed, 172 insertions(+), 41 deletions(-) create mode 100644 traincheck/parallel.py diff --git a/traincheck/checker.py b/traincheck/checker.py index 249945c7..2ed4222b 100644 --- a/traincheck/checker.py +++ b/traincheck/checker.py @@ -3,11 +3,13 @@ import json import logging import os +from concurrent.futures import ProcessPoolExecutor, as_completed from tqdm import tqdm import traincheck.utils as _tc_utils from traincheck.invariant import CheckerResult, Invariant, read_inv_file +from traincheck.parallel import config_snapshot, get_worker_traces, worker_init from traincheck.reporting import ( ReportEmitter, build_offline_report_data, @@ -36,9 +38,32 @@ def parse_checker_results(file_name: str): return all_results +def _check_worker_task(task: tuple): + """Check a single invariant against the worker's trace in a worker process. + + Returns (inv_idx, CheckerResult); the parent re-attaches its own invariant + object and computes the detection-time percentage, mirroring the serial path. + """ + inv_idx, inv, check_relation_first = task + trace = get_worker_traces()[0] + return inv_idx, inv.check(trace, check_relation_first) + + def check_engine( + trace: Trace, + invariants: list[Invariant], + check_relation_first: bool, + num_workers: int = 1, +) -> list[CheckerResult]: + if num_workers <= 1: + return _check_engine_serial(trace, invariants, check_relation_first) + return _check_engine_parallel(trace, invariants, check_relation_first, num_workers) + + +def _check_engine_serial( trace: Trace, invariants: list[Invariant], check_relation_first: bool ) -> list[CheckerResult]: + """Serial invariant checking (preserves the original behavior exactly).""" logger = logging.getLogger(__name__) results = [] total = len(invariants) @@ -77,6 +102,71 @@ def check_engine( return results +def _check_engine_parallel( + trace: Trace, + invariants: list[Invariant], + check_relation_first: bool, + num_workers: int, +) -> list[CheckerResult]: + """Parallel invariant checking. + + Each invariant check is an independent, CPU-bound scan over a read-only + trace, so we fan one task per invariant out across processes (threads would + serialize on the GIL). Workers load the trace once via the pool initializer; + results are placed by invariant index so the returned list is in the same + order as the serial path regardless of completion order. + """ + logger = logging.getLogger(__name__) + total = len(invariants) + print(f" Using {num_workers} worker processes") + results: list[CheckerResult | None] = [None] * total + n_violated = 0 + n_done = 0 + bar_fmt = ( + "{desc} {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]" + ) + with tqdm( + total=total, + bar_format=bar_fmt, + unit="inv", + desc=f"0 checked · {total} left · 0 violated", + ) as pbar: + with ProcessPoolExecutor( + max_workers=num_workers, + initializer=worker_init, + initargs=([trace], config_snapshot()), + ) as executor: + future_to_idx = { + executor.submit( + _check_worker_task, (inv_idx, inv, check_relation_first) + ): inv_idx + for inv_idx, inv in enumerate(invariants) + } + for future in as_completed(future_to_idx): + inv_idx, res = future.result() + # the result carries the worker's pickled copy of the invariant; + # point it back at the parent's object + res.invariant = invariants[inv_idx] + res.calc_and_set_time_precentage( + trace.get_start_time(), trace.get_end_time() + ) + logger.info("Invariant %s on trace %s: %s", res.invariant, trace, res) + results[inv_idx] = res + n_done += 1 + if not res.check_passed: + n_violated += 1 + pbar.set_description( + f"{n_done} checked · {total - n_done} left · {n_violated} violated" + ) + pbar.update(1) + + final_results: list[CheckerResult] = [] + for inv_idx, res_or_none in enumerate(results): + assert res_or_none is not None, f"Missing result for invariant {inv_idx}" + final_results.append(res_or_none) + return final_results + + def main(): parser = argparse.ArgumentParser( description="(Offline) Invariant Checker for ML Pipelines in Python" @@ -130,6 +220,16 @@ def main(): type=str, help="Output folder to store the results, defaulted to traincheck_checker_results_{timestamp}/", ) + parser.add_argument( + "-j", + "--num-workers", + type=int, + default=1, + help="Number of worker processes for invariant checking. 1 (default) runs " + "serially. Use a value > 1 to parallelize across cores (0 = use all " + "available cores). Note: each worker keeps its own copy of the trace, so " + "memory scales with the number of workers.", + ) parser.add_argument( "--no-html-report", action="store_true", @@ -274,10 +374,16 @@ def main(): ) logger.addHandler(stream_handler) + num_workers = args.num_workers + if num_workers == 0: + num_workers = os.cpu_count() or 1 + results_by_trace: list[tuple[str, list[CheckerResult]]] = [] for trace, trace_parent_folder in zip(traces, trace_parent_folders): - results_per_trace = check_engine(trace, invs, args.check_relation_first) + results_per_trace = check_engine( + trace, invs, args.check_relation_first, num_workers=num_workers + ) results_per_trace_failed = [ res for res in results_per_trace if not res.check_passed ] diff --git a/traincheck/infer_engine.py b/traincheck/infer_engine.py index 73b095be..09ff0411 100644 --- a/traincheck/infer_engine.py +++ b/traincheck/infer_engine.py @@ -19,6 +19,7 @@ find_precondition, relation_pool, ) +from traincheck.parallel import config_snapshot, get_worker_traces, worker_init from traincheck.trace import MDNONEJSONEncoder, select_trace_implementation from traincheck.utils import register_custom_excepthook @@ -30,39 +31,10 @@ random.seed(0) -# === parallel inference workers === +# === parallel inference worker tasks === # Used by the parallel paths in both generate_hypothesis and infer_precondition. -# Each worker process keeps its own copy of the traces in this module global so that -# tasks reuse them instead of re-pickling the (potentially large) traces per task. -_WORKER_TRACES: list | None = None - - -def _config_snapshot() -> dict: - """CLI-overridable config values that must be propagated to workers. - - These are set in main() at runtime; workers started with the 'spawn' method - begin from a fresh interpreter and would otherwise see only the module - defaults (harmless under 'fork'). Add any future CLI-overridable config here. - """ - return { - "ENABLE_PRECOND_SAMPLING": config.ENABLE_PRECOND_SAMPLING, - "PRECOND_SAMPLING_THRESHOLD": config.PRECOND_SAMPLING_THRESHOLD, - } - - -def _worker_init(traces: list, config_snapshot: dict) -> None: - """Initializer run once per worker process. - - Stores the traces in a module global (one copy per worker, not per task) and - restores CLI-overridable config values (see _config_snapshot). Also suppresses - the per-example inner progress bars, otherwise every worker would fight over - the terminal. - """ - global _WORKER_TRACES - _WORKER_TRACES = traces - for name, value in config_snapshot.items(): - setattr(config, name, value) - _tc_utils._suppress_inner_progress = True +# Workers hold the traces in traincheck.parallel's module global (loaded once per +# worker via worker_init, not re-pickled per task). def _genhypo_worker_task(task: tuple): @@ -72,8 +44,8 @@ def _genhypo_worker_task(task: tuple): deterministic order so the result matches the serial path. """ trace_idx, relation = task - assert _WORKER_TRACES is not None, "Worker traces were not initialized" - return trace_idx, relation, relation.generate_hypothesis(_WORKER_TRACES[trace_idx]) + traces = get_worker_traces() + return trace_idx, relation, relation.generate_hypothesis(traces[trace_idx]) def _precond_worker_task(hypothesis: Hypothesis): @@ -82,8 +54,7 @@ def _precond_worker_task(hypothesis: Hypothesis): Returns just the precondition (picklable); the parent applies it to its own copy of the hypothesis. """ - assert _WORKER_TRACES is not None, "Worker traces were not initialized" - return find_precondition(hypothesis, _WORKER_TRACES) + return find_precondition(hypothesis, get_worker_traces()) class InferEngine: @@ -247,8 +218,8 @@ def _generate_hypothesis_parallel( pbar.set_description("generating") with ProcessPoolExecutor( max_workers=self.num_workers, - initializer=_worker_init, - initargs=(self.traces, _config_snapshot()), + initializer=worker_init, + initargs=(self.traces, config_snapshot()), ) as executor: future_to_task = { executor.submit(_genhypo_worker_task, task): task for task in tasks @@ -359,8 +330,8 @@ def _record_result(hypothesis: Hypothesis, precondition): pbar.set_description("0 done · 0 failed") with ProcessPoolExecutor( max_workers=self.num_workers, - initializer=_worker_init, - initargs=(self.traces, _config_snapshot()), + initializer=worker_init, + initargs=(self.traces, config_snapshot()), ) as executor: future_to_hypo = { executor.submit(_precond_worker_task, hypothesis): hypothesis diff --git a/traincheck/parallel.py b/traincheck/parallel.py new file mode 100644 index 00000000..4d8bcb45 --- /dev/null +++ b/traincheck/parallel.py @@ -0,0 +1,54 @@ +"""Shared helpers for process-based parallelism in the infer engine and checker. + +Both traincheck-infer and traincheck-check fan CPU-bound, per-invariant work out +across a ProcessPoolExecutor (threads would serialize on the GIL). The pattern is: + + - The parent passes the traces to each worker exactly once, via the pool's + initializer (``worker_init``). Workers keep them in a module global so tasks + reuse them instead of re-pickling the (potentially large) traces per task. + - Workers only *compute* and return picklable results; all mutation of shared + state stays in the parent process. + +Task functions live next to their call sites (infer_engine.py / checker.py) and +access the traces through ``get_worker_traces()``. +""" + +import traincheck.config.config as config +import traincheck.utils as _tc_utils + +_WORKER_TRACES: list | None = None + + +def config_snapshot() -> dict: + """CLI-overridable config values that must be propagated to workers. + + These are set in the entry points' main() at runtime; workers started with + the 'spawn' method begin from a fresh interpreter and would otherwise see + only the module defaults (harmless under 'fork'). Add any future + CLI-overridable config here. + """ + return { + "ENABLE_PRECOND_SAMPLING": config.ENABLE_PRECOND_SAMPLING, + "PRECOND_SAMPLING_THRESHOLD": config.PRECOND_SAMPLING_THRESHOLD, + } + + +def worker_init(traces: list, snapshot: dict) -> None: + """Initializer run once per worker process. + + Stores the traces in a module global (one copy per worker, not per task) and + restores CLI-overridable config values (see config_snapshot). Also suppresses + the per-invariant inner progress bars, otherwise every worker would fight + over the terminal. + """ + global _WORKER_TRACES + _WORKER_TRACES = traces + for name, value in snapshot.items(): + setattr(config, name, value) + _tc_utils._suppress_inner_progress = True + + +def get_worker_traces() -> list: + """Return the traces stored by worker_init; only valid inside a worker.""" + assert _WORKER_TRACES is not None, "Worker traces were not initialized" + return _WORKER_TRACES From 0bb988fea690ed3b9010460cd381dc4368d206d5 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Thu, 16 Jul 2026 12:26:32 -0400 Subject: [PATCH 12/17] fixed pickle code objects APIContainsRelation signature issue --- traincheck/invariant/base_cls.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/traincheck/invariant/base_cls.py b/traincheck/invariant/base_cls.py index 944f57e6..3fbcdf62 100644 --- a/traincheck/invariant/base_cls.py +++ b/traincheck/invariant/base_cls.py @@ -181,6 +181,28 @@ def __init__( and param.default != inspect.Parameter.empty ): self.arguments[param_name] = var_to_serializable(param.default) + + def __getstate__(self) -> dict: + # ``signature`` is a cached ``inspect.Signature`` derived from + # ``func_name``. It is not picklable -- its ``_parameters`` is a + # ``mappingproxy`` and its return annotation may be a + # ``types.UnionType`` or a ``typing.ForwardRef`` that wraps a code + # object -- and it is redundant state. Drop it here and re-derive it in + # ``__setstate__`` so that ``Arguments`` (and the ``APIParam`` / + # ``Hypothesis`` objects that embed it, e.g. from APIContainRelation) + # can cross the ProcessPoolExecutor boundary used by the parallel + # infer/check paths. + state = self.__dict__.copy() + state["signature"] = None + return state + + def __setstate__(self, state: dict) -> None: + self.__dict__.update(state) + # Re-derive the dropped signature so an unpickled Arguments is identical + # to a freshly constructed one. ``load_function_signature`` is cached and + # returns None when the function cannot be resolved, matching __init__. + if self.__dict__.get("signature") is None: + self.signature = load_function_signature(self.func_name) def to_dict(self) -> dict: return { From 5219f19f9a4b72962dd11b24076e06dc5138f5b2 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Thu, 30 Jul 2026 13:22:30 -0400 Subject: [PATCH 13/17] added verbose flag to online checker --- traincheck/checker_online.py | 82 +++++++++++++++++++++++++++++++++++- 1 file changed, 80 insertions(+), 2 deletions(-) diff --git a/traincheck/checker_online.py b/traincheck/checker_online.py index d28c2a8e..b7205b4e 100644 --- a/traincheck/checker_online.py +++ b/traincheck/checker_online.py @@ -3,6 +3,7 @@ import json import logging import os +import queue import signal import sys import time @@ -35,6 +36,7 @@ REPORTER: ReportEmitter | None = None REPORT_GENERATED_AT = "" REPORT_OUTPUT_DIR = "" +VERBOSE = False ORIGINAL_SIGINT_HANDLER = signal.getsignal(signal.SIGINT) ORIGINAL_SIGTERM_HANDLER = signal.getsignal(signal.SIGTERM) @@ -217,8 +219,15 @@ def _emit_report(force: bool = False): def check( - invariants, traces, trace_folders, output_dir: str, check_relation_first: bool + invariants, + traces, + trace_folders, + output_dir: str, + check_relation_first: bool, + verbose: bool = False, ): + global VERBOSE + VERBOSE = verbose global OBSERVER global NUM_VIOLATIONS global FAILED_INV @@ -254,12 +263,53 @@ def check( _emit_report(force=True) + # --- verbose progress reporting --------------------------------------- + # The online checker is designed for a *live* trace and never terminates on + # its own (the queue blocks forever once drained). Verbose mode prints + # periodic progress and, crucially, tells the user when the trace appears + # drained so they know it is safe to Ctrl-C to finalize the report. + VERBOSE_PROGRESS_EVERY_S = 5.0 + VERBOSE_IDLE_TIMEOUT_S = 5.0 + records_processed = 0 + last_progress = time.monotonic() + idle_notified = False + + def _vprint(msg: str) -> None: + print(f"[online-checker] {msg}", flush=True) + + if verbose: + n_folders = len(trace_folders) if trace_folders else 0 + n_files = len(traces) if traces else 0 + _vprint( + f"watching {n_folders} folder(s) / {n_files} file(s); " + f"loaded {TOTAL_INVARIANTS} invariants. " + f"Processing... (Ctrl-C to finalize the report)" + ) + while True: - trace_record = checker_data.check_queue.get() + try: + trace_record = checker_data.check_queue.get( + timeout=VERBOSE_IDLE_TIMEOUT_S if verbose else None + ) + except queue.Empty: + # verbose-only: queue drained with no new data for a while. + if verbose and not idle_notified: + _vprint( + f"idle — no new trace data for {VERBOSE_IDLE_TIMEOUT_S:.0f}s. " + f"Processed {records_processed} records, " + f"step={CURRENT_STEP}, stage={CURRENT_STAGE}, " + f"{NUM_VIOLATIONS} violation(s), " + f"{len(TRIGGERED_INV)}/{TOTAL_INVARIANTS} invariants triggered. " + f"Trace appears drained; press Ctrl-C to finish." + ) + idle_notified = True + continue if checker_data.check_queue.empty(): logger.debug("Check queue is empty") if trace_record is None: continue + # new data arrived after an idle period -> re-arm the idle notice + idle_notified = False with checker_data.cond: while True: @@ -367,6 +417,16 @@ def check( _emit_report() + records_processed += 1 + if verbose and (time.monotonic() - last_progress) >= VERBOSE_PROGRESS_EVERY_S: + _vprint( + f"processed {records_processed} records | " + f"step={CURRENT_STEP} stage={CURRENT_STAGE} | " + f"{NUM_VIOLATIONS} violation(s) | " + f"{len(TRIGGERED_INV)}/{TOTAL_INVARIANTS} invariants triggered" + ) + last_progress = time.monotonic() + def stop_checker(): global OBSERVER @@ -391,6 +451,15 @@ def stop_checker(): if REPORTER is not None: REPORTER.close() + if VERBOSE: + print( + f"[online-checker] stopped. " + f"{NUM_VIOLATIONS} violation(s) across " + f"{len(FAILED_INV)} invariant(s). " + f"Report written to {REPORT_OUTPUT_DIR}", + flush=True, + ) + def main(): parser = argparse.ArgumentParser( @@ -422,6 +491,14 @@ def main(): action="store_true", help="Enable debug logging", ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="Print live progress to stdout (records processed, current " + "step/stage, violations) and notify when the trace appears drained " + "so you know when it is safe to Ctrl-C to finalize the report.", + ) parser.add_argument( "--check-relation-first", action="store_true", @@ -570,6 +647,7 @@ def main(): args.trace_folders, args.output_dir, args.check_relation_first, + args.verbose, ) From 9f8212f3c48788e6888fe7a7671f83d879caab36 Mon Sep 17 00:00:00 2001 From: danielryckman Date: Tue, 4 Aug 2026 03:55:47 -0400 Subject: [PATCH 14/17] online checker optimization, added automatic termination once trace drains --- traincheck/checker_online.py | 89 +++++++++++++++---- .../instrumentor/proxy_wrapper/subclass.py | 4 +- .../onlinechecker/streamhandler_filesystem.py | 4 + traincheck/onlinechecker/utils.py | 26 ++++++ traincheck/reporting/checker_report.py | 5 ++ 5 files changed, 108 insertions(+), 20 deletions(-) diff --git a/traincheck/checker_online.py b/traincheck/checker_online.py index b7205b4e..280acf11 100644 --- a/traincheck/checker_online.py +++ b/traincheck/checker_online.py @@ -199,6 +199,14 @@ def _read_sampling_config( def _emit_report(force: bool = False): if REPORTER is None: return + # Gate the (expensive) report build on the emitter's throttle decision. + # _emit_report is called on every processed record, but build_online_report_data + # re-serializes and re-hashes every invariant + violation; building it only to + # have the emitter throttle the write away dominated total runtime (~94% in + # profiling). report_state is cheap to compute, so decide first, build only if needed. + report_state = (NUM_VIOLATIONS, len(FAILED_INV)) + if not REPORTER.should_emit(report_state, force): + return report_data = build_online_report_data( generated_at=REPORT_GENERATED_AT, output_dir=REPORT_OUTPUT_DIR, @@ -214,7 +222,6 @@ def _emit_report(force: bool = False): sampling_interval=SAMPLING_INTERVAL, warm_up_steps=WARM_UP_STEPS, ) - report_state = (NUM_VIOLATIONS, len(FAILED_INV)) REPORTER.emit(report_data, force=force, report_state=report_state) @@ -225,6 +232,7 @@ def check( output_dir: str, check_relation_first: bool, verbose: bool = False, + exit_on_idle: bool = False, ): global VERBOSE VERBOSE = verbose @@ -263,13 +271,20 @@ def check( _emit_report(force=True) - # --- verbose progress reporting --------------------------------------- - # The online checker is designed for a *live* trace and never terminates on - # its own (the queue blocks forever once drained). Verbose mode prints - # periodic progress and, crucially, tells the user when the trace appears - # drained so they know it is safe to Ctrl-C to finalize the report. + # --- progress reporting / idle handling ------------------------------- + # The online checker is built for a *live* trace and by default never + # terminates on its own (the queue blocks forever once drained). Two opt-in + # behaviors change that: + # -v / verbose -> print periodic progress and a one-time notice when the + # trace looks drained (but keep waiting). + # --exit-on-idle -> when the trace is drained (no new data for + # IDLE_TIMEOUT_S) finalize the report and exit, so a + # finished / replayed trace runs to completion without a + # manual Ctrl-C. Works with or without -v. + # A finite get() timeout (to notice idle) is needed if either is set. VERBOSE_PROGRESS_EVERY_S = 5.0 - VERBOSE_IDLE_TIMEOUT_S = 5.0 + IDLE_TIMEOUT_S = 5.0 + use_idle_timeout = verbose or exit_on_idle records_processed = 0 last_progress = time.monotonic() idle_notified = False @@ -280,27 +295,46 @@ def _vprint(msg: str) -> None: if verbose: n_folders = len(trace_folders) if trace_folders else 0 n_files = len(traces) if traces else 0 + drain_behavior = ( + "will finalize and exit automatically once the trace drains" + if exit_on_idle + else "Ctrl-C to finalize the report (or use --exit-on-idle)" + ) _vprint( f"watching {n_folders} folder(s) / {n_files} file(s); " f"loaded {TOTAL_INVARIANTS} invariants. " - f"Processing... (Ctrl-C to finalize the report)" + f"Processing... ({drain_behavior})" ) while True: try: trace_record = checker_data.check_queue.get( - timeout=VERBOSE_IDLE_TIMEOUT_S if verbose else None + timeout=IDLE_TIMEOUT_S if use_idle_timeout else None ) except queue.Empty: - # verbose-only: queue drained with no new data for a while. + # Reachable only when a finite get() timeout is set (i.e. -v and/or + # --exit-on-idle). A timeout means no new trace data for + # IDLE_TIMEOUT_S -> the trace looks drained. + drained = ( + f"idle — no new trace data for {IDLE_TIMEOUT_S:.0f}s. " + f"Processed {records_processed} records, " + f"step={CURRENT_STEP}, stage={CURRENT_STAGE}, " + f"{NUM_VIOLATIONS} violation(s), " + f"{len(TRIGGERED_INV)}/{TOTAL_INVARIANTS} invariants triggered. " + ) + if exit_on_idle: + # Finished / replayed trace: treat the drain as end-of-input, + # finalize the report and exit instead of blocking forever on a + # (nonexistent) live stream. + if verbose: + _vprint(drained + "Trace drained; finalizing report and exiting.") + break + # -v without --exit-on-idle: announce the apparent drain once, then + # keep waiting for a live stream (re-armed when new data arrives). if verbose and not idle_notified: _vprint( - f"idle — no new trace data for {VERBOSE_IDLE_TIMEOUT_S:.0f}s. " - f"Processed {records_processed} records, " - f"step={CURRENT_STEP}, stage={CURRENT_STAGE}, " - f"{NUM_VIOLATIONS} violation(s), " - f"{len(TRIGGERED_INV)}/{TOTAL_INVARIANTS} invariants triggered. " - f"Trace appears drained; press Ctrl-C to finish." + drained + "Trace appears drained; press Ctrl-C to finish " + "(or re-run with --exit-on-idle to exit automatically)." ) idle_notified = True continue @@ -427,6 +461,12 @@ def _vprint(msg: str) -> None: ) last_progress = time.monotonic() + # Only reachable via the --exit-on-idle `break` above: the trace has drained, + # so finalize the report and stop the observer. stop_checker() emits the final + # report (force=True) and closes the reporter. Without --exit-on-idle the loop + # never falls through here; it is terminated by SIGINT/SIGTERM instead. + stop_checker() + def stop_checker(): global OBSERVER @@ -496,8 +536,20 @@ def main(): "--verbose", action="store_true", help="Print live progress to stdout (records processed, current " - "step/stage, violations) and notify when the trace appears drained " - "so you know when it is safe to Ctrl-C to finalize the report.", + "step/stage, violations) and a one-time notice when the trace appears " + "drained. Progress only -- does not by itself terminate the checker; " + "pair it with --exit-on-idle to auto-finalize a finished trace.", + ) + parser.add_argument( + "--exit-on-idle", + action="store_true", + help="Terminate automatically once the trace is drained (no new trace " + "data for 5s): finalize the report and exit instead of blocking forever " + "waiting for a live stream. Use this to run the checker to completion " + "against a finished / replayed trace folder without a manual Ctrl-C. Can " + "be combined with -v. Not recommended for a live run whose data can " + "pause for more than 5s (e.g. checkpointing), which would trigger an " + "early exit.", ) parser.add_argument( "--check-relation-first", @@ -648,6 +700,7 @@ def main(): args.output_dir, args.check_relation_first, args.verbose, + args.exit_on_idle, ) diff --git a/traincheck/instrumentor/proxy_wrapper/subclass.py b/traincheck/instrumentor/proxy_wrapper/subclass.py index 4c79207d..1e059fb7 100644 --- a/traincheck/instrumentor/proxy_wrapper/subclass.py +++ b/traincheck/instrumentor/proxy_wrapper/subclass.py @@ -279,8 +279,8 @@ def subclass_setattr_hook(self, name, value): @functools.wraps(orig_setattr) def wrapped_setattr(self, name, value): - hook = getattr(self, SUBCLASS_HOOK_KEY, None) - #hook = self.__dict__.get(SUBCLASS_HOOK_KEY) + #hook = getattr(self, SUBCLASS_HOOK_KEY, None) + hook = self.__dict__.get(SUBCLASS_HOOK_KEY) if hook is not None: # If hook returns True, skip the original setattr; otherwise continue. hook(self, name, value) diff --git a/traincheck/onlinechecker/streamhandler_filesystem.py b/traincheck/onlinechecker/streamhandler_filesystem.py index 5268e71f..f90f8bd4 100644 --- a/traincheck/onlinechecker/streamhandler_filesystem.py +++ b/traincheck/onlinechecker/streamhandler_filesystem.py @@ -302,6 +302,8 @@ def _set_context_map(self, trace_record): input=binded_args_and_kwargs, ) ) + # invalidate the meta-vars cache (see Checker_data.meta_vars_cache) + self.checker_data.context_version += 1 elif ".__exit__" in function_name and trace_type == TraceLineType.FUNC_CALL_PRE: context_manager_name = function_name.removesuffix(".__exit__") contextmanagerstate = None @@ -315,6 +317,8 @@ def _set_context_map(self, trace_record): contextmanagerstate = state if contextmanagerstate is not None: contextmanagerstate.liveness.end_time = trace_record["time"] + # invalidate the meta-vars cache (see Checker_data.meta_vars_cache) + self.checker_data.context_version += 1 def _set_read_time(self, trace_record): with self.cond: diff --git a/traincheck/onlinechecker/utils.py b/traincheck/onlinechecker/utils.py index 6da09e22..fa6b881b 100644 --- a/traincheck/onlinechecker/utils.py +++ b/traincheck/onlinechecker/utils.py @@ -25,6 +25,18 @@ def __init__(self, needed_data): self.context_map = {} self.init_map = {} + # --- meta-vars memoization --------------------------------------- + # get_meta_vars_online(time, pid, tid) is a pure function of its args + # plus the contents of context_map. context_version is bumped (under + # self.lock) on every context_map mutation; get_meta_vars_online caches + # results keyed by (time, pid, tid) and discards the whole cache when + # context_version changes. Every set_meta_vars_online call site and every + # context_map mutation holds self.lock, so cache access needs no extra + # synchronization. + self.context_version = 0 + self.meta_vars_cache: dict = {} + self.meta_vars_cache_version = -1 + self.read_time_map = {} self.min_read_time = None self.min_read_path = None @@ -237,11 +249,24 @@ def get_meta_vars_online( thread_id, checker_data, ): + # Memoize on (time, pid, tid). Callers all hold checker_data.lock, and the + # cache is invalidated whenever context_map changes (context_version), so a + # hit is guaranteed to be identical to a fresh computation. The cached value + # (a meta_vars dict or None) is only ever read by callers, never mutated. + cache = checker_data.meta_vars_cache + if checker_data.meta_vars_cache_version != checker_data.context_version: + cache.clear() + checker_data.meta_vars_cache_version = checker_data.context_version + cache_key = (time, process_id, thread_id) + if cache_key in cache: + return cache[cache_key] + ptid = PTID(process_id, thread_id) active_context_managers = [] meta_vars = {} if ptid not in checker_data.context_map: + cache[cache_key] = None return None context_managers = checker_data.context_map[ptid] for context_manager_name, context_manager_states in context_managers.items(): @@ -258,6 +283,7 @@ def get_meta_vars_online( "input" ] + cache[cache_key] = meta_vars return meta_vars diff --git a/traincheck/reporting/checker_report.py b/traincheck/reporting/checker_report.py index 1ea3c4a7..c2a95bd5 100644 --- a/traincheck/reporting/checker_report.py +++ b/traincheck/reporting/checker_report.py @@ -1344,6 +1344,11 @@ def __init__( self._wandb_run = None self._mlflow_active = False + def should_emit(self, report_state: tuple[int, int], force: bool = False) -> bool: + """Public throttle check so callers can skip the (expensive) report build + entirely when nothing would be emitted. Does not mutate emitter state.""" + return self._should_emit(report_state, force) + def _should_emit(self, report_state: tuple[int, int], force: bool) -> bool: if force: return True From 9eb4f320d3db34d372145261704b805b88e39bda Mon Sep 17 00:00:00 2001 From: danielryckman Date: Tue, 4 Aug 2026 04:03:54 -0400 Subject: [PATCH 15/17] temporary memory bound for infer engine --- traincheck/config/config.py | 2 ++ traincheck/invariant/base_cls.py | 22 +++++++++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/traincheck/config/config.py b/traincheck/config/config.py index 4aee9942..51b570fa 100644 --- a/traincheck/config/config.py +++ b/traincheck/config/config.py @@ -198,6 +198,8 @@ def should_disable_proxy_dumping() -> bool: MAX_PRECOND_DEPTH = 8 # the maximum depth of the precondition inference ENABLE_PRECOND_SAMPLING = True # whether to enable sampling of positive and negative examples for precondition inference, can be overridden by the command line argument PRECOND_SAMPLING_THRESHOLD = 10000 # the number of samples to take for precondition inference, if the number of samples is larger than this threshold, we will sample this number of samples +MAX_EXAMPLE_LIST_SIZE = 5000 +EXAMPLE_SAMPLING_SEED = 42 NOT_USE_AS_CLAUSE_FIELDS = [ "func_call_id", "process_id", diff --git a/traincheck/invariant/base_cls.py b/traincheck/invariant/base_cls.py index 3fbcdf62..588e7260 100644 --- a/traincheck/invariant/base_cls.py +++ b/traincheck/invariant/base_cls.py @@ -6,6 +6,7 @@ import json import logging import math +import random from enum import Enum from typing import Any, Hashable, Iterable, Optional, Type @@ -1679,6 +1680,15 @@ class ExampleList: def __init__(self, group_names: set[str]): self.group_names = group_names self.examples: list[Example] = [] + # Bound memory: some relations produce an unbounded number of examples + # (see config.MAX_EXAMPLE_LIST_SIZE). Once at capacity we keep a uniform + # random subset via reservoir sampling; num_seen tracks the true number + # of examples added so the retained subset stays uniform. The RNG is + # seeded so results are deterministic and identical across the serial + # and parallel code paths. + self.max_size = config.MAX_EXAMPLE_LIST_SIZE + self.num_seen = 0 + self._rng = random.Random(config.EXAMPLE_SAMPLING_SEED) def add_example(self, example: Example): if len(self.group_names) == 0: @@ -1688,7 +1698,16 @@ def add_example(self, example: Example): assert ( example.get_group_names() == self.group_names ), f"Example groups do not match the expected group names, expected: {self.group_names}, got: {set(example.trace_groups.keys())}" - self.examples.append(example) + #self.examples.append(example) + self.num_seen += 1 + if self.max_size is None or len(self.examples) < self.max_size: + self.examples.append(example) + else: + # reservoir sampling: replace an existing example with probability + # max_size / num_seen, keeping self.examples a uniform random sample + j = self._rng.randrange(self.num_seen) + if j < self.max_size: + self.examples[j] = example def get_group_from_examples(self, group_name: str) -> list[list[dict]]: return [example.get_group(group_name) for example in self.examples] @@ -1713,6 +1732,7 @@ def from_iterable_of_examples( example_list = ExampleList(preset_group_names) example_list.examples = examples + example_list.num_seen = len(examples) return example_list From 64cfcb00d021d4d5c370e02b916588f3c9cba01a Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 13 Oct 2026 14:05:24 +0800 Subject: [PATCH 16/17] Make inference sampling independent of worker scheduling --- tests/test_inference_sampling.py | 360 ++++++++++++++++++ traincheck/infer_engine.py | 70 +++- traincheck/inference_sampling.py | 54 +++ .../invariant/consistency_transient_vars.py | 17 +- traincheck/invariant/contain_relation.py | 14 +- traincheck/invariant/precondition.py | 7 +- 6 files changed, 485 insertions(+), 37 deletions(-) create mode 100644 tests/test_inference_sampling.py create mode 100644 traincheck/inference_sampling.py diff --git a/tests/test_inference_sampling.py b/tests/test_inference_sampling.py new file mode 100644 index 00000000..bfd87977 --- /dev/null +++ b/tests/test_inference_sampling.py @@ -0,0 +1,360 @@ +"""Regression coverage for task-local sampling, including actual relation paths.""" + +import importlib +import json +import multiprocessing +import os +import random +import subprocess +import sys +import threading +import unittest +from concurrent.futures import ProcessPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from unittest.mock import patch + +from traincheck.infer_engine import ( + InferEngine, + _find_precondition_task, + _generate_hypothesis_task, + _genhypo_worker_task, + _hypothesis_sampling_key, + _precond_worker_task, +) +from traincheck.inference_sampling import sampling_rng, sampling_task +from traincheck.invariant import relation_pool +from traincheck.invariant.base_cls import ( + APIParam, + Example, + ExampleList, + Hypothesis, + Invariant, + UnconditionalPrecondition, +) +from traincheck.invariant.contain_relation import ( + PARENT_GROUP_NAME, + APIContainRelation, + cache_events_scanner, + prune_func_call, +) +from traincheck.parallel import config_snapshot, worker_init +from traincheck.trace.trace_pandas import TracePandas +from traincheck.trace.types import FuncCallEvent + +PARENT = "torch.nn.modules.activation.ReLU.forward" +CHILD = "torch.nn.modules.dropout.Dropout.forward" + + +def nested_call_trace(prefix, calls=1005): + """Complete synthetic events, not training: >1000 calls exercises both samplers.""" + rows = [] + for index in range(calls): + events = [ + (PARENT, "parent", "pre"), + (CHILD, "child", "pre"), + (CHILD, "child", "post"), + (PARENT, "parent", "post"), + ] + for offset, (function, name, kind) in enumerate(events): + row = { + "time": index * 10 + offset + 1, + "func_call_id": f"{prefix}-{name}-{index}", + "function": function, + "type": f"function_call ({kind})", + "process_id": 1, + "thread_id": 1, + "is_bound_method": True, + "obj_id": 1, + "meta_vars": {}, + "args": {"0": {"object": {}}, "1": {"float": 1.0}}, + "kwargs": {}, + } + if kind == "post": + row["return_values"] = {"float": 1.0} + rows.append(row) + return TracePandas(rows) + + +def hypothesis_evidence(hypotheses): + """Compare exact retained call IDs/order as well as relation and parameters.""" + results = [] + for hypothesis in hypotheses: + + def examples(items): + return [ + { + name: [row["func_call_id"] for row in rows] + for name, rows in sorted(example.trace_groups.items()) + } + for example in items.examples + ] + + results.append( + { + "key": _hypothesis_sampling_key(hypothesis), + "positive": examples(hypothesis.positive_examples), + "negative": examples(hypothesis.negative_examples), + } + ) + return sorted(results, key=lambda item: item["key"]) + + +def simple_hypothesis(positive_count=3, negative_count=0): + def examples(count, prefix): + # Direct construction deliberately reaches the >5000 precondition branch; + # add_example's separate reservoir would otherwise cap the input first. + return ExampleList.from_iterable_of_examples( + [ + Example( + {PARENT_GROUP_NAME: [{"id": f"{prefix}-{i}", "function": PARENT}]} + ) + for i in range(count) + ], + {PARENT_GROUP_NAME}, + ) + + return Hypothesis( + Invariant( + APIContainRelation, [APIParam(PARENT), APIParam(CHILD)], None, "fixture" + ), + examples(positive_count, "positive"), + examples(negative_count, "negative"), + ) + + +class SamplingScopeTests(unittest.TestCase): + def test_reproducible_keys_and_untouched_global_rng(self): + before = random.getstate() + self.assertIs(sampling_rng(), random) + with sampling_task("generate", 0, "relation"): + first = [sampling_rng().random() for _ in range(20)] + self.assertEqual(random.getstate(), before) + random.random() + after_noise = random.getstate() + with sampling_task("generate", 0, "relation"): + second = [sampling_rng().random() for _ in range(20)] + with sampling_task("generate", 1, "relation"): + other = [sampling_rng().random() for _ in range(20)] + self.assertEqual(first, second) + self.assertNotEqual(first, other) + self.assertEqual(random.getstate(), after_noise) + + def test_nested_scopes_restore_on_frozen_exception(self): + @dataclass(frozen=True) + class FrozenError(Exception): + reason: str + + with sampling_task("outer") as expected_rng: + expected = [expected_rng.random() for _ in range(2)] + error = FrozenError("expected") + before = random.getstate() + with sampling_task("outer") as outer: + first = sampling_rng().random() + try: + with sampling_task("inner"): + sampling_rng().random() + raise error + except FrozenError as caught: + self.assertIs(caught, error) + self.assertIs(sampling_rng(), outer) + second = sampling_rng().random() + self.assertEqual([first, second], expected) + self.assertEqual(random.getstate(), before) + self.assertIs(sampling_rng(), random) + + def test_context_is_thread_local(self): + errors = [] + values = [] + + def worker(): + try: + self.assertIs(sampling_rng(), random) + with sampling_task("worker"): + values.append(sampling_rng().random()) + except BaseException as error: + errors.append(error) + + with sampling_task("parent") as parent: + thread = threading.Thread(target=worker) + thread.start() + thread.join(5) + self.assertFalse(thread.is_alive()) + self.assertIs(sampling_rng(), parent) + self.assertEqual(errors, []) + self.assertEqual(len(values), 1) + + def test_seed_does_not_depend_on_python_hash_seed(self): + script = ( + "import json; from traincheck.inference_sampling import sampling_task,sampling_rng; " + "\nwith sampling_task('generate', 4, {'b': 2, 'a': 1}): " + "print(json.dumps([sampling_rng().random() for _ in range(8)]))" + ) + root = str(Path(__file__).resolve().parents[1]) + outputs = [] + for seed in ("1", "98765"): + env = dict(os.environ, PYTHONHASHSEED=seed, PYTHONPATH=root) + outputs.append( + subprocess.check_output( + [sys.executable, "-c", script], env=env, text=True + ) + ) + self.assertEqual(outputs[0], outputs[1]) + + def test_actual_pruning_is_nontrivial_and_task_deterministic(self): + def retained(identity): + sizes, types, kept = [], [], [] + with sampling_task("pruning", identity): + for index in range(1500): + if not prune_func_call(1500, FuncCallEvent, sizes, types): + kept.append(index) + sizes.append(1) + types.append(FuncCallEvent) + return kept + + first = retained("same") + for _ in range(101): + random.random() + self.assertEqual(first, retained("same")) + self.assertNotEqual(first, retained("different")) + self.assertGreater(len(first), 10) + self.assertLess(len(first), 1500) + + +class InferenceTaskTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.traces = [nested_call_trace("trace0"), nested_call_trace("trace1")] + + def setUp(self): + cache_events_scanner.clear() + + def test_hypothesis_key_ignores_mutable_results_but_tracks_semantics(self): + hypothesis = simple_hypothesis() + key = _hypothesis_sampling_key(hypothesis) + hypothesis.invariant.num_positive_examples = 900 + hypothesis.invariant.num_negative_examples = 8 + hypothesis.invariant.text_description = "different display text" + hypothesis.invariant.precondition = object() + self.assertEqual(_hypothesis_sampling_key(hypothesis), key) + hypothesis.valid_only_for_stages = True + hypothesis.valid_stages = {"training", "evaluation"} + staged = _hypothesis_sampling_key(hypothesis) + self.assertNotEqual(staged, key) + hypothesis.valid_stages = set(["evaluation", "training"]) + self.assertEqual(_hypothesis_sampling_key(hypothesis), staged) + hypothesis.invariant.params.reverse() + self.assertNotEqual(_hypothesis_sampling_key(hypothesis), staged) + + def test_real_relation_retained_examples_repeat_despite_ambient_rng(self): + before = random.getstate() + hypotheses = _generate_hypothesis_task(0, APIContainRelation, self.traces[0]) + self.assertEqual(random.getstate(), before) + first = hypothesis_evidence(hypotheses) + self.assertEqual(len(first), 1) + self.assertGreater(len(first[0]["positive"]), 10) + self.assertLess(len(first[0]["positive"]), 120) + for _ in range(77): + random.random() + second = hypothesis_evidence( + _generate_hypothesis_task(0, APIContainRelation, self.traces[0]) + ) + self.assertEqual(first, second) + + def test_engine_serial_dispatches_each_relation_through_shared_task(self): + disabled = [ + relation for relation in relation_pool if relation is not APIContainRelation + ] + engine = InferEngine(self.traces, disabled_relations=disabled, num_workers=1) + with patch( + "traincheck.infer_engine._generate_hypothesis_task", + wraps=_generate_hypothesis_task, + ) as task: + hypotheses = engine.generate_hypothesis() + self.assertEqual(task.call_count, 2) + self.assertEqual([call.args[0] for call in task.call_args_list], [0, 1]) + self.assertEqual(len(hypotheses), 1) + self.assertEqual(next(iter(hypotheses.values())), [0, 1]) + + def test_actual_relation_matches_process_counts_orders_and_start_methods(self): + expected = { + index: hypothesis_evidence( + _generate_hypothesis_task(index, APIContainRelation, trace) + ) + for index, trace in enumerate(self.traces) + } + methods = [ + method + for method in ("spawn", "fork") + if method in multiprocessing.get_all_start_methods() + ] + self.assertTrue(methods) + for method in methods: + for workers in (1, 2, 4): + with self.subTest(start_method=method, workers=workers): + # Repeated tasks consume distinct worker histories. Reversing + # dispatch changes assignment without changing semantic keys. + tasks = [(index, APIContainRelation) for index in (1, 0, 1, 0)] + with ProcessPoolExecutor( + max_workers=workers, + mp_context=multiprocessing.get_context(method), + initializer=worker_init, + initargs=(self.traces, config_snapshot()), + ) as pool: + futures = [ + pool.submit(_genhypo_worker_task, task) for task in tasks + ] + for future in futures: + index, relation, hypotheses = future.result(timeout=90) + self.assertIs(relation, APIContainRelation) + self.assertEqual( + hypothesis_evidence(hypotheses), expected[index] + ) + + def test_precondition_over_5000_preserves_legacy_samples_without_global_seed(self): + precondition = importlib.import_module("traincheck.invariant.precondition") + hypothesis = simple_hypothesis(5007, 5011) + captured = [] + + def solve(positive, negative, traces, keys_to_skip): + captured.append( + ( + [item[0]["id"] for item in positive], + [item[0]["id"] for item in negative], + ) + ) + return [UnconditionalPrecondition()] + + legacy_rng = random.Random(42) + expected = ( + legacy_rng.sample([f"positive-{i}" for i in range(5007)], 5000), + legacy_rng.sample([f"negative-{i}" for i in range(5011)], 5000), + ) + before = random.getstate() + with patch.object( + precondition, "find_precondition_from_single_group", side_effect=solve + ): + result = _find_precondition_task( + hypothesis, [], _hypothesis_sampling_key(hypothesis) + ) + self.assertIsNotNone(result) + self.assertEqual(captured, [expected]) + self.assertEqual(random.getstate(), before) + + def test_precondition_worker_dispatch_matches_serial_key(self): + hypothesis = simple_hypothesis() + key = _hypothesis_sampling_key(hypothesis) + worker_init([], config_snapshot()) + + def result(*args): + return [sampling_rng().random() for _ in range(6)] + + with patch("traincheck.infer_engine.find_precondition", side_effect=result): + expected = _find_precondition_task(hypothesis, [], key) + random.random() + actual = _precond_worker_task((key, hypothesis)) + self.assertEqual(actual, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/infer_engine.py b/traincheck/infer_engine.py index 09ff0411..f5958efb 100644 --- a/traincheck/infer_engine.py +++ b/traincheck/infer_engine.py @@ -3,7 +3,6 @@ import json import logging import os -import random import time from concurrent.futures import ProcessPoolExecutor, as_completed @@ -11,6 +10,7 @@ import traincheck.config.config as config import traincheck.utils as _tc_utils +from traincheck.inference_sampling import sampling_task from traincheck.invariant import ( FailedHypothesis, Hypothesis, @@ -27,8 +27,35 @@ logger = logging.getLogger(__name__) -# set random seed -random.seed(0) + +def _generate_hypothesis_task(trace_idx, relation, trace): + """Use the same logical task RNG in serial and process-based generation.""" + with sampling_task( + "generate", trace_idx, relation.__module__, relation.__qualname__ + ): + return relation.generate_hypothesis(trace) + + +def _hypothesis_sampling_key(hypothesis): + """Identify a hypothesis before dispatch, excluding mutable results/counts.""" + relation = hypothesis.invariant.relation + identity = { + "relation": [relation.__module__, relation.__qualname__], + "params": [param.to_dict() for param in hypothesis.invariant.params], + "valid_only_for_stages": hypothesis.valid_only_for_stages, + "valid_stages": sorted(hypothesis.valid_stages), + } + # Normalize JSON object keys and MD_NONE with the invariant-output encoder, + # then sort keys for insertion-order independence. + normalized = json.loads(json.dumps(identity, cls=MDNONEJSONEncoder)) + return json.dumps( + normalized, sort_keys=True, separators=(",", ":"), ensure_ascii=False + ) + + +def _find_precondition_task(hypothesis, traces, key): + with sampling_task("precondition", key): + return find_precondition(hypothesis, traces) # === parallel inference worker tasks === @@ -45,16 +72,18 @@ def _genhypo_worker_task(task: tuple): """ trace_idx, relation = task traces = get_worker_traces() - return trace_idx, relation, relation.generate_hypothesis(traces[trace_idx]) + hypotheses = _generate_hypothesis_task(trace_idx, relation, traces[trace_idx]) + return trace_idx, relation, hypotheses -def _precond_worker_task(hypothesis: Hypothesis): +def _precond_worker_task(task): """Run precondition inference for a single hypothesis in a worker process. Returns just the precondition (picklable); the parent applies it to its own copy of the hypothesis. """ - return find_precondition(hypothesis, get_worker_traces()) + key, hypothesis = task + return _find_precondition_task(hypothesis, get_worker_traces(), key) class InferEngine: @@ -152,7 +181,7 @@ def _generate_hypothesis_serial( hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], hypo_lookup: dict[Hypothesis, Hypothesis], ) -> None: - """Serial hypothesis generation (preserves the original behavior exactly).""" + """Serial generation with the same per-task sampling as the worker path.""" n_traces = len(self.traces) rel_bar_fmt = "{desc} {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]" for trace_idx, trace in enumerate(self.traces): @@ -169,7 +198,9 @@ def _generate_hypothesis_serial( f"Generating hypotheses for relation: {relation.__name__}" ) t0 = time.time() - inferred_hypos = relation.generate_hypothesis(trace) + inferred_hypos = _generate_hypothesis_task( + trace_idx, relation, trace + ) elapsed = time.time() - t0 tqdm.write( f" {relation.__name__}: {len(inferred_hypos)} hypotheses ({elapsed:.1f}s)" @@ -197,7 +228,8 @@ def _generate_hypothesis_parallel( is CPU-bound and pure w.r.t. its (per-worker) trace, so we fan them out across processes. Workers only PRODUCE hypothesis lists; the merge into shared state stays in the parent. Results are merged in trace-then-relation order (identical - to the serial path) regardless of completion order, so the output is unchanged. + to the serial path) regardless of completion order. Each task also has + its own stable RNG. """ n_traces = len(self.traces) tasks = [ @@ -255,13 +287,15 @@ def collect_examples(self, hypotheses: dict[Hypothesis, list[int]]): f"\nCollecting examples for {len(cross_trace_hypos)} cross-trace hypotheses" ) for hypo, trace_idxs in cross_trace_hypos: + key = _hypothesis_sampling_key(hypo) for trace_idx, trace in enumerate(self.traces): if trace_idx in trace_idxs: continue logger.info( f"Collecting examples for hypothesis: {hypo} on trace {trace_idx + 1}/{len(self.traces)}" ) - hypo.invariant.relation.collect_examples(trace, hypo) + with sampling_task("collect", trace_idx, key): + hypo.invariant.relation.collect_examples(trace, hypo) def prune_incorrect_hypos(self, hypotheses: dict[Hypothesis, list[int]]): """Prune incorrect hypotheses based on the collected examples""" @@ -283,6 +317,8 @@ def infer_precondition(self, hypotheses: dict[Hypothesis, list[int]]): """TODO: move the precondition inference driving code into Hypothesis.get_invariant()""" logger.info("============= INFERING PRECONDITIONS =============") all_hypotheses = list(hypotheses.keys()) + # Build keys in the parent, before result mutation or worker pickling. + tasks = [(_hypothesis_sampling_key(hypo), hypo) for hypo in all_hypotheses] total = len(all_hypotheses) print(f"\nInferring preconditions for {total} hypotheses") @@ -301,17 +337,19 @@ def _record_result(hypothesis: Hypothesis, precondition): invariants.append(hypothesis.get_invariant(self.all_stages)) if self.num_workers <= 1: - # Serial path (preserves the original behavior exactly). + # Serial path uses the same task scopes as the workers. _tc_utils._suppress_inner_progress = True try: with tqdm(total=total, bar_format=bar_fmt, unit="hypo") as pbar: pbar.set_description("0 done · 0 failed") - for hypo_idx, hypothesis in enumerate(all_hypotheses): + for hypo_idx, (key, hypothesis) in enumerate(tasks): logger.info( f"Inferring precondition for hypothesis {hypo_idx + 1}/{total}: " f"{hypothesis.invariant.text_description}" ) - precondition = find_precondition(hypothesis, self.traces) + precondition = _find_precondition_task( + hypothesis, self.traces, key + ) _record_result(hypothesis, precondition) pbar.set_description( f"{len(invariants)} done · {len(failed_hypos)} failed" @@ -334,8 +372,10 @@ def _record_result(hypothesis: Hypothesis, precondition): initargs=(self.traces, config_snapshot()), ) as executor: future_to_hypo = { - executor.submit(_precond_worker_task, hypothesis): hypothesis - for hypothesis in all_hypotheses + executor.submit( + _precond_worker_task, (key, hypothesis) + ): hypothesis + for key, hypothesis in tasks } for future in as_completed(future_to_hypo): hypothesis = future_to_hypo[future] diff --git a/traincheck/inference_sampling.py b/traincheck/inference_sampling.py new file mode 100644 index 00000000..a4e26949 --- /dev/null +++ b/traincheck/inference_sampling.py @@ -0,0 +1,54 @@ +"""Task-local randomness for inference, independent of process scheduling. + +The engine supplies stable semantic task keys in both its serial and process +paths. Direct relation callers outside an engine task keep the legacy stdlib +random API, including control via random.seed(). Sampling rates and reservoirs +are unchanged; the chosen examples can differ from the old global RNG stream. +""" + +import hashlib +import json +import random +from contextvars import ContextVar + +_TASK_RNG: ContextVar[random.Random | None] = ContextVar( + "inference_task_rng", default=None +) + + +class sampling_task: + """Give one logical task a fresh reproducible RNG without touching global state. + + Keys must be JSON values with stable ordering, never Python hash(), object + addresses, worker IDs, or arrival order. Trace indexes refer to the engine's + ordered input list. A versioned digest keeps seeds stable across processes + and PYTHONHASHSEED values. ContextVar restores enclosing scopes on exceptions + and isolates concurrently executing threads. + """ + + def __init__(self, *identity): + self.identity = identity + + def __enter__(self): + payload = json.dumps( + ["traincheck-inference-sampling-v1", 0, *self.identity], + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + rng = random.Random(int.from_bytes(hashlib.sha256(payload).digest(), "big")) + self.token = _TASK_RNG.set(rng) + return rng + + def __exit__(self, exc_type, exc_value, traceback): + # A class context manager also preserves frozen exception objects; the + # generator-based contextlib helper may assign exception.__traceback__. + _TASK_RNG.reset(self.token) + return False + + +def sampling_rng(): + """Return the task's RNG, or the legacy module for standalone relation calls.""" + rng = _TASK_RNG.get() + return random if rng is None else rng diff --git a/traincheck/invariant/consistency_transient_vars.py b/traincheck/invariant/consistency_transient_vars.py index f02e65bb..eb9196ed 100644 --- a/traincheck/invariant/consistency_transient_vars.py +++ b/traincheck/invariant/consistency_transient_vars.py @@ -5,6 +5,7 @@ import pandas as pd from traincheck.config.config import ANALYSIS_SKIP_FUNC_NAMES +from traincheck.inference_sampling import sampling_rng from traincheck.instrumentor.tracer import TraceLineType from traincheck.invariant.base_cls import ( _NOT_SET, @@ -262,11 +263,9 @@ def get_events_of_funcs_with_tensors( } # sampling 1000 if more than 1000 - import random - all_func_call_ids = { func_name: ( - random.sample(func_call_ids, 1000) + sampling_rng().sample(func_call_ids, 1000) if len(func_call_ids) > 1000 else func_call_ids ) @@ -452,10 +451,8 @@ def collect_examples(trace, hypothesis): func_call_ids = trace.get_func_call_ids(func_name) # down sample to 1000 - import random - if len(func_call_ids) > 1000: - func_call_ids = random.sample(func_call_ids, 1000) + func_call_ids = sampling_rng().sample(func_call_ids, 1000) for func_call_id in tqdm( func_call_ids, desc=f"Adding examples for {inv.text_description}" @@ -867,10 +864,8 @@ def collect_examples(trace, hypothesis): func_name = api_param.api_full_name func_call_ids = trace.get_func_call_ids(func_name) - import random - if len(func_call_ids) > 1000: - func_call_ids = random.sample(func_call_ids, 1000) + func_call_ids = sampling_rng().sample(func_call_ids, 1000) for func_call_id in tqdm( func_call_ids, desc=f"Checking invariant {inv.text_description}" @@ -1360,10 +1355,8 @@ def collect_examples(trace, hypothesis): # get all function calls for the function func_call_ids = trace.get_func_call_ids(func_name) - import random - if len(func_call_ids) > 1000: - func_call_ids = random.sample(func_call_ids, 1000) + func_call_ids = sampling_rng().sample(func_call_ids, 1000) for func_call_id in tqdm( func_call_ids, desc=f"Checking invariant {inv.text_description}" diff --git a/traincheck/invariant/contain_relation.py b/traincheck/invariant/contain_relation.py index 84fa76d2..cdc82829 100644 --- a/traincheck/invariant/contain_relation.py +++ b/traincheck/invariant/contain_relation.py @@ -1,12 +1,12 @@ import copy import logging -import random import time from typing import Type import numpy as np from traincheck.config.config import ANALYSIS_SKIP_FUNC_NAMES +from traincheck.inference_sampling import sampling_rng from traincheck.instrumentor.tracer import TraceLineType from traincheck.invariant.base_cls import ( _NOT_SET, @@ -169,10 +169,10 @@ def prune_func_call( - whether the function call should be pruned or not The pruning logic is as follows: - - If the number of function calls total present in the trace is less than 1000, no pruning is done - - Pruning is probabilistic, with initial probability of (1 - 1000/total_func_calls) - - If `list_num_events_scanned` indicates that the last 10 function calls have the same number of events scanned, the pruning probability is increased by 100 times for the next function call - - We use random.random() to determine whether the function call should be pruned or not + - Initially keep all calls when there are at most 1000; otherwise keep 1000/total. + - After more than 10 scans, equal event counts in the last 10 reduce the keep + probability by 100, even below 1000 total calls. + - We use sampling_rng().random() to determine whether the function call should be pruned or not """ logger = logging.getLogger(__name__) @@ -197,7 +197,7 @@ def prune_func_call( f"Same number of events observed in the last 10 attempts: {list_num_events_scanned[0]}, increasing the pruning probability to: {1 - not_pruning_prob}" ) - is_skipping = random.random() > not_pruning_prob + is_skipping = sampling_rng().random() > not_pruning_prob if is_skipping: logger.debug( f"Skipping the function call due to sampling with a pruning probability of: {1 - not_pruning_prob}" @@ -650,7 +650,7 @@ def _infer( # down sample the function calls, keep the first 10, last 10, and randomly sample 100 from the rest parent_func_call_ids = ( parent_func_call_ids[:10] - + random.sample(parent_func_call_ids[10:-10], 100) # type: ignore + + sampling_rng().sample(parent_func_call_ids[10:-10], 100) # type: ignore + parent_func_call_ids[-10:] ) diff --git a/traincheck/invariant/precondition.py b/traincheck/invariant/precondition.py index f86ffb95..45c17f5e 100644 --- a/traincheck/invariant/precondition.py +++ b/traincheck/invariant/precondition.py @@ -306,17 +306,18 @@ def find_precondition( import random - random.seed(42) + # Preserve this group's existing sample while leaving the enclosing task RNG alone. + group_rng = random.Random(42) if len(positive_examples) > 5000: logger.warning( f"Too many positive examples found for group {group_name}, downsampling to 100000" ) - positive_examples = random.sample(positive_examples, 5000) + positive_examples = group_rng.sample(positive_examples, 5000) if len(negative_examples) > 5000: logger.warning( f"Too many negative examples found for group {group_name}, downsampling to 100000" ) - negative_examples = random.sample(negative_examples, 5000) + negative_examples = group_rng.sample(negative_examples, 5000) grouped_preconditions[group_name] = Preconditions( find_precondition_from_single_group( From 122e76937457cda9472804c9ad22e94b2b0980d4 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 21:23:33 +0800 Subject: [PATCH 17/17] Type deterministic sampling and inference task contracts --- traincheck/infer_engine.py | 27 +++++++++++++------ traincheck/inference_sampling.py | 45 +++++++++++++++++++++++++++----- 2 files changed, 58 insertions(+), 14 deletions(-) diff --git a/traincheck/infer_engine.py b/traincheck/infer_engine.py index f5958efb..d98182a5 100644 --- a/traincheck/infer_engine.py +++ b/traincheck/infer_engine.py @@ -19,8 +19,10 @@ find_precondition, relation_pool, ) +from traincheck.invariant.base_cls import GroupedPreconditions from traincheck.parallel import config_snapshot, get_worker_traces, worker_init from traincheck.trace import MDNONEJSONEncoder, select_trace_implementation +from traincheck.trace.trace import Trace from traincheck.utils import register_custom_excepthook register_custom_excepthook() @@ -28,7 +30,14 @@ logger = logging.getLogger(__name__) -def _generate_hypothesis_task(trace_idx, relation, trace): +GenerationTask = tuple[int, type[Relation]] +GenerationResult = tuple[int, type[Relation], list[Hypothesis]] +PreconditionTask = tuple[str, Hypothesis] + + +def _generate_hypothesis_task( + trace_idx: int, relation: type[Relation], trace: Trace +) -> list[Hypothesis]: """Use the same logical task RNG in serial and process-based generation.""" with sampling_task( "generate", trace_idx, relation.__module__, relation.__qualname__ @@ -36,7 +45,7 @@ def _generate_hypothesis_task(trace_idx, relation, trace): return relation.generate_hypothesis(trace) -def _hypothesis_sampling_key(hypothesis): +def _hypothesis_sampling_key(hypothesis: Hypothesis) -> str: """Identify a hypothesis before dispatch, excluding mutable results/counts.""" relation = hypothesis.invariant.relation identity = { @@ -47,13 +56,15 @@ def _hypothesis_sampling_key(hypothesis): } # Normalize JSON object keys and MD_NONE with the invariant-output encoder, # then sort keys for insertion-order independence. - normalized = json.loads(json.dumps(identity, cls=MDNONEJSONEncoder)) + normalized: object = json.loads(json.dumps(identity, cls=MDNONEJSONEncoder)) return json.dumps( normalized, sort_keys=True, separators=(",", ":"), ensure_ascii=False ) -def _find_precondition_task(hypothesis, traces, key): +def _find_precondition_task( + hypothesis: Hypothesis, traces: list[Trace], key: str +) -> GroupedPreconditions | None: with sampling_task("precondition", key): return find_precondition(hypothesis, traces) @@ -64,7 +75,7 @@ def _find_precondition_task(hypothesis, traces, key): # worker via worker_init, not re-pickled per task). -def _genhypo_worker_task(task: tuple): +def _genhypo_worker_task(task: GenerationTask) -> GenerationResult: """Generate hypotheses for one (trace_idx, relation) pair in a worker process. Returns (trace_idx, relation, inferred_hypos); the parent merges them in a @@ -76,7 +87,7 @@ def _genhypo_worker_task(task: tuple): return trace_idx, relation, hypotheses -def _precond_worker_task(task): +def _precond_worker_task(task: PreconditionTask) -> GroupedPreconditions | None: """Run precondition inference for a single hypothesis in a worker process. Returns just the precondition (picklable); the parent applies it to its own @@ -177,7 +188,7 @@ def _merge_hypotheses( def _generate_hypothesis_serial( self, - active_relations: list, + active_relations: list[type[Relation]], hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], hypo_lookup: dict[Hypothesis, Hypothesis], ) -> None: @@ -218,7 +229,7 @@ def _generate_hypothesis_serial( def _generate_hypothesis_parallel( self, - active_relations: list, + active_relations: list[type[Relation]], hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], hypo_lookup: dict[Hypothesis, Hypothesis], ) -> None: diff --git a/traincheck/inference_sampling.py b/traincheck/inference_sampling.py index a4e26949..1ca94b52 100644 --- a/traincheck/inference_sampling.py +++ b/traincheck/inference_sampling.py @@ -9,7 +9,32 @@ import hashlib import json import random -from contextvars import ContextVar +from collections.abc import Sequence +from contextvars import ContextVar, Token +from types import TracebackType +from typing import Literal, Protocol, TypeAlias, TypeVar, Union + +# The JSON encoder accepts tuples as arrays too; keep that existing contract. +SamplingIdentity: TypeAlias = Union[ + None, + bool, + int, + float, + str, + list["SamplingIdentity"], + tuple["SamplingIdentity", ...], + dict[str, "SamplingIdentity"], +] +_Sample = TypeVar("_Sample") + + +class SamplingRNG(Protocol): + """The shared sampling surface of random.Random and the random module.""" + + def random(self) -> float: ... + + def sample(self, population: Sequence[_Sample], k: int) -> list[_Sample]: ... + _TASK_RNG: ContextVar[random.Random | None] = ContextVar( "inference_task_rng", default=None @@ -26,10 +51,11 @@ class sampling_task: and isolates concurrently executing threads. """ - def __init__(self, *identity): + def __init__(self, *identity: SamplingIdentity) -> None: self.identity = identity + self.token: Token[random.Random | None] - def __enter__(self): + def __enter__(self) -> random.Random: payload = json.dumps( ["traincheck-inference-sampling-v1", 0, *self.identity], sort_keys=True, @@ -41,14 +67,21 @@ def __enter__(self): self.token = _TASK_RNG.set(rng) return rng - def __exit__(self, exc_type, exc_value, traceback): + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> Literal[False]: # A class context manager also preserves frozen exception objects; the # generator-based contextlib helper may assign exception.__traceback__. _TASK_RNG.reset(self.token) return False -def sampling_rng(): +def sampling_rng() -> SamplingRNG: """Return the task's RNG, or the legacy module for standalone relation calls.""" rng = _TASK_RNG.get() - return random if rng is None else rng + if rng is None: + return random + return rng