diff --git a/tests/test_selective_argument_kinds.py b/tests/test_selective_argument_kinds.py new file mode 100644 index 00000000..8f7da61d --- /dev/null +++ b/tests/test_selective_argument_kinds.py @@ -0,0 +1,321 @@ +"""Unsupported argument selectors must retain complete checker evidence.""" + +from __future__ import annotations + +import copy +import inspect +import json +import tempfile +import unittest +from collections.abc import Callable +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +from tests.test_selective_context_requirements import ( + JsonTraceDestination, + policy_wrapper, + string_map, +) +from tools.fresh_validation.check_coverage import ObservationReport, observe +from traincheck.collect_trace import get_per_func_instr_opts +from traincheck.config.config import InstrOpt +from traincheck.instrumentor import dumper, tracer +from traincheck.invariant.base_cls import ( + APIParam, + Arguments, + GroupedPreconditions, + Invariant, + Preconditions, + UnconditionalPrecondition, + load_function_signature, +) +from traincheck.invariant.contain_relation import ( + APIContainRelation, + cache_events_scanner, +) +from traincheck.trace.trace_pandas import read_trace_file_Pandas +from traincheck.utils import typename + + +def positional_child(value: int, /, other: int = 0) -> None: + pass + + +def mixed_child(value: int, /, other: int = 0, *, flag: bool = False) -> None: + pass + + +def variadic_child(receiver: str, /, *values: int, **options: int) -> None: + pass + + +def keyword_child(**options: int) -> None: + pass + + +def unknown_child(value: int) -> None: + pass + + +def parent_call( + child: Callable[..., object], + arguments: tuple[object, ...], + keywords: dict[str, object], +) -> None: + child(*arguments, **keywords) + + +# Discovery with `-s tests` imports this module without the `tests.` prefix. +# TrainCheck resolves fixture signatures from their package-qualified API names, +# so use the canonical module's actual callables in both unittest invocations. +# This leaves the production loader and the functions' metadata unchanged. +if __name__ != "tests.test_selective_argument_kinds": + from tests import test_selective_argument_kinds as _fixtures + + positional_child = _fixtures.positional_child + mixed_child = _fixtures.mixed_child + variadic_child = _fixtures.variadic_child + keyword_child = _fixtures.keyword_child + unknown_child = _fixtures.unknown_child + parent_call = _fixtures.parent_call + + +def containment_rule(child: Callable[..., object], arguments: Arguments) -> Invariant: + return Invariant( + APIContainRelation, + [ + APIParam(typename(parent_call)), + APIParam(typename(child), arguments=arguments), + ], + GroupedPreconditions( + {"parent_func_call_pre": Preconditions([UnconditionalPrecondition()])} + ), + "parent calls child with the required argument metadata", + ) + + +def required_arguments(rule: Invariant) -> Arguments: + child = rule.params[1] + assert isinstance(child, APIParam) + arguments = child.arguments + assert isinstance(arguments, Arguments) + return arguments + + +def collect_and_check( + rule: Invariant, + child: Callable[..., object], + arguments: tuple[object, ...], + keywords: dict[str, object], +) -> tuple[ObservationReport, list[dict[str, object]]]: + loaded: object = InstrOpt.from_json( + InstrOpt(get_per_func_instr_opts([rule]), None, True).to_json() + ) + assert isinstance(loaded, InstrOpt) + policy = string_map(loaded.funcs_instr_opts) + destination = JsonTraceDestination() + with ( + patch.object(tracer.config, "DISABLE_WRAPPER", False), + patch.object( + dumper, "THREAD_DATA", SimpleNamespace(api_trace_buffer=destination) + ), + patch.object( + tracer, "get_meta_vars", return_value={"step": 1, "stage": "training"} + ), + ): + wrapped_child = policy_wrapper( + child, typename(child), string_map(policy[typename(child)]) + ) + wrapped_parent = policy_wrapper( + parent_call, + typename(parent_call), + string_map(policy[typename(parent_call)]), + ) + wrapped_parent(wrapped_child, arguments, keywords) + records = [string_map(json.loads(line)) for line in destination.lines] + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "trace_API_arguments.jsonl" + path.write_text("\n".join(destination.lines) + "\n") + trace = read_trace_file_Pandas(str(path), truncate_incomplete_func_calls=False) + cache_events_scanner.clear() + try: + report = observe(trace, rule) + finally: + cache_events_scanner.clear() + return report, records + + +class SelectiveArgumentKindsTests(unittest.TestCase): + def test_fixture_signatures_resolve_under_both_unittest_invocations(self) -> None: + fixtures: tuple[Callable[..., object], ...] = ( + positional_child, + mixed_child, + variadic_child, + keyword_child, + unknown_child, + parent_call, + ) + for function in fixtures: + with self.subTest(function=typename(function)): + name = typename(function) + self.assertTrue(name.startswith("tests.test_selective_argument_kinds.")) + self.assertEqual( + load_function_signature(name), inspect.signature(function) + ) + + def test_positional_only_trace_passes_and_detects_changed_argument(self) -> None: + child = typename(positional_child) + rule = containment_rule( + positional_child, Arguments({"0": {"int": 7}}, {}, child) + ) + original = copy.deepcopy(rule.to_dict()) + for value, expected in ((7, True), (8, False)): + with self.subTest(value=value): + report, records = collect_and_check( + rule, positional_child, (value,), {} + ) + self.assertEqual(report["status"], "completed", report) + self.assertIs(report["active"], True, report) + self.assertEqual(report["decisions_completed"], 1) + self.assertIs(report["check_passed"], expected) + before = [ + record + for record in records + if record["function"] == child + and record["type"] == tracer.TraceLineType.FUNC_CALL_PRE + ] + self.assertEqual(len(before), 1) + self.assertEqual(string_map(before[0]["args"])["0"], {"int": value}) + self.assertEqual(rule.to_dict(), original) + + def test_variadic_positional_trace_passes_and_detects_changed_argument( + self, + ) -> None: + child = typename(variadic_child) + rule = containment_rule(variadic_child, Arguments({"1": {"int": 7}}, {}, child)) + for value, expected in ((7, True), (8, False)): + with self.subTest(value=value): + report, records = collect_and_check( + rule, variadic_child, ("receiver", value), {} + ) + self.assertIs(report["active"], True, report) + self.assertIs(report["check_passed"], expected, report) + before = [ + record + for record in records + if record["function"] == child + and record["type"] == tracer.TraceLineType.FUNC_CALL_PRE + ] + self.assertEqual(string_map(before[0]["args"])["1"], {"int": value}) + + def test_variadic_keyword_trace_passes_and_detects_changed_argument(self) -> None: + child = typename(keyword_child) + rule = containment_rule( + keyword_child, Arguments({}, {"options": {"int": 7}}, child) + ) + for value, expected in ((7, True), (8, False)): + with self.subTest(value=value): + report, records = collect_and_check( + rule, keyword_child, (), {"options": value} + ) + self.assertIs(report["active"], True, report) + self.assertIs(report["check_passed"], expected, report) + before = [ + record + for record in records + if record["function"] == child + and record["type"] == tracer.TraceLineType.FUNC_CALL_PRE + ] + self.assertEqual( + string_map(before[0]["kwargs"])["options"], {"int": value} + ) + + def test_full_fallback_dominates_supported_subset_in_either_rule_order( + self, + ) -> None: + name = typename(positional_child) + full = containment_rule( + positional_child, Arguments({"0": {"int": 7}}, {}, name) + ) + narrow = containment_rule( + positional_child, Arguments({}, {"other": {"int": 9}}, name) + ) + expected = get_per_func_instr_opts([full, narrow]) + self.assertNotIn("dump_args_config", expected[name]) + self.assertIs(expected[name]["dump_args"], True) + self.assertEqual(get_per_func_instr_opts([narrow, full]), expected) + + def test_required_supported_kinds_remain_selective(self) -> None: + name = typename(mixed_child) + arguments = Arguments({}, {"other": {"int": 9}, "flag": {"bool": True}}, name) + rule = containment_rule(mixed_child, arguments) + options = get_per_func_instr_opts([rule])[name] + self.assertEqual( + options["dump_args_config"], {"other": None, 1: None, "flag": None} + ) + for flag, expected in ((True, True), (False, False)): + report, _ = collect_and_check(rule, mixed_child, (7, 9), {"flag": flag}) + self.assertIs(report["active"], True, report) + self.assertIs(report["check_passed"], expected, report) + + def test_unknown_signature_requests_complete_arguments(self) -> None: + name = typename(unknown_child) + arguments = Arguments({"0": {"int": 7}}, {}, name) + arguments.signature = None + options = get_per_func_instr_opts([containment_rule(unknown_child, arguments)])[ + name + ] + self.assertTrue(options["dump_args"]) + self.assertNotIn("dump_args_config", options) + + def test_unresolved_required_name_requests_complete_arguments(self) -> None: + name = typename(unknown_child) + arguments = Arguments({}, {"unexpected": {"int": 7}}, name) + options = get_per_func_instr_opts([containment_rule(unknown_child, arguments)])[ + name + ] + self.assertTrue(options["dump_args"]) + self.assertNotIn("dump_args_config", options) + + def test_malformed_signature_or_argument_metadata_is_rejected(self) -> None: + name = typename(positional_child) + cases: list[tuple[str, object]] = [ + ("signature", "not a signature"), + ("arguments", []), + ("arguments", {7: {"int": 7}}), + ("arguments", {"": {"int": 7}}), + ("arguments", {"value": []}), + ("arguments", {"value": {}}), + ("arguments", {"value": {"int": 7, "str": "7"}}), + ("func_name", ""), + ("func_name", typename(unknown_child)), + ] + for field, malformed in cases: + with self.subTest(field=field, value=malformed): + arguments = Arguments({"0": {"int": 7}}, {}, name) + setattr(arguments, field, malformed) + with self.assertRaises((TypeError, ValueError)): + get_per_func_instr_opts( + [containment_rule(positional_child, arguments)] + ) + + def test_malformed_later_requirement_is_not_hidden_by_fallback(self) -> None: + name = typename(positional_child) + arguments = Arguments({"0": {"int": 7}}, {}, name) + setattr(arguments, "arguments", {"value": {"int": 7}, "other": []}) + with self.assertRaises(TypeError): + get_per_func_instr_opts([containment_rule(positional_child, arguments)]) + + def test_no_required_arguments_does_not_force_full_collection(self) -> None: + name = typename(unknown_child) + arguments = Arguments({}, {}, name) + arguments.signature = None + options = get_per_func_instr_opts([containment_rule(unknown_child, arguments)])[ + name + ] + self.assertEqual(options["dump_args_config"], {}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_selective_context_requirements.py b/tests/test_selective_context_requirements.py new file mode 100644 index 00000000..2fa6691d --- /dev/null +++ b/tests/test_selective_context_requirements.py @@ -0,0 +1,340 @@ +"""Context-dependent selective policies must preserve real checker activation.""" + +from __future__ import annotations + +import copy +import json +import sys +import tempfile +import unittest +from collections.abc import Callable, Mapping +from contextlib import ExitStack +from pathlib import Path +from types import SimpleNamespace, TracebackType +from typing import Literal +from unittest.mock import patch + +from tools.fresh_validation.check_coverage import ObservationReport, observe +from traincheck.collect_trace import get_per_func_instr_opts +from traincheck.config.config import InstrOpt +from traincheck.instrumentor import dumper, tracer +from traincheck.invariant.base_cls import ( + PT, + APIParam, + Arguments, + GroupedPreconditions, + Invariant, + Precondition, + PreconditionClause, + Preconditions, + UnconditionalPrecondition, +) +from traincheck.invariant.contain_relation import ( + APIContainRelation, + cache_events_scanner, +) +from traincheck.trace.trace_pandas import read_trace_file_Pandas + +CONTEXT = f"{__name__}.DeviceContext" +PROPERTY = f"meta_vars.context_managers.{CONTEXT}" +PARENT = f"{__name__}.context_parent" +CHILD = f"{__name__}.context_child" + + +class DeviceContext: + def __init__(self, device: int) -> None: + self.device = device + + def __enter__(self) -> DeviceContext: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> Literal[False]: + return False + + +def context_child() -> None: + return None + + +def context_parent(include_child: bool) -> None: + if include_child: + context_child() + + +def context_rule(*, inverted: bool = False, prop_name: str = PROPERTY) -> Invariant: + return Invariant( + APIContainRelation, + [APIParam(PARENT), APIParam(CHILD)], + GroupedPreconditions( + { + "parent_func_call_pre": Preconditions( + [ + Precondition( + [ + PreconditionClause( + prop_name, None, PT.EXIST, None, None + ), + PreconditionClause( + prop_name, int, PT.CONSTANT, ("device",), {0} + ), + ] + ) + ], + inverted=inverted, + ) + } + ), + "parent calls child while device context is zero", + ) + + +def string_map(value: object) -> dict[str, object]: + if not isinstance(value, Mapping): + raise TypeError("Expected string-keyed object") + result: dict[str, object] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError("Expected string-keyed object") + result[key] = item + return result + + +def policy_wrapper( + original: Callable[..., object], + name: str, + options: Mapping[str, object], + *, + bound: bool = False, + constructor: bool = False, +) -> Callable[..., object]: + dump_args = options["dump_args"] + dump_ret = options["dump_ret"] + if not isinstance(dump_args, bool) or not isinstance(dump_ret, bool): + raise TypeError("Dump options must be boolean") + # Absent selector is the documented request for complete inputs. + selector = ( + dict(string_map(options["dump_args_config"])) + if "dump_args_config" in options + else None + ) + + def wrapped(*args: object, **kwargs: object) -> object: + return tracer.function_wrapper( + original, + name, + bound, + False, + False, + dump_args, + selector, + dump_ret, + None, + False, + {}, + False, + constructor, + "self" if constructor else None, + *args, + **kwargs, + ) + + return wrapped + + +class JsonTraceDestination: + """Capture the production dumper's serialized bytes without background IO.""" + + def __init__(self) -> None: + self.lines: list[str] = [] + + def add_trace(self, record: dict[str, object]) -> None: + self.lines.append(dumper.serialize(record)) + + +def collect_and_check( + rule: Invariant, *, include_child: bool, device: int = 0 +) -> tuple[ObservationReport, list[dict[str, object]]]: + saved: object = InstrOpt.from_json( + InstrOpt(get_per_func_instr_opts([rule]), None, True).to_json() + ) + assert isinstance(saved, InstrOpt) + policy = string_map(saved.funcs_instr_opts) + destination = JsonTraceDestination() + module = sys.modules[__name__] + with ExitStack() as stack: + stack.enter_context(patch.object(tracer.config, "DISABLE_WRAPPER", False)) + stack.enter_context( + patch.object( + dumper, "THREAD_DATA", SimpleNamespace(api_trace_buffer=destination) + ) + ) + stack.enter_context( + patch.object( + tracer, + "get_meta_vars", + return_value={"step": 1, "stage": "training"}, + ) + ) + for method_name in ("__init__", "__enter__", "__exit__"): + name = f"{CONTEXT}.{method_name}" + if name in policy: + original = getattr(DeviceContext, method_name) + stack.enter_context( + patch.object( + DeviceContext, + method_name, + policy_wrapper( + original, + name, + string_map(policy[name]), + bound=True, + constructor=method_name == "__init__", + ), + ) + ) + stack.enter_context( + patch.object( + module, + "context_child", + policy_wrapper(context_child, CHILD, string_map(policy[CHILD])), + ) + ) + stack.enter_context( + patch.object( + module, + "context_parent", + policy_wrapper(context_parent, PARENT, string_map(policy[PARENT])), + ) + ) + with DeviceContext(device): + context_parent(include_child) + # No timestamp, identity, payload, or row-order rewriting: use exactly the + # production dumper output through the ordinary JSONL parser. + records = [dict(string_map(json.loads(line))) for line in destination.lines] + with tempfile.TemporaryDirectory() as directory: + trace_file = Path(directory) / "trace_API_context.jsonl" + trace_file.write_text("\n".join(destination.lines) + "\n") + trace = read_trace_file_Pandas( + str(trace_file), truncate_incomplete_func_calls=False + ) + cache_events_scanner.clear() + try: + report = observe(trace, rule) + finally: + cache_events_scanner.clear() + return report, records + + +class SelectiveContextRequirementsTests(unittest.TestCase): + def test_real_context_trace_activates_healthy_rule(self) -> None: + rule = context_rule() + original = copy.deepcopy(rule.to_dict()) + report, records = collect_and_check(rule, include_child=True) + self.assertEqual(report["status"], "completed", report) + self.assertIs(report["active"], True, report) + self.assertEqual(report["decisions_completed"], 1) + self.assertIs(report["check_passed"], True) + constructor = [ + record + for record in records + if record["function"] == f"{CONTEXT}.__init__" + and record["type"] == tracer.TraceLineType.FUNC_CALL_PRE + ] + self.assertEqual(len(constructor), 1) + arguments = string_map(constructor[0]["args"]) + self.assertEqual(arguments["1"], {"int": 0}) + self.assertEqual(rule.to_dict(), original) + + def test_real_context_trace_detects_missing_child(self) -> None: + report, _ = collect_and_check(context_rule(), include_child=False) + self.assertEqual(report["status"], "completed", report) + self.assertIs(report["active"], True, report) + self.assertEqual(report["decisions_completed"], 1) + self.assertIs(report["check_passed"], False) + + def test_inverted_context_requires_same_inputs_and_truth_is_preserved(self) -> None: + report, _ = collect_and_check(context_rule(inverted=True), include_child=False) + self.assertEqual(report["status"], "completed", report) + self.assertIs(report["dormant"], True, report) + self.assertEqual(report["decisions_completed"], 0) + active, _ = collect_and_check( + context_rule(inverted=True), include_child=False, device=1 + ) + self.assertIs(active["active"], True, active) + self.assertIs(active["check_passed"], False) + + def test_literal_thread_condition_remains_dormant(self) -> None: + rule = context_rule() + assert rule.precondition is not None + rule.precondition.grouped_preconditions["parent_func_call_pre"].preconditions[ + 0 + ].clauses.append(PreconditionClause("thread_id", int, PT.CONSTANT, None, {-1})) + original = copy.deepcopy(rule.to_dict()) + report, _ = collect_and_check(rule, include_child=False) + self.assertEqual(report["status"], "completed", report) + self.assertIs(report["dormant"], True, report) + self.assertEqual(report["decisions_completed"], 0) + self.assertEqual(rule.to_dict(), original) + + def test_complete_constructor_inputs_override_existing_subset(self) -> None: + constructor = f"{CONTEXT}.__init__" + subset = Invariant( + APIContainRelation, + [ + APIParam(PARENT), + APIParam( + constructor, + arguments=Arguments({"0": {CONTEXT: {}}}, {}, constructor), + ), + ], + GroupedPreconditions( + {"parent_func_call_pre": Preconditions([UnconditionalPrecondition()])} + ), + ) + for rules in ([context_rule(), subset], [subset, context_rule()]): + with self.subTest(order=[rule.text_description for rule in rules]): + options = get_per_func_instr_opts(rules)[constructor] + self.assertIs(options["dump_args"], True) + self.assertNotIn("dump_args_config", options) + for method in ("__enter__", "__exit__"): + self.assertIn(f"{CONTEXT}.{method}", get_per_func_instr_opts(rules)) + + def test_non_context_preconditions_do_not_add_instrumentation(self) -> None: + rule = context_rule(prop_name="meta_vars.other_metadata") + self.assertEqual(set(get_per_func_instr_opts([rule])), {PARENT, CHILD}) + + def test_missing_preconditions_fail_explicitly(self) -> None: + rule = context_rule() + rule.precondition = None + with self.assertRaisesRegex(ValueError, "precondition"): + get_per_func_instr_opts([rule]) + + def test_malformed_context_names_fail_explicitly(self) -> None: + for prop in ( + "meta_vars.context_managers", + "meta_vars.context_managers.", + "meta_vars.context_managers.torch..device", + ): + with self.subTest(property=prop), self.assertRaises(ValueError): + get_per_func_instr_opts([context_rule(prop_name=prop)]) + + def test_malformed_property_type_fails_explicitly(self) -> None: + rule = context_rule() + assert rule.precondition is not None + clause = ( + rule.precondition.grouped_preconditions["parent_func_call_pre"] + .preconditions[0] + .clauses[0] + ) + setattr(clause, "prop_name", 7) + with self.assertRaisesRegex(TypeError, "property"): + get_per_func_instr_opts([rule]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_selective_output_policy.py b/tests/test_selective_output_policy.py new file mode 100644 index 00000000..d6bca57d --- /dev/null +++ b/tests/test_selective_output_policy.py @@ -0,0 +1,333 @@ +"""Selective collection retains the evidence required by API invariants.""" + +import itertools +import json +import logging +import sys +import tempfile +import unittest +from collections.abc import Callable, Sequence +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import traincheck.collect_trace as collect_trace +import traincheck.config.config as config +from traincheck.config.config import InstrOpt +from traincheck.instrumentor import dumper, tracer +from traincheck.invariant.base_cls import ( + APIParam, + Arguments, + GroupedPreconditions, + InputOutputParam, + Invariant, + Param, + Preconditions, + UnconditionalPrecondition, + VarTypeParam, + read_inv_file, +) +from traincheck.invariant.consistency_transient_vars import ( + ConsistentInputOutputRelation, + ConsistentOutputRelation, +) +from traincheck.invariant.contain_relation import APIContainRelation +from traincheck.invariant.DistinctArgumentRelation import ( + EXP_GROUP_NAME, + DistinctArgumentRelation, +) +from traincheck.trace.trace_pandas import TracePandas, read_trace_file_Pandas +from traincheck.utils import typename + + +def tensor_result(value: torch.Tensor) -> torch.Tensor: + return value + 1 + + +def reference_tensor(value: torch.Tensor) -> torch.Tensor: + return value.clone() + + +def distinct_input(shared: str, key: int) -> int: + return key + + +def unconditional(group: str = "pre_event") -> GroupedPreconditions: + return GroupedPreconditions({group: Preconditions([UnconditionalPrecondition()])}) + + +def output_rule(expected: object = "torch.float32") -> Invariant: + return Invariant( + ConsistentOutputRelation, + [ + APIParam(typename(tensor_result)), + VarTypeParam("torch.Tensor", "dtype", const_value=expected), + ], + unconditional(), + "returned tensors have the expected dtype", + ) + + +def io_rule(function: str) -> Invariant: + return Invariant( + ConsistentInputOutputRelation, + [ + InputOutputParam(None, 0, "torch.Tensor", ("dtype",), function, True), + APIParam(function), + InputOutputParam(None, 0, "torch.Tensor", ("dtype",), function, False), + ], + unconditional(), + "input and output dtypes match", + ) + + +def distinct_rule() -> Invariant: + return Invariant( + DistinctArgumentRelation, + [APIParam(typename(distinct_input))], + unconditional(EXP_GROUP_NAME), + "calls in one step have distinct arguments", + ) + + +def argument_rule(function: str) -> Invariant: + return Invariant( + APIContainRelation, + [ + APIParam( + function, + arguments=Arguments({"0": {"str": "shared"}}, {}, function), + ), + APIParam(function), + ], + unconditional("parent_func_call_pre"), + "an independent rule requests only the first argument", + ) + + +class JsonTraceDestination: + def __init__(self) -> None: + self.lines: list[str] = [] + + def add_trace(self, record: dict[str, object]) -> None: + # Replace only the final buffer destination. Keep real metadata + # conversion, timestamping, and production JSON serialization. + self.lines.append(dumper.serialize(record)) + + +def collect_calls( + root: Path, + invariants: list[Invariant], + function: Callable[..., object], + calls: Sequence[tuple[object, ...]], +) -> TracePandas: + reference = io_rule(typename(reference_tensor)) + invariant_file = root / "invariants.jsonl" + invariant_file.write_text( + "".join(json.dumps(inv.to_dict()) + "\n" for inv in [*invariants, reference]) + ) + loaded_invariants = read_inv_file(str(invariant_file)) + options = collect_trace.get_per_func_instr_opts(loaded_invariants) + saved: object = InstrOpt.from_json(InstrOpt(options, None, True).to_json()) + assert isinstance(saved, InstrOpt) + destination = JsonTraceDestination() + with ( + patch.object( + dumper, "THREAD_DATA", SimpleNamespace(api_trace_buffer=destination) + ), + patch.object(config, "DISABLE_WRAPPER", False), + patch.object(config, "INSTR_OPTS", saved), + patch.object( + tracer, + "get_instrumentation_logger_for_process", + return_value=logging.getLogger(__name__), + ), + patch.object( + dumper, + "get_instrumentation_logger_for_process", + return_value=logging.getLogger(__name__), + ), + patch.dict(config.INSTRUMENTATION_POLICY, dict(config.INSTRUMENTATION_POLICY)), + patch.dict(tracer.META_VARS, {"step": 0}, clear=True), + patch.dict( + dumper.tensor_dump_format, + dump_tensor_hash=True, + dump_tensor_stats=False, + dump_tensor_full=False, + ), + ): + instrumentor = tracer.Instrumentor( + sys.modules[__name__], + scan_proxy_in_args=False, + use_full_instr=False, + funcs_to_instr=list(options), + ) + # An ordinary IO-instrumented API ensures missing output values take + # the same normal-loader path as a real trace containing other returns. + instrumentor.get_wrapped_function(reference_tensor)(torch.ones(2)) + wrapped = instrumentor.get_wrapped_function(function) + for arguments in calls: + wrapped(*arguments) + + trace_file = root / "trace_API_fixture.jsonl" + trace_file.write_text("\n".join(destination.lines) + "\n") + return read_trace_file_Pandas(str(trace_file), truncate_incomplete_func_calls=False) + + +class SelectiveOutputPolicyTests(unittest.TestCase): + def test_output_only_rule_requests_full_returns_without_model_tracking( + self, + ) -> None: + invariant = output_rule() + options = collect_trace.get_per_func_instr_opts([invariant])[ + typename(tensor_result) + ] + self.assertTrue(options["dump_ret"]) + self.assertFalse(options["dump_args"]) + self.assertFalse(options["scan_proxy_in_args"]) + self.assertEqual(options["var_types_to_track"], {}) + self.assertNotIn("dump_ret_config", options) + self.assertIsNone( + collect_trace.get_model_tracker_instr_opts([invariant], "subclass") + ) + + def test_return_requirement_is_an_order_independent_union(self) -> None: + function = typename(tensor_result) + rules = [output_rule(), output_rule("torch.float64"), argument_rule(function)] + expected = collect_trace.get_per_func_instr_opts(rules) + for order in itertools.permutations(rules): + with self.subTest(order=[rule.relation.__name__ for rule in order]): + self.assertEqual( + collect_trace.get_per_func_instr_opts(list(order)), expected + ) + self.assertTrue(expected[function]["dump_args"]) + self.assertTrue(expected[function]["dump_ret"]) + self.assertIn("dump_args_config", expected[function]) + self.assertNotIn("dump_ret_config", expected[function]) + + def test_full_arguments_dominate_narrow_selectors_in_both_orders(self) -> None: + function = typename(distinct_input) + for full_rule in (distinct_rule(), io_rule(function)): + for rules in ( + [full_rule, argument_rule(function)], + [argument_rule(function), full_rule], + ): + with self.subTest(relation=full_rule.relation.__name__, order=rules): + options = collect_trace.get_per_func_instr_opts(rules)[function] + self.assertTrue(options["dump_args"]) + self.assertNotIn("dump_args_config", options) + self.assertEqual( + options["dump_ret"], + full_rule.relation == ConsistentInputOutputRelation, + ) + + def test_output_rule_rejects_malformed_selector_schema(self) -> None: + cases: list[list[Param]] = [ + [], + [APIParam("torch.fake")], + [VarTypeParam("torch.Tensor", "dtype"), APIParam("torch.fake")], + [APIParam("torch.fake"), APIParam("torch.other")], + [ + APIParam("torch.fake"), + VarTypeParam("torch.Tensor", "dtype"), + APIParam("torch.extra"), + ], + ] + for params in cases: + with self.subTest(params=params): + invariant = Invariant(ConsistentOutputRelation, params, unconditional()) + with self.assertRaises((TypeError, ValueError)): + collect_trace.get_per_func_instr_opts([invariant]) + + def test_output_rule_rejects_invalid_names_and_missing_constant(self) -> None: + for target, field in ((0, "api_full_name"), (1, "var_type"), (1, "attr_name")): + for value in ("", None, 7): + with self.subTest(field=field, value=value): + invariant = output_rule() + setattr(invariant.params[target], field, value) + with self.assertRaisesRegex(ValueError, field): + collect_trace.get_per_func_instr_opts([invariant]) + invariant = Invariant( + ConsistentOutputRelation, + [APIParam("torch.fake"), VarTypeParam("torch.Tensor", "dtype")], + unconditional(), + ) + with self.assertRaisesRegex(ValueError, "const_value"): + collect_trace.get_per_func_instr_opts([invariant]) + # None is a legitimate expected value, distinct from an absent value. + self.assertTrue( + collect_trace.get_per_func_instr_opts([output_rule(None)])[ + typename(tensor_result) + ]["dump_ret"] + ) + + def test_distinct_rule_rejects_malformed_selectors(self) -> None: + cases: list[list[Param]] = [ + [], + [VarTypeParam("torch.Tensor", "dtype")], + [APIParam("one"), APIParam("two")], + ] + for params in cases: + invariant = Invariant( + DistinctArgumentRelation, params, unconditional(EXP_GROUP_NAME) + ) + with ( + self.subTest(params=params), + self.assertRaises((TypeError, ValueError)), + ): + collect_trace.get_per_func_instr_opts([invariant]) + for name in ("", None, 7): + invariant = distinct_rule() + setattr(invariant.params[0], "api_full_name", name) + with ( + self.subTest(name=name), + self.assertRaisesRegex(ValueError, "api_full_name"), + ): + collect_trace.get_per_func_instr_opts([invariant]) + + def test_real_return_trace_passes_healthy_rule_and_rejects_wrong_dtype( + self, + ) -> None: + invariant = output_rule() + with tempfile.TemporaryDirectory() as directory: + trace = collect_calls( + Path(directory), + [invariant], + tensor_result, + [(torch.ones(2),), (torch.ones(3),)], + ) + self.assertEqual(len(trace.get_func_call_ids(typename(tensor_result))), 2) + healthy = invariant.check(trace, check_relation_first=False) + wrong = output_rule("torch.float64").check( + trace, check_relation_first=False + ) + self.assertTrue(healthy.triggered) + self.assertTrue(healthy.check_passed) + self.assertTrue(wrong.triggered) + self.assertFalse(wrong.check_passed) + + def test_real_distinct_trace_preserves_key_and_detects_actual_duplicates( + self, + ) -> None: + invariant = distinct_rule() + rules = [invariant, argument_rule(typename(distinct_input))] + for keys, should_pass in (((1, 2), True), ((1, 1), False)): + with self.subTest(keys=keys), tempfile.TemporaryDirectory() as directory: + trace = collect_calls( + Path(directory), + rules, + distinct_input, + [("shared", key) for key in keys], + ) + self.assertEqual( + len(trace.get_func_call_ids(typename(distinct_input))), 2 + ) + result = invariant.check(trace, check_relation_first=False) + self.assertTrue(result.triggered) + self.assertEqual(result.check_passed, should_pass) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/collect_trace.py b/traincheck/collect_trace.py index cac7645b..adb66335 100644 --- a/traincheck/collect_trace.py +++ b/traincheck/collect_trace.py @@ -1,5 +1,6 @@ import argparse import datetime +import inspect import logging import os @@ -11,16 +12,168 @@ import traincheck.runner as runner from traincheck.config.config import InstrOpt from traincheck.invariant.base_cls import ( + _NOT_SET, APIParam, Arguments, + GroupedPreconditions, InputOutputParam, Invariant, + Precondition, + PreconditionClause, + Preconditions, VarNameParam, VarTypeParam, read_inv_file, ) from traincheck.invariant.consistency_relation import ConsistencyRelation +from traincheck.invariant.consistency_transient_vars import ConsistentOutputRelation from traincheck.invariant.contain_relation import VAR_GROUP_NAME, APIContainRelation +from traincheck.invariant.DistinctArgumentRelation import DistinctArgumentRelation + +FunctionInstrumentationOptions = dict[str, bool | dict[str | int, object]] + + +def _required_context_manager_names(invariant: Invariant) -> set[str]: + """Collect context dependencies without evaluating or rewriting predicates. + + The parser reconstructs context liveness from __init__/__enter__/__exit__, + then binds complete constructor inputs. A subset can omit required inputs + or violate the parser's signature-binding contract. + """ + conditions = invariant.precondition + if conditions is None: + if invariant.relation == APIContainRelation: + raise ValueError("APIContainRelation requires invariant preconditions") + # Other relations already permit compilation of partial hypotheses whose + # preconditions are not assigned yet. There are no clauses to inspect; + # this does not make the hypothesis valid for invariant checking. + return set() + if not isinstance(conditions, GroupedPreconditions): + raise TypeError("Invariant precondition must be GroupedPreconditions") + groups = conditions.grouped_preconditions + if not isinstance(groups, dict): + raise TypeError("Precondition groups must be a dictionary") + names: set[str] = set() + stem = "meta_vars.context_managers" + prefix = f"{stem}." + for group_name, alternatives in groups.items(): + if not isinstance(group_name, str) or not group_name: + raise ValueError("Precondition group names must be nonempty strings") + if not isinstance(alternatives, Preconditions): + raise TypeError("Each precondition group must contain Preconditions") + if not isinstance(alternatives.preconditions, list): + raise TypeError("Precondition alternatives must be a list") + for condition in alternatives.preconditions: + if not isinstance(condition, Precondition): + raise TypeError("Each alternative must be a Precondition") + if not isinstance(condition.clauses, list): + raise TypeError("Precondition clauses must be a list") + for clause in condition.clauses: + if not isinstance(clause, PreconditionClause): + raise TypeError("Each clause must be a PreconditionClause") + property_name = clause.prop_name + if not isinstance(property_name, str): + raise TypeError("Precondition property names must be strings") + if property_name == stem: + raise ValueError("Context precondition must name a context class") + if not property_name.startswith(prefix): + continue + class_name = property_name[len(prefix) :] + parts = class_name.split(".") + if len(parts) < 2 or not all(part.isidentifier() for part in parts): + raise ValueError(f"Invalid context class name: {class_name!r}") + names.add(class_name) + return names + + +def _consistent_output_api_name(invariant: Invariant) -> str: + """Validate the output relation's required selectors before compiling it.""" + if len(invariant.params) != 2: + raise ValueError("ConsistentOutputRelation requires exactly two parameters") + api, output = invariant.params + if not isinstance(api, APIParam) or not isinstance(output, VarTypeParam): + raise TypeError("ConsistentOutputRelation requires [APIParam, VarTypeParam]") + for name, value in ( + ("api_full_name", api.api_full_name), + ("var_type", output.var_type), + ("attr_name", output.attr_name), + ): + if not isinstance(value, str) or not value: + raise ValueError( + f"ConsistentOutputRelation {name} must be a nonempty string" + ) + if output.const_value is _NOT_SET: + raise ValueError("ConsistentOutputRelation requires a const_value") + return api.api_full_name + + +def _distinct_argument_api_name(invariant: Invariant) -> str: + """Validate the API whose complete inputs the distinctness checker compares.""" + if len(invariant.params) != 1: + raise ValueError("DistinctArgumentRelation requires exactly one parameter") + api = invariant.params[0] + if not isinstance(api, APIParam): + raise TypeError("DistinctArgumentRelation requires an APIParam") + if not isinstance(api.api_full_name, str) or not api.api_full_name: + raise ValueError( + "DistinctArgumentRelation api_full_name must be a nonempty string" + ) + return api.api_full_name + + +def _requires_full_argument_metadata(arguments: Arguments, api_name: str) -> bool: + """Retain complete inputs when the current selector cannot express a requirement. + + Missing signatures and unresolved required names are explicit unsupported + cases. Malformed metadata is an error, even if another requirement already + requires full collection. The caller unions this decision across all rules. + """ + if not isinstance(api_name, str) or not api_name: + raise ValueError("API argument requirements need a nonempty API name") + if not isinstance(arguments.func_name, str) or arguments.func_name != api_name: + raise ValueError("Argument function name must match its API parameter") + requirements: object = arguments.arguments + if not isinstance(requirements, dict): + raise TypeError("Required argument metadata must be a dictionary") + required_names: set[str] = set() + for name, metadata in requirements.items(): + if not isinstance(name, str) or not name: + raise ValueError("Required argument names must be nonempty strings") + if not isinstance(metadata, dict): + raise TypeError("Each argument requirement must contain typed metadata") + if len(metadata) != 1: + raise ValueError("Each argument requirement must contain exactly one type") + type_name = next(iter(metadata)) + if not isinstance(type_name, str) or not type_name: + raise ValueError("Argument metadata type names must be nonempty strings") + required_names.add(name) + + signature: object = arguments.signature + if signature is not None and not isinstance(signature, inspect.Signature): + raise TypeError("Argument signature must be inspect.Signature or None") + if signature is None: + return bool(required_names) + supported_kinds = ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + known_kinds = supported_kinds + ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ) + for name, parameter in signature.parameters.items(): + if not isinstance(name, str) or not name: + raise ValueError("Signature parameter names must be nonempty strings") + if not isinstance(parameter, inspect.Parameter) or parameter.name != name: + raise TypeError("Signature entries must be matching inspect.Parameters") + if parameter.kind not in known_kinds: + raise ValueError("Unknown signature parameter kind") + return any( + name not in signature.parameters + or signature.parameters[name].kind not in supported_kinds + for name in required_names + ) def get_list_of_funcs_from_invariants(invariants: list[Invariant]) -> list[str]: @@ -37,15 +190,25 @@ def get_list_of_funcs_from_invariants(invariants: list[Invariant]) -> list[str]: def get_per_func_instr_opts( invariants: list[Invariant], -) -> dict[str, dict[str, bool | dict]]: +) -> dict[str, FunctionInstrumentationOptions]: """ Get per function instrumentation options """ # TODO: for APIContainRelation that describes a variable, if the precondition is not unconditional on the variable and the API belongs to a class, then all class methods should be instrumented with `scan_proxy_in_args` set to True logger = logging.getLogger(__name__) - func_instr_opts: dict[str, dict[str, bool | dict]] = {} + func_instr_opts: dict[str, FunctionInstrumentationOptions] = {} + full_argument_apis: set[str] = set() + context_names: set[str] = set() for inv in invariants: + context_names.update(_required_context_manager_names(inv)) + output_api = ( + _consistent_output_api_name(inv) + if inv.relation == ConsistentOutputRelation + else None + ) + if inv.relation == DistinctArgumentRelation: + full_argument_apis.add(_distinct_argument_api_name(inv)) for param in inv.params: if isinstance(param, APIParam) or isinstance(param, InputOutputParam): func_name = ( @@ -58,27 +221,37 @@ def get_per_func_instr_opts( func_instr_opts[func_name] = { "scan_proxy_in_args": False, "dump_args": False, - "dump_ret": False, # not really used for now' + "dump_ret": False, "var_types_to_track": {}, # NOTE: do selective proxy dumping in APIs might interfere with correctness of the Consistency Relation Checking } if isinstance(param, APIParam): func_name = param.api_full_name - def merge(a: dict, b: dict, path=[]): + def merge( + a: dict[str | int, object], + b: dict[str | int, object], + path: tuple[str, ...] = (), + ) -> dict[str | int, object]: for key in b: if key in a: - if isinstance(a[key], dict) and isinstance(b[key], dict): - merge(a[key], b[key], path + [str(key)]) - elif a[key] != b[key]: - if a[key] is None: - a[key] = b[key] - elif b[key] is None: + existing_value = a[key] + incoming_value = b[key] + if isinstance(existing_value, dict) and isinstance( + incoming_value, dict + ): + merge( + existing_value, incoming_value, path + (str(key),) + ) + elif existing_value != incoming_value: + if existing_value is None: + a[key] = incoming_value + elif incoming_value is None: pass else: raise Exception( "Conflict at " - + ".".join(path + [str(key)]) + + ".".join(path + (str(key),)) + f" {a[key]} != {b[key]}" ) else: @@ -87,27 +260,41 @@ def merge(a: dict, b: dict, path=[]): if isinstance(param.arguments, Arguments): func_instr_opts[func_name]["dump_args"] = True - arg_instr_opt = param.arguments.to_instr_opts_dict() - if "dump_args_config" in func_instr_opts[func_name]: - existing_config = func_instr_opts[func_name]["dump_args_config"] - assert isinstance(existing_config, dict) - func_instr_opts[func_name]["dump_args_config"] = merge( - existing_config, - arg_instr_opt, - ) + if _requires_full_argument_metadata(param.arguments, func_name): + full_argument_apis.add(func_name) else: - func_instr_opts[func_name]["dump_args_config"] = arg_instr_opt + arg_instr_opt = param.arguments.to_instr_opts_dict() + if "dump_args_config" in func_instr_opts[func_name]: + existing_config = func_instr_opts[func_name][ + "dump_args_config" + ] + assert isinstance(existing_config, dict) + func_instr_opts[func_name]["dump_args_config"] = merge( + existing_config, + arg_instr_opt, + ) + else: + func_instr_opts[func_name][ + "dump_args_config" + ] = arg_instr_opt if isinstance(param, InputOutputParam): func_name = param.api_name assert isinstance(func_name, str) func_instr_opts[func_name]["dump_args"] = True func_instr_opts[func_name]["dump_ret"] = True + full_argument_apis.add(func_name) # TODO: convert the arguments to instr_opts_dict (currently not possible as the index indicates the index of the argument/ret value among other tensors not all arguments) logger.debug( "Currently not supporting fine-grained dumping of arguments and return values for InputOutputParam" ) + if output_api is not None: + # Output-only relations use APIParam + VarTypeParam, not the + # InputOutputParam branch above. Use the existing full return + # metadata converter; it has no per-field return selector. + func_instr_opts[output_api]["dump_ret"] = True + if inv.relation == APIContainRelation: assert isinstance(inv.params[0], APIParam) assert inv.precondition is not None @@ -134,6 +321,24 @@ def merge(a: dict, b: dict, path=[]): "dump_unchanged": False } + for context_name in sorted(context_names): + for method_name in ("__init__", "__enter__", "__exit__"): + function_name = f"{context_name}.{method_name}" + if function_name not in func_instr_opts: + func_instr_opts[function_name] = { + "scan_proxy_in_args": False, + "dump_args": False, + "dump_ret": False, + "var_types_to_track": {}, + } + full_argument_apis.add(f"{context_name}.__init__") + + for func_name in full_argument_apis: + func_instr_opts[func_name]["dump_args"] = True + # Absence requests full arguments. A narrower requirement from another + # invariant must never shrink this coverage, regardless of rule order. + func_instr_opts[func_name].pop("dump_args_config", None) + return func_instr_opts