diff --git a/docs/experiments/trace-integrity-results.json b/docs/experiments/trace-integrity-results.json new file mode 100644 index 00000000..d44c4508 --- /dev/null +++ b/docs/experiments/trace-integrity-results.json @@ -0,0 +1,125 @@ +{ + "measurement_date": "2026-10-13", + "traincheck_base": "9eb4f320d3db34d372145261704b805b88e39bda", + "benchmark_source": "ae6601978217bc0f06b9a78479cf00f0bc7242c9", + "controlled_reentry": { + "buggy": { + "application_calls": 100, + "events": 400, + "metadata_helper_events": 200, + "application_results_match": true, + "recursion_errors": 0, + "scope": "Controlled injected conversion helper, not a natural workload frequency" + }, + "fixed": { + "application_calls": 100, + "events": 200, + "metadata_helper_events": 0, + "application_results_match": true, + "recursion_errors": 0, + "scope": "Controlled injected conversion helper, not a natural workload frequency" + } + }, + "controlled_buffer_interleavings": { + "scope": "Controlled forced interleavings, not a natural workload frequency or performance benchmark", + "trials_per_mode": 50, + "records_preloaded_per_trial": 100, + "results": { + "buggy": { + "duplicate": { + "records": 10000, + "duplicates": 5000, + "missing": 0 + }, + "loss": { + "records": 5000, + "duplicates": 0, + "missing": 50 + } + }, + "fixed": { + "duplicate": { + "records": 5000, + "duplicates": 0, + "missing": 0 + }, + "loss": { + "records": 5050, + "duplicates": 0, + "missing": 0 + } + } + } + }, + "natural_workload_matrix": { + "runs_completed": 45, + "instrumented_runs": 36, + "conditions": [ + "baseline", + "buggy", + "guard", + "buffer", + "fixed" + ], + "seeds": [ + 42, + 43, + 44 + ], + "training_steps": 5, + "batch_size": 2, + "native_threads": 1, + "runtime": { + "python": "3.10.16", + "torch": "2.2.2+cu121", + "gpu": "RTX A2000 12 GB" + }, + "scope": "Adapted small model workloads with ordinary local tensors; not full published benchmark reproductions. All instrumented variants share the same dynamic-type-name compatibility patch.", + "instrumentation": "Default API coverage, use_full_instr=False, sampling interval one, no model/proxy tracking; training bodies only.", + "workloads": [ + { + "name": "mnist", + "parameters": 1199882, + "instrumented_runs": 12, + "records_per_run": 654, + "raw_bytes_per_run": 472863, + "instrumented_variants_and_seeds_equal": true + }, + { + "name": "transformer", + "parameters": 23522, + "instrumented_runs": 12, + "records_per_run": 2006, + "raw_bytes_per_run": 1564970, + "instrumented_variants_and_seeds_equal": true + }, + { + "name": "dcgan", + "parameters": 190672, + "instrumented_runs": 12, + "records_per_run": 4596, + "raw_bytes_per_run": 3401128, + "instrumented_variants_and_seeds_equal": true + } + ], + "total_instrumented_records": 87072, + "total_instrumented_raw_bytes": 65267532, + "all_recorded_losses_and_final_model_state_hashes_match_baseline": true, + "observed_duplicate_records": 0, + "observed_recursion_error_events": 0, + "natural_trace_size_reduction_observed": false + }, + "evidence_artifact_sha256": { + "matrix.csv": "f196211cc478ec532338de06247895cfa5bcd379770fd66583c5b80189127cfe", + "matrix-summary.json": "a6a18173e3a0550500160b0e0b3169553ee714bb0ce175048cc5998f19c2ba06", + "controlled-races.json": "6b0ab42d2d8f1ff37557627bd8b261610f7d84adf65f88e2025739dfd875c4aa", + "finite-reentry-buggy.json": "25f5dfeab0f68c080682f21f0fd6e0314e7c4254a2a9e538b9081aad3a382475", + "finite-reentry-fixed.json": "18470fb8d6c929296feecf3485242775dd63331e89c0669244d7ed763bdf6595" + }, + "limitations": [ + "Injected helpers and forced interleavings establish possible failure severity, not natural occurrence rates.", + "All 45 ordinary-tensor runs completed; none of the 36 instrumented runs showed trace inflation or a trace-size reduction from these fixes.", + "No general performance improvement or invariant-checking overhead claim is made.", + "Historical incomplete and completed TorchTitan traces are not comparable percentage-size-reduction measurements." + ] +} diff --git a/docs/experiments/trace-integrity.md b/docs/experiments/trace-integrity.md new file mode 100644 index 00000000..4cdff2aa --- /dev/null +++ b/docs/experiments/trace-integrity.md @@ -0,0 +1,68 @@ +# Recorder re-entry and buffer integrity + +Recorder-owned argument conversion, tensor metadata inspection, and JSON encoding can call instrumented APIs again. The wrapper previously did not honor the conversion helper's local suppression flag, so recorder work could generate more records or recurse. Independently, the producer and periodic flusher accessed the same buffer without synchronization: overlapping flushes could duplicate a batch, and a producer append could be cleared before delivery. + +The recorder now uses an exception-safe, thread-local suppression depth around its own work. It releases suppression before executing application code, preserving nested application calls and generator execution. A class context manager preserves frozen framework exceptions. The buffer snapshots mutable values synchronously, then locks append, queue handoff, and clear as one ordered operation. Buffers still belong to individual producer threads for trace routing. + +## Controlled impact + +These tests deliberately inject a conversion helper or force a thread interleaving; they do not estimate how often an ordinary training job triggers either defect. + +| Controlled trigger | Buggy | Fixed | +|---|---:|---:| +| 100 application calls with a wrapped serialization helper | 400 records, including 200 recorder-owned records | 200 records, zero recorder-owned records | +| Two overlapping flushes; 50 trials × 100 records | 10,000 emitted, including 5,000 duplicates | 5,000 emitted, zero duplicates | +| Append during paused flush; 50 trials, 5,050 expected records | 50 records missing | Zero records missing | + +The finite re-entry case preserves application return values and raises no `RecursionError`: self-tracing can inflate a trace without an obvious recursion failure. The buffer race can cause loss as well as duplication; valid pre/post pairing alone cannot detect a completely missing pair. + +## Ordinary workload control + +All 45 runs completed: MNIST CNN, a small Transformer, and reduced-width DCGAN × five conditions × seeds 42/43/44. Each run used five training steps, batch size two, one GPU, and one native CPU thread. Conditions were no instrumentation, the buggy snapshot, guard only, buffer lock only, and both fixes. + +None of the 36 instrumented runs triggered either defect. Each instrumented condition had the same record count and raw byte size at every seed: + +| Workload | Records per run | Raw bytes per run | Buggy → fixed size change | +|---|---:|---:|---:| +| MNIST CNN | 654 | 472,863 | 0% | +| Small Transformer | 2,006 | 1,564,970 | 0% | +| DCGAN | 4,596 | 3,401,128 | 0% | + +All recorded losses and canonical final model-state hashes matched their corresponding uninstrumented baselines. These are adapted micro-workloads with ordinary local tensors on PyTorch 2.2.2+cu121, not full published benchmark reproductions or DTensor coverage. The observations support these correctness repairs; they do not establish a general performance improvement or invariant-checking overhead result. + +## Validation and provenance + +Run the eleven portable regression tests from the checkout: + +```sh +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m unittest tests.test_trace_integrity -v +``` + +The suite covers recorder re-entry, JSON emission re-entry, real tensor snapshots, nested application calls, generators, original and frozen exceptions, thread-local suppression, overlapping flushes, and concurrent producer ordering. + +[Sanitized results](trace-integrity-results.json) retain controlled counts, workload settings, base/source revisions, and evidence-file hashes. Measurements used TrainCheck base `9eb4f320d3db34d372145261704b805b88e39bda` and adapted benchmark source `ae6601978217bc0f06b9a78479cf00f0bc7242c9`; all instrumented conditions shared the same dynamic-type-name compatibility patch. No private paths, raw training traces, or checkpoint data are included. + +## Contract hardening (2026-10-25) + +Runtime checkpoint `aba4ee8e164f4d24eedae97ea641cd9a6f69c343` makes the +thread-local suppression state, context-manager exception/traceback interface, +and buffer queue/record types explicit. Suppression starts at zero independently +in each thread; process/thread IDs retain their existing lazy initialization. +`__exit__` returns `Literal[False]`, preserving application exceptions. Buffer +serialization and queue failures propagate; a failed enqueue retains its batch +for a later successful flush without dropping or duplicating records. + +Validation used Python 3.10.16, PyTorch 2.2.2, and one OpenMP/MKL/OpenBLAS thread: + +```sh +python -m unittest discover -s tests -p 'test_trac*.py' +pre-commit run --files traincheck/instrumentor/tracer.py traincheck/instrumentor/dumper.py tests/test_tracing_contracts.py +mypy --follow-imports=silent --ignore-missing-imports --check-untyped-defs --disallow-untyped-defs --disallow-any-generics tests/test_tracing_contracts.py +``` + +All 15 tests passed (11 existing regressions and 4 state/write-failure contract +tests). All repository-pinned hooks passed; the stronger check of the new test +module used pinned mypy 1.9. Legacy serializer and proxy-wrapper annotations are +outside this narrowly scoped change; this does not claim repository-wide strict +typing. No benchmark timings were rerun for this typing/contract checkpoint; +the experiment numbers above remain evidence for their original revisions. diff --git a/tests/test_trace_integrity.py b/tests/test_trace_integrity.py new file mode 100644 index 00000000..26d9e0b2 --- /dev/null +++ b/tests/test_trace_integrity.py @@ -0,0 +1,272 @@ +"""Regression coverage for tracer self-reentry and concurrent buffer flushing.""" + +import copy +import json +import threading +import unittest +from queue import Queue +from unittest.mock import patch + +from traincheck.instrumentor import dumper, tracer + + +def wrapped(fn, name="outer", scan_proxy=False): + def call(*args, **kwargs): + return tracer.function_wrapper( + fn, + name, + False, + scan_proxy, + False, + True, + None, + True, + None, + False, + {}, + False, + *args, + **kwargs, + ) + + return call + + +class TracingTests(unittest.TestCase): + def setUp(self): + self.records = [] + self.patches = [ + patch.object(tracer.config, "DISABLE_WRAPPER", False), + patch.object( + tracer, + "dump_trace_API", + side_effect=lambda r: self.records.append(copy.deepcopy(r)), + ), + ] + for p in self.patches: + p.start() + self.addCleanup(p.stop) + + def test_serializer_reentry_does_not_trace_itself(self): + helper = wrapped(lambda x: x, "serializer_helper") + with patch.object( + tracer, "var_to_serializable", side_effect=lambda x: helper(x) + ): + self.assertEqual(wrapped(lambda x: x + 1)(3), 4) + self.assertEqual([r["function"] for r in self.records], ["outer", "outer"]) + + def test_encoding_reentry_and_nested_guard_restoration(self): + helper = wrapped(lambda: "encoded", "encoder_helper") + + def emit(record): + helper() + self.records.append(copy.deepcopy(record)) + + with ( + patch.object(tracer, "dump_trace_API", side_effect=emit), + patch.object(tracer, "var_to_serializable", side_effect=copy.deepcopy), + ): + wrapped(lambda x: x)(1) + with tracer.suppress_tracing(): + tracer.to_dict_return_value(1) + self.assertTrue(tracer.tracing_suppressed()) + self.assertFalse(tracer.tracing_suppressed()) + self.assertEqual([r["function"] for r in self.records], ["outer", "outer"]) + + def test_nested_user_calls_and_mutable_snapshots(self): + inner = wrapped(lambda x: x.append(2), "inner") + + def outer(x): + inner(x) + return x + + with patch.object(tracer, "var_to_serializable", side_effect=copy.deepcopy): + self.assertEqual(wrapped(outer)([1]), [1, 2]) + self.assertEqual( + [r["function"] for r in self.records], ["outer", "inner", "inner", "outer"] + ) + self.assertEqual(self.records[0]["args"][0], [1]) + self.assertEqual(self.records[-1]["return_values"], [1, 2]) + self.assertEqual( + self.records[0]["func_call_id"], self.records[-1]["func_call_id"] + ) + self.assertEqual( + self.records[1]["func_call_id"], self.records[2]["func_call_id"] + ) + + def test_real_tensor_conversion_preserves_pre_post_shape(self): + import torch + + value = torch.ones(2) + + def resize(tensor): + tensor.resize_(3) + return tensor + + self.assertIs(wrapped(resize)(value), value) + before = next(iter(self.records[0]["args"][0].values())) + after = next(iter(self.records[-1]["return_values"].values())) + self.assertEqual(before["shape"], (2,)) + self.assertEqual(after["shape"], (3,)) + self.assertEqual(self.records[0]["meta_vars"], self.records[-1]["meta_vars"]) + + def test_generator_application_calls_remain_visible(self): + inner = wrapped(lambda: 7, "inner") + + def values(): + yield inner() + + with patch.object(tracer, "var_to_serializable", side_effect=lambda x: repr(x)): + self.assertEqual(wrapped(lambda x: list(x), scan_proxy=True)(values()), [7]) + self.assertEqual( + [r["function"] for r in self.records], ["inner", "inner", "outer", "outer"] + ) + self.assertFalse(tracer.tracing_suppressed()) + + def test_conversion_exception_restores_guard(self): + with patch.object( + tracer, "var_to_serializable", side_effect=ValueError("conversion") + ): + with self.assertRaisesRegex(ValueError, "conversion"): + tracer.to_dict_args_kwargs((1,), {}) + with patch.object(tracer, "var_to_serializable", side_effect=copy.deepcopy): + wrapped(lambda x: x)(1) + self.assertEqual(len(self.records), 2) + + def test_frozen_framework_exception_preserved(self): + from dataclasses import dataclass + + @dataclass(frozen=True, slots=True) + class FrameworkSignal(Exception): + reason: str + + error = FrameworkSignal("cache bypass") + + def fail(): + raise error + + try: + wrapped(fail)() + except FrameworkSignal as caught: + self.assertIs(caught, error) + else: + self.fail("Original exception was swallowed") + self.assertFalse(tracer.tracing_suppressed()) + + def test_original_exception_preserved(self): + error = ValueError("original") + + def fail(): + raise error + + with self.assertRaises(ValueError) as caught: + wrapped(fail)() + self.assertIs(caught.exception, error) + self.assertEqual( + self.records[-1]["type"], tracer.TraceLineType.FUNC_CALL_POST_EXCEPTION + ) + self.assertEqual( + self.records[0]["func_call_id"], self.records[-1]["func_call_id"] + ) + + def test_conversion_suppression_is_thread_local(self): + entered, release = threading.Event(), threading.Event() + + def convert(x): + if x == "blocked": + entered.set() + if not release.wait(5): + raise TimeoutError("test release") + return x + + with patch.object(tracer, "var_to_serializable", side_effect=convert): + thread = threading.Thread( + target=tracer.to_dict_args_kwargs, args=(("blocked",), {}) + ) + thread.start() + try: + self.assertTrue(entered.wait(5)) + self.assertEqual(wrapped(lambda x: x)(3), 3) + finally: + release.set() + thread.join(5) + self.assertEqual(len(self.records), 2) + + +class BufferTests(unittest.TestCase): + def make_buffer(self, queue): + # Drive the periodic thread's flush explicitly, so scheduling is controlled. + with patch.object(threading.Thread, "start"): + return dumper.TraceBuffer( + lambda: queue, buffer_size=100, flush_interval=3600 + ) + + def test_concurrent_flush_exactly_once(self): + entered, release, second_done = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + + class BlockingQueue(Queue): + def put(self, item): + if not entered.is_set(): + entered.set() + if not release.wait(5): + raise TimeoutError("test release") + super().put(item) + + queue = BlockingQueue() + buffer = self.make_buffer(queue) + buffer.add_trace({"id": 1}) + first = threading.Thread(target=buffer._flush) + + def flush_again(): + buffer._flush() + second_done.set() + + second = threading.Thread(target=flush_again) + first.start() + try: + self.assertTrue(entered.wait(5)) + second.start() + second_done.wait(0.1) + finally: + release.set() + first.join(5) + second.join(5) + rows = [] + while not queue.empty(): + rows.extend(json.loads(row) for row in queue.get_nowait().splitlines()) + self.assertEqual(rows, [{"id": 1}]) + + def test_concurrent_producer_and_flush_order_and_snapshot(self): + queue = Queue() + buffer = self.make_buffer(queue) + done = threading.Event() + + def flush(): + while not done.is_set(): + buffer._flush() + + thread = threading.Thread(target=flush) + thread.start() + try: + item = {"id": 0} + for i in range(2000): + item["id"] = i + buffer.add_trace(item) + finally: + done.set() + thread.join(5) + buffer._flush() + rows = [] + while not queue.empty(): + rows.extend( + json.loads(row)["id"] for row in queue.get_nowait().splitlines() + ) + self.assertEqual(rows, list(range(2000))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_tracing_contracts.py b/tests/test_tracing_contracts.py new file mode 100644 index 00000000..1564692d --- /dev/null +++ b/tests/test_tracing_contracts.py @@ -0,0 +1,73 @@ +"""Failure contracts for recorder state and buffered writes.""" + +import threading +import unittest +from queue import Queue +from unittest.mock import patch + +from traincheck.instrumentor import dumper, tracer + + +class TracingContractTests(unittest.TestCase): + def test_application_scope_restores_suppression_after_exception(self) -> None: + error = ValueError("application failure") + with tracer.suppress_tracing(): + with self.assertRaises(ValueError) as caught: + with tracer.application_tracing(): + self.assertFalse(tracer.tracing_suppressed()) + raise error + self.assertIs(caught.exception, error) + self.assertTrue(tracer.tracing_suppressed()) + self.assertFalse(tracer.tracing_suppressed()) + + def test_new_thread_initializes_independent_state(self) -> None: + observations: list[bool] = [] + + def observe() -> None: + observations.append(tracer.tracing_suppressed()) + with tracer.suppress_tracing(): + observations.append(tracer.tracing_suppressed()) + observations.append(tracer.tracing_suppressed()) + + with tracer.suppress_tracing(): + thread = threading.Thread(target=observe) + thread.start() + thread.join(5) + self.assertFalse(thread.is_alive()) + self.assertTrue(tracer.tracing_suppressed()) + self.assertEqual(observations, [False, True, False]) + self.assertFalse(tracer.tracing_suppressed()) + + def test_failed_enqueue_retains_batch_and_propagates(self) -> None: + queue: Queue[str] = Queue() + with patch.object(threading.Thread, "start"): + buffer = dumper.TraceBuffer(lambda: queue, buffer_size=100) + buffer.add_trace({"id": 1}) + error = RuntimeError("queue unavailable") + with patch.object(queue, "put", side_effect=error): + with self.assertRaises(RuntimeError) as caught: + buffer._flush() + self.assertIs(caught.exception, error) + buffer._flush() + self.assertEqual(queue.get_nowait(), '{"id":1}') + self.assertTrue(queue.empty()) + buffer._flush() + self.assertTrue(queue.empty()) + + def test_serialization_failure_preserves_existing_batch(self) -> None: + queue: Queue[str] = Queue() + with patch.object(threading.Thread, "start"): + buffer = dumper.TraceBuffer(lambda: queue, buffer_size=100) + buffer.add_trace({"id": 1}) + error = TypeError("unsupported record") + with patch.object(dumper, "serialize", side_effect=error): + with self.assertRaises(TypeError) as caught: + buffer.add_trace({"id": object()}) + self.assertIs(caught.exception, error) + buffer._flush() + self.assertEqual(queue.get_nowait(), '{"id":1}') + self.assertTrue(queue.empty()) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/instrumentor/dumper.py b/traincheck/instrumentor/dumper.py index b808e986..5a852d33 100644 --- a/traincheck/instrumentor/dumper.py +++ b/traincheck/instrumentor/dumper.py @@ -4,7 +4,7 @@ import threading import time from queue import Empty, Queue -from typing import Hashable +from typing import Callable, Hashable import orjson import torch @@ -181,43 +181,53 @@ def get_trace_VAR_dumper_queue(): class TraceBuffer: """Buffer traces and flush them to the trace log file periodically or when the buffer is full. - Caveat: **This class must be instantiated PER THREAD*, as the buffer is not thread-safe. + Caveat: **This class must be instantiated PER THREAD*, to preserve per-thread trace routing. Also, the queue must be created at the same time as the buffer (as did in the __init__ method). Previously, the queue was dynamic getted from the queue_getter method but it caused a mysterious bug, where a trace file contains traces of different thread ids and process ids. """ def __init__( - self, queue_getter, buffer_size=BUFFER_SIZE, flush_interval=FLUSH_INTERVAL - ): + self, + queue_getter: Callable[[], Queue[str]], + buffer_size: int = BUFFER_SIZE, + flush_interval: float = FLUSH_INTERVAL, + ) -> None: self.queue = queue_getter() self.buffer_size = buffer_size self.flush_interval = flush_interval - self.buffer = [] + self.buffer: list[str] = [] + self._lock = threading.Lock() self.last_flush_time = time.time() self.flush_thread = threading.Thread(target=self._flush_periodically) self.flush_thread.daemon = True self.flush_thread.start() - def add_trace(self, trace): - self.buffer.append( - serialize(trace) - ) # TODO: serialization step cannot be buffered rn as trace dicts might get modified - if ( - len(self.buffer) >= self.buffer_size - or (time.time() - self.last_flush_time) >= self.flush_interval - ): - self._flush() + def add_trace(self, trace: dict[str, object]) -> None: + # Snapshot mutable values in the producer, before returning to the caller. + serialized = serialize(trace) + with self._lock: + self.buffer.append(serialized) + if ( + len(self.buffer) >= self.buffer_size + or (time.time() - self.last_flush_time) >= self.flush_interval + ): + self._flush_locked() + + def _flush(self) -> None: + with self._lock: + self._flush_locked() - def _flush(self): + def _flush_locked(self) -> None: if not self.buffer: return - # serialize all traces in the buffer + # Keep enqueue and clear atomic with respect to the producer and timer. + # Enqueue under the lock also preserves batch ordering. self.queue.put("\n".join(self.buffer)) self.buffer.clear() self.last_flush_time = time.time() - def _flush_periodically(self): + def _flush_periodically(self) -> None: while not stop_event.is_set(): time.sleep(self.flush_interval) self._flush() diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 2808ba30..6a6f11c0 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -8,8 +8,9 @@ import time import traceback import types +from contextlib import ContextDecorator from importlib.util import find_spec -from typing import Callable, Optional +from typing import Callable, Literal, Optional import torch @@ -46,8 +47,6 @@ IS_INSTRUMENTING = False -DISABLE_WRAPPER = False - # for prompt generation tasks using the transformers library (see traincheck/developer/instr_stage_annotation.py:annotate_answer_start_token_ids) GENERATE_START_TOKEN_ID: None | int = None GENERATE_START_TOKEN_ID_INCLUDE_START_TOKEN = False @@ -55,7 +54,58 @@ COLLECT_OVERHEAD_METRICS = os.environ.get("COLLECT_OVERHEAD_METRICS", "0") == "1" -THREAD_DATA = threading.local() +class _TracingThreadState(threading.local): + # IDs remain lazy, matching get_process_thread_id's existing behavior. + thread_id: int + process_id: int + + def __init__(self) -> None: + self.suppression_depth: int = 0 + + +THREAD_DATA = _TracingThreadState() + + +class suppress_tracing(ContextDecorator): + """Suppress recorder-owned calls independently per thread, exception-safely. + + Use a class rather than a generator context manager: Python 3.11+ contextlib + assigns exception.__traceback__, which breaks frozen framework exceptions. + Never hold this scope over the original application call. + """ + + depth: int + + def _recreate_cm(self) -> "suppress_tracing": + # Decorated conversions may nest or run concurrently on different threads. + return type(self)() + + def __enter__(self) -> "suppress_tracing": + self.depth = THREAD_DATA.suppression_depth + THREAD_DATA.suppression_depth = self.depth + 1 + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: types.TracebackType | None, + ) -> Literal[False]: + THREAD_DATA.suppression_depth = self.depth + return False + + +class application_tracing(suppress_tracing): + """Restore application tracing while advancing a user-supplied generator.""" + + def __enter__(self) -> "application_tracing": + super().__enter__() + THREAD_DATA.suppression_depth = 0 + return self + + +def tracing_suppressed() -> bool: + return THREAD_DATA.suppression_depth > 0 class TraceLineType: @@ -103,9 +153,8 @@ def get_owner_class(func): return owner +@suppress_tracing() def to_dict_args_kwargs(args, kwargs, dump_args_config=None) -> dict: - global DISABLE_WRAPPER - DISABLE_WRAPPER = True if dump_args_config is None: result = { "args": {i: var_to_serializable(arg) for i, arg in enumerate(args)}, @@ -125,20 +174,17 @@ def to_dict_args_kwargs(args, kwargs, dump_args_config=None) -> dict: result["args"] = args_dicts result["kwargs"] = kwargs_dicts - DISABLE_WRAPPER = False return result +@suppress_tracing() def to_dict_return_value(result) -> dict | list[dict]: - global DISABLE_WRAPPER - DISABLE_WRAPPER = True result_dict: dict | list[dict] if isinstance(result, tuple): result_dict = [var_to_serializable(r) for r in result] else: result_dict = var_to_serializable(result) - DISABLE_WRAPPER = False return result_dict @@ -167,7 +213,7 @@ def function_wrapper( global PROCESS_ID - if config.DISABLE_WRAPPER: + if config.DISABLE_WRAPPER or tracing_suppressed(): # TODO: all meta vars update should be done outside the function_wrapper (e.g. step increment) by applying a separate wrapper return original_function(*args, **kwargs) @@ -179,79 +225,81 @@ def function_wrapper( # tracing here and just delegate to the original function. return original_function(*args, **kwargs) - if COLLECT_OVERHEAD_METRICS: - ENTER_PERF_TIME = time.perf_counter() - - func_call_id = get_unique_id() - process_id, thread_id = get_process_thread_id() - - pre_meta_vars = get_meta_vars() - if IS_INSTRUMENTING: # during instrumentation, skip the dumping to avoid infinite recursion # and interference with the import system return original_function(*args, **kwargs) - pre_record = { - "func_call_id": func_call_id, - "thread_id": thread_id, - "process_id": process_id, - "meta_vars": pre_meta_vars, - "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]), - } - - if dump_stack_trace: - pre_record["stack_trace"] = traceback.format_stack() - - if scan_proxy_in_args: - proxy_in_args = [] - - def find_proxy_in_args(args): - for i, arg in enumerate(args): - if is_proxied(arg) or is_proxyparameter(arg): - proxy_in_args.append(arg) - elif type(arg) in [list, tuple]: - find_proxy_in_args(arg) - elif isinstance(arg, types.GeneratorType) and not isinstance( - arg, tuple - ): - arg_list = list(arg) - 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) - - if proxy_in_args: - if "proxy_obj_names" not in pre_record: - pre_record["proxy_obj_names"] = [] - for proxy in proxy_in_args: - if is_proxyparameter(proxy): - pre_record["proxy_obj_names"].append( - [proxy.__dict__["var_name"], "Parameter"] - ) - else: - pre_record["proxy_obj_names"].append( - [proxy.__dict__["var_name"], type(proxy._obj).__name__] - ) - if dump_args: - dict_args_kwargs = to_dict_args_kwargs(args, kwargs, dump_args_config) - pre_record["args"] = dict_args_kwargs["args"] - pre_record["kwargs"] = dict_args_kwargs["kwargs"] - dump_trace_API(pre_record) - - if trigger_var_dump: - """Mimicking the behavior the observer wrapper: pre-observe""" - get_global_registry().dump_modified( - dump_loc=original_function_name, dump_config=var_dump_config - ) + with suppress_tracing(): + if COLLECT_OVERHEAD_METRICS: + ENTER_PERF_TIME = time.perf_counter() + + func_call_id = get_unique_id() + process_id, thread_id = get_process_thread_id() + + pre_meta_vars = get_meta_vars() + + pre_record = { + "func_call_id": func_call_id, + "thread_id": thread_id, + "process_id": process_id, + "meta_vars": pre_meta_vars, + "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]), + } - if need_unproxy_args_kwargs: - args, kwargs = unproxy_args_kwargs(args, kwargs) + if dump_stack_trace: + pre_record["stack_trace"] = traceback.format_stack() + + if scan_proxy_in_args: + proxy_in_args = [] + + def find_proxy_in_args(args): + for i, arg in enumerate(args): + if is_proxied(arg) or is_proxyparameter(arg): + proxy_in_args.append(arg) + elif type(arg) in [list, tuple]: + find_proxy_in_args(arg) + elif isinstance(arg, types.GeneratorType) and not isinstance( + arg, tuple + ): + with application_tracing(): + arg_list = list(arg) + 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) + + if proxy_in_args: + if "proxy_obj_names" not in pre_record: + pre_record["proxy_obj_names"] = [] + for proxy in proxy_in_args: + if is_proxyparameter(proxy): + pre_record["proxy_obj_names"].append( + [proxy.__dict__["var_name"], "Parameter"] + ) + else: + pre_record["proxy_obj_names"].append( + [proxy.__dict__["var_name"], type(proxy._obj).__name__] + ) + if dump_args: + dict_args_kwargs = to_dict_args_kwargs(args, kwargs, dump_args_config) + pre_record["args"] = dict_args_kwargs["args"] + pre_record["kwargs"] = dict_args_kwargs["kwargs"] + dump_trace_API(pre_record) + + if trigger_var_dump: + """Mimicking the behavior the observer wrapper: pre-observe""" + get_global_registry().dump_modified( + dump_loc=original_function_name, dump_config=var_dump_config + ) + + if need_unproxy_args_kwargs: + args, kwargs = unproxy_args_kwargs(args, kwargs) try: if COLLECT_OVERHEAD_METRICS: @@ -260,134 +308,138 @@ def find_proxy_in_args(args): if COLLECT_OVERHEAD_METRICS: ORIG_EXIT_PERF_TIME = time.perf_counter() except Exception as e: - if COLLECT_OVERHEAD_METRICS: - ORIG_EXIT_PERF_TIME = time.perf_counter() + with suppress_tracing(): + if COLLECT_OVERHEAD_METRICS: + ORIG_EXIT_PERF_TIME = time.perf_counter() + + if trigger_var_dump: + """Mimicking the behavior the observer wrapper: post-observe""" + get_global_registry().dump_modified( + dump_loc=original_function_name, dump_config=var_dump_config + ) + + dump_trace_API( + { + "func_call_id": func_call_id, + "thread_id": thread_id, + "process_id": process_id, + "meta_vars": pre_meta_vars, + "type": TraceLineType.FUNC_CALL_POST_EXCEPTION, + "function": original_function_name, + "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]), + }, + ) + if COLLECT_OVERHEAD_METRICS: + EXIT_PERF_TIME = time.perf_counter() + print( + f"WRAPPER TIME: {original_function_name},{ORIG_EXIT_PERF_TIME - ORIG_ENTER_PERF_TIME},{EXIT_PERF_TIME - ENTER_PERF_TIME}" + ) + raise e + + with suppress_tracing(): if trigger_var_dump: - """Mimicking the behavior the observer wrapper: post-observe""" get_global_registry().dump_modified( dump_loc=original_function_name, dump_config=var_dump_config ) - dump_trace_API( - { - "func_call_id": func_call_id, - "thread_id": thread_id, - "process_id": process_id, - "meta_vars": pre_meta_vars, - "type": TraceLineType.FUNC_CALL_POST_EXCEPTION, - "function": original_function_name, - "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]), - }, - ) + post_record = { + "func_call_id": func_call_id, + "thread_id": thread_id, + "process_id": process_id, + "meta_vars": pre_meta_vars, + "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]), + } + + result_to_dump = result + + # if the current function name is transformers.generate, then we will dump the response tokens only, let's see. + # a concrete name: "transformers.models.whisper.modeling_whisper.WhisperForConditionalGeneration.generate" + # we want a pattern that abstracts the specific model name + pattern = "transformers.models.*.generate" + # find matches in the pattern + import re + + if ( + GENERATE_START_TOKEN_ID is not None + and re.match(pattern, original_function_name) + and isinstance(result, torch.Tensor) + ): + print(f"Found match for {original_function_name}") + # the first dimension is the batch size, and each corresponds to a separate response, let's try to match the batch size with the start token ids first + response_starting_indices = [] + for i in range(result.size(0)): + # try to find the match of the start token ids in the response + response = result[i] + # Find all indices where the start_token_id matches + matches = (response == GENERATE_START_TOKEN_ID).nonzero(as_tuple=True)[ + 0 + ] + indexes = matches.tolist() + if len(indexes) == 0: + # No occurrences found + print( + f"start_token_id ({GENERATE_START_TOKEN_ID}) not found in response {i}" + ) + start_index = -1 # Handle case where token is not found + elif len(indexes) > 1: + # Multiple occurrences found, raise an error + raise ValueError( + f"Multiple occurrences of start_token_id ({GENERATE_START_TOKEN_ID}) found in response {i}: {matches.tolist()}" + ) + else: + # Single occurrence found, get the index + start_index = indexes[0] + if not GENERATE_START_TOKEN_ID_INCLUDE_START_TOKEN: + start_index += 1 + + response_starting_indices.append(start_index) + + # compute the length of each response + response_lengths = [] + for i in range(result.size(0)): + response = result[i] + start_index = response_starting_indices[i] + if start_index == -1: + response_lengths.append(0) + else: + response_lengths.append(response.size(0) - start_index) + + result_to_dump = result.detach() + setattr( + result_to_dump, + "_TRAINCHECK_RESPONSE_STARTING_INDICES", + response_starting_indices, + ) + setattr(result_to_dump, "_TRAINCHECK_RESPONSE_LENGTHS", response_lengths) + + print(response_starting_indices) + print(response_lengths) + if dump_ret: + post_record["return_values"] = to_dict_return_value(result_to_dump) + dump_trace_API(post_record) if COLLECT_OVERHEAD_METRICS: EXIT_PERF_TIME = time.perf_counter() print( f"WRAPPER TIME: {original_function_name},{ORIG_EXIT_PERF_TIME - ORIG_ENTER_PERF_TIME},{EXIT_PERF_TIME - ENTER_PERF_TIME}" ) - raise e - - if trigger_var_dump: - get_global_registry().dump_modified( - dump_loc=original_function_name, dump_config=var_dump_config - ) - - post_record = { - "func_call_id": func_call_id, - "thread_id": thread_id, - "process_id": process_id, - "meta_vars": pre_meta_vars, - "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]), - } - - result_to_dump = result - - # if the current function name is transformers.generate, then we will dump the response tokens only, let's see. - # a concrete name: "transformers.models.whisper.modeling_whisper.WhisperForConditionalGeneration.generate" - # we want a pattern that abstracts the specific model name - pattern = "transformers.models.*.generate" - # find matches in the pattern - import re - - if ( - GENERATE_START_TOKEN_ID is not None - and re.match(pattern, original_function_name) - and isinstance(result, torch.Tensor) - ): - print(f"Found match for {original_function_name}") - # the first dimension is the batch size, and each corresponds to a separate response, let's try to match the batch size with the start token ids first - response_starting_indices = [] - for i in range(result.size(0)): - # try to find the match of the start token ids in the response - response = result[i] - # Find all indices where the start_token_id matches - matches = (response == GENERATE_START_TOKEN_ID).nonzero(as_tuple=True)[0] - indexes = matches.tolist() - if len(indexes) == 0: - # No occurrences found - print( - f"start_token_id ({GENERATE_START_TOKEN_ID}) not found in response {i}" - ) - start_index = -1 # Handle case where token is not found - elif len(indexes) > 1: - # Multiple occurrences found, raise an error - raise ValueError( - f"Multiple occurrences of start_token_id ({GENERATE_START_TOKEN_ID}) found in response {i}: {matches.tolist()}" - ) - else: - # Single occurrence found, get the index - start_index = indexes[0] - if not GENERATE_START_TOKEN_ID_INCLUDE_START_TOKEN: - start_index += 1 - - response_starting_indices.append(start_index) - - # compute the length of each response - response_lengths = [] - for i in range(result.size(0)): - response = result[i] - start_index = response_starting_indices[i] - if start_index == -1: - response_lengths.append(0) - else: - response_lengths.append(response.size(0) - start_index) - - result_to_dump = result.detach() - setattr( - result_to_dump, - "_TRAINCHECK_RESPONSE_STARTING_INDICES", - response_starting_indices, - ) - setattr(result_to_dump, "_TRAINCHECK_RESPONSE_LENGTHS", response_lengths) - - print(response_starting_indices) - print(response_lengths) - if dump_ret: - post_record["return_values"] = to_dict_return_value(result_to_dump) - dump_trace_API(post_record) - - if COLLECT_OVERHEAD_METRICS: - EXIT_PERF_TIME = time.perf_counter() - print( - f"WRAPPER TIME: {original_function_name},{ORIG_EXIT_PERF_TIME - ORIG_ENTER_PERF_TIME},{EXIT_PERF_TIME - ENTER_PERF_TIME}" - ) - return result + return result def core_wrapper_proxy(original_function, *args, **kwargs): """Core wrapper that only handles unproxying for built-in functions.""" - global DISABLE_WRAPPER - if DISABLE_WRAPPER: + if tracing_suppressed(): return original_function(*args, **kwargs) - args, kwargs = unproxy_args_kwargs(args, kwargs) + with suppress_tracing(): + args, kwargs = unproxy_args_kwargs(args, kwargs) return original_function(*args, **kwargs)