fix(router): predict_batch hook follow-ups to #289 # Conflicts: # docs/hooks/lifecycle.md # laya/router.py
12 KiB
Lifecycle
This page is the precise order of events for each entry point. If you only read one diagram, read the Router one; it is the superset.
- Agent.predict_batch
- Agent.system_one / predict
- Router.predict
- Router.predict_batch
- Model lifecycle
- Caching with skip
- Empty inputs
- Ordering rules
- Concurrency
Agent.predict_batch
predict_batch is the single implementation; system_one and predict call it with one state.
predict_batch(states, questions, batch_size=..., hooks=..., ...)
│
├─ active = installed hooks + per-call hooks (installed first)
├─ ctx = PredictContext(states, questions, model=self.model_id, agent=self)
│
├─ try:
│ │
│ ├─ on_predict_start ─────────────────────────────┐
│ │ a hook may: │
│ │ • rewrite ctx.states / ctx.questions │
│ │ • set ctx.max_len / ctx.head_max_len │
│ │ • ctx.skip(results) ─────────────┐ │
│ │ • raise (aborts; see errors) │ │
│ │ │ │
│ ├─ if ctx.results is not None: ◄─────────┘ │ cache hit
│ │ skip tokenization and forward │
│ ├─ else: │
│ │ validate states is a list │
│ │ for each batch chunk: │
│ │ _encode_state ─► collate ─► _forward │
│ │ _decode_answers │
│ │ ctx.results = [...] │
│ │ │
│ └─ (any failure here) ──► except BaseException: │
│ ctx.error = exc │
│ on_error │
│ re-raise │
│ │
│ finally: │
│ ctx.elapsed_ms = now - ctx.started_at │
│ if ctx.results: ctx.usage = aggregate_usage(...)│
│ on_predict_end ─────────────────────────────────┘
│
└─ return ctx.results
The same ctx object flows through start, error and end, so run_id correlates them and an end
hook can read ctx.error.
Agent.system_one / predict
system_one(state, questions, hooks=..., ...)
└─ predict_batch([state], questions, hooks=..., ...)[0]
So system_one inherits every hook and the same lifecycle, with ctx.states == [state].
Router.predict
Router.predict(state, questions, model=..., hooks=..., on_predict_start=..., on_predict_end=...)
│
├─ active = installed hooks + per-call hooks
│
├─ route(state, questions, ..., hooks=per-call, hooks_raise=...)
│ │
│ ├─ _route(...) detect script / language / workflow
│ ├─ on_route ──► ctx.decision a hook may replace the decision
│ └─ return ctx.decision
│
├─ load(decision["model"])
│ │
│ ├─ already resident? ──► return it
│ ├─ else build Agent(...) ──► on_load (after the Router lock is released)
│ └─ evict LRU checkpoints ──► on_evict (after the Router lock is released)
│
├─ ctx = PredictContext(states=[state], questions, decision, model=decision.model,
│ agent=agent, router=self)
├─ try:
│ ├─ on_predict_start
│ ├─ if ctx.results is None:
│ │ result = agent.system_one(ctx.states[0], ctx.questions)
│ │ └─ the Agent's own hooks run here (start / forward / end)
│ │ result["routing"] = decision
│ │ ctx.results = [result]
│ └─ else:
│ for each cached result: result.setdefault("routing", decision)
│ └─ (any failure) ──► except: on_error, re-raise
│ └─ finally: elapsed_ms, usage, on_predict_end
│
└─ return ctx.results[0]
Key points:
on_routeruns before the model is loaded, so a hook can pin a checkpoint and avoid loading another one.- Router-level predict hooks wrap the whole call. They are not forwarded into the Agent; an attached Agent with its own hooks runs those too, which is expected.
- A Router-level
ctx.skip()still addsrouting, so the return shape is stable.
Router.predict_batch
Each result is what predict returns for that request, so Router-level predict hooks run per
request here too: every request gets its own PredictContext, run_id and elapsed_ms.
Router.predict_batch(requests, batch_size=...)
│
├─ route_batch(requests) ──► on_route, once per request (no checkpoint loaded yet)
│
└─ for each checkpoint, in order of first appearance:
│
├─ load(checkpoint) ──► on_load / on_evict
├─ for each request of this checkpoint, in input order:
│ ctx = PredictContext(states=[state], questions, decision, model, agent, router,
│ max_len=request.get("max_len"),
│ head_max_len=request.get("head_max_len"))
│ on_predict_start a hook may redact, rewrite, set a token budget or skip
├─ group the requests left to infer by (questions, ctx.max_len, ctx.head_max_len)
│ agent.predict_batch(states, questions, ...) ──► one shared forward pass per group
│ result["routing"] = decision; ctx.results = [result]
├─ ctx.usage, once per request
├─ (any failure) ──► for every started request, in reverse input order:
│ ctx.error = exc, on_error, on_predict_end; then re-raise
└─ on_predict_end, once per request of this checkpoint, in reverse input order
Key points:
- Every start hook of a checkpoint's requests runs before any of their end hooks, because they
share forward passes. A cache that fills in
on_predict_endtherefore cannot serve a duplicate state within the same checkpoint group; it can across calls. - For the same reason the requests end in reverse of the order they started, so a hook that sets
something in start and resets it in end (a
contextvarsvalue, an OpenTelemetrycontext.attach/detach) unwinds to the value it found. - A start hook that replaces
ctx.states,ctx.questionsor the token budget changes its own request only: requests are grouped for the forward pass after their start hooks have run. Mutating a shared questions dict in place is different, and not whatpredictdoes: nothing of the group is inferred until all of its start hooks have run, so the change reaches every request that shares the dict, including those whose start hooks ran earlier, and the caller too. Assign a new dict toctx.questionsinstead. - The checkpoint name is resolved before grouping, so an
on_routehook that pins a request by an alias ("ml") shares that checkpoint's forward pass, andctx.modelis the resolved name. - A request may carry its own
max_len/head_max_len, the per-request form of the token budgetpredicttakes as call arguments. A start hook that setsctx.max_lenoverrides it, because the hook runs after the context is built. Requests that ask for different budgets cannot share a forward pass, so a batch that mixes budgets makes oneagent.predict_batchcall per budget. - Every started request gets exactly one
on_predict_end, even when another request's end hook raises; the first such error is raised after all of them have run. - If a checkpoint group fails, every one of its started requests fails with the exception: each
gets
on_errorwithctx.errorset to it, thenon_predict_end. That includes a cache hit and a request whose question group had already run, because the caller gets the exception and no result for any of them, and it meansctx.errormay be another request's failure (a start hook that raises for one request fails its group). Requests of checkpoint groups that already finished have ended with their results, as the earlier calls of[router.predict(...) for ...]would have.
Model lifecycle
on_load fires when a checkpoint is built; on_evict when one is freed. Both run after the
Router's internal lock is released, so a hook may safely call back into the Router.
load("multilingual")
│
├─ [lock]
│ build Agent(...) (seconds: download + weights)
│ register in _agents / _order
│ evict LRU if over max_loaded ──► evicted = ["english"]
├─ [unlock]
├─ on_evict("english")
└─ on_load("multilingual")
unload("english")
├─ [lock] remove from _agents / _order
├─ [unlock]
└─ on_evict("english")
attach(name, agent) registers an existing agent and does not fire on_load, because no
checkpoint was built.
Caching with skip
on_predict_start
├─ cache hit? ctx.skip([cached_result])
│ └─ forward pass skipped
│ └─ on_predict_end still runs
│ └─ Router adds `routing` if missing
└─ cache miss? nothing
└─ forward pass runs
└─ on_predict_end can store the result
See examples/hooks/cache.py for a working cache.
Empty inputs
Hooks still fire so audit sees every call:
| input | ctx.results at end |
|---|---|
predict_batch([]) |
[] |
predict_batch(states, {}) (no questions) |
one empty-answer payload per state |
system_one(state, {}) |
a single empty-answer payload |
No tokenization or forward pass happens in these cases, but on_predict_start and
on_predict_end run.
Ordering rules
- Installed hooks run before per-call hooks, always.
- Within a list, hooks run in list order.
- For one event, every hook that implements it runs, in that order, before the next event.
on_errorruns beforeon_predict_endon the failure path.on_evictruns beforeon_loadwhen a singleloadboth evicts and builds.
installed: [A, B] per-call: [C]
on_predict_start: A, B, C
on_predict_end: A, B, C
Concurrency
Agent and Router are safe to call from many threads. Each call creates its own
PredictContext, so contexts never leak across requests. The only shared state is the hook
objects themselves, so a hook that is not thread-safe must either guard its own state or be
installed with hooks_concurrent=False.
hooks_concurrent=True (default) hooks_concurrent=False
thread 1 ─┐ thread 1 ─┐
thread 2 ─┼─ hooks run in parallel thread 2 ─┼─ one hook at a time
thread 3 ─┘ thread 3 ─┘ (RLock)
hooks_concurrent=False serialises each hook invocation, not whole calls: two calls can still
interleave between events. It uses a re-entrant lock, so a hook may call back into the same
Agent/Router without deadlocking.
See also
- Errors: what happens when a hook raises, per event.
- Patterns and anti-patterns: how to use the lifecycle well.