From 34ac80e530bc72e5a4f36936d05db59f08909a3a Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 13 Oct 2026 14:05:24 +0800 Subject: [PATCH 1/5] 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 6c1a5d414137c96013acd2fc90e0e90088919de1 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 15:05:47 +0800 Subject: [PATCH 2/5] Fix bounded reservoir construction and population-aware merges --- tests/test_example_reservoir.py | 528 +++++++++++++++++++++++ traincheck/infer_engine.py | 40 +- traincheck/invariant/base_cls.py | 138 ++++-- traincheck/invariant/contain_relation.py | 51 ++- traincheck/parallel.py | 14 +- 5 files changed, 705 insertions(+), 66 deletions(-) create mode 100644 tests/test_example_reservoir.py diff --git a/tests/test_example_reservoir.py b/tests/test_example_reservoir.py new file mode 100644 index 00000000..d632ab0a --- /dev/null +++ b/tests/test_example_reservoir.py @@ -0,0 +1,528 @@ +"""Reservoir invariants, weighted population merges, and actual merge call sites.""" + +import multiprocessing +import random +import unittest +from collections import Counter +from concurrent.futures import ProcessPoolExecutor +from dataclasses import dataclass +from functools import partial +from unittest.mock import patch + +import traincheck.config.config as config +from traincheck.infer_engine import InferEngine +from traincheck.invariant.base_cls import ( + APIParam, + Example, + ExampleList, + GroupedPreconditions, + Hypothesis, + Invariant, + Preconditions, + UnconditionalPrecondition, +) +from traincheck.invariant.contain_relation import PARENT_GROUP_NAME, APIContainRelation +from traincheck.parallel import config_snapshot, worker_init + +GROUP = PARENT_GROUP_NAME + + +def example(value, group=GROUP): + return Example({group: [{"id": value}]}) + + +def values(reservoir): + return [item.get_group(GROUP)[0]["id"] for item in reservoir.examples] + + +def reservoir(items, cap, seed=42, groups=None): + with ( + patch.object(config, "MAX_EXAMPLE_LIST_SIZE", cap), + patch.object(config, "EXAMPLE_SAMPLING_SEED", seed), + ): + result = ExampleList({GROUP} if groups is None else groups) + for item in items: + result.add_example(example(item)) + return result + + +def snapshot(reservoir): + return ( + reservoir.num_seen, + reservoir.max_size, + set(reservoir.group_names), + list(reservoir.examples), + reservoir._rng.getstate(), + ) + + +def hypothesis(prefix, positive, negative, cap): + precondition = GroupedPreconditions( + {GROUP: Preconditions([UnconditionalPrecondition()])} + ) + invariant = Invariant( + APIContainRelation, + [ + APIParam("torch.nn.modules.activation.ReLU.forward"), + APIParam("torch.nn.modules.dropout.Dropout.forward"), + ], + precondition, + "ReLU.forward() contains Dropout.forward()", + ) + return Hypothesis( + invariant, + reservoir([f"{prefix}-positive-{i}" for i in range(positive)], cap), + reservoir([f"{prefix}-negative-{i}" for i in range(negative)], cap), + ) + + +def worker_reservoir_result(count): + result = ExampleList({GROUP}) + for index in range(count): + result.add_example(example(index)) + return ( + result.max_size, + config.EXAMPLE_SAMPLING_SEED, + result.num_seen, + values(result), + ) + + +@dataclass +class ReservoirFixtureTrace: + """A picklable trace carrying populations, without any trace parser work.""" + + prefix: str + positive_count: int + negative_count: int + + def is_stage_annotated(self): + return False + + +class ReservoirFixtureRelation(APIContainRelation): + @classmethod + def generate_hypothesis(cls, trace): + # Read config directly here: patching construction inside the fixture + # would mask a missing config_snapshot/worker_init propagation field. + precondition = GroupedPreconditions( + {GROUP: Preconditions([UnconditionalPrecondition()])} + ) + invariant = Invariant( + cls, + [APIParam("fixture.parent"), APIParam("fixture.child")], + precondition, + "fixture.parent contains fixture.child", + ) + positive, negative = ExampleList({GROUP}), ExampleList({GROUP}) + for kind, count, examples in [ + ("positive", trace.positive_count, positive), + ("negative", trace.negative_count, negative), + ]: + for index in range(count): + examples.add_example(example(f"{trace.prefix}-{kind}-{index}")) + return [Hypothesis(invariant, positive, negative)] + + +class ExampleReservoirConstructionTests(unittest.TestCase): + def test_from_iterable_streams_through_capped_add_and_counts_every_item(self): + applied = [] + original = ExampleList.add_example + + def observed_add(result, item): + original(result, item) + applied.append(result.num_seen) + self.assertLessEqual(len(result), 3) + + def items(): + for index in range(37): + self.assertEqual( + len(applied), index, "input must be consumed incrementally" + ) + yield example(index) + + with ( + patch.object(config, "MAX_EXAMPLE_LIST_SIZE", 3), + patch.object(ExampleList, "add_example", observed_add), + ): + result = ExampleList.from_iterable_of_examples(items(), {GROUP}) + self.assertEqual(len(result), 3) + self.assertEqual(result.num_seen, 37) + self.assertEqual(applied, list(range(1, 38))) + + def test_unlimited_zero_and_empty_iterables(self): + for cap, expected in [(None, list(range(12))), (0, []), (20, list(range(12)))]: + with ( + self.subTest(cap=cap), + patch.object(config, "MAX_EXAMPLE_LIST_SIZE", cap), + ): + result = ExampleList.from_iterable_of_examples( + (example(i) for i in range(12)), {GROUP} + ) + self.assertEqual(values(result), expected) + self.assertEqual(result.num_seen, 12) + empty = ExampleList.from_iterable_of_examples(iter(()), {GROUP}) + self.assertEqual(len(empty), 0) + self.assertEqual(empty.num_seen, 0) + + def test_invalid_caps_fail_at_construction(self): + for cap in [-1, True, False, 1.5, "3"]: + with ( + self.subTest(cap=cap), + patch.object(config, "MAX_EXAMPLE_LIST_SIZE", cap), + ): + with self.assertRaises(ValueError): + ExampleList({GROUP}) + + def test_from_iterable_validates_group_membership(self): + with self.assertRaises((AssertionError, ValueError)): + ExampleList.from_iterable_of_examples([example(1, "different")], {GROUP}) + + +class ExampleReservoirMergeTests(unittest.TestCase): + def assert_rejected_without_mutation(self, left, right): + prior_left, prior_right = snapshot(left), snapshot(right) + rng = random.Random(761) + rng_before = rng.getstate() + with self.assertRaises(ValueError): + left.merge(right, rng=rng) + self.assertEqual(snapshot(left), prior_left) + self.assertEqual(snapshot(right), prior_right) + self.assertEqual(rng.getstate(), rng_before) + + def test_exact_unlimited_and_under_cap_order_without_rng_consumption(self): + for cap in [None, 10]: + with self.subTest(cap=cap): + left = reservoir(["a", "a", "b"], cap) + right = reservoir(["c", "a"], cap) + before = snapshot(right) + rng = random.Random(54) + rng_before = rng.getstate() + self.assertIsNone(left.merge(right, rng=rng)) + self.assertEqual(values(left), ["a", "a", "b", "c", "a"]) + self.assertEqual(left.num_seen, 5) + self.assertEqual(snapshot(right), before) + self.assertEqual(rng.getstate(), rng_before) + + def test_empty_destination_adopts_groups_and_empty_source_preserves_order(self): + left = reservoir([], 3, groups=set()) + right = reservoir([1, 2], 3) + left.merge(right, rng=random.Random(1)) + self.assertEqual(left.group_names, {GROUP}) + self.assertEqual(values(left), [1, 2]) + self.assertEqual(left.num_seen, 2) + + uninitialized = reservoir([], 3, groups=set()) + rng = random.Random(3) + before = rng.getstate() + uninitialized.merge(reservoir([], 3), rng=rng) + self.assertEqual(uninitialized.group_names, {GROUP}) + self.assertEqual(uninitialized.num_seen, 0) + self.assertEqual(rng.getstate(), before) + left.merge(reservoir([], 3), rng=random.Random(2)) + self.assertEqual(values(left), [1, 2]) + self.assertEqual(left.num_seen, 2) + + def test_zero_capacity_accumulates_counts_and_preserves_no_items(self): + left, right = reservoir(range(20), 0), reservoir(range(7), 0) + left.merge(right, rng=random.Random(7)) + left.add_example(example("next")) + self.assertEqual(left.num_seen, 28) + self.assertEqual(len(left), 0) + + def test_unequal_populations_are_counted_and_source_remains_unchanged(self): + left, right = reservoir(range(10), 4, 11), reservoir(range(10, 100), 4, 19) + prior_right = snapshot(right) + left.merge(right, rng=random.Random(37)) + self.assertEqual(left.num_seen, 100) + self.assertEqual(len(left), 4) + self.assertTrue(set(values(left)).issubset(range(100))) + self.assertEqual(snapshot(right), prior_right) + for value in range(100, 117): + left.add_example(example(value)) + self.assertEqual(left.num_seen, 117) + self.assertEqual(len(left), 4) + + def test_duplicate_examples_are_a_multiset_not_a_set(self): + left, right = reservoir(["same"] * 7, 3), reservoir(["same"] * 11, 3) + left.merge(right, rng=random.Random(42)) + self.assertEqual(left.num_seen, 18) + self.assertEqual(values(left), ["same", "same", "same"]) + + def test_different_caps_are_allowed_when_source_has_enough_information(self): + left, right = reservoir(range(10), 5), reservoir([10, 11], 2) + left.merge(right, rng=random.Random(8)) + self.assertEqual((len(left), left.num_seen), (5, 12)) + left, right = reservoir(range(7), 2), reservoir(range(7, 17), None) + left.merge(right, rng=random.Random(9)) + self.assertEqual((len(left), left.num_seen), (2, 17)) + + def test_rejects_self_group_mismatch_and_incompatible_sampled_sources(self): + left = reservoir(range(3), 3) + self.assert_rejected_without_mutation(left, left) + right = reservoir([], 3, groups={"other"}) + self.assert_rejected_without_mutation(left, right) + self.assert_rejected_without_mutation( + reservoir(range(7), 5), reservoir(range(9), 2) + ) + self.assert_rejected_without_mutation( + reservoir(range(3), None), reservoir(range(9), 2) + ) + + def test_rejects_underfilled_and_malformed_reservoir_states(self): + for corrupt_left in [False, True]: + left, right = reservoir(range(9), 3), reservoir(range(12), 3) + (left if corrupt_left else right).examples.pop() + self.assert_rejected_without_mutation(left, right) + for malformed in [-1, True, 2.5]: + left, right = reservoir(range(3), 3), reservoir(range(3), 3) + right.num_seen = malformed + self.assert_rejected_without_mutation(left, right) + + def test_weighted_merge_and_continued_add_have_expected_marginals(self): + # Fixed seed enumeration is deterministic, not a fresh statistical draw. + # Concatenating retained samples gives 'left' ~75%, not the required30%. + before_add = after_add = 0 + trials = 2000 + for seed in range(trials): + left = reservoir(["left"], 3, 2 * seed + 1) + right = reservoir(range(9), 3, 2 * seed + 2) + left.merge(right, rng=random.Random(100000 + seed)) + before_add += "left" in values(left) + for index in range(10): + left.add_example(example(f"new-{index}")) + after_add += "left" in values(left) + self.assertAlmostEqual(before_add / trials, 3 / 10, delta=0.04) + self.assertAlmostEqual(after_add / trials, 3 / 20, delta=0.04) + + def test_nested_merge_tree_preserves_marginals_with_shared_or_distinct_source_seeds( + self, + ): + # An independent sequential merge RNG must be threaded through the entire + # uneven tree. Replayed per-list merge streams can over-select 'e' (~27%). + trials = 4000 + for same_source_seed in [True, False]: + merged, continued = Counter(), Counter() + for seed in range(trials): + source_seeds = ( + [seed] * 5 + if same_source_seed + else [31 * seed + i for i in range(5)] + ) + a, b, c, d, e = [ + reservoir([name], 1, value) + for name, value in zip("abcde", source_seeds) + ] + rng = random.Random(200000 + seed) + a.merge(b, rng=rng) + a.merge(c, rng=rng) + d.merge(e, rng=rng) + a.merge(d, rng=rng) + self.assertEqual(a.num_seen, 5) + merged.update(values(a)) + a.add_example(example("f")) + self.assertEqual(a.num_seen, 6) + continued.update(values(a)) + with self.subTest(same_source_seed=same_source_seed): + for name in "abcde": + self.assertAlmostEqual(merged[name] / trials, 1 / 5, delta=0.025) + for name in "abcdef": + self.assertAlmostEqual(continued[name] / trials, 1 / 6, delta=0.025) + + def test_sampled_unequal_sources_then_add_preserve_every_occurrence_marginal(self): + # Both sources have already sampled, using the same construction seed. + # Continuing the destination's old RNG replays draws used by the longer + # source and biases some occurrence marginals to ~14% or ~20%, versus + # 2/12. Checking just aggregate source weights would hide this effect. + merged, continued = Counter(), Counter() + trials = 6000 + for seed in range(trials): + left = reservoir(range(3), 2, seed) + right = reservoir(range(3, 10), 2, seed) + left.merge(right, rng=random.Random(300000 + seed)) + self.assertEqual((left.num_seen, len(left)), (10, 2)) + merged.update(values(left)) + for index in (10, 11): + left.add_example(example(index)) + self.assertEqual((left.num_seen, len(left)), (12, 2)) + continued.update(values(left)) + for occurrence in range(10): + self.assertAlmostEqual(merged[occurrence] / trials, 2 / 10, delta=0.02) + for occurrence in range(12): + self.assertAlmostEqual(continued[occurrence] / trials, 2 / 12, delta=0.02) + + +class ReservoirIntegrationTests(unittest.TestCase): + def test_actual_generate_hypothesis_matches_workers_1_2_4_fork_and_spawn(self): + traces = [ + ReservoirFixtureTrace("a", 17, 11), + ReservoirFixtureTrace("b", 31, 19), + ReservoirFixtureTrace("c", 53, 29), + ] + + def run(workers): + generated = InferEngine(traces, num_workers=workers).generate_hypothesis() + self.assertEqual(len(generated), 1) + item, trace_indexes = next(iter(generated.items())) + self.assertEqual(trace_indexes, [0, 1, 2]) + result = [] + for examples, count in [ + (item.positive_examples, 101), + (item.negative_examples, 59), + ]: + self.assertEqual((examples.max_size, len(examples)), (5, 5)) + self.assertEqual(examples.num_seen, count) + # Include the continuation state: pickling must preserve it and + # the parent must consume the same merge stream in both paths. + result.append((values(examples), count, examples._rng.getstate())) + return result + + with ( + patch("traincheck.infer_engine.relation_pool", [ReservoirFixtureRelation]), + patch.object(config, "MAX_EXAMPLE_LIST_SIZE", 5), + patch.object(config, "EXAMPLE_SAMPLING_SEED", 907), + ): + expected = run(1) + for method in [ + value + for value in ("spawn", "fork") + if value in multiprocessing.get_all_start_methods() + ]: + executor = partial( + ProcessPoolExecutor, mp_context=multiprocessing.get_context(method) + ) + for workers in (1, 2, 4): + with ( + self.subTest(method=method, workers=workers), + patch("traincheck.infer_engine.ProcessPoolExecutor", executor), + ): + self.assertEqual(run(workers), expected) + + def test_actual_engine_cross_trace_merge_keeps_caps_and_total_counts(self): + first, second = hypothesis("a", 17, 8, 3), hypothesis("b", 29, 11, 3) + engine, trace_indexes, lookup = InferEngine([]), {}, {} + rng = random.Random(712) + engine._merge_hypotheses([first], 0, trace_indexes, lookup, rng) + engine._merge_hypotheses([second], 1, trace_indexes, lookup, rng) + self.assertEqual(len(trace_indexes), 1) + self.assertEqual(next(iter(trace_indexes.values())), [0, 1]) + self.assertEqual( + (len(first.positive_examples), first.positive_examples.num_seen), (3, 46) + ) + self.assertEqual( + (len(first.negative_examples), first.negative_examples.num_seen), (3, 19) + ) + + def test_actual_stage_merge_keeps_caps_and_total_counts(self): + first, second = hypothesis("a", 17, 8, 3), hypothesis("b", 29, 11, 3) + + class Stages: + def is_stage_annotated(self): + return True + + def get_traces_for_stage(self): + return {"training": "training", "evaluation": "evaluation"} + + def get_all_stages(self): + return {"training", "evaluation"} + + by_stage = {"training": first, "evaluation": second} + + def inferred(stage): + item = by_stage[stage] + return [item.invariant], [], [item] + + with patch.object(APIContainRelation, "_infer", side_effect=inferred): + invariants, failures, successful = APIContainRelation.infer(Stages(), True) + self.assertEqual((len(invariants), len(failures), len(successful)), (1, 0, 1)) + self.assertEqual( + ( + len(successful[0].positive_examples), + successful[0].positive_examples.num_seen, + ), + (3, 46), + ) + self.assertEqual( + ( + len(successful[0].negative_examples), + successful[0].negative_examples.num_seen, + ), + (3, 19), + ) + + def test_sampled_generalization_does_not_treat_retained_union_as_population(self): + from traincheck.invariant.contain_relation import _merge_hypotheses + + hypotheses = [hypothesis("a", 20, 9, 3), hypothesis("b", 40, 15, 3)] + with patch( + "traincheck.invariant.contain_relation.calc_likelihood", + side_effect=AssertionError("sampled union must not be generalized"), + ): + result = _merge_hypotheses(hypotheses) + self.assertEqual( + [id(item) for item in result], [id(item) for item in hypotheses] + ) + + def test_complete_generalization_deduplicates_known_union_then_applies_cap(self): + from traincheck.invariant.contain_relation import _merge_hypotheses + + first, second = hypothesis("a", 0, 0, 10), hypothesis("b", 0, 0, 10) + first.positive_examples = reservoir(range(10), 10) + second.positive_examples = reservoir(range(10, 20), 10) + first.negative_examples = reservoir(range(10, 20), 10) + second.negative_examples = reservoir(range(10), 10) + expected = values(reservoir(range(20), 10)) + # Repeating one complete input must not double-count the known union. + for inputs in ([first, second], [first, second, first]): + with ( + self.subTest(inputs=len(inputs)), + patch.object(config, "MAX_EXAMPLE_LIST_SIZE", 10), + patch.object(config, "EXAMPLE_SAMPLING_SEED", 42), + ): + merged = _merge_hypotheses(inputs) + self.assertEqual(len(merged), 1) + positive, negative = ( + merged[0].positive_examples, + merged[0].negative_examples, + ) + self.assertEqual((len(positive), positive.num_seen), (10, 20)) + self.assertEqual((len(negative), negative.num_seen), (0, 0)) + self.assertEqual(values(positive), expected) + self.assertEqual(negative.group_names, {GROUP}) + self.assertEqual(first.positive_examples.num_seen, 10) + self.assertEqual(second.positive_examples.num_seen, 10) + + def test_worker_snapshot_carries_reservoir_settings_under_fork_and_spawn(self): + with ( + patch.object(config, "MAX_EXAMPLE_LIST_SIZE", 7), + patch.object(config, "EXAMPLE_SAMPLING_SEED", 907), + ): + settings = config_snapshot() + expected = worker_reservoir_result(29) + self.assertEqual(settings["MAX_EXAMPLE_LIST_SIZE"], 7) + self.assertEqual(settings["EXAMPLE_SAMPLING_SEED"], 907) + for method in [ + value + for value in ("spawn", "fork") + if value in multiprocessing.get_all_start_methods() + ]: + with ( + self.subTest(method=method), + ProcessPoolExecutor( + max_workers=2, + mp_context=multiprocessing.get_context(method), + initializer=worker_init, + initargs=([], settings), + ) as executor, + ): + futures = [ + executor.submit(worker_reservoir_result, 29) for _ in range(2) + ] + self.assertEqual( + [future.result(timeout=60) for future in futures], + [expected, expected], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/infer_engine.py b/traincheck/infer_engine.py index 09ff0411..ffc27963 100644 --- a/traincheck/infer_engine.py +++ b/traincheck/infer_engine.py @@ -6,6 +6,7 @@ import random import time from concurrent.futures import ProcessPoolExecutor, as_completed +from random import Random from tqdm import tqdm @@ -92,14 +93,19 @@ def generate_hypothesis(self) -> dict[Hypothesis, list[int]]: active_relations = [ r for r in relation_pool if r not in self.disabled_relations ] + # One stream for the fixed parent merge order, independent of reservoir + # construction and stage merging. Never reset it per hypothesis. + merge_rng = Random( + f"traincheck-example-merge-v1:cross-trace:{config.EXAMPLE_SAMPLING_SEED!r}" + ) if self.num_workers <= 1: self._generate_hypothesis_serial( - active_relations, hypotheses_and_trace_idxs, hypo_lookup + active_relations, hypotheses_and_trace_idxs, hypo_lookup, merge_rng ) else: self._generate_hypothesis_parallel( - active_relations, hypotheses_and_trace_idxs, hypo_lookup + active_relations, hypotheses_and_trace_idxs, hypo_lookup, merge_rng ) total = len(hypotheses_and_trace_idxs) @@ -113,6 +119,7 @@ def _merge_hypotheses( trace_idx: int, hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], hypo_lookup: dict[Hypothesis, Hypothesis], + merge_rng: Random, ) -> None: """Merge one trace's inferred hypotheses into the running collection. @@ -126,33 +133,21 @@ def _merge_hypotheses( 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.positive_examples.merge( + hypo.positive_examples, rng=merge_rng ) - original_hypo.negative_examples.examples.extend( - hypo.negative_examples.examples + original_hypo.negative_examples.merge( + hypo.negative_examples, rng=merge_rng ) - 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], + merge_rng: Random, ) -> None: - """Serial hypothesis generation (preserves the original behavior exactly).""" + """Generate hypotheses serially, merging in trace-then-relation order.""" 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): @@ -183,6 +178,7 @@ def _generate_hypothesis_serial( trace_idx, hypotheses_and_trace_idxs, hypo_lookup, + merge_rng, ) def _generate_hypothesis_parallel( @@ -190,6 +186,7 @@ def _generate_hypothesis_parallel( active_relations: list, hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], hypo_lookup: dict[Hypothesis, Hypothesis], + merge_rng: Random, ) -> None: """Parallel hypothesis generation. @@ -197,7 +194,7 @@ 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. """ n_traces = len(self.traces) tasks = [ @@ -241,6 +238,7 @@ def _generate_hypothesis_parallel( trace_idx, hypotheses_and_trace_idxs, hypo_lookup, + merge_rng, ) def collect_examples(self, hypotheses: dict[Hypothesis, list[int]]): diff --git a/traincheck/invariant/base_cls.py b/traincheck/invariant/base_cls.py index 588e7260..59ddd2c3 100644 --- a/traincheck/invariant/base_cls.py +++ b/traincheck/invariant/base_cls.py @@ -182,7 +182,7 @@ 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 @@ -1677,18 +1677,47 @@ def __eq__(self, value: object) -> bool: class ExampleList: + """A bounded sample of example occurrences, including duplicate values. + + ``num_seen`` counts the underlying stream, while ``len(self)`` counts the + retained examples. Different lists retain the historical equal-seed RNG + behavior, so their samples can be correlated; merging does not promise a + uniform joint sample across lists. + """ + 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) + self.max_size = config.MAX_EXAMPLE_LIST_SIZE + self._validate_cap(self.max_size) + self.num_seen = 0 + self._rng = random.Random(config.EXAMPLE_SAMPLING_SEED) + + @staticmethod + def _validate_cap(cap: int | None) -> None: + # Reject bool and int subclasses rather than accepting isinstance(int). + if cap is not None and (type(cap) is not int or cap < 0): # noqa: E721 + raise ValueError( + "Example reservoir capacity must be a nonnegative integer or None" + ) + + def _validate_state(self) -> None: + self._validate_cap(self.max_size) + # A population count must be a builtin integer. + if type(self.num_seen) is not int or self.num_seen < 0: # noqa: E721 + raise ValueError( + "Example reservoir population must be a nonnegative integer" + ) + expected = ( + self.num_seen + if self.max_size is None + else min(self.max_size, self.num_seen) + ) + if len(self.examples) != expected: + raise ValueError( + "Example reservoir is underfilled or has inconsistent population accounting: " + f"retained {len(self.examples)}, expected {expected} for {self.num_seen} seen" + ) def add_example(self, example: Example): if len(self.group_names) == 0: @@ -1698,17 +1727,80 @@ 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.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.num_seen += 1 + if self.max_size is None or len(self.examples) < self.max_size: + self.examples.append(example) + else: + # Algorithm R retains each occurrence with probability cap / num_seen. + j = self._rng.randrange(self.num_seen) + if j < self.max_size: self.examples[j] = example + def merge(self, other: ExampleList, *, rng: random.Random) -> None: + """Append another stream's reservoir without exceeding this list's cap. + + The caller's RNG must be independent of both input selections. Share one + RNG sequentially within a merge task; do not replay its state in nested + merge operations. This preserves each occurrence's marginal inclusion + probability, even when the input reservoirs are correlated. A uniform + joint sample additionally requires independent input reservoirs. + + A sampled source cannot be expanded beyond its retained capacity or + merged into an uncapped destination: discarded examples cannot be + recovered. Invalid inputs are rejected before changing either list or + consuming the supplied RNG. Complete streams preserve concatenation + order whenever they fit in the destination. + """ + if self is other: + raise ValueError("Cannot merge an example reservoir into itself") + self._validate_state() + other._validate_state() + adopt_groups = not self.group_names and self.num_seen == 0 + source_uninitialized = not other.group_names and other.num_seen == 0 + if ( + not adopt_groups + and not source_uninitialized + and self.group_names != other.group_names + ): + raise ValueError( + "Cannot merge example reservoirs with different group names" + ) + + total = self.num_seen + other.num_seen + target = total if self.max_size is None else min(self.max_size, total) + for source in (self, other): + if len(source.examples) < min(target, source.num_seen): + raise ValueError( + "Source reservoir retains too few examples for the destination capacity" + ) + if other.num_seen == 0: + if adopt_groups: + self.group_names = other.group_names.copy() + return + + if target == total: + # Both inputs are complete: no sampling or RNG consumption is needed. + examples = self.examples + other.examples + continuation_rng = self._rng + else: + # Sampling population positions gives an exact hypergeometric split + # in O(cap) space, without expanding either represented population. + take_self = sum( + index < self.num_seen for index in rng.sample(range(total), target) + ) + examples = rng.sample(self.examples, take_self) + rng.sample( + other.examples, target - take_self + ) + # Further additions must not replay RNG draws used by a longer source + # with the same initial reservoir seed. + continuation_rng = random.Random(rng.getrandbits(256)) + + self.examples = examples + self.num_seen = total + self._rng = continuation_rng + if adopt_groups: + self.group_names = other.group_names.copy() + def get_group_from_examples(self, group_name: str) -> list[list[dict]]: return [example.get_group(group_name) for example in self.examples] @@ -1725,14 +1817,10 @@ def __len__(self): def from_iterable_of_examples( input: Iterable[Example], preset_group_names: set[str] ) -> ExampleList: - examples = [] + example_list = ExampleList(preset_group_names) for exp in input: assert preset_group_names == exp.get_group_names() - examples.append(exp) - - example_list = ExampleList(preset_group_names) - example_list.examples = examples - example_list.num_seen = len(examples) + example_list.add_example(exp) return example_list diff --git a/traincheck/invariant/contain_relation.py b/traincheck/invariant/contain_relation.py index 84fa76d2..e1a9de13 100644 --- a/traincheck/invariant/contain_relation.py +++ b/traincheck/invariant/contain_relation.py @@ -2,10 +2,12 @@ import logging import random import time +from random import Random from typing import Type import numpy as np +import traincheck.config.config as config from traincheck.config.config import ANALYSIS_SKIP_FUNC_NAMES from traincheck.instrumentor.tracer import TraceLineType from traincheck.invariant.base_cls import ( @@ -209,6 +211,15 @@ def _merge_hypotheses(hypotheses: list[Hypothesis]) -> list[Hypothesis]: assert len(hypotheses) > 0, "Expected at least one hypotheses to be merged" logger = logging.getLogger(__name__) + if any( + examples.num_seen > len(examples) + for hypothesis in hypotheses + for examples in (hypothesis.positive_examples, hypothesis.negative_examples) + ): + # This optional generalization uses set unions and differences. Reservoirs + # cannot recover the discarded values or their overlap, so keep the + # original hypotheses instead of treating retained samples as full sets. + return hypotheses pos_exp_group_names = hypotheses[0].positive_examples.get_group_names() neg_exp_group_names = hypotheses[0].negative_examples.get_group_names() logger.debug( @@ -221,9 +232,10 @@ def _merge_hypotheses(hypotheses: list[Hypothesis]) -> list[Hypothesis]: all_likelihoods = [hypo.calc_likelihood() for hypo in hypotheses] # calculate the number of positive and negative examples after merge - all_positive_examples = set() + # Insertion-ordered sets keep constructor sampling independent of hash order. + all_positive_examples: dict[Example, None] = {} all_positive_parent_examples = set() - all_negative_examples = set() + all_negative_examples: dict[Example, None] = {} for idx in range(len(hypotheses)): @@ -234,16 +246,22 @@ def _merge_hypotheses(hypotheses: list[Hypothesis]) -> list[Hypothesis]: Example({PARENT_GROUP_NAME: pre_record}) for pre_record in pre_records ] all_positive_parent_examples.update(exps_with_only_parent) - all_positive_examples.update(hypotheses[idx].positive_examples.examples) - all_negative_examples.update(hypotheses[idx].negative_examples.examples) + all_positive_examples.update( + (example, None) for example in hypotheses[idx].positive_examples.examples + ) + all_negative_examples.update( + (example, None) for example in hypotheses[idx].negative_examples.examples + ) # let's remove those negative examples that are present in the positive examples of other hypotheses - if used_dynamic_analysis: - all_negative_examples = all_negative_examples.difference(all_positive_examples) - else: - all_negative_examples = all_negative_examples.difference( - all_positive_parent_examples - ) + positive_comparison = ( + all_positive_examples if used_dynamic_analysis else all_positive_parent_examples + ) + all_negative_examples = { + example: None + for example in all_negative_examples + if example not in positive_comparison + } # calculate the likelihood of the merged hypotheses now merged_likelihood = calc_likelihood( @@ -575,15 +593,20 @@ def infer( merged_invariants = [] merged_successful_hypotheses: list[Hypothesis] = [] all_stages = trace.get_all_stages() + # Reuse one independent stream for this stage-merge task. Cross-trace + # merging uses a different domain, so nested merges do not replay draws. + merge_rng = Random( + f"traincheck-example-merge-v1:stage:{config.EXAMPLE_SAMPLING_SEED!r}" + ) for invariant in invariants_by_stage: supported_stages = set(invariants_by_stage[invariant]) hypothesis = invariants_to_hypos[invariant][0] for other_hypothesis in invariants_to_hypos[invariant][1:]: - hypothesis.positive_examples.examples.extend( - other_hypothesis.positive_examples.examples + hypothesis.positive_examples.merge( + other_hypothesis.positive_examples, rng=merge_rng ) - hypothesis.negative_examples.examples.extend( - other_hypothesis.negative_examples.examples + hypothesis.negative_examples.merge( + other_hypothesis.negative_examples, rng=merge_rng ) hypothesis.add_stage_info(supported_stages, all_stages) merged_invariants.append(hypothesis.get_invariant(all_stages)) diff --git a/traincheck/parallel.py b/traincheck/parallel.py index 4d8bcb45..b1c35dac 100644 --- a/traincheck/parallel.py +++ b/traincheck/parallel.py @@ -20,16 +20,18 @@ def config_snapshot() -> dict: - """CLI-overridable config values that must be propagated to workers. + """Runtime inference 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. + These include CLI sampling options and programmatic reservoir settings. + Workers started with 'spawn' begin from a fresh interpreter and would + otherwise see only the module defaults. Add future runtime-overridable + inference configuration here. """ return { "ENABLE_PRECOND_SAMPLING": config.ENABLE_PRECOND_SAMPLING, "PRECOND_SAMPLING_THRESHOLD": config.PRECOND_SAMPLING_THRESHOLD, + "MAX_EXAMPLE_LIST_SIZE": config.MAX_EXAMPLE_LIST_SIZE, + "EXAMPLE_SAMPLING_SEED": config.EXAMPLE_SAMPLING_SEED, } @@ -37,7 +39,7 @@ 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 + restores runtime inference config values (see config_snapshot). Also suppresses the per-invariant inner progress bars, otherwise every worker would fight over the terminal. """ From c6385665083550ac7f3750c932cfabe54a4c0906 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 21:23:33 +0800 Subject: [PATCH 3/5] 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 From ba5268668f46db580f002643262f55495be97308 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 21:25:00 +0800 Subject: [PATCH 4/5] Define reservoir worker configuration and mutation contracts --- traincheck/invariant/base_cls.py | 6 +++--- traincheck/parallel.py | 21 +++++++++++++++++---- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/traincheck/invariant/base_cls.py b/traincheck/invariant/base_cls.py index 59ddd2c3..d2f340d2 100644 --- a/traincheck/invariant/base_cls.py +++ b/traincheck/invariant/base_cls.py @@ -1685,10 +1685,10 @@ class ExampleList: uniform joint sample across lists. """ - def __init__(self, group_names: set[str]): + def __init__(self, group_names: set[str]) -> None: self.group_names = group_names self.examples: list[Example] = [] - self.max_size = config.MAX_EXAMPLE_LIST_SIZE + self.max_size: int | None = config.MAX_EXAMPLE_LIST_SIZE self._validate_cap(self.max_size) self.num_seen = 0 self._rng = random.Random(config.EXAMPLE_SAMPLING_SEED) @@ -1719,7 +1719,7 @@ def _validate_state(self) -> None: f"retained {len(self.examples)}, expected {expected} for {self.num_seen} seen" ) - def add_example(self, example: Example): + def add_example(self, example: Example) -> None: if len(self.group_names) == 0: assert len(self.examples) == 0 self.group_names = example.get_group_names() diff --git a/traincheck/parallel.py b/traincheck/parallel.py index b1c35dac..943922ef 100644 --- a/traincheck/parallel.py +++ b/traincheck/parallel.py @@ -13,13 +13,26 @@ access the traces through ``get_worker_traces()``. """ +from typing import TypedDict + import traincheck.config.config as config import traincheck.utils as _tc_utils +from traincheck.trace.trace import Trace + + +class InferenceConfigSnapshot(TypedDict): + """Required runtime settings transferred between our parent and workers.""" + + ENABLE_PRECOND_SAMPLING: bool + PRECOND_SAMPLING_THRESHOLD: int + MAX_EXAMPLE_LIST_SIZE: int | None + EXAMPLE_SAMPLING_SEED: int + -_WORKER_TRACES: list | None = None +_WORKER_TRACES: list[Trace] | None = None -def config_snapshot() -> dict: +def config_snapshot() -> InferenceConfigSnapshot: """Runtime inference config values that must be propagated to workers. These include CLI sampling options and programmatic reservoir settings. @@ -35,7 +48,7 @@ def config_snapshot() -> dict: } -def worker_init(traces: list, snapshot: dict) -> None: +def worker_init(traces: list[Trace], snapshot: InferenceConfigSnapshot) -> None: """Initializer run once per worker process. Stores the traces in a module global (one copy per worker, not per task) and @@ -50,7 +63,7 @@ def worker_init(traces: list, snapshot: dict) -> None: _tc_utils._suppress_inner_progress = True -def get_worker_traces() -> list: +def get_worker_traces() -> list[Trace]: """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 fa4e498e5bcfba9844081ca9b077a9b9ce533be2 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 21:32:42 +0800 Subject: [PATCH 5/5] Validate complete worker snapshots before changing process state --- tests/test_parallel_config.py | 77 +++++++++++++++++++++++++++++++++++ traincheck/parallel.py | 63 +++++++++++++++++++++++----- 2 files changed, 130 insertions(+), 10 deletions(-) create mode 100644 tests/test_parallel_config.py diff --git a/tests/test_parallel_config.py b/tests/test_parallel_config.py new file mode 100644 index 00000000..6e534669 --- /dev/null +++ b/tests/test_parallel_config.py @@ -0,0 +1,77 @@ +"""Worker settings are required and validated before any process state changes.""" + +import unittest +from contextlib import ExitStack +from unittest.mock import patch + +import traincheck.config.config as config +import traincheck.parallel as parallel +import traincheck.utils as utils +from traincheck.trace.trace import Trace + + +class WorkerConfigTests(unittest.TestCase): + def setUp(self) -> None: + stack = ExitStack() + self.addCleanup(stack.close) + stack.enter_context(patch.multiple(config, **parallel.config_snapshot())) + self.original_traces: list[Trace] = [] + stack.enter_context( + patch.object(parallel, "_WORKER_TRACES", self.original_traces) + ) + stack.enter_context(patch.object(utils, "_suppress_inner_progress", False)) + + def assert_rejected_unchanged( + self, snapshot: dict[str, object], error: type[Exception] + ) -> None: + before = parallel.config_snapshot() + with self.assertRaises(error): + parallel.worker_init([], snapshot) + self.assertEqual(parallel.config_snapshot(), before) + self.assertIs(parallel.get_worker_traces(), self.original_traces) + self.assertFalse(utils._suppress_inner_progress) + + def test_all_fields_are_required(self) -> None: + for field in parallel.config_snapshot(): + snapshot: dict[str, object] = dict(parallel.config_snapshot()) + del snapshot[field] + with self.subTest(field=field): + self.assert_rejected_unchanged(snapshot, KeyError) + + def test_unknown_fields_are_rejected(self) -> None: + snapshot: dict[str, object] = dict(parallel.config_snapshot()) + snapshot["UNRECOGNIZED_SETTING"] = 1 + self.assert_rejected_unchanged(snapshot, ValueError) + + def test_malformed_fields_leave_all_state_unchanged(self) -> None: + invalid: dict[str, list[object]] = { + "ENABLE_PRECOND_SAMPLING": [1, None, "true"], + "PRECOND_SAMPLING_THRESHOLD": [True, 1.5, "100"], + "MAX_EXAMPLE_LIST_SIZE": [True, -1, 2.5, "10"], + "EXAMPLE_SAMPLING_SEED": [object(), {}, []], + } + for field, values in invalid.items(): + for value in values: + snapshot: dict[str, object] = dict(parallel.config_snapshot()) + snapshot[field] = value + with self.subTest(field=field, value=value): + self.assert_rejected_unchanged(snapshot, ValueError) + + def test_supported_seeds_and_uncapped_or_zero_capacity(self) -> None: + for seed in (None, 9, -8, 1.5, "seed", b"seed", bytearray(b"seed")): + for capacity in (None, 0, 7): + snapshot: dict[str, object] = { + "ENABLE_PRECOND_SAMPLING": False, + "PRECOND_SAMPLING_THRESHOLD": 0, + "MAX_EXAMPLE_LIST_SIZE": capacity, + "EXAMPLE_SAMPLING_SEED": seed, + } + traces: list[Trace] = [] + parallel.worker_init(traces, snapshot) + self.assertEqual(parallel.config_snapshot(), snapshot) + self.assertIs(parallel.get_worker_traces(), traces) + self.assertTrue(utils._suppress_inner_progress) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/parallel.py b/traincheck/parallel.py index 943922ef..218d40d9 100644 --- a/traincheck/parallel.py +++ b/traincheck/parallel.py @@ -13,12 +13,15 @@ access the traces through ``get_worker_traces()``. """ -from typing import TypedDict +from collections.abc import Mapping +from typing import TypeAlias, TypedDict import traincheck.config.config as config import traincheck.utils as _tc_utils from traincheck.trace.trace import Trace +RandomSeed: TypeAlias = int | float | str | bytes | bytearray | None + class InferenceConfigSnapshot(TypedDict): """Required runtime settings transferred between our parent and workers.""" @@ -26,7 +29,7 @@ class InferenceConfigSnapshot(TypedDict): ENABLE_PRECOND_SAMPLING: bool PRECOND_SAMPLING_THRESHOLD: int MAX_EXAMPLE_LIST_SIZE: int | None - EXAMPLE_SAMPLING_SEED: int + EXAMPLE_SAMPLING_SEED: RandomSeed _WORKER_TRACES: list[Trace] | None = None @@ -48,18 +51,58 @@ def config_snapshot() -> InferenceConfigSnapshot: } -def worker_init(traces: list[Trace], snapshot: InferenceConfigSnapshot) -> None: - """Initializer run once per worker process. +def _validate_snapshot(snapshot: Mapping[str, object]) -> InferenceConfigSnapshot: + """Reject incomplete/corrupt worker settings before changing process state.""" + required = { + "ENABLE_PRECOND_SAMPLING", + "PRECOND_SAMPLING_THRESHOLD", + "MAX_EXAMPLE_LIST_SIZE", + "EXAMPLE_SAMPLING_SEED", + } + unexpected = set(snapshot) - required + if unexpected: + raise ValueError(f"Unexpected inference settings: {sorted(unexpected)}") + # Direct access makes every field required; missing fields retain KeyError. + enabled = snapshot["ENABLE_PRECOND_SAMPLING"] + threshold = snapshot["PRECOND_SAMPLING_THRESHOLD"] + capacity = snapshot["MAX_EXAMPLE_LIST_SIZE"] + seed = snapshot["EXAMPLE_SAMPLING_SEED"] + if type(enabled) is not bool: # noqa: E721 - booleans are not integer settings + raise ValueError("ENABLE_PRECOND_SAMPLING must be a bool") + if type(threshold) is not int: # noqa: E721 - match the integer CLI contract + raise ValueError("PRECOND_SAMPLING_THRESHOLD must be an integer") + if capacity is not None and ( + type(capacity) is not int or capacity < 0 # noqa: E721 - reservoir contract + ): + raise ValueError("MAX_EXAMPLE_LIST_SIZE must be a nonnegative integer or None") + # Preserve the documented random.Random seed types, including None and bytes. + if seed is not None and not isinstance(seed, (int, float, str, bytes, bytearray)): + raise ValueError("EXAMPLE_SAMPLING_SEED must be a supported random.Random seed") + return { + "ENABLE_PRECOND_SAMPLING": enabled, + "PRECOND_SAMPLING_THRESHOLD": threshold, + "MAX_EXAMPLE_LIST_SIZE": capacity, + "EXAMPLE_SAMPLING_SEED": seed, + } + + +def worker_init(traces: list[Trace], snapshot: Mapping[str, object]) -> None: + """Validate required settings, then initialize this worker atomically. - Stores the traces in a module global (one copy per worker, not per task) and - restores runtime inference config values (see config_snapshot). Also suppresses - the per-invariant inner progress bars, otherwise every worker would fight - over the terminal. + A malformed snapshot must not install some settings or replace worker traces. + The parent supplies config_snapshot(); accepting a Mapping also lets this + process boundary validate deserialized or otherwise untrusted payloads. """ + settings = _validate_snapshot(snapshot) global _WORKER_TRACES + config.ENABLE_PRECOND_SAMPLING = settings["ENABLE_PRECOND_SAMPLING"] + config.PRECOND_SAMPLING_THRESHOLD = settings["PRECOND_SAMPLING_THRESHOLD"] + # The legacy config module infers int from defaults; programmatic reservoir + # options also support None and the other stdlib seed types. Fixed attribute + # names avoid pretending arbitrary keys are supported configuration fields. + setattr(config, "MAX_EXAMPLE_LIST_SIZE", settings["MAX_EXAMPLE_LIST_SIZE"]) + setattr(config, "EXAMPLE_SAMPLING_SEED", settings["EXAMPLE_SAMPLING_SEED"]) _WORKER_TRACES = traces - for name, value in snapshot.items(): - setattr(config, name, value) _tc_utils._suppress_inner_progress = True