Skip to content

fix(profiling): patch sub-component methods at class level to prevent instance-isolation bug#13468

Open
Anai-Guo wants to merge 3 commits intohuggingface:mainfrom
Anai-Guo:fix-profiling-bound-method-isolation
Open

fix(profiling): patch sub-component methods at class level to prevent instance-isolation bug#13468
Anai-Guo wants to merge 3 commits intohuggingface:mainfrom
Anai-Guo:fix-profiling-bound-method-isolation

Conversation

@Anai-Guo
Copy link
Copy Markdown

Fixes #13462

Root cause

annotate_pipeline called setattr(component, method_name, annotate(method, label)) where method is a bound method (retrieved via getattr(component, method_name)). The resulting wrapper is a plain function that closes over the original bound method — and thus retains a hard reference to the original component instance.

When a pipeline internally copies a sub-component (e.g. audio_scheduler = copy.copy(self.scheduler) in pipeline_ltx2.py), the copy's instance __dict__ inherits the wrapper. Calling audio_scheduler.step(...) then silently dispatches to self.scheduler.step(...), sharing internal state (_step_index) between the two schedulers and producing index-out-of-bounds errors.

Fix

Detect bound methods with inspect.ismethod and patch at the class level instead:

if inspect.ismethod(method):
    setattr(type(component), method_name, annotate(method.__func__, label))
else:
    setattr(component, method_name, annotate(method, label))

Because the wrapper is now set as a class attribute (an unbound function), Python's descriptor protocol re-binds it to whichever instance accesses it — including copies. Both the original scheduler and audio_scheduler will have properly isolated step calls.

Impact

  • examples/profiling/profiling_utils.py only (no production pipeline code changed)
  • Purely additive: adds import inspect and a 5-line if/else guard
  • Pipelines without internal component copies are unaffected

@github-actions github-actions Bot added examples size/S PR with diff < 50 LOC labels Apr 14, 2026
Comment thread examples/profiling/profiling_utils.py Outdated
continue
setattr(component, method_name, annotate(method, label))
if inspect.ismethod(method):
# Wrap the underlying function and patch at the class level so
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Class-level patches should be avoided as they're not transient.

…() callable

Class-level patches are saved before application and the function now
returns a restore() callable that undoes them, making the patches
explicitly transient. This addresses reviewer feedback that class-level
patches without cleanup are non-transient.
@github-actions github-actions Bot added size/S PR with diff < 50 LOC and removed size/S PR with diff < 50 LOC labels Apr 22, 2026
Store the restore callable from annotate_pipeline() and call it at
the end of run() so class-level method patches are cleaned up.
@Anai-Guo
Copy link
Copy Markdown
Author

Thanks for the feedback @sayakpaul! Updated the PR to make the patches explicitly transient:

  1. annotate_pipeline() now returns a restore() callable that undoes all class-level (and instance-level) patches by restoring saved original attributes.

  2. PipelineProfiler stores the restore callable as self._restore_annotations and calls it at the end of run() — before del pipe — so the class is restored to its original state after each profiling run.

Class-level patches are still needed to fix the instance-isolation bug (the copy issue where audio_scheduler = copy.copy(self.scheduler) would otherwise dispatch to the original scheduler's state), but they're now explicitly cleaned up rather than being permanent.

@github-actions github-actions Bot added size/M PR with diff < 200 LOC and removed size/S PR with diff < 50 LOC labels Apr 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

examples size/M PR with diff < 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Decorating bound methods in profiling utils breaks instance isolation

2 participants