diff --git a/docs/experiments/reservoir-cap-validation.json b/docs/experiments/reservoir-cap-validation.json new file mode 100644 index 00000000..b777c692 --- /dev/null +++ b/docs/experiments/reservoir-cap-validation.json @@ -0,0 +1,262 @@ +{ + "baseline_commit": "9eb4f320d3db34d372145261704b805b88e39bda", + "checkpoint_verified_date": "2026-10-25", + "combined_changed_production_source_sha256": { + "traincheck/infer_engine.py": "476a3ec99bef0ab823b23f94bc95842c20b6d9b7b4a57ab4a2efa2987f30f286", + "traincheck/invariant/base_cls.py": "2a2092626065dd871549dd7c969f085c4bc38fe5fffd92469d2ed775bda084ae", + "traincheck/invariant/contain_relation.py": "b32f819db4d5a11a11318bf7e38c64148f4c9360137ac5f09475b480c34c16a5", + "traincheck/parallel.py": "6843319dd8c5c397e93f361a29fe564161e7307fbe4f0f8368090d252d4a7bc4" + }, + "combined_sampling_pr": 139, + "combined_source_manifest_sha256": "60f431cd032532f6b3521169c29ce3ed3943c30fb5b55449ef7ebb608d6cc811", + "combined_test_source_sha256": { + "tests/test_example_reservoir.py": "0381a81e714980f5cff16699db951d984c4f42cd95b6fa47b1611d4d38c2f636", + "tests/test_inference_sampling.py": "046554588613da74e5dbad524acd9ca9ed7eee9c96611d1341eab3cc6cd16e8e" + }, + "controlled_probe": { + "after": { + "cap": 2, + "constructed_num_seen": 5, + "constructed_retained": 2, + "merged_num_seen": 20, + "merged_retained": 2 + }, + "before": { + "cap": 2, + "constructed_num_seen": 5, + "constructed_retained": 5, + "merged_num_seen": 10, + "merged_retained": 4 + } + }, + "limits": [ + "Per-occurrence marginal inclusion only for correlated source reservoirs; no joint uniformity claim.", + "Optional containment generalization is skipped for already-sampled inputs; generalized hypotheses can change above the cap.", + "Retained reservoirs are bounded; temporary unions and total hypothesis/process memory are not.", + "Saved raw benchmark traces are external to the checked-in regression tests." + ], + "recorded_date": "2026-10-13", + "regression_tests": { + "reservoir_standalone_passed": 21, + "reservoir_with_sampling_passed": 21, + "sampling_with_reservoir_passed": 11, + "start_methods": [ + "fork", + "spawn" + ], + "worker_counts": [ + 1, + 2, + 4 + ] + }, + "replay": { + "comparison": "Full structural JSON multiset; preserves multiplicity and array order", + "complete_runs": 9, + "disabled_relations": [ + "FunctionLeadRelation", + "FunctionCoverRelation" + ], + "input_hashes_unchanged": true, + "native_threads": 1, + "repeats_per_workload_and_workers": 1, + "runs": [ + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "8e39cf710d943898ca007478953584d363db135ce3cb458619e71836ae179820", + "distinct_records": 21, + "records": 21 + }, + "invariants": { + "canonical_multiset_sha256": "5691e461a1fba5bf5314b18a62bf72f82fc95e56f71b26eca87218e53f321e84", + "distinct_records": 193, + "records": 193 + }, + "matches_sampling_only": true, + "workers": 1, + "workload": "mnist" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "8e39cf710d943898ca007478953584d363db135ce3cb458619e71836ae179820", + "distinct_records": 21, + "records": 21 + }, + "invariants": { + "canonical_multiset_sha256": "5691e461a1fba5bf5314b18a62bf72f82fc95e56f71b26eca87218e53f321e84", + "distinct_records": 193, + "records": 193 + }, + "matches_sampling_only": true, + "workers": 2, + "workload": "mnist" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "8e39cf710d943898ca007478953584d363db135ce3cb458619e71836ae179820", + "distinct_records": 21, + "records": 21 + }, + "invariants": { + "canonical_multiset_sha256": "5691e461a1fba5bf5314b18a62bf72f82fc95e56f71b26eca87218e53f321e84", + "distinct_records": 193, + "records": 193 + }, + "matches_sampling_only": true, + "workers": 4, + "workload": "mnist" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "b3424febf9358e01c3e079362e17bc68a939a933e14f1cfab954884482073bd7", + "distinct_records": 10, + "records": 10 + }, + "invariants": { + "canonical_multiset_sha256": "21332736553d2fe20c1fdb1e2c852597950e3b799bc399d47007a2b1f7a71b09", + "distinct_records": 518, + "records": 518 + }, + "matches_sampling_only": true, + "workers": 1, + "workload": "transformer" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "b3424febf9358e01c3e079362e17bc68a939a933e14f1cfab954884482073bd7", + "distinct_records": 10, + "records": 10 + }, + "invariants": { + "canonical_multiset_sha256": "21332736553d2fe20c1fdb1e2c852597950e3b799bc399d47007a2b1f7a71b09", + "distinct_records": 518, + "records": 518 + }, + "matches_sampling_only": true, + "workers": 2, + "workload": "transformer" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "b3424febf9358e01c3e079362e17bc68a939a933e14f1cfab954884482073bd7", + "distinct_records": 10, + "records": 10 + }, + "invariants": { + "canonical_multiset_sha256": "21332736553d2fe20c1fdb1e2c852597950e3b799bc399d47007a2b1f7a71b09", + "distinct_records": 518, + "records": 518 + }, + "matches_sampling_only": true, + "workers": 4, + "workload": "transformer" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "dfba6a7ed565ce66163b9eb98921ad05857649e37f78075b48fdf9341b912204", + "distinct_records": 86, + "records": 86 + }, + "invariants": { + "canonical_multiset_sha256": "596f82fcfb946a7fe2fe138fa9e61302bafb93465fdc529317ccf93d2da1022a", + "distinct_records": 512, + "records": 512 + }, + "matches_sampling_only": true, + "workers": 1, + "workload": "dcgan" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "dfba6a7ed565ce66163b9eb98921ad05857649e37f78075b48fdf9341b912204", + "distinct_records": 86, + "records": 86 + }, + "invariants": { + "canonical_multiset_sha256": "596f82fcfb946a7fe2fe138fa9e61302bafb93465fdc529317ccf93d2da1022a", + "distinct_records": 512, + "records": 512 + }, + "matches_sampling_only": true, + "workers": 2, + "workload": "dcgan" + }, + { + "exit_code": 0, + "failed_hypotheses": { + "canonical_multiset_sha256": "dfba6a7ed565ce66163b9eb98921ad05857649e37f78075b48fdf9341b912204", + "distinct_records": 86, + "records": 86 + }, + "invariants": { + "canonical_multiset_sha256": "596f82fcfb946a7fe2fe138fa9e61302bafb93465fdc529317ccf93d2da1022a", + "distinct_records": 512, + "records": 512 + }, + "matches_sampling_only": true, + "workers": 4, + "workload": "dcgan" + } + ], + "source_hashes_unchanged": true + }, + "runtime": { + "packages": { + "numpy": "1.26.4", + "pandas": "2.3.3", + "torch": "2.2.2+cu121" + }, + "platform": "Linux", + "python": "3.10.16" + }, + "saved_evidence_sha256": { + "baseline-probe.json": "69d232dbeae11972ddbb9d27abb153d5f88dd0e2025417c041800583740d9289", + "cap-source-manifest.json": "6910fabc4a993976f9cf90cd98d68d2797e70ff306fb5ff9d4052fab2fb2f8b1", + "cap-tests.log": "7f512ee57813789d5bb907cdccbceaeea7782c95cf95b7dac62fccf9448e3a82", + "combined-reservoir-tests.log": "6b0f38d6f17bafb808ba505c604ff60306a6362452233d22b6c77700a0210dbd", + "combined-sampling-tests.log": "dac8a7daaff443800210eeecbadfa8b257e9cbbe7ccd30859f8951c7c0a9246c", + "combined-source-manifest.json": "60f431cd032532f6b3521169c29ce3ed3943c30fb5b55449ef7ebb608d6cc811", + "fixed-probe.json": "72161a6b9fe85d4b0df03f9ae4b0921bbe27def313da4ffa738fef10ca8a7304", + "replay/summary.json": "aa081af685c92f9e62c3ca7f72b9082cd2ed61af7c4f1a37b543eb8e709d914c", + "sampling-only-comparison.json": "e526d5cc2c9f8fa37c2aeef21a05e9481338dc142789f27edff129d3121d0a15" + }, + "schema_version": 1, + "standalone_source_sha256": { + "tests/test_example_reservoir.py": "0381a81e714980f5cff16699db951d984c4f42cd95b6fa47b1611d4d38c2f636", + "traincheck/infer_engine.py": "741ac7701ad4fef3089fce778ddd4d422c5245432cf582bec06095e0ed6ec71c", + "traincheck/invariant/base_cls.py": "2a2092626065dd871549dd7c969f085c4bc38fe5fffd92469d2ed775bda084ae", + "traincheck/invariant/contain_relation.py": "67a6613f0048542174f6655fcb1f984fabf5c5f2fa9303eab13aed93103b410d", + "traincheck/parallel.py": "6843319dd8c5c397e93f361a29fe564161e7307fbe4f0f8368090d252d4a7bc4" + }, + "trace_inputs": { + "dcgan": { + "file_count": 1, + "file_sha256": [ + "da74105c67ff699b91d96eac1fac89f3626ba9199594c9c373eb1ff6a69b15d1" + ], + "total_bytes": 3401128 + }, + "mnist": { + "file_count": 1, + "file_sha256": [ + "6e3dd86b5ed0d7b5d556a167eaf95226852dd816a33f872668fbaf1b26638c20" + ], + "total_bytes": 472863 + }, + "transformer": { + "file_count": 1, + "file_sha256": [ + "467e587463e72e874f5abfd14452cabdf966fd804f448eecf85f0155ea49ffb6" + ], + "total_bytes": 1564970 + } + } +} diff --git a/docs/experiments/reservoir-cap-validation.md b/docs/experiments/reservoir-cap-validation.md new file mode 100644 index 00000000..4e71de54 --- /dev/null +++ b/docs/experiments/reservoir-cap-validation.md @@ -0,0 +1,46 @@ +# Example reservoir cap validation + +Iterable construction and stage/cross-trace merging previously appended retained examples directly. This bypassed `MAX_EXAMPLE_LIST_SIZE` and left the represented population count stale. The fix uses bounded insertion and merges according to the original stream populations, preserving duplicate occurrences and exact concatenation order when complete inputs fit. + +Results below were recorded on 2026-10-13. The pending source and saved evidence hashes were independently rechecked on 2026-10-25. The baseline is `9eb4f320d3db34d372145261704b805b88e39bda`; combined validation includes the deterministic sampling change in #139. + +## Controlled fixture + +The fixture sets capacity to 2, merges two streams containing ten occurrences each, and separately constructs a reservoir from five occurrences. + +| Measurement | Before | After | +|---|---:|---:| +| Retained after merge | 4 | 2 | +| `num_seen` after merge | 10 | 20 | +| Retained from five-item iterable | 5 | 2 | +| `num_seen` for iterable | 5 | 5 | + +These measure reservoir correctness. They do not measure whole-process memory savings. + +## Regression and replay checks + +Run the portable regression suite from the repository root: + +```bash +python -m unittest discover -s tests -p test_example_reservoir.py -v +``` + +All 21 reservoir tests passed on Linux with Python 3.10.16, PyTorch 2.2.2+cu121, pandas 2.3.3, and NumPy 1.26.4. Coverage includes cap/count accounting, duplicates, atomic rejection, continuation RNG state, stage/cross-trace call sites, actual inference with 1/2/4 workers under fork and spawn, and fixed-seed marginal-inclusion sweeps. With #139, all 21 reservoir and 11 sampling tests passed; its sampling fixture explicitly disables the cap to exercise precondition downsampling separately. + +Nine complete saved-trace inference replays used one pass per workload and worker count, with native thread counts set to 1 and `FunctionLeadRelation`/`FunctionCoverRelation` disabled. Every invariant and failed-hypothesis JSON multiset matched across workers and against sampling-only outputs; multiplicity and array order were retained in comparisons. + +| Workload | Workers tested | Invariants | Failed hypotheses | +|---|---|---:|---:| +| MNIST | 1, 2, 4 | 193 | 21 | +| Small transformer | 1, 2, 4 | 518 | 10 | +| DCGAN | 1, 2, 4 | 512 | 86 | + +Input and source hashes remained unchanged. Compact counts, canonical output hashes, and source/evidence hashes are recorded in [reservoir-cap-validation.json](reservoir-cap-validation.json). Its combined-source manifest and test hashes identify the validated combination with #139, including the uncapped sampling fixture. Raw saved traces are external to this note; the synthetic regression suite runs directly from the checkout. + +Publication validation on 2026-10-25 reran all 32 combined tests successfully on macOS. Production Python ASTs match the previously validated combined source; publication changes are formatting and narrow lint comments for intentional exact-integer checks. + +## Scope + +The merge RNG must advance independently of source selection and must not reset across nested merges. Existing equal-seed source reservoirs can still be correlated: the guarantee is marginal occurrence inclusion, not a uniform joint sample across correlated reservoirs. + +Optional containment generalization is skipped when an input has already been sampled, because discarded set overlap cannot be reconstructed. Complete inputs retain exact set deduplication before capping the known union. Consequently, generalized hypotheses can change above the cap. Temporary unions, total hypothesis counts, and total process memory are outside this per-reservoir bound. 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/tests/test_inference_sampling.py b/tests/test_inference_sampling.py index bfd87977..b861ebf5 100644 --- a/tests/test_inference_sampling.py +++ b/tests/test_inference_sampling.py @@ -102,17 +102,22 @@ def examples(items): 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}, - ) + # This fixture deliberately reaches the >5000 precondition branch; + # disable the independent reservoir cap while constructing its input. + with patch("traincheck.config.config.MAX_EXAMPLE_LIST_SIZE", None): + 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( 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/infer_engine.py b/traincheck/infer_engine.py index d98182a5..4e48a4c3 100644 --- a/traincheck/infer_engine.py +++ b/traincheck/infer_engine.py @@ -5,6 +5,7 @@ import os import time from concurrent.futures import ProcessPoolExecutor, as_completed +from random import Random from tqdm import tqdm @@ -132,14 +133,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) @@ -153,6 +159,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. @@ -166,31 +173,19 @@ 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[type[Relation]], hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], hypo_lookup: dict[Hypothesis, Hypothesis], + merge_rng: Random, ) -> None: """Serial generation with the same per-task sampling as the worker path.""" n_traces = len(self.traces) @@ -225,6 +220,7 @@ def _generate_hypothesis_serial( trace_idx, hypotheses_and_trace_idxs, hypo_lookup, + merge_rng, ) def _generate_hypothesis_parallel( @@ -232,6 +228,7 @@ def _generate_hypothesis_parallel( active_relations: list[type[Relation]], hypotheses_and_trace_idxs: dict[Hypothesis, list[int]], hypo_lookup: dict[Hypothesis, Hypothesis], + merge_rng: Random, ) -> None: """Parallel hypothesis generation. @@ -284,6 +281,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..d2f340d2 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,20 +1677,49 @@ def __eq__(self, value: object) -> bool: class ExampleList: - def __init__(self, group_names: set[str]): + """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]) -> None: 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): + 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) + + @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) -> None: if len(self.group_names) == 0: assert len(self.examples) == 0 self.group_names = example.get_group_names() @@ -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 cdc82829..b68231fc 100644 --- a/traincheck/invariant/contain_relation.py +++ b/traincheck/invariant/contain_relation.py @@ -1,10 +1,12 @@ import copy import logging 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.inference_sampling import sampling_rng from traincheck.instrumentor.tracer import TraceLineType @@ -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..218d40d9 100644 --- a/traincheck/parallel.py +++ b/traincheck/parallel.py @@ -13,42 +13,100 @@ access the traces through ``get_worker_traces()``. """ +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.""" + + ENABLE_PRECOND_SAMPLING: bool + PRECOND_SAMPLING_THRESHOLD: int + MAX_EXAMPLE_LIST_SIZE: int | None + EXAMPLE_SAMPLING_SEED: RandomSeed + -_WORKER_TRACES: list | None = None +_WORKER_TRACES: list[Trace] | None = None -def config_snapshot() -> dict: - """CLI-overridable config values that must be propagated to workers. +def config_snapshot() -> InferenceConfigSnapshot: + """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, + } + + +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, snapshot: dict) -> None: - """Initializer run once per worker process. +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 CLI-overridable 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 -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