layer_numerics: make stage observations producer-ordered

[oyazdanb]

Summary

• order copy/sparse tensor observations after TorchRec's real producer streams and track each batch externally across copy, sparse, and forward
• add trusted/overlap/skip evidence, a bounded nonblocking drain, and optional post-step observations for separating producer stages from overlapping compute
• harden shutdown/reset/error paths and expand the public docs, recipe examples, CPU/GPU coverage, and customer self-checks

Test plan

• PYTHONPATH=src conda run -n ort_312 python -m pytest tests/instrumentation/test_layer_numerics.py tests/instrumentation/test_layer_numerics_docs.py tests/instrumentation/test_layer_numerics_spec.py tests/instrumentation/test_layer_numerics_runtime.py tests/instrumentation/test_layer_numerics_smoke_gpu.py -q (168 passed, 9 CUDA/Docker skips)
• python -m py_compile for the logger and modified tests
• git diff --check
• real ROCm/TorchRec recom-repro smoke: Pass A, C, and C+ completed 20/20 steps; all 2,604 C and 3,844 C+ records were trusted with no missing phases or skip reasons
• rerun the ROCm smoke after the summary-counter follow-up to confirm pipeline_observer_count=1 and pipeline_ticks=20 in the emitted summary

[Copilot]

Pull request overview

This PR updates the layer_numerics NaN logger’s pipeline-stage “follow” mode to produce timing-safe, producer-ordered stage observations (copy/sparse via TorchRec producer streams + side streams; forward entry ordered on compute stream), adds an optional post-step (C+) observation path, and hardens draining/shutdown behavior with expanded documentation and tests.

Changes:

• Reworked pipeline tracking to avoid mutating customer batch objects by introducing external per-pipeline batch tokens and richer stage evidence auditing.
• Implemented deferred, non-blocking draining for incomplete stage observations, plus bounded deferral and stricter “skip + audit” behavior when evidence can’t be trusted.
• Expanded docs/recipes and added extensive CPU/GPU/runtime test coverage for v6 pipeline wrapper behavior, stage evidence validity, and post-step observations.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/aorta/instrumentation/layer_numerics/instrument_nan_logger.py	Major v6 refactor: producer-ordered stage reads, external batch tokens, deferred draining, post-step observations, and expanded summary audit fields.
src/aorta/instrumentation/layer_numerics/README.md	Updated behavior description and added Buck .par integration guidance + customer self-check jq snippet.
docs/layer-numerics.md	Updated spec schema/docs for stage_reads + new post_step, added operational guidance and expanded output/audit semantics.
recipes/README.md	Updated recipe example NANLOG_SPEC to the new stage-read schema.
tests/instrumentation/test_layer_numerics_spec.py	Added spec validation/rollback tests for post_step and stale flat-var clearing.
tests/instrumentation/test_layer_numerics_smoke_gpu.py	Updated smoke scenarios to new follow schema and added assertions for stage-read audit fields.
tests/instrumentation/test_layer_numerics_runtime.py	Added extensive runtime tests covering wrapper signatures, token lifecycle, evidence invalidation, deferred drain behavior, and post-step semantics.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.