fix(router): process-wide default hooks never reached predict_batch - #379
Conversation
`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
left a comment
There was a problem hiding this comment.
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.
|
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. |
fix(router): process-wide default hooks never reached predict_batch
set_default_hooksdocuments that defaults "run before installed and per-call hooks for everyAgent,RouterandONNXAgentin the process", andpredict_batchpromises each request itsown
PredictContext. Neither held on the batched path, because its dispatch site was the onlyone in the library that did not compose with the registry:
compose_hooksis what mergesdefault_hooks()in;list(self.hooks)reads the instance listalone. Reproduced on
23a1752with one recorder registered as a default hook, recording whethereach event carried a Router:
The same hook passed to
Router(hooks=[...])does fire router-level onpredict_batch, whichisolates it to the composition of
activerather than the dispatch around it:After the change,
predict_batchemits router-level start/end per request, and the event listfor one request is identical to what
predictproduces 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 indocs/hooks/api.md("Process-wide defaults"),docs/hooks/patterns.mdanddocs/hooks/examples.md— and serves throughRouter.predict_batch.Those hooks saw no Router-level context for batched requests: no
on_predict_start,on_predict_endoron_error,ctx.routerunset, andctx.modelthe repo id rather than therouted checkpoint. A per-request redaction hook in that position silently does not run, which is
the dangerous direction.
Deliberate choices tested:
active = compose_hooks(self.hooks)is whatpredictalready does atrouter.py:525, androuter.py:266,:309and:407compose 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.compose_hooks(self.hooks)with an empty registry returns the same listlist(self.hooks)did, so a process that never callsset_default_hookssees identical dispatch, ordering and forward-pass count. The instance-hook case is asserted too, since it was already working and must stay unchanged.defaults/run before installed and per-callcheck requires; the new checks only add the Router level that was missing.predictwas already correct, so "same events aspredictfor the same request" is the property that was actually violated and the one worth pinning.The new block in
tests/test_hooks.pyfails without the production line:The existing default-hook checks in that file (around line 918) use
make_fake(), which is anAgent, so nothing there ever exercised aRouter— which is why this survived. The new checksuse the file's existing
batch_router()helper.Known limitation: open PR
#277rewrites the Router's hook-dispatch call sites for async hooksand per-hook timeouts, so the two will conflict in this function.
#277does not change thisline, so whichever lands second needs a trivial rebase.