[feat]: capture FX graph metadata during optimization profiles - #19
[feat]: capture FX graph metadata during optimization profiles#19aryan5v wants to merge 10 commits into
Conversation
Add a model-independent FX capture hook that runs alongside the existing worker-side optimization profile. Repeated module stacks are found structurally (any nn.ModuleList whose children share one class), so no architecture is named anywhere in the capture path. - forward hooks record only tensor signatures and call counters - symbolic tracing runs after the profiler window closes, so captured graphs never contaminate the exported timings - export gains optional regions / graph_breaks / unsupported keys plus a separately versioned capture block; readers that only understand rows are unaffected - capture stays off unless the optimization profile is requested and FASTVIDEO_OPTIMIZATION_PROFILE_CAPTURE_FX is set - trace and finalize failures are recorded as data; generation and the timing export continue - payloads are asserted metadata-only: no tensor values, weights, or prompts
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if graph is None: | ||
| raise RuntimeError("trace result has no FX graph") | ||
| operations, dependencies, constants, notes = _extract_graph(graph) |
There was a problem hiding this comment.
no_tensor_inputs graph-break appended on every call, not just the first
Each invocation of _observe that finds no tensor inputs unconditionally appends a new entry to self._graph_breaks. For a model that calls a non-tensor-input scope N times during the profiling window, N entries accumulate before _coalesce reduces them at finalize(). While _coalesce handles this correctly in the final output, the in-flight list is unbounded. A small per-scope dedup set would prevent repeated appends while keeping the semantics identical.
Prompt To Fix With AI
This is a comment left during a code review.
Path: fastvideo/optimization/fx_capture.py
Line: 429-431
Comment:
**`no_tensor_inputs` graph-break appended on every call, not just the first**
Each invocation of `_observe` that finds no tensor inputs unconditionally appends a new entry to `self._graph_breaks`. For a model that calls a non-tensor-input scope N times during the profiling window, N entries accumulate before `_coalesce` reduces them at `finalize()`. While `_coalesce` handles this correctly in the final output, the in-flight list is unbounded. A small per-scope dedup set would prevent repeated appends while keeping the semantics identical.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
1 participant
Stacked on #18 (
universal-ws1-generation-launcher), which adds the worker-side optimization profile this builds on. Fork-only — not for upstream.What this does
Adds a generic FX-capture hook to the existing dedicated profiling generation. Repeated transformer/module calls are captured with no Wan- or LTX-specific code anywhere in the path.
fastvideo/optimization/fx_capture.py(new, torch-only):default_capture_targets()selects anynn.ModuleListwith >=2 identically-typed children — the model-independent signature of a block stack. Every member is hooked so call/shape frequency covers the whole stack; one trace is taken per stack.symbolic_traceruns infinalize(), after the profiler window closes, so captured graphs never contaminate the exported timings.calls/shape_frequency, graph breaks and unsupported ops, with a stable content fingerprint the consumer can recompute as an integrity check.Export schema
The profile JSON gains optional, additive keys —
capture,regions,graph_breaks,unsupported— only when capture ran.capture.capture_schema_versionis versioned independently of the export'sschema_version, so a capture-format change does not invalidate timing-only readers, and older exports (rows only) still parse unchanged.Off by default
Capture requires both
FASTVIDEO_OPTIMIZATION_PROFILE_OUTPUT(existing) andFASTVIDEO_OPTIMIZATION_PROFILE_CAPTURE_FX=1(new, default off). Clean timing runs and normal generation attach no hooks at all. New knobs:..._FX_TRACER,..._FX_MAX_SCOPES(64),..._FX_MAX_SHAPES(8) to bound the export.Failure handling
Target selection, hooking, per-call observation, tracing and finalize are each guarded. Failures are recorded as
graph_breaksentries orcapture.errorsand the run continues — the timing export is still written.Privacy
assert_metadata_only()walks the finished payload before it is written and raises on any forbidden key (prompt,weights,values,activations, ...) or live tensor. Parameter/bufferget_attrreads are recorded as a note, never a value.Testing
pytest fastvideo/tests/optimization/ -q-> 13 passed. CPU-only, no GPU needed; a fake pipeline drives a synthetic block stack through the hook.Covered: region/frequency contents, keyword tensor inputs, capture off when not requested, capture off when no profile output, hook removal plus bit-identical outputs vs. an unhooked run, untraceable block recorded as
fx_trace_failedwith rows still exported, forcedfinalizeexception recorded incapture.errors, shape-variant bounding, no forbidden keys in the written file, and fingerprint stability.Consumer side round-trips: an export produced by this branch loads through the downstream discovery loader, passing its independent fingerprint recomputation and its pure-tensor safety check.
Note
The
mypypre-commit hook fails withFastVideo-main is not a valid Python package name— pre-existing, reproduces on untouched files (checkout directory name), unrelated to this change.yapf,ruffandcodespellpass on all touched files.Greptile Summary
This PR adds a generic FX graph metadata capture hook layered on top of the existing optimization profiler. It introduces
FXCaptureSession, which hooks repeatednn.ModuleListblock stacks, records tensor shape/dtype signatures during the profiling window, then runs FX traces (symbolic→export→dynamofallback) after the window closes to avoid contaminating profiler timings.fastvideo/optimization/fx_capture.py(new, 958 lines): self-contained session class, privacy guard (assert_metadata_only), stable content fingerprinting, and bounded export with graceful failure recording throughout.fastvideo/optimization/profiler.py:optimization_profile()now accepts an optionalmodulesmapping; FX capture is attached before profiling and finalized (including hook removal) in afinallyblock; profiler rows emitted by the FX range hooks gainparent_moduleandscope_kindfields.fastvideo/tests/optimization/test_fx_capture.py(new, 565 lines): 13 CPU-only tests covering region/frequency data, hook cleanup, untraceable blocks, forced finalize failure, shape-variant bounding, privacy guards, and fingerprint stability.Confidence Score: 5/5
Safe to merge; the two findings are edge-case quality issues in best-effort telemetry that degrade gracefully without affecting generation or timing output.
The core generation path is fully protected — capture is opt-in, all hook attachment and finalization is guarded with broad try/except, and failures are recorded as data rather than raised. The two observations are: (1) assert_metadata_only would produce a false-positive finalize_failed error if a model's nn.ModuleList happens to be named exactly a forbidden key (e.g. 'data'), silently dropping capture metadata while leaving timing intact; (2) the 'motionkernel::' prefix string is duplicated between fx_capture.py and profiler.py rather than imported from a single source. Neither affects production inference, timing correctness, or the written profile file in the common case.
Files Needing Attention: fastvideo/optimization/fx_capture.py — the assert_metadata_only call at line 984 and the hardcoded prefix string at line 652 are the two spots worth a second look before any future refactor.
Important Files Changed
Reviews (3): Last reviewed commit: "Merge pull request #20 from aryan5v/agen..." | Re-trigger Greptile