From 86cbfc297cb27bac4a99359224cf21d48ecdf581 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 27 Oct 2026 03:25:17 +0800 Subject: [PATCH 1/2] Snapshot uninitialized operator constructor receivers without getters --- tests/test_constructor_snapshot.py | 269 ++++++++++++++++++ .../instrumentor/constructor_snapshot.py | 126 ++++++++ traincheck/instrumentor/tracer.py | 44 ++- 3 files changed, 431 insertions(+), 8 deletions(-) create mode 100644 tests/test_constructor_snapshot.py create mode 100644 traincheck/instrumentor/constructor_snapshot.py diff --git a/tests/test_constructor_snapshot.py b/tests/test_constructor_snapshot.py new file mode 100644 index 00000000..2166dda8 --- /dev/null +++ b/tests/test_constructor_snapshot.py @@ -0,0 +1,269 @@ +"""Constructor tracing must not inspect an uninitialized receiver dynamically.""" + +import functools +import inspect +import json +import tempfile +import unittest +from contextlib import ExitStack +from pathlib import Path +from types import SimpleNamespace +from typing import Callable +from unittest.mock import patch + +import torch + +import traincheck.config.config as config +from traincheck.instrumentor import dumper, tracer +from traincheck.instrumentor.constructor_snapshot import constructor_receiver_snapshot +from traincheck.trace.trace_pandas import read_trace_file_Pandas +from traincheck.trace.utils import bind_args_kwargs_to_signature + +CONSTRUCTOR = "torch._ops.OpOverloadPacket.__init__" +PACKET_TYPE = "torch._ops.OpOverloadPacket" +SNAPSHOT_FIELD = "_TRAINCHECK_CONSTRUCTOR_SNAPSHOT" + + +class UnexpectedGetter(BaseException): + """Bypass legacy serializer exception swallowing to fail the regression fast.""" + + +class ConstructorSnapshotTests(unittest.TestCase): + def setUp(self) -> None: + self.records: list[dict[str, object]] = [] + destination = SimpleNamespace(add_trace=self.capture) + patches = ExitStack() + self.addCleanup(patches.close) + patches.enter_context( + patch.object( + dumper, + "THREAD_DATA", + SimpleNamespace( + api_trace_buffer=destination, var_trace_buffer=destination + ), + ) + ) + patches.enter_context(patch.object(config, "DISABLE_WRAPPER", False)) + patches.enter_context(patch.object(config, "INSTR_OPTS", None)) + patches.enter_context(patch.object(tracer, "IS_INSTRUMENTING", False)) + + def capture(self, record: dict[str, object]) -> None: + value: object = json.loads(dumper.serialize(record)) + assert isinstance(value, dict) + validated: dict[str, object] = {} + for key, item in value.items(): + assert isinstance(key, str) + validated[key] = item + self.records.append(validated) + + def wrap(self, constructor: Callable[..., object]) -> Callable[..., object]: + return tracer.wrapper( + constructor, + is_bound_method=True, + scan_proxy_in_args=False, + dump_stack_trace=False, + handle_proxy=False, + ) + + def snapshot(self, encoded: dict[str, object] | None) -> dict[str, object]: + assert encoded is not None + packet = encoded[PACKET_TYPE] + assert isinstance(packet, dict) + snapshot = packet[SNAPSHOT_FIELD] + assert isinstance(snapshot, dict) + result: dict[str, object] = {} + for key, value in snapshot.items(): + assert isinstance(key, str) + result[key] = value + return result + + def test_constructor_keeps_pre_post_and_other_arguments_without_getter_reads( + self, + ) -> None: + source = torch.ops.aten.add + packet = object.__new__(torch._ops.OpOverloadPacket) + object.__getattribute__(packet, "__dict__")["existing_value"] = 7 + + def forbidden_getter(receiver: object) -> object: + raise UnexpectedGetter("receiver property accessed before initialization") + + with patch.object( + torch._ops.OpOverloadPacket, "op", property(forbidden_getter) + ): + result = self.wrap(torch._ops.OpOverloadPacket.__init__)( + packet, "aten::add", "add", source._op, source._overload_names + ) + self.assertIsNone(result) + self.assertEqual(str(packet), "aten.add") + self.assertEqual(packet.existing_value, 7) + self.assertEqual( + [r["type"] for r in self.records], + [tracer.TraceLineType.FUNC_CALL_PRE, tracer.TraceLineType.FUNC_CALL_POST], + ) + pre, post = self.records + self.assertEqual(pre["obj_id"], id(packet)) + self.assertEqual(pre["func_call_id"], post["func_call_id"]) + args = pre["args"] + assert isinstance(args, dict) + self.assertEqual(args["1"], dumper.var_to_serializable("aten::add")) + self.assertEqual(args["2"], dumper.var_to_serializable("add")) + self.assertEqual(args["3"], dumper.var_to_serializable(source._op)) + self.assertEqual(args["4"], dumper.var_to_serializable(source._overload_names)) + snapshot = self.snapshot(args["0"]) + self.assertEqual(snapshot["phase"], "before_init") + self.assertEqual(snapshot["coverage"], "stored_attributes_only") + self.assertEqual(snapshot["object_id"], id(packet)) + self.assertEqual( + snapshot["attributes"], {"existing_value": {"kind": "value", "value": 7}} + ) + + # Existing readers preserve the marker and bind every original argument. + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "trace_API.jsonl" + path.write_text("\n".join(json.dumps(row) for row in self.records) + "\n") + trace = read_trace_file_Pandas( + str(path), truncate_incomplete_func_calls=False + ) + call_id = pre["func_call_id"] + assert isinstance(call_id, str) + loaded = trace.get_pre_func_call_record(call_id) + bound = bind_args_kwargs_to_signature( + loaded["args"], + loaded["kwargs"], + inspect.signature(torch._ops.OpOverloadPacket.__init__), + ) + self.assertEqual(bound.get_arg("qualified_op_name"), args["1"]) + self.assertEqual( + bound.get_arg("self")[PACKET_TYPE][SNAPSHOT_FIELD]["coverage"], + "stored_attributes_only", + ) + + def test_constructor_failure_retains_pre_exception_and_original_error(self) -> None: + packet = object.__new__(torch._ops.OpOverloadPacket) + original_error = RuntimeError("constructor failed") + calls: list[object] = [] + + @functools.wraps(torch._ops.OpOverloadPacket.__init__) + def failing(receiver: object, value: int) -> None: + calls.extend([receiver, value]) + raise original_error + + with self.assertRaises(RuntimeError) as result: + self.wrap(failing)(packet, 9) + self.assertIs(result.exception, original_error) + self.assertIs(calls[0], packet) + self.assertEqual(calls[1], 9) + self.assertEqual( + [r["type"] for r in self.records], + [ + tracer.TraceLineType.FUNC_CALL_PRE, + tracer.TraceLineType.FUNC_CALL_POST_EXCEPTION, + ], + ) + self.assertEqual(self.records[1]["exception_msg"], "constructor failed") + self.assertEqual( + self.records[0]["func_call_id"], self.records[1]["func_call_id"] + ) + + def test_initialized_packet_and_other_arguments_keep_existing_encoding( + self, + ) -> None: + packet = torch.ops.aten.add + self.assertIsNone(constructor_receiver_snapshot(CONSTRUCTOR, packet)) + arguments = (packet, "name", 3) + before = tracer.to_dict_args_kwargs(arguments, {"flag": True}) + after = tracer.to_dict_args_kwargs( + arguments, {"flag": True}, function_name=CONSTRUCTOR + ) + self.assertEqual(dumper.serialize(dict(before)), dumper.serialize(dict(after))) + self.assertIsNone(constructor_receiver_snapshot("different.__init__", object())) + self.assertIsNone(constructor_receiver_snapshot(CONSTRUCTOR, object())) + + def test_stored_none_false_and_opaque_references_are_distinct(self) -> None: + packet = object.__new__(torch._ops.OpOverloadPacket) + opaque = object() + object.__getattribute__(packet, "__dict__").update( + {"unset": None, "flag": False, "opaque": opaque} + ) + snapshot = self.snapshot(constructor_receiver_snapshot(CONSTRUCTOR, packet)) + self.assertEqual( + snapshot["attributes"], + { + "unset": {"kind": "value", "value": None}, + "flag": {"kind": "value", "value": False}, + "opaque": { + "kind": "object_reference", + "object_id": id(opaque), + "type": "builtins.object", + }, + }, + ) + + def test_selective_config_retains_explicit_unavailable_fields(self) -> None: + packet = object.__new__(torch._ops.OpOverloadPacket) + object.__getattribute__(packet, "__dict__").update({"first": 1, "second": 2}) + encoded = tracer.to_dict_args_kwargs( + (packet, "unchanged"), + {}, + {"0": {"first": True, "op": True}}, + function_name=CONSTRUCTOR, + ) + self.assertEqual(set(encoded["args"]), {0}) + snapshot = self.snapshot(encoded["args"][0]) + self.assertEqual(snapshot["requested_attributes"], ["first", "op"]) + self.assertEqual(snapshot["unavailable_attributes"], ["op"]) + self.assertEqual( + snapshot["attributes"], {"first": {"kind": "value", "value": 1}} + ) + + def test_malformed_state_and_selector_are_rejected(self) -> None: + packet = object.__new__(torch._ops.OpOverloadPacket) + for config_value in (1, {1: True}): + with self.subTest(config_value=config_value), self.assertRaises(TypeError): + constructor_receiver_snapshot(CONSTRUCTOR, packet, config_value) + object.__getattribute__(packet, "__dict__")[1] = "malformed name" + with self.assertRaisesRegex(TypeError, "attribute names must be strings"): + constructor_receiver_snapshot(CONSTRUCTOR, packet) + + def test_opaque_type_metadata_failure_propagates(self) -> None: + def unavailable_module(receiver: object) -> str: + raise RuntimeError("type metadata unavailable") + + metaclass = type( + "MetadataFailure", (type,), {"__module__": property(unavailable_module)} + ) + opaque_type = metaclass("Opaque", (), {}) + packet = object.__new__(torch._ops.OpOverloadPacket) + object.__getattribute__(packet, "__dict__")["opaque"] = opaque_type() + with self.assertRaisesRegex(RuntimeError, "type metadata unavailable"): + constructor_receiver_snapshot(CONSTRUCTOR, packet) + + def test_nonfinite_stored_values_survive_json_distinct_from_none(self) -> None: + packet = object.__new__(torch._ops.OpOverloadPacket) + object.__getattribute__(packet, "__dict__").update( + { + "missing": None, + "nan": float("nan"), + "positive": float("inf"), + "negative": float("-inf"), + "finite": 1.5, + } + ) + encoded = constructor_receiver_snapshot(CONSTRUCTOR, packet) + assert encoded is not None + decoded: object = json.loads(dumper.serialize(encoded)) + assert isinstance(decoded, dict) + self.assertEqual( + self.snapshot(decoded)["attributes"], + { + "missing": {"kind": "value", "value": None}, + "nan": {"kind": "nonfinite_float", "value": "nan"}, + "positive": {"kind": "nonfinite_float", "value": "+inf"}, + "negative": {"kind": "nonfinite_float", "value": "-inf"}, + "finite": {"kind": "value", "value": 1.5}, + }, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/instrumentor/constructor_snapshot.py b/traincheck/instrumentor/constructor_snapshot.py new file mode 100644 index 00000000..2132ab76 --- /dev/null +++ b/traincheck/instrumentor/constructor_snapshot.py @@ -0,0 +1,126 @@ +"""Explicit stored-state snapshots for receivers that are not initialized yet.""" + +import math +from collections.abc import Mapping +from typing import Literal, TypeAlias, TypedDict + +import torch + +_PACKET_CONSTRUCTOR = "torch._ops.OpOverloadPacket.__init__" +_PACKET_TYPE = "torch._ops.OpOverloadPacket" +_SNAPSHOT_FIELD = "_TRAINCHECK_CONSTRUCTOR_SNAPSHOT" + +Scalar: TypeAlias = str | int | float | bool | None + + +class StoredValue(TypedDict): + kind: Literal["value"] + value: Scalar + + +class StoredReference(TypedDict): + kind: Literal["object_reference"] + object_id: int + type: str + + +class StoredNonfiniteFloat(TypedDict): + kind: Literal["nonfinite_float"] + value: Literal["nan", "+inf", "-inf"] + + +StoredAttribute: TypeAlias = StoredValue | StoredReference | StoredNonfiniteFloat + + +class ConstructorSnapshot(TypedDict): + phase: Literal["before_init"] + coverage: Literal["stored_attributes_only"] + object_id: int + requested_attributes: list[str] | None + unavailable_attributes: list[str] + attributes: dict[str, StoredAttribute] + + +def _stored_value(value: object) -> StoredAttribute: + if value is None or type(value) in (str, int, float, bool): + # The explicit isinstance check narrows the exact built-in types above. + assert value is None or isinstance(value, (str, int, float, bool)) + if isinstance(value, float) and not math.isfinite(value): + label: Literal["nan", "+inf", "-inf"] = ( + "nan" if math.isnan(value) else "+inf" if value > 0 else "-inf" + ) + return {"kind": "nonfinite_float", "value": label} + return {"kind": "value", "value": value} + + value_type = type(value) + # Read class metadata without invoking an instance's dynamic lookup or repr. + module: object = type.__getattribute__(value_type, "__module__") + qualified_name: object = type.__getattribute__(value_type, "__qualname__") + if not isinstance(module, str) or not isinstance(qualified_name, str): + raise TypeError("Stored attribute type metadata must contain string names") + return { + "kind": "object_reference", + "object_id": id(value), + "type": f"{module}.{qualified_name}", + } + + +def constructor_receiver_snapshot( + function_name: str, receiver: object, dump_config: object = None +) -> dict[str, object] | None: + """Return an explicit partial snapshot only for an uninitialized packet. + + None means the normal serializer applies. Full instrumentation observes + OpOverloadPacket.__init__ before PyTorch sets _op and _qualified_op_name. + Its ``op`` property and dynamic name lookup recurse when either is absent. + This adapter keeps the receiver and its identity in the PRE record without + running those getters. Stored scalars retain their values; opaque values are + clearly marked identity references, not claims about their internal state. + Nonfinite floats use tagged names so JSON cannot turn them into null. + + Initialized packets, other arguments, other constructors, and subclasses + retain their existing serialization. This is not a fallback after a getter + fails and does not suppress application errors. The caller supplies the + positional receiver used by PyTorch. The tracer's existing inability to + identify a bound receiver supplied only as a keyword is outside this adapter. + """ + if function_name != _PACKET_CONSTRUCTOR: + return None + if type(receiver) is not torch._ops.OpOverloadPacket: + return None + + state: object = object.__getattribute__(receiver, "__dict__") + if not isinstance(state, dict): + raise TypeError("OpOverloadPacket stored state must be a dictionary") + if "_op" in state and "_qualified_op_name" in state: + return None + + requested: list[str] | None = None + if dump_config is not None: + if not isinstance(dump_config, Mapping): + raise TypeError("Constructor receiver dump configuration must be a mapping") + requested = [] + for name in dump_config: + if not isinstance(name, str): + raise TypeError("Requested constructor attribute names must be strings") + requested.append(name) + + attributes: dict[str, StoredAttribute] = {} + for name, value in state.items(): + if not isinstance(name, str): + raise TypeError("OpOverloadPacket stored attribute names must be strings") + if requested is None or name in requested: + attributes[name] = _stored_value(value) + snapshot: ConstructorSnapshot = { + "phase": "before_init", + "coverage": "stored_attributes_only", + "object_id": id(receiver), + "requested_attributes": requested, + "unavailable_attributes": ( + [name for name in requested if name not in state] + if requested is not None + else [] + ), + "attributes": attributes, + } + return {_PACKET_TYPE: {_SNAPSHOT_FIELD: snapshot}} diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 10bb81ab..b950c9d6 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -10,7 +10,7 @@ import types from contextlib import ContextDecorator from importlib.util import find_spec -from typing import Callable, Literal, Optional +from typing import Callable, Literal, Mapping, Optional, Sequence, TypedDict import torch @@ -22,6 +22,7 @@ WRAP_WITHOUT_DUMP_WHITELIST, ) from traincheck.instrumentor.caches import META_VARS +from traincheck.instrumentor.constructor_snapshot import constructor_receiver_snapshot from traincheck.instrumentor.dumper import ( convert_var_to_dict, dump_trace_API, @@ -157,20 +158,45 @@ def get_owner_class(func): return owner +class SerializedArguments(TypedDict): + args: dict[int, dict[str, object]] + kwargs: dict[str, dict[str, object]] + + @suppress_tracing() -def to_dict_args_kwargs(args, kwargs, dump_args_config=None) -> dict: +def to_dict_args_kwargs( + args: Sequence[object], + kwargs: Mapping[str, object], + dump_args_config: Mapping[str, object] | None = None, + *, + function_name: str | None = None, +) -> SerializedArguments: + def serialize_argument(index: int, argument: object) -> dict[str, object]: + argument_config = ( + dump_args_config[str(index)] if dump_args_config is not None else None + ) + if index == 0 and function_name is not None: + snapshot = constructor_receiver_snapshot( + function_name, argument, argument_config + ) + if snapshot is not None: + return snapshot + if dump_args_config is None: + return var_to_serializable(argument) + return var_to_serializable(argument, argument_config) + if dump_args_config is None: - result = { - "args": {i: var_to_serializable(arg) for i, arg in enumerate(args)}, + result: SerializedArguments = { + "args": {i: serialize_argument(i, arg) for i, arg in enumerate(args)}, "kwargs": {k: var_to_serializable(v) for k, v in kwargs.items()}, } else: result = {"args": {}, "kwargs": {}} - args_dicts = {} - kwargs_dicts = {} + args_dicts: dict[int, dict[str, object]] = {} + kwargs_dicts: dict[str, dict[str, object]] = {} for i, arg in enumerate(args): if str(i) in dump_args_config: - args_dicts[i] = var_to_serializable(arg, dump_args_config[str(i)]) + args_dicts[i] = serialize_argument(i, arg) for k, v in kwargs.items(): if k in dump_args_config: @@ -291,7 +317,9 @@ def find_proxy_in_args(args): [proxy.__dict__["var_name"], type(proxy._obj).__name__] ) if dump_args: - dict_args_kwargs = to_dict_args_kwargs(args, kwargs, dump_args_config) + dict_args_kwargs = to_dict_args_kwargs( + args, kwargs, dump_args_config, function_name=original_function_name + ) pre_record["args"] = dict_args_kwargs["args"] pre_record["kwargs"] = dict_args_kwargs["kwargs"] dump_trace_API(pre_record) From 98fed5058bfc0c712c0e9a717a52c1bc1c224666 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 27 Oct 2026 13:24:45 +0800 Subject: [PATCH 2/2] Generalize constructor PRE snapshots to native stored state --- docs/constructor-snapshots.md | 66 ++ tests/test_constructor_snapshot.py | 713 +++++++++++++++++- tests/test_trace_integrity.py | 9 +- .../instrumentor/constructor_snapshot.py | 199 +++-- traincheck/instrumentor/tracer.py | 140 +++- 5 files changed, 1010 insertions(+), 117 deletions(-) create mode 100644 docs/constructor-snapshots.md diff --git a/docs/constructor-snapshots.md b/docs/constructor-snapshots.md new file mode 100644 index 00000000..ac6ea8dd --- /dev/null +++ b/docs/constructor-snapshots.md @@ -0,0 +1,66 @@ +# Constructor PRE snapshots + +When argument dumping is enabled for a class's `__init__`, TrainCheck records +the receiver before the initializer body runs. The receiver may be empty, +partially initialized by a subclass, or being initialized again. The instrumenter +identifies initializers from their owning class and attribute name; a standalone +function named `__init__` keeps ordinary argument serialization. + +The receiver is encoded under its static, fully qualified type name with an +`_TRAINCHECK_CONSTRUCTOR_SNAPSHOT` marker. This is a deliberate change from +ordinary object introspection for **all selected initializers**, including +metaclass initializers. Other arguments and ordinary method calls keep their +existing representation. The function's PRE, POST, and exception records remain. + +## Stored-state contract + +The marker contains: + +| Field | Meaning | +| --- | --- | +| `phase` | `before_init`; this does not claim initialization succeeded. | +| `coverage` | `stored_attributes_only`; computed properties are not captured. | +| `object_id` | Identity of the actual receiver. | +| `dictionary_status` | `captured`, `absent`, or `unavailable` when storage exists but no native accessor remains accessible. | +| `attributes` | Selected entries from the instance dictionary or native class namespace. | +| `slots` | Selected populated native slots, each with its declaring `owner_type`, `name`, and tagged `value`. | +| `unset_slots` | Selected declared slots whose storage has not been assigned. | +| `requested_attributes` | The explicit selector, or `null` for all accessible stored fields. | +| `unavailable_attributes` | Requested names with no captured dictionary entry or populated slot. | + +Dictionary entries and slot records are separate. A base slot, a derived slot, +and a dictionary entry can all have the same name and hold different values; +none overwrites another. Slot selectors use the stored name, including Python's +name mangling for private slots. A requested name selects each matching storage +location. An unset location stays explicit even if another matching slot has a +value. An empty selector is distinct from unselected/all-field collection. + +Exact built-in scalars retain their values. Nonfinite floats have explicit +`nan`, `+inf`, and `-inf` tags. Other values, including containers and tensors, +are **identity/type references**, not recursive content snapshots or tensor +hashes. Missing storage is never represented as a stored `None` value. + +## Avoiding application execution + +The snapshot reads native dictionary and slot descriptors from raw class +namespaces. It bypasses receiver `__getattribute__`, `__getattr__`, `__dir__`, +properties, custom descriptors, and custom metaclass lookup/properties. It also +bypasses overridden dictionary iteration methods and custom metaclass equality. +Constructor receivers are excluded from proxy-marker scans and argument +unproxying on this dumping path; ordinary arguments keep those behaviors. + +Normal positional construction and explicit keyword-receiver calls are covered. +If a native callable exposes no signature, its positional receiver is supported; +no keyword receiver name is guessed. Missing-receiver errors are produced by +the original callable and retain an exception record. + +Native extension state without a recognized storage descriptor, replaced slot +descriptors, and computed attributes are outside this snapshot's coverage. +Dictionary storage hidden behind a replacement property is marked unavailable +unless an inherited native accessor still exposes it. Snapshots are synchronous, +do not cache application metadata, and do not promise an atomic view during +concurrent mutation of the same object's storage. + +Existing invariants that depended on introspected constructor receiver fields +need recollection and activation validation against this representation. A +requested property is explicitly unavailable, not a successfully observed value. diff --git a/tests/test_constructor_snapshot.py b/tests/test_constructor_snapshot.py index 2166dda8..6ba0c21d 100644 --- a/tests/test_constructor_snapshot.py +++ b/tests/test_constructor_snapshot.py @@ -3,12 +3,14 @@ import functools import inspect import json +import logging +import random import tempfile import unittest from contextlib import ExitStack from pathlib import Path -from types import SimpleNamespace -from typing import Callable +from types import ModuleType, SimpleNamespace +from typing import Callable, NoReturn from unittest.mock import patch import torch @@ -19,7 +21,6 @@ from traincheck.trace.trace_pandas import read_trace_file_Pandas from traincheck.trace.utils import bind_args_kwargs_to_signature -CONSTRUCTOR = "torch._ops.OpOverloadPacket.__init__" PACKET_TYPE = "torch._ops.OpOverloadPacket" SNAPSHOT_FIELD = "_TRAINCHECK_CONSTRUCTOR_SNAPSHOT" @@ -56,18 +57,21 @@ def capture(self, record: dict[str, object]) -> None: validated[key] = item self.records.append(validated) - def wrap(self, constructor: Callable[..., object]) -> Callable[..., object]: + def wrap( + self, constructor: Callable[..., object], *, scan_proxy: bool = False + ) -> Callable[..., object]: return tracer.wrapper( constructor, is_bound_method=True, - scan_proxy_in_args=False, + scan_proxy_in_args=scan_proxy, dump_stack_trace=False, handle_proxy=False, + is_constructor=True, ) - def snapshot(self, encoded: dict[str, object] | None) -> dict[str, object]: - assert encoded is not None - packet = encoded[PACKET_TYPE] + def snapshot(self, encoded: dict[str, object]) -> dict[str, object]: + self.assertEqual(len(encoded), 1) + packet = next(iter(encoded.values())) assert isinstance(packet, dict) snapshot = packet[SNAPSHOT_FIELD] assert isinstance(snapshot, dict) @@ -165,19 +169,27 @@ def failing(receiver: object, value: int) -> None: self.records[0]["func_call_id"], self.records[1]["func_call_id"] ) - def test_initialized_packet_and_other_arguments_keep_existing_encoding( - self, - ) -> None: - packet = torch.ops.aten.add - self.assertIsNone(constructor_receiver_snapshot(CONSTRUCTOR, packet)) + def test_initialized_receiver_is_snapshotted_on_reinitialization(self) -> None: + source = torch.ops.aten.add + packet = torch._ops.OpOverloadPacket( + "aten::add", "add", source._op, source._overload_names + ) + snapshot = self.snapshot(constructor_receiver_snapshot(packet)) + self.assertEqual(snapshot["object_id"], id(packet)) + attributes = snapshot["attributes"] + assert isinstance(attributes, dict) + self.assertEqual( + attributes["_qualified_op_name"], {"kind": "value", "value": "aten::add"} + ) arguments = (packet, "name", 3) - before = tracer.to_dict_args_kwargs(arguments, {"flag": True}) - after = tracer.to_dict_args_kwargs( - arguments, {"flag": True}, function_name=CONSTRUCTOR + ordinary = tracer.to_dict_args_kwargs(arguments, {"flag": True}) + constructor = tracer.to_dict_args_kwargs( + arguments, {"flag": True}, is_constructor=True ) - self.assertEqual(dumper.serialize(dict(before)), dumper.serialize(dict(after))) - self.assertIsNone(constructor_receiver_snapshot("different.__init__", object())) - self.assertIsNone(constructor_receiver_snapshot(CONSTRUCTOR, object())) + self.assertEqual(ordinary["kwargs"], constructor["kwargs"]) + for index in (1, 2): + self.assertEqual(ordinary["args"][index], constructor["args"][index]) + self.assertNotEqual(ordinary["args"][0], constructor["args"][0]) def test_stored_none_false_and_opaque_references_are_distinct(self) -> None: packet = object.__new__(torch._ops.OpOverloadPacket) @@ -185,7 +197,7 @@ def test_stored_none_false_and_opaque_references_are_distinct(self) -> None: object.__getattribute__(packet, "__dict__").update( {"unset": None, "flag": False, "opaque": opaque} ) - snapshot = self.snapshot(constructor_receiver_snapshot(CONSTRUCTOR, packet)) + snapshot = self.snapshot(constructor_receiver_snapshot(packet)) self.assertEqual( snapshot["attributes"], { @@ -206,7 +218,7 @@ def test_selective_config_retains_explicit_unavailable_fields(self) -> None: (packet, "unchanged"), {}, {"0": {"first": True, "op": True}}, - function_name=CONSTRUCTOR, + is_constructor=True, ) self.assertEqual(set(encoded["args"]), {0}) snapshot = self.snapshot(encoded["args"][0]) @@ -220,23 +232,43 @@ def test_malformed_state_and_selector_are_rejected(self) -> None: packet = object.__new__(torch._ops.OpOverloadPacket) for config_value in (1, {1: True}): with self.subTest(config_value=config_value), self.assertRaises(TypeError): - constructor_receiver_snapshot(CONSTRUCTOR, packet, config_value) + constructor_receiver_snapshot(packet, config_value) object.__getattribute__(packet, "__dict__")[1] = "malformed name" with self.assertRaisesRegex(TypeError, "attribute names must be strings"): - constructor_receiver_snapshot(CONSTRUCTOR, packet) + constructor_receiver_snapshot(packet) - def test_opaque_type_metadata_failure_propagates(self) -> None: - def unavailable_module(receiver: object) -> str: - raise RuntimeError("type metadata unavailable") + def test_metaclass_properties_are_not_evaluated_for_receiver_or_values( + self, + ) -> None: + def forbidden(receiver: object) -> object: + raise UnexpectedGetter("metaclass property executed") metaclass = type( - "MetadataFailure", (type,), {"__module__": property(unavailable_module)} + "MetadataFailure", + (type,), + { + "__module__": property(forbidden), + "__dict__": property(forbidden), + "__mro__": property(forbidden), + }, + ) + receiver_type = metaclass("Receiver", (), {"__module__": "fixture"}) + opaque_type = metaclass("Opaque", (), {"__module__": "fixture"}) + receiver = receiver_type() + opaque = opaque_type() + object.__setattr__(receiver, "opaque", opaque) + encoded = constructor_receiver_snapshot(receiver) + self.assertEqual(set(encoded), {"fixture.Receiver"}) + self.assertEqual( + self.snapshot(encoded)["attributes"], + { + "opaque": { + "kind": "object_reference", + "object_id": id(opaque), + "type": "fixture.Opaque", + } + }, ) - opaque_type = metaclass("Opaque", (), {}) - packet = object.__new__(torch._ops.OpOverloadPacket) - object.__getattribute__(packet, "__dict__")["opaque"] = opaque_type() - with self.assertRaisesRegex(RuntimeError, "type metadata unavailable"): - constructor_receiver_snapshot(CONSTRUCTOR, packet) def test_nonfinite_stored_values_survive_json_distinct_from_none(self) -> None: packet = object.__new__(torch._ops.OpOverloadPacket) @@ -249,8 +281,7 @@ def test_nonfinite_stored_values_survive_json_distinct_from_none(self) -> None: "finite": 1.5, } ) - encoded = constructor_receiver_snapshot(CONSTRUCTOR, packet) - assert encoded is not None + encoded = constructor_receiver_snapshot(packet) decoded: object = json.loads(dumper.serialize(encoded)) assert isinstance(decoded, dict) self.assertEqual( @@ -264,6 +295,620 @@ def test_nonfinite_stored_values_survive_json_distinct_from_none(self) -> None: }, ) + def test_proxy_scanning_does_not_inspect_constructor_receiver(self) -> None: + class Hostile: + def __init__(self, value: int) -> None: + object.__setattr__(self, "value", value) + + def __getattribute__(self, name: str) -> object: + raise UnexpectedGetter("receiver dynamically inspected") + + receiver = object.__new__(Hostile) + wrapped = tracer.wrapper( + Hostile.__init__, + is_bound_method=True, + scan_proxy_in_args=True, + dump_stack_trace=False, + handle_proxy=True, + is_constructor=True, + ) + wrapped(receiver, 5) + self.assertEqual(object.__getattribute__(receiver, "value"), 5) + self.assertEqual(len(self.records), 2) + args = self.records[0]["args"] + assert isinstance(args, dict) + self.assertEqual(self.snapshot(args["0"])["object_id"], id(receiver)) + + def test_generic_constructor_skips_properties_and_keeps_application_effects( + self, + ) -> None: + calls: list[int] = [] + + class Model: + def __init__(self, width: int) -> None: + calls.append(width) + self.width = width + + @property + def initialized_width(self) -> int: + raise UnexpectedGetter("constructor property was evaluated") + + receiver = object.__new__(Model) + result = self.wrap(Model.__init__)(receiver, 12) + self.assertIsNone(result) + self.assertEqual(receiver.width, 12) + self.assertEqual(calls, [12]) + self.assertEqual(len(self.records), 2) + args = self.records[0]["args"] + assert isinstance(args, dict) + snapshot = self.snapshot(args["0"]) + self.assertEqual(snapshot["attributes"], {}) + self.assertEqual(snapshot["dictionary_status"], "captured") + self.assertEqual(args["1"], dumper.var_to_serializable(12)) + + def test_snapshot_does_not_call_dynamic_lookup_dir_or_repr(self) -> None: + class Hostile: + def __getattribute__(self, name: str) -> object: + raise UnexpectedGetter("__getattribute__ executed") + + def __getattr__(self, name: str) -> object: + raise UnexpectedGetter("__getattr__ executed") + + def __dir__(self) -> list[str]: + raise UnexpectedGetter("__dir__ executed") + + def __repr__(self) -> str: + raise UnexpectedGetter("__repr__ executed") + + receiver = object.__new__(Hostile) + object.__setattr__(receiver, "value", 4) + snapshot = self.snapshot(constructor_receiver_snapshot(receiver)) + self.assertEqual(snapshot["object_id"], id(receiver)) + self.assertEqual( + snapshot["attributes"], {"value": {"kind": "value", "value": 4}} + ) + + def test_shadowed_dictionary_accessor_is_explicitly_unavailable(self) -> None: + def forbidden(receiver: object) -> object: + raise UnexpectedGetter("__dict__ property executed") + + receiver_type = type( + "ShadowedDictionary", (), {"__dict__": property(forbidden)} + ) + receiver = receiver_type() + object.__setattr__(receiver, "stored", 5) + snapshot = self.snapshot( + constructor_receiver_snapshot(receiver, {"stored": True}) + ) + self.assertEqual(snapshot["dictionary_status"], "unavailable") + self.assertEqual(snapshot["attributes"], {}) + self.assertEqual(snapshot["unavailable_attributes"], ["stored"]) + + def test_inherited_native_dictionary_accessor_bypasses_shadowing_property( + self, + ) -> None: + class Base: + pass + + def forbidden(receiver: object) -> object: + raise UnexpectedGetter("shadowing __dict__ property executed") + + receiver_type = type("Shadow", (Base,), {"__dict__": property(forbidden)}) + receiver = receiver_type() + object.__setattr__(receiver, "stored", 5) + snapshot = self.snapshot(constructor_receiver_snapshot(receiver)) + self.assertEqual(snapshot["dictionary_status"], "captured") + self.assertEqual( + snapshot["attributes"], {"stored": {"kind": "value", "value": 5}} + ) + + def test_native_dict_iteration_does_not_call_dict_subclass_override(self) -> None: + class HostileDict(dict[str, object]): + def items(self) -> NoReturn: + raise UnexpectedGetter("overridden dictionary items executed") + + class Model: + pass + + receiver = Model() + object.__setattr__(receiver, "__dict__", HostileDict({"value": 8})) + snapshot = self.snapshot(constructor_receiver_snapshot(receiver)) + self.assertEqual( + snapshot["attributes"], {"value": {"kind": "value", "value": 8}} + ) + + def test_metaclass_equality_is_not_used_for_scalar_classification(self) -> None: + class HostileMeta(type): + def __eq__(self, other: object) -> bool: + raise UnexpectedGetter("metaclass equality executed") + + class Opaque(metaclass=HostileMeta): + pass + + receiver = SimpleNamespace(value=Opaque()) + snapshot = self.snapshot(constructor_receiver_snapshot(receiver)) + attributes = snapshot["attributes"] + assert isinstance(attributes, dict) + value = attributes["value"] + assert isinstance(value, dict) + self.assertEqual(value["kind"], "object_reference") + self.assertEqual(value["object_id"], id(receiver.value)) + + def test_slots_preserve_populated_unset_inherited_and_private_storage(self) -> None: + class Base: + __slots__ = ("base", "__private") + + class Child(Base): + __slots__ = ("child", "pending") + + receiver = object.__new__(Child) + object.__setattr__(receiver, "base", 3) + object.__setattr__(receiver, "_Base__private", 9) + object.__setattr__(receiver, "child", False) + snapshot = self.snapshot(constructor_receiver_snapshot(receiver)) + self.assertEqual(snapshot["dictionary_status"], "absent") + self.assertEqual(snapshot["attributes"], {}) + slots = snapshot["slots"] + assert isinstance(slots, list) + self.assertCountEqual( + slots, + [ + { + "owner_type": f"{Base.__module__}.{Base.__qualname__}", + "name": "base", + "value": {"kind": "value", "value": 3}, + }, + { + "owner_type": f"{Base.__module__}.{Base.__qualname__}", + "name": "_Base__private", + "value": {"kind": "value", "value": 9}, + }, + { + "owner_type": f"{Child.__module__}.{Child.__qualname__}", + "name": "child", + "value": {"kind": "value", "value": False}, + }, + ], + ) + self.assertEqual( + snapshot["unset_slots"], + [ + { + "owner_type": f"{Child.__module__}.{Child.__qualname__}", + "name": "pending", + } + ], + ) + + def test_duplicate_slots_and_dictionary_field_are_not_overwritten(self) -> None: + class Base: + __slots__ = ("value",) + + class Child(Base): + __slots__ = ("value", "__dict__") + + receiver = object.__new__(Child) + Base.__dict__["value"].__set__(receiver, 1) + Child.__dict__["value"].__set__(receiver, 2) + object.__getattribute__(receiver, "__dict__")["value"] = 3 + snapshot = self.snapshot( + constructor_receiver_snapshot(receiver, {"value": True}) + ) + self.assertEqual( + snapshot["attributes"], {"value": {"kind": "value", "value": 3}} + ) + slots = snapshot["slots"] + assert isinstance(slots, list) + self.assertCountEqual( + slots, + [ + { + "owner_type": f"{Base.__module__}.{Base.__qualname__}", + "name": "value", + "value": {"kind": "value", "value": 1}, + }, + { + "owner_type": f"{Child.__module__}.{Child.__qualname__}", + "name": "value", + "value": {"kind": "value", "value": 2}, + }, + ], + ) + self.assertEqual(snapshot["unavailable_attributes"], []) + + def test_slot_storage_is_read_without_evaluating_shadowing_property(self) -> None: + class Base: + __slots__ = ("value",) + + def forbidden(receiver: object) -> object: + raise UnexpectedGetter("slot shadowing property executed") + + receiver_type = type( + "Shadow", (Base,), {"__slots__": (), "value": property(forbidden)} + ) + receiver: object = object.__new__(receiver_type) + Base.__dict__["value"].__set__(receiver, 7) + snapshot = self.snapshot(constructor_receiver_snapshot(receiver)) + self.assertEqual( + snapshot["slots"], + [ + { + "owner_type": f"{Base.__module__}.{Base.__qualname__}", + "name": "value", + "value": {"kind": "value", "value": 7}, + } + ], + ) + + def test_selective_unset_slot_is_distinct_from_missing_and_stored_none( + self, + ) -> None: + class Model: + __slots__ = ("present", "pending", "excluded") + + receiver = object.__new__(Model) + object.__setattr__(receiver, "present", None) + object.__setattr__(receiver, "excluded", 4) + snapshot = self.snapshot( + constructor_receiver_snapshot( + receiver, {"present": True, "pending": True, "unknown": True} + ) + ) + self.assertEqual(snapshot["unavailable_attributes"], ["pending", "unknown"]) + self.assertEqual( + snapshot["slots"], + [ + { + "owner_type": f"{Model.__module__}.{Model.__qualname__}", + "name": "present", + "value": {"kind": "value", "value": None}, + } + ], + ) + self.assertEqual( + snapshot["unset_slots"], + [ + { + "owner_type": f"{Model.__module__}.{Model.__qualname__}", + "name": "pending", + } + ], + ) + + def test_empty_selector_and_absent_dictionary_are_distinct(self) -> None: + receiver = SimpleNamespace(value=1) + selected = self.snapshot(constructor_receiver_snapshot(receiver, {})) + self.assertEqual(selected["dictionary_status"], "captured") + self.assertEqual(selected["requested_attributes"], []) + self.assertEqual(selected["attributes"], {}) + absent = self.snapshot(constructor_receiver_snapshot(object())) + self.assertEqual(absent["dictionary_status"], "absent") + self.assertIsNone(absent["requested_attributes"]) + self.assertEqual(absent["attributes"], {}) + self.assertEqual(absent["slots"], []) + + def test_constructor_preserves_rng_and_does_not_mutate_pre_snapshot(self) -> None: + class Model: + def __init__(self, initial: int) -> None: + self.initial = initial + self.python_draw = random.random() + self.torch_draw = float(torch.rand(())) + + @property + def destructive(self) -> float: + return random.random() + float(torch.rand(())) + + python_state = random.getstate() + torch_state = torch.get_rng_state().clone() + self.addCleanup(random.setstate, python_state) + self.addCleanup(torch.set_rng_state, torch_state) + plain = Model(2) + expected_python = random.getstate() + expected_torch = torch.get_rng_state().clone() + random.setstate(python_state) + torch.set_rng_state(torch_state) + receiver = object.__new__(Model) + object.__setattr__(receiver, "initial", 1) + self.wrap(Model.__init__)(receiver, 2) + self.assertEqual(receiver.__dict__, plain.__dict__) + self.assertEqual(random.getstate(), expected_python) + self.assertTrue(torch.equal(torch.get_rng_state(), expected_torch)) + self.assertEqual(len(self.records), 2) + args = self.records[0]["args"] + assert isinstance(args, dict) + self.assertEqual( + self.snapshot(args["0"])["attributes"], + {"initial": {"kind": "value", "value": 1}}, + ) + + def test_base_constructor_records_partially_initialized_subclass_state( + self, + ) -> None: + class Base: + def __init__(self) -> None: + self.base = 9 + + class Child(Base): + def __init__(self) -> None: + self.child = 7 + super().__init__() + + with ( + patch.object(Base, "__init__", self.wrap(Base.__init__)), + patch.object(Child, "__init__", self.wrap(Child.__init__)), + ): + receiver = Child() + self.assertEqual(receiver.child, 7) + self.assertEqual(receiver.base, 9) + self.assertEqual( + [row["type"] for row in self.records], + [ + tracer.TraceLineType.FUNC_CALL_PRE, + tracer.TraceLineType.FUNC_CALL_PRE, + tracer.TraceLineType.FUNC_CALL_POST, + tracer.TraceLineType.FUNC_CALL_POST, + ], + ) + outer_args = self.records[0]["args"] + inner_args = self.records[1]["args"] + assert isinstance(outer_args, dict) + assert isinstance(inner_args, dict) + self.assertEqual(self.snapshot(outer_args["0"])["attributes"], {}) + self.assertEqual( + self.snapshot(inner_args["0"])["attributes"], + {"child": {"kind": "value", "value": 7}}, + ) + self.assertEqual( + self.records[0]["func_call_id"], self.records[3]["func_call_id"] + ) + self.assertEqual( + self.records[1]["func_call_id"], self.records[2]["func_call_id"] + ) + + def test_malformed_static_type_metadata_is_rejected_without_formatting( + self, + ) -> None: + class MisleadingString(str): + def __str__(self) -> str: + raise UnexpectedGetter("module string subclass formatted") + + for module in (1, MisleadingString("fixture")): + with self.subTest(module_type=type(module).__name__): + receiver_type = type("BadModule", (), {"__module__": module}) + with self.assertRaises(TypeError): + constructor_receiver_snapshot(receiver_type()) + + def test_native_constructor_without_signature_still_snapshots_and_executes( + self, + ) -> None: + receiver: dict[str, int] = {} + with patch.object( + tracer.inspect, + "signature", + side_effect=ValueError("native signature unavailable"), + ): + wrapped = self.wrap(dict.__init__) + self.assertIsNone(wrapped(receiver, {"value": 7})) + self.assertEqual(receiver, {"value": 7}) + self.assertEqual(len(self.records), 2) + self.assertEqual(self.records[0]["obj_id"], id(receiver)) + args = self.records[0]["args"] + assert isinstance(args, dict) + snapshot = self.snapshot(args["0"]) + self.assertEqual(snapshot["object_id"], id(receiver)) + self.assertEqual(snapshot["dictionary_status"], "absent") + self.assertEqual( + self.records[0]["func_call_id"], self.records[1]["func_call_id"] + ) + + def test_malformed_python_constructor_signature_is_not_silently_ignored( + self, + ) -> None: + def initialize(receiver: object) -> None: + raise AssertionError("malformed signature must fail during wrapping") + + with patch.object(initialize, "__signature__", object(), create=True): + with self.assertRaises(TypeError): + self.wrap(initialize) + original_error = ValueError("malformed Python signature") + with patch.object(tracer.inspect, "signature", side_effect=original_error): + with self.assertRaises(ValueError) as caught: + self.wrap(initialize) + self.assertIs(caught.exception, original_error) + self.assertEqual(self.records, []) + + def test_keyword_receiver_preserves_argument_order_and_identity(self) -> None: + class Model: + def __init__(receiver, value: int, *, enabled: bool) -> None: + object.__setattr__(receiver, "value", value) + object.__setattr__(receiver, "enabled", enabled) + + def __getattribute__(self, name: str) -> object: + raise UnexpectedGetter("keyword receiver dynamically inspected") + + receiver = object.__new__(Model) + with patch.object(tracer, "is_funcs_to_be_unproxied", return_value=True): + wrapped = tracer.wrapper( + Model.__init__, + is_bound_method=True, + scan_proxy_in_args=True, + dump_stack_trace=False, + handle_proxy=True, + is_constructor=True, + ) + wrapped(value=5, receiver=receiver, enabled=False) + self.assertEqual(object.__getattribute__(receiver, "value"), 5) + self.assertIs(object.__getattribute__(receiver, "enabled"), False) + self.assertEqual(len(self.records), 2) + pre = self.records[0] + self.assertEqual(pre["obj_id"], id(receiver)) + self.assertEqual(pre["args"], {}) + kwargs = pre["kwargs"] + assert isinstance(kwargs, dict) + self.assertEqual(list(kwargs), ["value", "receiver", "enabled"]) + self.assertEqual(kwargs["value"], dumper.var_to_serializable(5)) + self.assertEqual(kwargs["enabled"], dumper.var_to_serializable(False)) + self.assertEqual(self.snapshot(kwargs["receiver"])["object_id"], id(receiver)) + + def test_keyword_metaclass_receiver_captures_class_namespace(self) -> None: + class Meta(type): + def __init__( + cls, + name: str, + bases: tuple[type[object], ...], + namespace: dict[str, object], + ) -> None: + type.__init__(cls, name, bases, namespace) + type.__setattr__(cls, "initialized", True) + + @property + def dangerous(cls) -> object: + raise UnexpectedGetter("metaclass receiver property evaluated") + + receiver = type.__new__( + Meta, "Subject", (), {"__module__": "fixture", "seed": 3} + ) + self.wrap(Meta.__init__, scan_proxy=True)( + cls=receiver, name="Subject", bases=(), namespace={"seed": 3} + ) + self.assertIs(type.__getattribute__(receiver, "initialized"), True) + self.assertEqual(len(self.records), 2) + self.assertEqual(self.records[0]["obj_id"], id(receiver)) + kwargs = self.records[0]["kwargs"] + assert isinstance(kwargs, dict) + snapshot = self.snapshot(kwargs["cls"]) + self.assertEqual(snapshot["dictionary_status"], "captured") + attributes = snapshot["attributes"] + assert isinstance(attributes, dict) + self.assertEqual(attributes["seed"], {"kind": "value", "value": 3}) + self.assertNotIn("initialized", attributes) + + def test_missing_receiver_keeps_original_typeerror_and_exception_pair(self) -> None: + class Model: + def __init__(receiver, value: int) -> None: + raise AssertionError("missing receiver must prevent constructor body") + + # The callable interface deliberately permits an invalid call so its + # original runtime TypeError can be compared with the traced call. + original: Callable[..., object] = Model.__init__ + with self.assertRaises(TypeError) as plain: + original(value=5) + with self.assertRaises(TypeError) as traced: + self.wrap(original, scan_proxy=True)(value=5) + self.assertEqual(str(traced.exception), str(plain.exception)) + self.assertEqual( + [row["type"] for row in self.records], + [ + tracer.TraceLineType.FUNC_CALL_PRE, + tracer.TraceLineType.FUNC_CALL_POST_EXCEPTION, + ], + ) + self.assertIsNone(self.records[0]["obj_id"]) + self.assertEqual( + self.records[0]["func_call_id"], self.records[1]["func_call_id"] + ) + self.assertEqual(self.records[1]["exception_msg"], str(plain.exception)) + + def test_metaclass_initializer_snapshots_class_namespace_without_getters( + self, + ) -> None: + class Meta(type): + def __init__( + cls, + name: str, + bases: tuple[type[object], ...], + namespace: dict[str, object], + ) -> None: + type.__init__(cls, name, bases, namespace) + type.__setattr__(cls, "initialized", True) + + def __getattribute__(cls, name: str) -> object: + raise UnexpectedGetter("class receiver dynamically inspected") + + @property + def dangerous(cls) -> object: + raise UnexpectedGetter("class receiver property evaluated") + + with ( + patch.object(config, "MODEL_TRACKER_STYLE", "subclass"), + patch.object( + tracer, + "get_instrumentation_logger_for_process", + return_value=logging.getLogger(__name__), + ), + ): + instrumentor = tracer.Instrumentor( + Meta, + scan_proxy_in_args=True, + use_full_instr=True, + sampling_interval=0, + warm_up_steps=0, + ) + instrumentor._instrument_module(Meta, set(), True, 0) + receiver = Meta("Subject", (), {"__module__": "fixture", "seed": 8}) + self.assertIs(type.__getattribute__(receiver, "initialized"), True) + self.assertEqual(len(self.records), 2) + self.assertEqual(self.records[0]["obj_id"], id(receiver)) + args = self.records[0]["args"] + assert isinstance(args, dict) + snapshot = self.snapshot(args["0"]) + self.assertEqual(snapshot["dictionary_status"], "captured") + attributes = snapshot["attributes"] + assert isinstance(attributes, dict) + self.assertEqual(attributes["seed"], {"kind": "value", "value": 8}) + self.assertNotIn("initialized", attributes) + self.assertEqual( + self.records[0]["func_call_id"], self.records[1]["func_call_id"] + ) + + def test_instrumenter_distinguishes_class_initializer_from_same_named_function( + self, + ) -> None: + class Model: + def __init__(receiver, value: int) -> None: + receiver.value = value + + @property + def dangerous(self) -> object: + raise UnexpectedGetter("class constructor property read") + + def __init__(self: object) -> None: + return None + + module = ModuleType(__name__) + setattr(module, "Model", Model) + setattr(module, "__init__", __init__) + with ( + patch.object(config, "MODEL_TRACKER_STYLE", "subclass"), + patch.object( + tracer, + "get_instrumentation_logger_for_process", + return_value=logging.getLogger(__name__), + ), + ): + instrumentor = tracer.Instrumentor( + module, + scan_proxy_in_args=False, + use_full_instr=True, + sampling_interval=0, + warm_up_steps=0, + ) + instrumentor._instrument_module(module, set(), True, 0) + receiver = Model(5) + standalone = vars(module)["__init__"] + assert callable(standalone) + standalone(SimpleNamespace(value=6)) + self.assertEqual(receiver.value, 5) + self.assertEqual(len(self.records), 4) + constructor_args = self.records[0]["args"] + standalone_args = self.records[2]["args"] + assert isinstance(constructor_args, dict) + assert isinstance(standalone_args, dict) + self.assertEqual(self.snapshot(constructor_args["0"])["attributes"], {}) + self.assertEqual( + standalone_args["0"], dumper.var_to_serializable(SimpleNamespace(value=6)) + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_trace_integrity.py b/tests/test_trace_integrity.py index 26d9e0b2..4a88055a 100644 --- a/tests/test_trace_integrity.py +++ b/tests/test_trace_integrity.py @@ -5,13 +5,16 @@ import threading import unittest from queue import Queue +from typing import Callable from unittest.mock import patch from traincheck.instrumentor import dumper, tracer -def wrapped(fn, name="outer", scan_proxy=False): - def call(*args, **kwargs): +def wrapped( + fn: Callable[..., object], name: str = "outer", scan_proxy: bool = False +) -> Callable[..., object]: + def call(*args: object, **kwargs: object) -> object: return tracer.function_wrapper( fn, name, @@ -25,6 +28,8 @@ def call(*args, **kwargs): False, {}, False, + False, + None, *args, **kwargs, ) diff --git a/traincheck/instrumentor/constructor_snapshot.py b/traincheck/instrumentor/constructor_snapshot.py index 2132ab76..1d5115b4 100644 --- a/traincheck/instrumentor/constructor_snapshot.py +++ b/traincheck/instrumentor/constructor_snapshot.py @@ -1,14 +1,13 @@ -"""Explicit stored-state snapshots for receivers that are not initialized yet.""" +"""Constructor PRE snapshots without evaluating application attributes.""" import math -from collections.abc import Mapping +import types +from collections.abc import Iterable, Mapping from typing import Literal, TypeAlias, TypedDict -import torch - -_PACKET_CONSTRUCTOR = "torch._ops.OpOverloadPacket.__init__" -_PACKET_TYPE = "torch._ops.OpOverloadPacket" _SNAPSHOT_FIELD = "_TRAINCHECK_CONSTRUCTOR_SNAPSHOT" +# Immutable built-in type namespace; no application class metadata is cached. +_TYPE_NAMESPACE: Mapping[str, object] = vars(type) Scalar: TypeAlias = str | int | float | bool | None @@ -32,18 +31,58 @@ class StoredNonfiniteFloat(TypedDict): StoredAttribute: TypeAlias = StoredValue | StoredReference | StoredNonfiniteFloat +class SlotIdentity(TypedDict): + owner_type: str + name: str + + +class StoredSlot(SlotIdentity): + value: StoredAttribute + + class ConstructorSnapshot(TypedDict): phase: Literal["before_init"] coverage: Literal["stored_attributes_only"] object_id: int + dictionary_status: Literal["captured", "absent", "unavailable"] requested_attributes: list[str] | None unavailable_attributes: list[str] attributes: dict[str, StoredAttribute] + slots: list[StoredSlot] + unset_slots: list[SlotIdentity] + + +def _type_metadata(cls: type[object], name: str) -> object: + descriptor = _TYPE_NAMESPACE[name] + assert isinstance( + descriptor, (types.GetSetDescriptorType, types.MemberDescriptorType) + ) + return descriptor.__get__(cls) + + +def _type_name(cls: type[object]) -> str: + # Invoke type's own native descriptors, bypassing custom metaclass + # __getattribute__ AND properties with these names. type.__getattribute__ + # alone still invokes such properties. + module: object = _type_metadata(cls, "__module__") + qualified_name: object = _type_metadata(cls, "__qualname__") + module_type, name_type = type(module), type(qualified_name) + if module_type is not str or name_type is not str: + raise TypeError("Stored attribute type metadata must contain string names") + assert isinstance(module, str) and isinstance(qualified_name, str) + return f"{module}.{qualified_name}" def _stored_value(value: object) -> StoredAttribute: - if value is None or type(value) in (str, int, float, bool): - # The explicit isinstance check narrows the exact built-in types above. + value_type = type(value) + # Identity comparisons also avoid custom metaclass __eq__ implementations. + if ( + value is None + or value_type is str + or value_type is int + or value_type is float + or value_type is bool + ): assert value is None or isinstance(value, (str, int, float, bool)) if isinstance(value, float) and not math.isfinite(value): label: Literal["nan", "+inf", "-inf"] = ( @@ -51,76 +90,132 @@ def _stored_value(value: object) -> StoredAttribute: ) return {"kind": "nonfinite_float", "value": label} return {"kind": "value", "value": value} - - value_type = type(value) - # Read class metadata without invoking an instance's dynamic lookup or repr. - module: object = type.__getattribute__(value_type, "__module__") - qualified_name: object = type.__getattribute__(value_type, "__qualname__") - if not isinstance(module, str) or not isinstance(qualified_name, str): - raise TypeError("Stored attribute type metadata must contain string names") return { "kind": "object_reference", "object_id": id(value), - "type": f"{module}.{qualified_name}", + "type": _type_name(value_type), } def constructor_receiver_snapshot( - function_name: str, receiver: object, dump_config: object = None -) -> dict[str, object] | None: - """Return an explicit partial snapshot only for an uninitialized packet. - - None means the normal serializer applies. Full instrumentation observes - OpOverloadPacket.__init__ before PyTorch sets _op and _qualified_op_name. - Its ``op`` property and dynamic name lookup recurse when either is absent. - This adapter keeps the receiver and its identity in the PRE record without - running those getters. Stored scalars retain their values; opaque values are - clearly marked identity references, not claims about their internal state. - Nonfinite floats use tagged names so JSON cannot turn them into null. - - Initialized packets, other arguments, other constructors, and subclasses - retain their existing serialization. This is not a fallback after a getter - fails and does not suppress application errors. The caller supplies the - positional receiver used by PyTorch. The tracer's existing inability to - identify a bound receiver supplied only as a keyword is outside this adapter. + receiver: object, dump_config: object = None +) -> dict[str, object]: + """Capture stored state at entry to any selected instance initializer. + + The caller identifies __init__ using the owning class, not a name heuristic. + Reinitialization and base-class initializers use the same PRE representation. + Read native dictionary/slot storage only: never evaluate properties, custom + descriptors, dir(), repr(), or receiver/metaclass attribute lookup hooks. + Opaque values are identity references, not recursive snapshots. In particular + tensor references do not claim tensor-content or hash coverage. + + Dictionary entries and slots are separate because inherited slots (and the + dictionary) can contain distinct values with the same name. Unset slots and + inaccessible dictionary storage are explicit. Native extension storage not + exposed by these descriptors is outside stored_attributes_only coverage. + No metadata is cached: each call observes current class and instance storage. """ - if function_name != _PACKET_CONSTRUCTOR: - return None - if type(receiver) is not torch._ops.OpOverloadPacket: - return None - - state: object = object.__getattribute__(receiver, "__dict__") - if not isinstance(state, dict): - raise TypeError("OpOverloadPacket stored state must be a dictionary") - if "_op" in state and "_qualified_op_name" in state: - return None - requested: list[str] | None = None if dump_config is not None: if not isinstance(dump_config, Mapping): raise TypeError("Constructor receiver dump configuration must be a mapping") requested = [] for name in dump_config: - if not isinstance(name, str): + name_type = type(name) + if name_type is not str: raise TypeError("Requested constructor attribute names must be strings") requested.append(name) + receiver_type = type(receiver) + mro: object = _type_metadata(receiver_type, "__mro__") + assert isinstance(mro, tuple) + dictionary_offset: object = _type_metadata(receiver_type, "__dictoffset__") + assert isinstance(dictionary_offset, int) + dictionary_status: Literal["captured", "absent", "unavailable"] = ( + "unavailable" if dictionary_offset else "absent" + ) attributes: dict[str, StoredAttribute] = {} - for name, value in state.items(): - if not isinstance(name, str): - raise TypeError("OpOverloadPacket stored attribute names must be strings") - if requested is None or name in requested: - attributes[name] = _stored_value(value) + slots: list[StoredSlot] = [] + unset_slots: list[SlotIdentity] = [] + available: set[str] = set() + + for owner in mro: + assert isinstance(owner, type) + namespace: object = _type_metadata(owner, "__dict__") + assert isinstance(namespace, types.MappingProxyType) + for name, descriptor in namespace.items(): + name_type = type(name) + if name_type is not str: + raise TypeError("Class attribute names must be strings") + descriptor_type = type(descriptor) + if ( + name == "__dict__" + and dictionary_status != "captured" + and ( + descriptor_type is types.GetSetDescriptorType + or descriptor_type is types.MemberDescriptorType + ) + and descriptor.__name__ == name + and descriptor.__objclass__ is owner + ): + state: object = descriptor.__get__(receiver, receiver_type) + # A metaclass initializer receives a class: its native stored + # namespace is a mappingproxy. Instance dictionaries may be dict + # subclasses, whose overridden methods must not run here. + entries: Iterable[tuple[object, object]] + if isinstance(state, dict): + entries = dict.items(state) + elif isinstance(state, types.MappingProxyType): + entries = state.items() + else: + raise TypeError( + "Constructor stored state must be a dictionary or class namespace" + ) + for field, value in entries: + field_type = type(field) + if field_type is not str: + raise TypeError("Stored attribute names must be strings") + assert isinstance(field, str) + available.add(field) + if requested is None or field in requested: + attributes[field] = _stored_value(value) + dictionary_status = "captured" + elif ( + descriptor_type is types.MemberDescriptorType + and descriptor.__name__ == name + and descriptor.__objclass__ is owner + and name not in ("__dict__", "__weakref__") + and (requested is None or name in requested) + ): + owner_name = _type_name(owner) + try: + slot_value: object = descriptor.__get__(receiver, receiver_type) + except AttributeError: + # Native slot descriptors raise AttributeError for unset storage. + unset_slots.append({"owner_type": owner_name, "name": name}) + else: + available.add(name) + slots.append( + { + "owner_type": owner_name, + "name": name, + "value": _stored_value(slot_value), + } + ) + snapshot: ConstructorSnapshot = { "phase": "before_init", "coverage": "stored_attributes_only", "object_id": id(receiver), + "dictionary_status": dictionary_status, "requested_attributes": requested, "unavailable_attributes": ( - [name for name in requested if name not in state] + [name for name in requested if name not in available] if requested is not None else [] ), "attributes": attributes, + "slots": slots, + "unset_slots": unset_slots, } - return {_PACKET_TYPE: {_SNAPSHOT_FIELD: snapshot}} + return {_type_name(receiver_type): {_SNAPSHOT_FIELD: snapshot}} diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index b950c9d6..59b2658b 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -169,18 +169,25 @@ def to_dict_args_kwargs( kwargs: Mapping[str, object], dump_args_config: Mapping[str, object] | None = None, *, - function_name: str | None = None, + is_constructor: bool = False, + constructor_receiver_name: str | None = None, ) -> SerializedArguments: def serialize_argument(index: int, argument: object) -> dict[str, object]: argument_config = ( dump_args_config[str(index)] if dump_args_config is not None else None ) - if index == 0 and function_name is not None: - snapshot = constructor_receiver_snapshot( - function_name, argument, argument_config - ) - if snapshot is not None: - return snapshot + if index == 0 and is_constructor: + return constructor_receiver_snapshot(argument, argument_config) + if dump_args_config is None: + return var_to_serializable(argument) + return var_to_serializable(argument, argument_config) + + def serialize_keyword(name: str, argument: object) -> dict[str, object]: + argument_config = ( + dump_args_config[name] if dump_args_config is not None else None + ) + if is_constructor and name == constructor_receiver_name: + return constructor_receiver_snapshot(argument, argument_config) if dump_args_config is None: return var_to_serializable(argument) return var_to_serializable(argument, argument_config) @@ -188,7 +195,7 @@ def serialize_argument(index: int, argument: object) -> dict[str, object]: if dump_args_config is None: result: SerializedArguments = { "args": {i: serialize_argument(i, arg) for i, arg in enumerate(args)}, - "kwargs": {k: var_to_serializable(v) for k, v in kwargs.items()}, + "kwargs": {k: serialize_keyword(k, v) for k, v in kwargs.items()}, } else: result = {"args": {}, "kwargs": {}} @@ -200,7 +207,7 @@ def serialize_argument(index: int, argument: object) -> dict[str, object]: for k, v in kwargs.items(): if k in dump_args_config: - kwargs_dicts[k] = var_to_serializable(v, dump_args_config[k]) + kwargs_dicts[k] = serialize_keyword(k, v) result["args"] = args_dicts result["kwargs"] = kwargs_dicts @@ -219,21 +226,23 @@ def to_dict_return_value(result) -> dict | list[dict]: def function_wrapper( - original_function: Callable, + original_function: Callable[..., object], original_function_name: str, is_bound_method: bool, scan_proxy_in_args: bool, dump_stack_trace: bool, dump_args: bool, - dump_args_config, + dump_args_config: Mapping[str, object] | None, dump_ret: bool, - dump_ret_config, + dump_ret_config: Mapping[str, object] | None, trigger_var_dump: bool, var_dump_config: dict[str, object] | None, need_unproxy_args_kwargs: bool, - *args, - **kwargs, -): + is_constructor: bool, + constructor_receiver_name: str | None, + *args: object, + **kwargs: object, +) -> object: """Instrumentation for Function When using this wrapper, pass in the control parameters as positional arguments, as any kwargs are passed to the original function. @@ -268,6 +277,20 @@ def function_wrapper( process_id, thread_id = get_process_thread_id() pre_meta_vars = get_meta_vars() + if is_constructor: + # Missing receivers are left to the original initializer to reject. + receiver_id = ( + id(args[0]) + if args + else ( + id(kwargs[constructor_receiver_name]) + if constructor_receiver_name is not None + and constructor_receiver_name in kwargs + else None + ) + ) + else: + receiver_id = None if not is_bound_method else id(args[0]) pre_record = { "func_call_id": func_call_id, @@ -277,7 +300,7 @@ def function_wrapper( "type": TraceLineType.FUNC_CALL_PRE, "function": original_function_name, "is_bound_method": is_bound_method, - "obj_id": None if not is_bound_method else id(args[0]), + "obj_id": receiver_id, } if dump_stack_trace: @@ -300,9 +323,12 @@ def find_proxy_in_args(args): args[i] = iter(arg_list) find_proxy_in_args(arg_list) - args = list(args) # type: ignore[assignment] - find_proxy_in_args(args) - args = tuple(args) + # An initializer receiver may not support even proxy-marker lookup. + # Preserve it untouched while scanning the ordinary arguments. + receiver_prefix = args[:1] if is_constructor else () + scanned_args = list(args[1:] if is_constructor else args) + find_proxy_in_args(scanned_args) + args = receiver_prefix + tuple(scanned_args) if proxy_in_args: if "proxy_obj_names" not in pre_record: @@ -318,7 +344,11 @@ def find_proxy_in_args(args): ) if dump_args: dict_args_kwargs = to_dict_args_kwargs( - args, kwargs, dump_args_config, function_name=original_function_name + args, + kwargs, + dump_args_config, + is_constructor=is_constructor, + constructor_receiver_name=constructor_receiver_name, ) pre_record["args"] = dict_args_kwargs["args"] pre_record["kwargs"] = dict_args_kwargs["kwargs"] @@ -331,7 +361,28 @@ def find_proxy_in_args(args): ) if need_unproxy_args_kwargs: - args, kwargs = unproxy_args_kwargs(args, kwargs) + # Construct the actual receiver, never an object obtained by reading + # proxy metadata from its partially initialized state. + if is_constructor: + ordinary_kwargs = { + key: value + for key, value in kwargs.items() + if key != constructor_receiver_name + } + ordinary_args, converted_kwargs = unproxy_args_kwargs( + args[1:], ordinary_kwargs + ) + args = args[:1] + tuple(ordinary_args) + kwargs = { + key: ( + value + if key == constructor_receiver_name + else converted_kwargs[key] + ) + for key, value in kwargs.items() + } + else: + args, kwargs = unproxy_args_kwargs(args, kwargs) try: if COLLECT_OVERHEAD_METRICS: @@ -361,7 +412,7 @@ def find_proxy_in_args(args): "exception": typename(e, is_runtime=True), "exception_msg": str(e), "is_bound_method": is_bound_method, - "obj_id": None if not is_bound_method else id(args[0]), + "obj_id": receiver_id, }, ) @@ -386,7 +437,7 @@ def find_proxy_in_args(args): "type": TraceLineType.FUNC_CALL_POST, "function": original_function_name, "is_bound_method": is_bound_method, - "obj_id": None if not is_bound_method else id(args[0]), + "obj_id": receiver_id, } result_to_dump = result @@ -488,7 +539,27 @@ def wrapper( handle_proxy: bool = True, trigger_var_dump: bool = False, var_dump_config: dict[str, object] | None = None, + is_constructor: bool = False, ) -> Callable[..., object]: + # Supplied by class traversal, not guessed from a function name or the + # spelling of its first parameter. Direct wrapper callers opt in explicitly. + is_bound_method = is_bound_method or is_constructor + constructor_receiver_name: str | None = None + if is_constructor: + try: + parameters = tuple(inspect.signature(original_function).parameters.values()) + except ValueError: + if isinstance(original_function, types.FunctionType): + raise + # Native callables may lack a signature; Python's normal construction + # still supplies the receiver positionally. No keyword name is guessed. + pass + else: + if ( + parameters + and parameters[0].kind == inspect.Parameter.POSITIONAL_OR_KEYWORD + ): + constructor_receiver_name = parameters[0].name is_builtin = is_c_level_function(original_function) need_unproxy_args_kwargs = handle_proxy and ( is_builtin or is_funcs_to_be_unproxied(original_function) @@ -513,6 +584,8 @@ def wrapped(*args: object, **kwargs: object) -> object: trigger_var_dump, var_dump_config, need_unproxy_args_kwargs, + is_constructor, + constructor_receiver_name, *args, **kwargs, ) @@ -967,7 +1040,9 @@ def should_disable_dump(self, attr) -> bool: return True return False - def get_wrapped_function(self, func_obj: Callable) -> Callable: + def get_wrapped_function( + self, func_obj: Callable[..., object], *, is_constructor: bool = False + ) -> Callable[..., object]: """Get the wrapped function for the provided function object, based on the instrumentation options provided in instr_opts.json. """ @@ -986,6 +1061,7 @@ def get_wrapped_function(self, func_obj: Callable) -> Callable: disable_dump=self.should_disable_dump(func_obj), dump_stack_trace=self.API_dump_stack_trace, handle_proxy=used_proxy, + is_constructor=is_constructor, ) else: # checking stage instrumentation @@ -998,6 +1074,7 @@ def get_wrapped_function(self, func_obj: Callable) -> Callable: dump_stack_trace=False, disable_dump=True, handle_proxy=used_proxy, + is_constructor=is_constructor, ) func_instr_opt = self.instr_opts.funcs_instr_opts[func_name] @@ -1020,6 +1097,7 @@ def get_wrapped_function(self, func_obj: Callable) -> Callable: else None ), handle_proxy=used_proxy, + is_constructor=is_constructor, trigger_var_dump=self.instr_opts.disable_proxy_dumping and len(func_instr_opt["var_types_to_track"]) > 0, var_dump_config=func_instr_opt["var_types_to_track"], @@ -1027,11 +1105,11 @@ def get_wrapped_function(self, func_obj: Callable) -> Callable: def _instrument_module( self, - pymodule: types.ModuleType | type, - visited_file_paths: set, + pymodule: types.ModuleType | type[object], + visited_file_paths: set[str], recurse_into_sub_module: bool, - depth, - ): + depth: int, + ) -> int: target_name = pymodule.__name__ if not recurse_into_sub_module and inspect.ismodule(pymodule): @@ -1113,7 +1191,11 @@ def _instrument_module( ) attr = funcs_to_be_replaced[typename(attr)] - wrapped = self.get_wrapped_function(attr) + wrapped = self.get_wrapped_function( + attr, + is_constructor=inspect.isclass(pymodule) + and attr_name == "__init__", + ) try: setattr(pymodule, attr_name, wrapped) except Exception as e: