两个任务为什么不能共用一个 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_context
slime/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_sample
slime/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 result

ContextVar 隔离的是上下文中的绑定,不会把绑定到的可变对象自动深拷贝。如果多个上下文都持有同一个列表,它们仍可能修改同一份列表;这时还要考虑状态设计和同步。

复述检查

锁决定谁可以同时操作共享状态;ContextVar 决定这次任务读取哪一个上下文值。两者解决不同问题。

语义核对:contextvars、isawaitable。

本节对应的全部源码位置

下面把同类用法在 Slime 中的每个扫描位置接回本节讲解;点击文件名可查看完整源码快照。重复出现的 API 共用解释,具体输入与清理由所在函数决定。

展开本节 18 条源码记录
记录 / 讲解源码位置原代码
occ-000856slime/observability/trace_utils.py:3
<module>
import contextvars
occ-000858slime/observability/trace_utils.py:47
<module>
_TRACE_STACK: contextvars.ContextVar[tuple[tuple[str, str], ...]] = contextvars.ContextVar(
occ-000859slime/observability/trace_utils.py:47
<module>
_TRACE_STACK: contextvars.ContextVar[tuple[tuple[str, str], ...]] = contextvars.ContextVar(
occ-000860slime/observability/trace_utils.py:51
<module>
_TRACE_HANDLE_STACK: contextvars.ContextVar[tuple[tuple[TraceHandle, ...], ...]] = contextvars.ContextVar(
occ-000861slime/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-001104slime/rollout/sample_hooks.py:5
<module>
from contextvars import ContextVar
occ-001105slime/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 · 进程池、子进程与系统边界 · 源码与教程对照 · 全部用法