mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
c4b93e6579
#153622 introduced a hook for getting the relevant code objects after frame tracing. The idea is to have vLLM use this instead of monkey-patching `inline_call_()` to determine the source code files to hash. Unfortunately, the hook runs too late; the vLLM backend needs access to the set of source code filenames while it's running. This PR replaces the newly-added hook with a utility function that a backend can call to get this information. I've made the change in vLLM and can verify that this allows the information to be queried at the right time. Pull Request resolved: https://github.com/pytorch/pytorch/pull/155249 Approved by: https://github.com/zou3519
25 lines
867 B
Python
25 lines
867 B
Python
"""Hook system for Dynamo's guard functionality.
|
|
|
|
This module provides a way to register callback functions that are triggered during
|
|
guard-related operations.
|
|
|
|
The Hooks class manages two types of hook functions:
|
|
- guard_export_fn: Called when guards need to be exported, taking a GuardsSet as input
|
|
- guard_fail_fn: Called when a guard check fails, taking a GuardFail object as input
|
|
These hooks enable customization of guard export and failure handling behaviors.
|
|
"""
|
|
|
|
import dataclasses
|
|
from typing import Callable, Optional
|
|
|
|
from torch._guards import GuardsSet
|
|
|
|
from .types import GuardFail, GuardFilterEntry
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Hooks:
|
|
guard_export_fn: Optional[Callable[[GuardsSet], None]] = None
|
|
guard_fail_fn: Optional[Callable[[GuardFail], None]] = None
|
|
guard_filter_fn: Optional[Callable[[list[GuardFilterEntry]], list[bool]]] = None
|