两个任务为什么不能共用一个 current_id
如果两个任务都把自己的编号写进同一个普通全局变量,某个任务暂停后再恢复时,读到的可能已经是另一个任务的编号。这里需要的是上下文隔离。
下面保留 Slime 的 _task_rollout_id、rollout_context() 和 token 恢复结构,删去 hook 和 Sample,只观察任务编号。
import asyncio
from contextlib import contextmanager
from contextvars import ContextVar
_task_rollout_id = ContextVar("rollout_sample_hook_id", default=None)
@contextmanager
def rollout_context(rollout_id):
token = _task_rollout_id.set(rollout_id)
try:
yield
finally:
_task_rollout_id.reset(token)
async def generate(rollout_id):
with rollout_context(rollout_id):
await asyncio.sleep(0)
return _task_rollout_id.get()
async def exercise():
print("task values", await asyncio.gather(generate(10), generate(11)))
print("outer value", _task_rollout_id.get())
if __name__ == "__main__":
asyncio.run(exercise())
下载并运行本节 example.py。只使用 Python 标准库;推荐 Python 3.12,最低 Python 3.11。
运行结果:
task values [10, 11]
outer value None
set 的返回值为什么还要保留
ContextVar.set() 返回 token。reset(token) 把变量恢复到设置之前的状态,比一律设成 None 更适合嵌套调用。
实验中的两个任务都在设置编号后暂停一次。它们恢复后仍读到各自的编号;退出上下文以后,外层仍是默认值 None。
Slime 怎样把它交给 hook
rollout_context 保存和恢复编号;apply_rollout_sample_hooks 读取当前任务编号。hook 的返回值还可能是 awaitable,因此先用 inspect.isawaitable() 判断,再决定是否 await。
展开 Slime 源码:rollout_context
rollout_contextslime/rollout/sample_hooks.py · L11–L27
快照 2f2318653f6f794dddd321eff7c9d4b7b174643f· 查看固定版本源码
_current_rollout_id: int | None = None
_task_rollout_id = ContextVar("rollout_sample_hook_id", default=None)
def set_current_rollout_id(rollout_id: int | None) -> None:
global _current_rollout_id
_current_rollout_id = rollout_id
@contextmanager
def rollout_context(rollout_id):
"""Keep concurrent worker requests' hook contexts independent."""
token = _task_rollout_id.set(rollout_id)
try:
yield
finally:
_task_rollout_id.reset(token)展开 Slime 源码:_apply_to_sample
_apply_to_sampleslime/rollout/sample_hooks.py · L37–L42
快照 2f2318653f6f794dddd321eff7c9d4b7b174643f· 查看固定版本源码
async def _apply_to_sample(args, sample: Sample, paths: list[str], **kwargs) -> Sample:
for path in paths:
hook = load_function(path)
result = hook(args, sample, **_accepted_kwargs(hook, kwargs))
if inspect.isawaitable(result):
result = await resultContextVar 隔离的是上下文中的绑定,不会把绑定到的可变对象自动深拷贝。如果多个上下文都持有同一个列表,它们仍可能修改同一份列表;这时还要考虑状态设计和同步。
复述检查
锁决定谁可以同时操作共享状态;ContextVar 决定这次任务读取哪一个上下文值。两者解决不同问题。
语义核对:contextvars、isawaitable。
本节对应的全部源码位置
下面把同类用法在 Slime 中的每个扫描位置接回本节讲解;点击文件名可查看完整源码快照。重复出现的 API 共用解释,具体输入与清理由所在函数决定。
展开本节 18 条源码记录
| 记录 / 讲解 | 源码位置 | 原代码 |
|---|---|---|
| occ-000856 | slime/observability/trace_utils.py:3 <module> | import contextvars |
| occ-000858 | slime/observability/trace_utils.py:47 <module> | _TRACE_STACK: contextvars.ContextVar[tuple[tuple[str, str], ...]] = contextvars.ContextVar( |
| occ-000859 | slime/observability/trace_utils.py:47 <module> | _TRACE_STACK: contextvars.ContextVar[tuple[tuple[str, str], ...]] = contextvars.ContextVar( |
| occ-000860 | slime/observability/trace_utils.py:51 <module> | _TRACE_HANDLE_STACK: contextvars.ContextVar[tuple[tuple[TraceHandle, ...], ...]] = contextvars.ContextVar( |
| occ-000861 | slime/observability/trace_utils.py:51 <module> | _TRACE_HANDLE_STACK: contextvars.ContextVar[tuple[tuple[TraceHandle, ...], ...]] = contextvars.ContextVar( |
| occ-000877 有接收者线索 | slime/observability/trace_utils.py:363 trace_span | stack_before = _TRACE_STACK.get() |
| occ-000878 有接收者线索 | slime/observability/trace_utils.py:364 trace_span | handle_stack_before = _TRACE_HANDLE_STACK.get() |
| occ-000879 有接收者线索 | slime/observability/trace_utils.py:383 trace_span | token = _TRACE_STACK.set(stack_before + tuple(new_entries)) |
| occ-000880 有接收者线索 | slime/observability/trace_utils.py:384 trace_span | handle_token = _TRACE_HANDLE_STACK.set(handle_stack_before + (tuple(handles),)) |
| occ-000881 有接收者线索 | slime/observability/trace_utils.py:419 trace_span | _TRACE_STACK.reset(token) |
| occ-000882 有接收者线索 | slime/observability/trace_utils.py:423 trace_span | _TRACE_HANDLE_STACK.reset(handle_token) |
| occ-000894 有接收者线索 | slime/observability/trace_utils.py:635 _get_current_parent_span_id | stack = _TRACE_STACK.get() if stack is None else stack |
| occ-000896 有接收者线索 | slime/observability/trace_utils.py:755 _get_current_trace_target | handle_stack = _TRACE_HANDLE_STACK.get() |
| occ-001104 | slime/rollout/sample_hooks.py:5 <module> | from contextvars import ContextVar |
| occ-001105 | slime/rollout/sample_hooks.py:12 <module> | _task_rollout_id = ContextVar("rollout_sample_hook_id", default=None) |
| occ-001106 有接收者线索 | slime/rollout/sample_hooks.py:23 rollout_context | token = _task_rollout_id.set(rollout_id) |
| occ-001107 有接收者线索 | slime/rollout/sample_hooks.py:27 rollout_context | _task_rollout_id.reset(token) |
| occ-001112 有接收者线索 | slime/rollout/sample_hooks.py:58 apply_rollout_sample_hooks | rollout_id = _task_rollout_id.get() |
专题首页 · 上一节 · 下一节:12 · 进程池、子进程与系统边界 · 源码与教程对照 · 全部用法