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(