Skip to content

fix(router): process-wide default hooks never reached predict_batch - #379

Merged
NandhaKishorM merged 1 commit into
NandhaKishorM:mainfrom
PerryLink:fix/router-batch-default-hooks
Sep 24, 2026
Merged

NandhaKishorM merged 1 commit into
NandhaKishorM:mainfrom
PerryLink:fix/router-batch-default-hooks

Conversation

@PerryLink

Copy link
Copy Markdown
Contributor

fix(router): process-wide default hooks never reached predict_batch

set_default_hooks documents that defaults "run before installed and per-call hooks for every
Agent, Router and ONNXAgent in the process", and predict_batch promises each request its
own PredictContext. Neither held on the batched path, because its dispatch site was the only
one in the library that did not compose with the registry:

# laya/router.py, predict_batch
results: List[Optional[Dict[str, Any]]] = [None] * len(requests)
active = list(self.hooks)                                    # <- the defect

# laya/router.py:525, predict
active = compose_hooks(self.hooks, hooks, on_predict_start, on_predict_end)

compose_hooks is what merges default_hooks() in; list(self.hooks) reads the instance list
alone. Reproduced on 23a1752 with one recorder registered as a default hook, recording whether
each event carried a Router:

Router.predict        -> [('default', 'start', 'router'), ('default', 'start', 'agent'),
                          ('default', 'end', 'agent'),   ('default', 'end', 'router')]
Router.predict_batch  -> [('default', 'start', 'agent'),  ('default', 'end', 'agent')]
                          ^ no router-level event at all, for either request

The same hook passed to Router(hooks=[...]) does fire router-level on predict_batch, which
isolates it to the composition of active rather than the dispatch around it:

Router(hooks=[recorder]).predict_batch -> [('instance', 'start', 'router'),
                                           ('instance', 'end', 'router')]

After the change, predict_batch emits router-level start/end per request, and the event list
for one request is identical to what predict produces for the same request.

Who this affected: a process that installs an audit, metrics or tracing hook once with
set_default_hooks — the documented pattern in docs/hooks/api.md ("Process-wide defaults"),
docs/hooks/patterns.md and docs/hooks/examples.md — and serves through Router.predict_batch.
Those hooks saw no Router-level context for batched requests: no on_predict_start,
on_predict_end or on_error, ctx.router unset, and ctx.model the repo id rather than the
routed checkpoint. A per-request redaction hook in that position silently does not run, which is
the dangerous direction.

Deliberate choices tested:

  • One line, at the composition site. active = compose_hooks(self.hooks) is what predict already does at router.py:525, and router.py:266, :309 and :407 compose the same way for the load and route events. This is the only call site that did not, so the fix makes it match rather than introducing a new pattern.
  • No behaviour change without default hooks. compose_hooks(self.hooks) with an empty registry returns the same list list(self.hooks) did, so a process that never calls set_default_hooks sees identical dispatch, ordering and forward-pass count. The instance-hook case is asserted too, since it was already working and must stay unchanged.
  • Ordering is unchanged where it already worked. Defaults run before installed and per-call hooks, as the existing defaults/run before installed and per-call check requires; the new checks only add the Router level that was missing.
  • The test asserts agreement between the two entry points, not just the presence of an event. predict was already correct, so "same events as predict for the same request" is the property that was actually violated and the one worth pinning.

The new block in tests/test_hooks.py fails without the production line:

# with `active = list(self.hooks)` restored
139 passed, 2 failed
  FAIL defaults/reach the Router on predict_batch: got [], want [('default','start','router'), ...]
  FAIL defaults/predict and predict_batch give the same event shape: got [], want [...]

# restored
141 passed, 0 failed

The existing default-hook checks in that file (around line 918) use make_fake(), which is an
Agent, so nothing there ever exercised a Router — which is why this survived. The new checks
use the file's existing batch_router() helper.

$ python tests/test_hooks.py                   all hook tests passed   (139 before)
$ python tests/test_hooks_api.py               all hook API tests passed
$ python tests/test_router.py                  all routing tests passed
$ python tests/test_router_memory.py           OK
$ python tests/test_batch.py                   49 passed, 0 failed
$ python -m pytest tests/test_router_batch.py tests/test_predict_batch.py \
                   tests/test_system_one_lang.py tests/test_truncation_direction.py -q
31 passed
$ python -m ruff check laya/ --select=E9,F63,F7,F82,F401,F811 --line-length=120
All checks passed!
$ python -m compileall -q laya/ tests/

Known limitation: open PR #277 rewrites the Router's hook-dispatch call sites for async hooks
and per-hook timeouts, so the two will conflict in this function. #277 does not change this
line, so whichever lands second needs a trivial rebase.

`set_default_hooks` documents that defaults "run before installed and per-call hooks for every
`Agent`, `Router` and `ONNXAgent` in the process", and `predict_batch` promises each request its
own `PredictContext`. Neither held on the batched path, because its dispatch site was the only
one in the library that did not compose with the registry:

```python
# laya/router.py, predict_batch
results: List[Optional[Dict[str, Any]]] = [None] * len(requests)
active = list(self.hooks)                                    # <- the defect

# laya/router.py:525, predict
active = compose_hooks(self.hooks, hooks, on_predict_start, on_predict_end)
```

`compose_hooks` is what merges `default_hooks()` in; `list(self.hooks)` reads the instance list
alone. Reproduced on `23a1752` with one recorder registered as a default hook, recording whether
each event carried a Router:

```
Router.predict        -> [('default', 'start', 'router'), ('default', 'start', 'agent'),
                          ('default', 'end', 'agent'),   ('default', 'end', 'router')]
Router.predict_batch  -> [('default', 'start', 'agent'),  ('default', 'end', 'agent')]
                          ^ no router-level event at all, for either request
```

The same hook passed to `Router(hooks=[...])` does fire router-level on `predict_batch`, which
isolates it to the composition of `active` rather than the dispatch around it:

```
Router(hooks=[recorder]).predict_batch -> [('instance', 'start', 'router'),
                                           ('instance', 'end', 'router')]
```

After the change, `predict_batch` emits router-level start/end per request, and the event list
for one request is identical to what `predict` produces for the same request.

Who this affected: a process that installs an audit, metrics or tracing hook once with
`set_default_hooks` — the documented pattern in `docs/hooks/api.md` ("Process-wide defaults"),
`docs/hooks/patterns.md` and `docs/hooks/examples.md` — and serves through `Router.predict_batch`.
Those hooks saw no Router-level context for batched requests: no `on_predict_start`,
`on_predict_end` or `on_error`, `ctx.router` unset, and `ctx.model` the repo id rather than the
routed checkpoint. A per-request redaction hook in that position silently does not run, which is
the dangerous direction.

Deliberate choices tested:

- **One line, at the composition site.** `active = compose_hooks(self.hooks)` is what `predict` already does at `router.py:525`, and `router.py:266`, `:309` and `:407` compose the same way for the load and route events. This is the only call site that did not, so the fix makes it match rather than introducing a new pattern.
- **No behaviour change without default hooks.** `compose_hooks(self.hooks)` with an empty registry returns the same list `list(self.hooks)` did, so a process that never calls `set_default_hooks` sees identical dispatch, ordering and forward-pass count. The instance-hook case is asserted too, since it was already working and must stay unchanged.
- **Ordering is unchanged where it already worked.** Defaults run before installed and per-call hooks, as the existing `defaults/run before installed and per-call` check requires; the new checks only add the Router level that was missing.
- **The test asserts agreement between the two entry points**, not just the presence of an event. `predict` was already correct, so "same events as `predict` for the same request" is the property that was actually violated and the one worth pinning.

The new block in `tests/test_hooks.py` fails without the production line:

```
# with `active = list(self.hooks)` restored
139 passed, 2 failed
  FAIL defaults/reach the Router on predict_batch: got [], want [('default','start','router'), ...]
  FAIL defaults/predict and predict_batch give the same event shape: got [], want [...]

# restored
141 passed, 0 failed
```

The existing default-hook checks in that file (around line 918) use `make_fake()`, which is an
`Agent`, so nothing there ever exercised a `Router` — which is why this survived. The new checks
use the file's existing `batch_router()` helper.

```
$ python tests/test_hooks.py                   all hook tests passed   (139 before)
$ python tests/test_hooks_api.py               all hook API tests passed
$ python tests/test_router.py                  all routing tests passed
$ python tests/test_router_memory.py           OK
$ python tests/test_batch.py                   49 passed, 0 failed
$ python -m pytest tests/test_router_batch.py tests/test_predict_batch.py \
                   tests/test_system_one_lang.py tests/test_truncation_direction.py -q
31 passed
$ python -m ruff check laya/ --select=E9,F63,F7,F82,F401,F811 --line-length=120
All checks passed!
$ python -m compileall -q laya/ tests/
```

Known limitation: open PR `NandhaKishorM#277` rewrites the Router's hook-dispatch call sites for async hooks
and per-hook timeouts, so the two will conflict in this function. `NandhaKishorM#277` does not change this
line, so whichever lands second needs a trivial rebase.

@emiliano-go emiliano-go left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Confirmed against main (0.3.20): predict_batch reads active = list(self.hooks) while route and predict use compose_hooks, so a process-wide default hook sees the Agent-level events and none of the Router-level ones. The fix is the right one, and since predict_batch takes no per-call hooks=, compose_hooks(self.hooks) is exactly the composition it needs.

The tests are shaped the way I would want: they assert the two entry points produce the same event shape and that an instance hook still fires once alongside defaults, which pins the property that was violated rather than just the patch. I do not see an interaction with my hooks PR (#277); that branch composes in predict/route and leaves predict_batch as main has it, so the two should rebase cleanly either way.

@NandhaKishorM

Copy link
Copy Markdown
Owner

Thank you, this is an important one. A default audit or redaction hook installed with set_default_hooks now sees Router-level events for batched requests too, and testing that predict and predict_batch produce the same event shape pins exactly the property that was broken. Merging.

@NandhaKishorM
NandhaKishorM merged commit 623b4a6 into NandhaKishorM:main Sep 24, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants