文档导航

生命周期

生命周期

这一页讲每个入口点的事件确切顺序。如果只看一张图,就看 Router 那张;它是 超集。

Agent.predict_batch

predict_batch 是唯一的实现;system_one 和 predict 用一个状态调用它。

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

同一个 ctx 对象贯穿 start、error 和 end,所以 run_id 把它们关联起来,end 钩子可以读 ctx.error。

Agent.system_one / predict

system_one(state, questions, hooks=..., ...)
  └─ predict_batch([state], questions, hooks=..., ...)[0]

所以 system_one 继承每个钩子和同样的生命周期,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]

要点:

  • on_route 在模型加载之前运行,所以钩子可以钉住一个 checkpoint,避免加载另一个。
  • Router 级的 predict 钩子包住整次调用。它们不会被转发进 Agent;一个挂载的 Agent 若有自己 的钩子,也会运行它们,这是预期内的。
  • Router 级的 ctx.skip() 仍然会加上 routing,所以返回形状稳定。

Router.predict_batch

每个结果就是 predict 对该请求返回的东西,所以 Router 级的 predict 钩子在这里也是每请求运行一次: 每个请求都得到自己的 PredictContext、run_id 和 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

要点:

  • 一个 checkpoint 的各个请求,其每个 start 钩子都在它们任何一个 end 钩子之前运行,因为它们共享 前向传播。所以在 on_predict_end 里填充的缓存无法在同一个 checkpoint 组内服务一个重复的状态; 跨调用可以。
  • 同样的原因,请求结束的顺序和它们开始的顺序相反,所以一个在 start 里设置、在 end 里重置东西的 钩子(一个 contextvars 值,一次 OpenTelemetry 的 context.attach / detach)会展开回它发现 时的值。
  • 一个替换 ctx.states、ctx.questions 或 token 预算的 start 钩子只改变它自己的请求:请求是在 它们的 start 钩子运行之后才为前向传播分组的。原地改动一个共享的 questions dict 则不同,也不是 predict 的做法:这个组里没有任何东西会被推理,直到它的所有 start 钩子都运行过,所以改动会抵达 每一个共享该 dict 的请求,包括那些 start 钩子更早运行的请求,也包括调用方。请改为给 ctx.questions 赋一个新的 dict。
  • checkpoint 名在分组之前解析,所以一个用别名("ml")钉住某个请求的 on_route 钩子会共享那个 checkpoint 的前向传播,而 ctx.model 是解析后的名字。
  • 一个请求可以携带它自己的 max_len / head_max_len,那是 predict 作为调用参数接受的 token 预算的逐请求形式。一个设置 ctx.max_len 的 start 钩子会覆盖它,因为钩子在上下文构建之后运行。 要求不同预算的请求无法共享一次前向传播,所以一个混用预算的批次会按预算各做一次 agent.predict_batch 调用。
  • 每个已开始的请求都恰好得到一个 on_predict_end,即便另一个请求的 end 钩子抛异常;第一个这样的 错误在所有它们都运行完之后才抛出。
  • 如果一个 checkpoint 组失败,它每一个已开始的请求都以那个异常失败:每个都拿到 on_error, ctx.error 设为它,然后是 on_predict_end。这包括一次缓存命中,以及一个其问题组已经跑过的 请求,因为调用方为它们中任何一个都拿到异常、拿不到结果;这也意味着 ctx.error 可能是另一个请求 的失败(一个为某个请求抛异常的 start 钩子会让它所在的组失败)。已经完成的 checkpoint 组的请求 已经带着它们的结果结束了,就像更早的 [router.predict(...) for ...] 调用本会做的那样。

模型生命周期

on_load 在一个 checkpoint 被构建时触发;on_evict 在一个被释放时触发。两者都在 Router 内部 锁释放之后运行,所以钩子可以安全地回调进 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) 注册一个已有的 agent,且不会触发 on_load,因为没有构建任何 checkpoint。

用 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

一个能用的缓存见 examples/hooks/cache.py。

空输入

钩子仍会触发,所以审计能看到每一次调用:

输入 end 时的 ctx.results
predict_batch([]) []
predict_batch(states, {})(没有 questions) 每个状态一个空答案载荷
system_one(state, {}) 单个空答案载荷

这些情况下不发生 tokenization 或前向传播,但 on_predict_start 和 on_predict_end 会运行。

顺序规则

  1. 已安装的钩子总在逐调用钩子之前运行。
  2. 在一个列表内,钩子按列表顺序运行。
  3. 对于一个事件,实现它的每个钩子都按那个顺序运行,然后才轮到下一个事件。
  4. 在失败路径上 on_error 在 on_predict_end 之前运行。
  5. 当一次 load 同时驱逐和构建时,on_evict 在 on_load 之前运行。
installed: [A, B]   per-call: [C]
on_predict_start: A, B, C
on_predict_end:   A, B, C

并发

Agent 和 Router 可以从很多线程安全地调用。每次调用创建它自己的 PredictContext,所以上下文 永不跨请求泄漏。唯一共享的状态就是钩子对象本身,所以一个非线程安全的钩子必须要么自己守住它的 状态,要么用 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 串行化的是每次钩子调用,不是整次调用:两次调用仍然可以在事件之间交错。 它用一把可重入锁,所以一个钩子可以回调进同一个 Agent/Router 而不会死锁。

另见