[XPU]fix cuda event used in XPU model runner #23708
Conversation
jikunshang
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to fix a compatibility issue with torch.cuda.Event on XPU devices by introducing a _torch_cuda_wrapper context manager to monkey-patch torch.cuda.Event during the initialization of XPUModelRunner.
While the approach of using a context manager to temporarily patch the torch.cuda.Event is sound, the current implementation has a critical flaw. It permanently replaces torch.cuda.Event with a placeholder, which can cause side effects across the application. My review includes a critical comment with a suggested fix to properly restore the original torch.cuda.Event after use, ensuring the change is safe and contained.
| class _EventPlaceholder: | ||
|
|
||
| def __init__(self, *args, **kwargs) -> None: | ||
| self.record = lambda: None | ||
| self.synchronize = lambda: None | ||
|
|
||
| try: | ||
| # replace cuda Event with xpu Event, this should work by default | ||
| torch.cuda.Event = torch.xpu.Event | ||
| yield | ||
| finally: | ||
| # if anything goes wrong, just patch it with a placeholder | ||
| torch.cuda.Event = _EventPlaceholder |
There was a problem hiding this comment.
The current implementation of _torch_cuda_wrapper permanently alters the global torch.cuda.Event by setting it to _EventPlaceholder in the finally block. This is a dangerous side effect that can break functionality in other parts of the codebase that rely on the original torch.cuda.Event, especially in mixed-device environments or if other operations use torch.cuda.Event after this runner is initialized.
The finally block should restore the original state of torch.cuda.Event instead of replacing it with a placeholder. This ensures that the monkey-patching is contained within the with statement and does not leak. The _EventPlaceholder class is also no longer needed with this change.
| class _EventPlaceholder: | |
| def __init__(self, *args, **kwargs) -> None: | |
| self.record = lambda: None | |
| self.synchronize = lambda: None | |
| try: | |
| # replace cuda Event with xpu Event, this should work by default | |
| torch.cuda.Event = torch.xpu.Event | |
| yield | |
| finally: | |
| # if anything goes wrong, just patch it with a placeholder | |
| torch.cuda.Event = _EventPlaceholder | |
| original_event = getattr(torch.cuda, "Event", None) | |
| try: | |
| # replace cuda Event with xpu Event, this should work by default | |
| torch.cuda.Event = torch.xpu.Event | |
| yield | |
| finally: | |
| # Restore the original torch.cuda.Event. | |
| if original_event is not None: | |
| torch.cuda.Event = original_event | |
| elif hasattr(torch.cuda, "Event"): | |
| # If it didn't exist before, remove our monkey-patched attribute. | |
| delattr(torch.cuda, "Event") |
github-actions
bot
added
the
ready
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> Signed-off-by: Xiao Yu <xiao.yu@amd.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2 participants
Purpose
#22760 introduced
torch.cuda.Eventin GPUModelRunner init method. this cause some compatible issue in xpu. this PR fix it by adding a api wrapper.Test Plan
CI.
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.