Skip to content

Hook Manager

hook_manager

Classes:

  • HookManager –

    Manages hooks (DGHooks) for a DGraph, supporting both shared and key-specific hooks.

HookManager 

HookManager(keys: List[str])

Manages hooks (DGHooks) for a DGraph, supporting both shared and key-specific hooks.

This class allows you to register hooks that modify or enrich batches of graph events via transformation on the current, and optionally past history of the temporal graph. The hook manager executed these transformation transparently to the user during data iteration. Hooks can be shared across all keys or specific to a key. Dependencies between hooks are automatically resolved using topological sorting based on their requires and produces attributes.

Parameters:

  • keys (List[str]) –

    List of valid keys for key-specific hooks. Each key can have its own set of hooks, in addition to shared hooks.

Raises:

Methods:

  • activate –

    Context manager to temporarily set a key as active for hook execution.

  • execute_active_hooks –

    Executes all hooks (shared + key-specific) for the active key on a batch.

  • register –

    Registers a key-specific hook.

  • register_shared –

    Registers a shared hook that runs for all keys.

  • reset_state –

    Resets the internal state of all stateful hooks.

  • resolve_hooks –

    Resolves hook execution order by topologically sorting them based on dependencies.

  • set_active_hooks –

    Sets the currently active key for executing hooks.

  • validate_requirement –

    Validate that the registered hooks satisfy the requirements of a given NNModule.

Source code in tgm/hooks/hook_manager.py 
56
57
58
59
60
61
62
63
64
def __init__(self, keys: List[str]) -> None:
    if not len(keys):
        raise ValueError('HookManager keys list must be non-empty')

    self._dirty: Dict[str, bool] = {k: False for k in keys}
    self._key_to_hooks: Dict[str, List[DGHook]] = {k: [] for k in keys}
    self._shared_hooks: List[DGHook] = []
    self._active_key: str | None = None
    self._registered_key = keys
activate 
activate(key: str) -> Iterator[None]

Context manager to temporarily set a key as active for hook execution.

Parameters:

  • key (str) –

    The key to activate.

Source code in tgm/hooks/hook_manager.py 
214
215
216
217
218
219
220
221
222
223
224
225
226
@contextmanager
def activate(self, key: str) -> Iterator[None]:
    """Context manager to temporarily set a key as active for hook execution.

    Args:
        key (str): The key to activate.
    """
    prev_key = self._active_key
    self.set_active_hooks(key)
    try:
        yield
    finally:
        self._active_key = prev_key  # Restore previous active key
execute_active_hooks 
execute_active_hooks(dg: DGraph, batch: DGBatch) -> DGBatch

Executes all hooks (shared + key-specific) for the active key on a batch.

Parameters:

  • dg (DGraph) –

    The graph on which hooks operate.

  • batch (DGBatch) –

    The batch of events to modify.

Returns:

  • DGBatch ( DGBatch ) –

    The modified batch after all hooks are applied.

Raises:

  • RuntimeError –

    If no active key is set.

  • UnresolvableHookDependenciesError –

    If required attributes are missing or hooks form a cyclic dependency.

Source code in tgm/hooks/hook_manager.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@log_latency(level=logging.DEBUG)
def execute_active_hooks(self, dg: DGraph, batch: DGBatch) -> DGBatch:
    """Executes all hooks (shared + key-specific) for the active key on a batch.

    Args:
        dg (DGraph): The graph on which hooks operate.
        batch (DGBatch): The batch of events to modify.

    Returns:
        DGBatch: The modified batch after all hooks are applied.

    Raises:
        RuntimeError: If no active key is set.
        UnresolvableHookDependenciesError: If required attributes are missing or hooks form a cyclic dependency.
    """
    if self._active_key is None:
        raise RuntimeError('No active key set. Use activate() context manager.')

    # Lazily validate and topological sort if needed
    key = self._active_key
    if self._dirty[key]:
        self.resolve_hooks(key)

    for hook in self._key_to_hooks[key]:
        start_time = time.perf_counter()
        batch = hook(dg, batch)
        latency = time.perf_counter() - start_time
        logger.debug('%s hook executed in %.4fs', hook.__class__.__name__, latency)

    return batch
register 
register(key: str, hook: DGHook) -> None

Registers a key-specific hook.

Parameters:

  • key (str) –

    The key to associate the hook with.

  • hook (DGHook) –

    The hook to register.

Raises:

  • KeyError –

    If key is not a declared key.

  • BadHookProtocolError –

    If hook does not implement the DGHook protocol.

  • RuntimeError –

    If called while a key is active.

Source code in tgm/hooks/hook_manager.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def register(self, key: str, hook: DGHook) -> None:
    """Registers a key-specific hook.

    Args:
        key (str): The key to associate the hook with.
        hook (DGHook): The hook to register.

    Raises:
        KeyError: If `key` is not a declared key.
        BadHookProtocolError: If `hook` does not implement the `DGHook` protocol.
        RuntimeError: If called while a key is active.
    """
    self._ensure_valid_key(key)
    self._ensure_valid_hook(hook)
    self._ensure_no_active_key()
    self._key_to_hooks[key].append(hook)
    self._dirty[key] = True  # Mark registered key as 'dirty'
    logger.debug('Registered hook: %s for key: %s', hook.__class__.__name__, key)
register_shared 
register_shared(hook: DGHook) -> None

Registers a shared hook that runs for all keys.

Parameters:

  • hook (DGHook) –

    The hook to register.

Raises:

  • BadHookProtocolError –

    If hook does not implement the DGHook protocol.

  • RuntimeError –

    If called while a key is active.

Source code in tgm/hooks/hook_manager.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def register_shared(self, hook: DGHook) -> None:
    """Registers a shared hook that runs for all keys.

    Args:
        hook (DGHook): The hook to register.

    Raises:
        BadHookProtocolError: If `hook` does not implement the `DGHook` protocol.
        RuntimeError: If called while a key is active.
    """
    self._ensure_valid_hook(hook)
    self._ensure_no_active_key()
    self._shared_hooks.append(hook)
    for k in self._dirty:  # Mark all keys as 'dirty'
        self._dirty[k] = True
    logger.debug('Registered shared hook: %s', hook.__class__.__name__)
reset_state 
reset_state(key: str | None = None) -> None

Resets the internal state of all stateful hooks.

Parameters:

  • key (str | None, default: None ) –

    If specified, resets only hooks for this key. Otherwise resets all keys and shared hooks.

Source code in tgm/hooks/hook_manager.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def reset_state(self, key: str | None = None) -> None:
    """Resets the internal state of all stateful hooks.

    Args:
        key (str | None): If specified, resets only hooks for this key.
            Otherwise resets all keys and shared hooks.
    """
    logger.debug('Resetting hooks for key: %s', key)
    if key is not None:
        self._ensure_valid_key(key)

    for hook in self._shared_hooks:
        logger.debug('Resetting state for shared hook: %s', hook.__class__.__name__)
        hook.reset_state()

    keys_to_reset = [key] if key is not None else list(self._key_to_hooks.keys())
    for k in keys_to_reset:
        for h in self._key_to_hooks[k]:
            logger.debug('Resetting state for keyed hook: %s', h.__class__.__name__)
            h.reset_state()
resolve_hooks 
resolve_hooks(key: str | None = None) -> None

Resolves hook execution order by topologically sorting them based on dependencies.

Parameters:

  • key (str | None, default: None ) –

    If specified, resolves hooks only for this key. Otherwise resolves all keys.

Raises:

  • KeyError –

    If key is invalid.

  • UnresolvableHookDependenciesError –

    If required attributes are missing or hooks form a cyclic dependency.

Source code in tgm/hooks/hook_manager.py 
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
def resolve_hooks(self, key: str | None = None) -> None:
    """Resolves hook execution order by topologically sorting them based on dependencies.

    Args:
        key (str | None): If specified, resolves hooks only for this key.
            Otherwise resolves all keys.

    Raises:
        KeyError: If `key` is invalid.
        UnresolvableHookDependenciesError: If required attributes are missing or hooks form a cyclic dependency.
    """
    if key is not None:
        self._ensure_valid_key(key)

    keys_to_validate = [key] if key else list(self._key_to_hooks.keys())
    for k in keys_to_validate:
        hooks = self._shared_hooks + [
            h for h in self._key_to_hooks[k] if h not in self._shared_hooks
        ]
        logger.debug('Running topological sort to resolve hooks for key: %s', k)
        self._key_to_hooks[k] = self._topological_sort_hooks(hooks)
        self._dirty[k] = False
set_active_hooks 
set_active_hooks(key: str) -> None

Sets the currently active key for executing hooks.

Parameters:

  • key (str) –

    The key to activate.

Raises:

  • KeyError –

    If key is not a declared key.

Source code in tgm/hooks/hook_manager.py 
126
127
128
129
130
131
132
133
134
135
136
137
def set_active_hooks(self, key: str) -> None:
    """Sets the currently active key for executing hooks.

    Args:
        key (str): The key to activate.

    Raises:
        KeyError: If `key` is not a declared key.
    """
    self._ensure_valid_key(key)
    self._active_key = key
    logger.debug('Set active hooks to key: %s', key)
validate_requirement 
validate_requirement(
    module: EncoderModule, key: str | None = None
) -> None

Validate that the registered hooks satisfy the requirements of a given NNModule.

Checks whether the hooks registered under the specified key (or all keys if none is provided) fulfill the module's declared requirements. Shared hooks are always included in the validation. Logs a confirmation message if all requirements are met. If one or more of the module's requirements are not satisfied by the registered hooks, This method will return error with suggestion to fix. For example:

Cannot resolve the following requirements {'foo', 'neighbour', 'num_edge_eventss'} from any hook registered under key 'train'. Suggestions: - 'foo': Can not find any existing hooks that satisfy this requirement. - 'neighbour': Found keyword 'neighbour' in 'NeighborSamplerHook' documentation. If this hook produces what you are looking for, update the module requirement with the correct name and register 'NeighborSamplerHook' with key 'train'. - 'num_edge_eventss': Do you mean 'num_edge_events' or 'num_node_events'?. If so, please update the module requirement with the correct name and register 'BatchAnalyticsHook' with key 'train' to resolve this.

Parameters:

  • module (NNModule) –

    The module whose requirements are to be validated.

  • key (str | None, default: None ) –

    The key to validate against. If provided, only hooks registered under that key are checked. If None, all registered keys are validated. Defaults to None.

Raises:

  • ValueError –

    If the provided key is not registered in the HookManager.

  • UnresolvableHookDependenciesError –

    If one or more of the module's requirements are not satisfied by the registered hooks.

  • BadEncoderProtocolError –

    If provided encoder doesn't follow the protocol EncoderModule

Source code in tgm/hooks/hook_manager.py 
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
def validate_requirement(
    self, module: EncoderModule, key: str | None = None
) -> None:
    r"""Validate that the registered hooks satisfy the requirements of a given NNModule.

    Checks whether the hooks registered under the specified key (or all keys if none
    is provided) fulfill the module's declared requirements. Shared hooks are always
    included in the validation. Logs a confirmation message if all requirements are met.
    If one or more of the module's requirements are not satisfied by the registered hooks,
    This method will return error with suggestion to fix. For example:

    Cannot resolve the following requirements {'foo', 'neighbour', 'num_edge_eventss'} from any hook registered under key 'train'.
    Suggestions:
            - 'foo': Can not find any existing hooks that satisfy this requirement.
            - 'neighbour': Found keyword 'neighbour' in 'NeighborSamplerHook' documentation. If this hook produces what you are looking for, update the module requirement with the correct name and register 'NeighborSamplerHook' with key 'train'.
            - 'num_edge_eventss': Do you mean 'num_edge_events' or 'num_node_events'?. If so, please update the module requirement with the correct name and register 'BatchAnalyticsHook' with key 'train' to resolve this.

    Args:
        module (NNModule): The module whose requirements are to be validated.
        key (str | None): The key to validate against. If provided, only hooks
            registered under that key are checked. If None, all registered keys
            are validated. Defaults to None.

    Raises:
        ValueError: If the provided key is not registered in the HookManager.
        UnresolvableHookDependenciesError: If one or more of the module's requirements are not
            satisfied by the registered hooks.

        BadEncoderProtocolError: If provided encoder doesn't follow the protocol `EncoderModule`
    """
    if not isinstance(module, EncoderModule):
        raise BadEncoderProtocolError(
            f'Cannot validate {type(module).__name__}: must implement __call__(self, batch: DGBatch, *args: Any, **kwargs: Any) and have `requires` attribute'
        )

    keys = []
    if key is not None:
        # In the case user provide key, we validate only provided key
        self._ensure_valid_key(key)
        keys.append(key)
    else:
        # In the case user provide no key, we validdate all registered keys
        keys.extend(self._key_to_hooks.keys())

    for k in keys:
        hooks_by_key = self._key_to_hooks[k] + self._shared_hooks
        self._resolve_requirements_against_hooks(module.requires, hooks_by_key, k)

    logger.info(
        f'HookManager satisfies all requirements for {module.__class__.__name__}.'
    )