MMCV学习——基础篇4(Hook)| 八千字:从设计模式到源码解读

MMCV学习——基础篇4(Hook)

Hook 机制在MMCV的各个开源库中应用的十分广泛,它主要用于管理和扩展Runner的整个生命周期。通过Hook机制在Runner的各个生命周期节点调用自定义的函数,可以实现丰富的定制功能

1. 从观察者模式谈起

Hook本身是一种程序设计的机制,并不是某种语言或者框架独有的。在程序设计模式中,有一种模式叫观察者模式就可以通过Hook机制去实现。观察者模式描述的是被观察者(Subject)观察者(Observer)之间的一对多关系,在观察者不需要知道被观察者是谁的情况下,将被观察者的状态改变推送到观察者这里。

1.1 课程更改,学生如何知道?

 下面我们通过一个简单的场景代码来介绍观察者模式:假设一群人订阅了一门课程,如果课程有内容更新,如何比较优雅地让这群人知道这门课程更新了呢?我们天然的方案可能是,每个人定期查询一下课程看看是否更新,但是这样做会导致每隔一段时间大量用户去访问一个课程,显然这样是不合理并且浪费资源的。
 所以我们在这里引入观察者模式,仅仅在被观察者(Subject)和观察者(Observer)之间建立一个抽象的耦合关系,就可以做到让观察者及时感知到被观察者的变化。首先,我们先定义一个Subject(课程)类:

class Subject(object):

    def __init__(self, state: int) -> None:
        self._state = state
        self._observers = []
    
    @property
    def state(self) -> int:
        return self._state
    
    @state.setter
    def state(self, state: int):
        print('===== start change state! =====')
        old_state = self._state
        self._state = state
        self.notify_observers(old_state)
        print('===== end change state! =====')
    
    @property
    def observers(self):
        # read-only property
        return self._observers
    
    def attach(self, observer):
        self._observers.append(observer)
    
    def detach(self, observer):
        self._observers.remove(observer)
    
    def notify_observers(self, old_state):
        for observer in self._observers:
            observer.update(old_state)

    def __str__(self) -> str:
        return f'Subject(state: {self._state})'
  • Subject类有一个state(状态)和observers(观察者列表)实例属性。
  • 要实现观察者模式Subject就需要实现notify_observers方法,在state发生改变时通知所有观察者。
  • attachdetach方法用来管理观察者列表。

 然后再定义一个Observer(观察者)类:

class Observer(object):

    def __init__(self, name: str, subject: Subject) -> None:
        self._name = name
        self._subject = subject
        self._subject.attach(self)
    
    @property
    def name(self):
        # read-only
        return self._name
    
    @property
    def subject(self):
        # read-only
        return self._subject
    
    def update(self, old_state):
        print(f'{self._name}: subject from {old_state} to {self._subject.state}')
    
    def __str__(self) -> str:
        return f'Observer(name: {self._name}, subject: {self._subject})'
  • Observer类有一个name(名字)和subject(课程)只读实例属性。
  • 要实现观察者模式Observer就需要实现update方法以供Subjectnotify_observers中调用。

接下来我们来写一段Running script运行一下观察者模式的示例代码:

if __name__ == '__main__':
    subject = Subject(1)
    observers = [Observer(name, subject) for name in ['Tom', 'Ben', 'Jerry']]
    subject.state = 2
    print('Now detach Tom and change the state!')
    subject.detach(observers[0])
    subject.state = 3
'''
Output:
===== start change state! =====
Tom: subject from 1 to 2
Ben: subject from 1 to 2
Jerry: subject from 1 to 2
===== end change state! =====
Now detach Tom and change the state!
===== start change state! =====
Ben: subject from 2 to 3
Jerry: subject from 2 to 3
===== end change state! =====
'''

1.2 观察者模式到底与Hook有什么关系?

 上面的观察者设计模式的实现依赖于SubjectObserver两个类,但是如果有些时候我们只是想在Subject状态改变的时候唤起某个自定义的函数,而不想费这么大功夫去专门去写个Observer类的时候该怎么做呢?对于Python来说,函数是里面的一等公民,所以我们可以按照下面的方式去实现观察者设计模式

from functools import partial


class Subject(object):

    def __init__(self, state: int) -> None:
        self._state = state
        self._hooks = []
    
    @property
    def state(self) -> int:
        return self._state
    
    @state.setter
    def state(self, state: int):
        print('===== start change state! =====')
        old_state = self._state
        self._state = state
        self.notify_hooks(old_state)
        print('===== end change state! =====')
    
    def attach(self, hook):
        self._hooks.append(hook)
    
    def detach(self, hook):
        self._hooks.remove(hook)
    
    def notify_hooks(self, old_state):
        for hook in self._hooks:
            hook(old_state)

    def __str__(self) -> str:
        return f'Subject(state: {self._state})'


def obs_hook(old_state, name, subject):
    print(f'{name}: subject from {old_state} to {subject.state}')


if __name__ == '__main__':
    subject = Subject(1)
    hook0 = partial(obs_hook, name='Tom', subject=subject)
    print(type(hook0))
    subject.attach(hook0)
    # Now change the state of subject
    subject.state = 2
'''
Output:
<class 'functools.partial'>
===== start change state! =====
Tom: subject from 1 to 2
===== end change state! =====
Now detach Tom and change the state!
'''
  • 我们将Observer这个类在这里简化成了一个obs_hook函数,并通过partial工具给obs_hook函数绑定name和subject参数。
  • Subject的state的setter方法中去激活所有注册的Hook函数并执行。

2. MMCV中的Hook

MMCV这类第三方框架都会按照工作流程进行一定程度地抽象并归纳出一套通用的执行流程(Runner),但是对于第三方框架的开发者来说,并不知道我们用户在使用这个框架时碰到的具体问题,所以既要保证开发时框架的通用性,又要保证使用时用户可以定制化地修改框架的部分逻辑,就需要用到Hook函数了。

2.1 MMCV Runner的生命周期与Hook

 如下图所示,MMCV Runner的生命周期大体上分为这6个阶段,每个阶段都可以插入Hook从而实现扩展功能。当然,Runner中还涉及到了训练(train)和验证(val)模式,可以按照不同的模式划分不同的阶段,具体可以参看MMCV Hook源码
在这里插入图片描述
 下面的代码以EpochBasedRunner为例展示了Runner调用hook的基本流程,用户要做的只是在各个点位注册好自己写的Hook函数就可以实现自定义的功能:

# 运行前准备工作
before_run()

while self.epoch < self._max_epochs:

    # 开始 epoch 迭代前调用
    before_train_epoch()

    for i, data_batch in enumerate(self.data_loader):
        # 开始 iter 迭代前调用
        before_train_iter()

        self.model.train_step()

        # 经过一次迭代后调用
        after_train_iter()

    # 经过一个 epoch 迭代后调用
    after_train_epoch()

# 运行完成后调用
after_run()

2.2 MMCV的Hook分类

 上图是在MMCV官方分享的知乎上拿来的一张Hook分类图,可以看到MMCV中Hook可以分为默认和定制Hook两类。其各自的相关功能图上也写的很清楚了。如果想知道更细节一点的各类Hook的功能实现可以参看MMCV常用 Hook 类简析
在这里插入图片描述

2.3 如何自定义Hook?

 用户如果想要自定义一个Hook就可以选择继承基类或者利用ClosureHook快速注册。如果是选择继承的方式自定义Hook,就需要创建子类去重写插入阶段的函数,下面是v1.7.0版本Hook基类部分源代码:

# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.utils import Registry, is_method_overridden

HOOKS = Registry('hook')


class Hook:
    stages = ('before_run', 'before_train_epoch', 'before_train_iter',
              'after_train_iter', 'after_train_epoch', 'before_val_epoch',
              'before_val_iter', 'after_val_iter', 'after_val_epoch',
              'after_run')

    def before_run(self, runner):
        pass

    def after_run(self, runner):
        pass

    def before_epoch(self, runner):
        pass

    def after_epoch(self, runner):
        pass

    def before_iter(self, runner):
        pass

    def after_iter(self, runner):
        pass

    def before_train_epoch(self, runner):
        self.before_epoch(runner)

    def before_val_epoch(self, runner):
        self.before_epoch(runner)

    def after_train_epoch(self, runner):
        self.after_epoch(runner)

    def after_val_epoch(self, runner):
        self.after_epoch(runner)

    def before_train_iter(self, runner):
        self.before_iter(runner)

    def before_val_iter(self, runner):
        self.before_iter(runner)

    def after_train_iter(self, runner):
        self.after_iter(runner)

    def after_val_iter(self, runner):
        self.after_iter(runner)
	...
  • 为了方便模块管理和从config构建自定义Hook,除了需要继承Hook基类,还需要用@HOOKS.register_module()去注册模块。

 写好了自定义的Hook,就可以创建实例并注册到runner实例中使用了,MMCV中BaseRunner提供了register_hook按照priority优先级方法注册自定义的Hook。

    def register_hook(self,
                      hook: Hook,
                      priority: Union[int, str, Priority] = 'NORMAL') -> None:
        """Register a hook into the hook list.
        The hook will be inserted into a priority queue, with the specified
        priority (See :class:`Priority` for details of priorities).
        For hooks with the same priority, they will be triggered in the same
        order as they are registered.
        Args:
            hook (:obj:`Hook`): The hook to be registered.
            priority (int or str or :obj:`Priority`): Hook priority.
                Lower value means higher priority.
        """
        assert isinstance(hook, Hook)
        if hasattr(hook, 'priority'):
            raise ValueError('"priority" is a reserved attribute for hooks')
        priority = get_priority(priority)
        hook.priority = priority  # type: ignore
        # insert the hook to a sorted list
        inserted = False
        for i in range(len(self._hooks) - 1, -1, -1):
            if priority >= self._hooks[i].priority:  # type: ignore
                self._hooks.insert(i + 1, hook)
                inserted = True
                break
        if not inserted:
            self._hooks.insert(0, hook)
  • 对于那几个训练过程中需要用到的默认Hook,BaseRunner也提供了register_training_hooks去修改它们。

3. 参考资料

DeepSeek 3FS解读源码分析专栏收录该内容5 文章订阅专栏客户端模式3FS 实现了两套客户端模式,FUSE Client和Native Client。前者更方便适配,性能相对较差。后者适合集成性能敏感的应用程序,适配成本较高。接下来做进一步分析。Fuse ClientFuse Client 模式的原理如下图所示。和传统 FUSE 应用类似,在 libfuse 中注册了 FUSE Daemon 实现的 fuse_lowlevel_ops,之后通过 FUSE 的所有的文件操作,都会通过 libfuse 回调到 FUSE Daemon 进行处理。同时在 libfuse 中实现了一个多线程模式来高效读取请求。这种模式对于业务逻辑影响较小,可以做到无感知。但是每次 I/O 单向需要经过两次“用户态-内核态”上下文切换,以及“用户空间-内核空间”之间数据拷贝。Native Client (USRBIO)介绍这个模式之前我们先了解一下 User Space Ring Based IO(简称 USRBIO)[1],它是一组构建在 3FS 上的高速 I/O 函数。用户应用程序能够通过 USRBIO API 直接提交 I/O 请求给 FUSE Daemon 进程中的 3FS I/O queue 来和 FUSE 进程通信,从而做到 kernel bypass。两者之间通过基于共享内存 ior/iov 的机制交换数据,这部分在后面章节介绍。Native Client 模式的原理如下图所示。使用这种模式能有效避免“用户态-内核态”上下文切换,同时做到零数据拷贝,全链路基本无锁化设计,性能上要比 Fuse Client 模式提升很多。(根据我们对 3FS 开源的 fio ioengine hf3fs_usrbio 压测结果看,在不进行参数调优的情况下,USRBIO 比 **Fuse Client **模式顺序写性能提升 20%-40%,其他场景性能还在进一步验证中)3FS 使用 Pybind 定义 Python 扩展模块 hf3fs_py_usrbio,这也方便 Python 能够访问 hf3fs 的功能。以此推测 USRBIO 模式适合在大模型训练和推理等对性能有极致需求的场景中使用。另外,从上面分析我们注意到 Fuse Daemon 在两种客户端模式下都起到核心作用的重要组件。在 Fuse Client 模式中,它通过 fuseMainLoop 创建 FuseClients,注册 fuse 相关 op 的 hook,并根据配置拉起单线程(或多线程) fuse session loop 处理 fuse op。在 USRBIO 模式中,与 I/O 读写链路相关的 USRBIO API 通过共享内存和 Fuse Daemon 通信,部分与 I/O 无关的控制路径请求例如 hardlink,punchhole 等,USRBIO API 则还是通过 ioctl 直接走了内核 FUSE 路径。这可能是一个 tradeoff 的设计,后面会做讨论。基础组件ServerLauncherFuse Daemon 也就是 FuseApplication, 通过 core::ServerLauncher 拉起。同样的还有 MgmtdServer,MetaServer,StorageServer 都是类似的 Daemon。Fuse Daemon 拉起之后就创建一个 FuseClients 进行核心功能操作。FuseClients一个 FuseApplication 包含 一个 FuseClients,一个 FuseClients 和一个挂载点对应。FuseClients 主要包括下图所示组件,其中包含与其他组件(meta,mgmtd,storage)打交道的 "client for client"。FuseClients 在启动时也会初始化 mgmtdClient,创建 StorageClient,metaClient,启动周期性 Sync Runner(用来更新文件长度等元数据),创建 notifyInvalExec 线程池等。同时还为每一个 FuseClients 创建一组 IOV 和 IOR。FuseClients 最重要的部分还是在和 USRBIO 协同设计。下面我们着重分析这部分。USRBIO 的设计和思考3FS USRBIO 设计思想借鉴了 io_uring 以及 SPDK NVMe 协议栈的设计。原生 io_uring [2] 由一组 Submission Queue 和 Completion Queue 组成,每个 queue 是一个 ring buffer。用户进程提交请求到 SQ,内核选择 polling 模式或事件驱动模式处理 SQ 中的请求,完成之后内核向 CQ 队尾 put 完成 entry,应用程序根据 polling 模式或者事件驱动模式处理 CQ 队首的请求。整个过程无锁,共享内存无内存拷贝。在 polling 模式下,io_uring 接近纯用户态 SPDK polling mode 性能,但是 io_uring 需要通过额外的 CPU cost 达到这个效果。3FS USRBIO 的核心设计围绕 ior 和 iov 来开展。 ior 是一个用于用户进程与 FUSE 进程之间通信的小型共享内存环。用户进程将读/写请求入队,而 FUSE 进程从队列中取出这些请求并完成执行。ior 记录的是读写操作的元数据,并不包含用户数据,其对应的用户数据 buffer 指向一个叫做 iov 的共享内存文件。这个 iov 映射到用户进程和 FUSE 进程共享的一段内存,其中 InfiniBand 内存注册由 FUSE 进程管理。在 USRBIO 中,所有的读取数据都会被读取到 iov,而所有的写入数据应由用户先写入 iov。ior 用来管理 op 操作任务,和 io_uring 不同的是这个 queue 中既包括提交 I/O 请求(sqe)又接收完成 I/O 结果(cqe),而且不通过 kernel,纯用户态操作。ior 中包含的 sqeSection 和 cqeSection 的地址范围由创建 ring 的时候计算出来的 entries 个数确定,用来查询 sqe 和 cqe 在 ring 中的 位置。ior 中还包含一个 ringSection,这个 section 用来帮助 sqe 定位 iov id 的索引和位置。如下图所示,sqe 里包含 idx 是 IOArgs* ringSection 这个数组的下标,索引后才是真正的 io 参数。例如:seq -> ringSection[idx] -> IovId -> Iov。USRBIO 中提供了一个 API hf3fs_iorwrap 用来创建和管理 ior,其中 Hf3fsIorHandle 用来管理 ior。之后 hf3fs_iorwrap 会通过 cqeSem 解析 submit-ios 信号量的路径,并通过 sem_open 打开关联信号量,用于 I/O 任务同步。这里的信号量根据优先级被放置在不同目录中。之后在提交 IO 过程中,会 post 信号量通知 cqe section 中 available 的 slots。在 ior 中,通过 IoRingJob 分配工作,任务被拆分成 IoRingJob,每个任务会处理一定数量的 I/O 请求做批处理。和 io_uring 一样,采用 shared memory 减少用户态与内核态切换。1. IoRing 初始化资源2. 提交 I/O 请求 addSqe3. 获取待处理的 I/O 任务 IoRing::jobsToProc4. 处理 I/O 任务 IoRing::process,如上图所示。IoRing::process() -->ioExec.addWrite() --> ioExec.executeWrite()--> ioExec.finishIo()IoRing 中的 ioExec 就是 PioV。PioV::executeWrite() 执行写操作中根据是否需要 truncate chunk,选择将 truncate WriteIO 包到一个 std::vectorstorage::client::WriteIO wios2中,或者直接传输std::vectorstorage::client::WriteIO wios_,最后通过 StorageClient::batchWrite() 将 Write IO 通过发送 RPC 写请求到 Storage 端。其中,写请求 WriteReq 包括 payload,tag,retryCount,userInfo,featureFlags 等段。FuseClients 中最核心的逻辑之一在 ioRingWorker 中。它负责从 FuseClients 的 ior job queue 中拿到一个 ior,并调用 process 处理它。在处理过程中考虑了取消任务的设计,这里使用了一个 co_withCancellation 来封装,它能够在异步操作中优雅地处理任务取消,避免不必要的计算或资源占用,并且支持嵌套任务的取消感知。有关 co_cancellation 的原理可以参考 [3]:另外,还支持可配置的对任务 job 分优先级,优先级高的 job 优先处理。这些优化都能在复杂的场景下让性能得到极致提升。值得提到的一点是,所有的 iovs 共享内存文件在挂载点 3fs-virt/iovs/ 目录下均建有 symlink,指向 /dev/shm 下的对应文件。USRBIO 代码逻辑错综复杂,偏差之处在所难免,在这里抛砖引玉一些阅读代码的思路和头绪,如有错误也请不吝批评指正。关于USRBIO的思考USRBIO 在共享内存设计上使用了映射到物理内存的一个文件上,而不是使用匿名映射到物理内存。这可能是因为用户进程和 FUSE Daemon 进程不是父子进程关系。实现非派生关系进程间的内存共享,只能使用基于文件的映射或 POSIX 共享内存对象。USRBIO 没有采用直接以 SDK 形式,放弃 Fuse Daemon,直接和元数据服务器与 Chunk Server 来通信的方式设计客户端,而采用了关键 I/O 路径使用纯用户态共享内存,非关键路径上依旧复用 libfuse 这种方式。这可能是简化控制链路设计,追求 FUSE 上的复用性,追求关键路径性能考虑。另外在 IoRing 的设计上并没有使用类似 io_uring 中的可配置的 polling 模式,而是采用信号量进行同步,这里暂时还没有理解背后的原因是什么。USRBIO 使用共享内存还是不可避免会带来一些开销和性能损耗,如此设计的本质原因还是所有核心逻辑都做在了 FUSE Daemon 进程中。如果提供重客户端 SDK,所有逻辑都实现在 SDK 中,以动态连接库形式发布给客户端,可能就不需要进行这样的 IoRing 设计,或者只需要保留 io_uring 这样的无锁设计,不再需要共享内存设计。这样的好处和坏处都很鲜明:好处是 SDK 的实现能避免跨进程的通信开销,性能能达到理想的极限;坏处是如果需要保留 FUSE 功能的话需要实现两套代码,逻辑还很雷同,带来较大的开发和维护成本。而且 SDK 的升级比较重,对客户端造成的影响相对较大。当然从工程角度上可以由 FUSE 抽象出公共函数库让 Native Client 直接调用也可以避免重复开发。
最新发布
04-08
### DeepSeek 3FS 架构设计与源码分析 #### 1. 概述 DeepSeek 3FS 是一种专门为大规模分布式存储系统设计的文件系统,其核心设计理念围绕着高性能读带宽展开。为了实现这一目标,3FS 在多个方面进行了权衡,例如牺牲随机写性能、元数据操作效率以及部分 POSIX 兼容性来换取更高的吞吐量和更低的延迟[^1]。 #### 2. 架构设计特点 3FS 的架构主要分为以下几个模块: - **FUSE Daemon**: 提供了一个用户空间接口,允许开发者通过标准文件系统 API 访问底层对象存储服务。 - **I/O Queue**: 用户应用可以直接向 FUSE Daemon 中的 I/O 队列提交请求,减少传统文件系统的开销。 - **Shared Memory Mechanism**: 利用共享内存技术在客户端和服务端之间传递数据,避免了不必要的上下文切换和数据复制。 #### 3. USRBIO 性能优化原理 User Space Ring Based IO (USRBIO) 是建立在 3FS 基础上的高效 I/O 子系统。它的主要优势在于绕过了内核层处理逻辑,实现了真正的 zero-copy 和 lock-free 设计。具体来说: - **Kernel Bypass**: 应用程序无需经过操作系统内核即可完成 I/O 操作,显著降低了 CPU 使用率并提高了响应速度[^2]。 - **Zero Copy**: 数据传输过程中不需要额外缓冲区参与,减少了内存占用和访问次数。 - **Lock-Free Design**: 整条路径几乎没有任何同步原语介入,极大提升了并发能力。 以下是 USRBIO 实现的一个简化伪代码示例: ```c struct usrbio_request { char *buffer; size_t length; }; void submit_io(struct usrbio_queue *queue, struct usrbio_request req) { spin_lock(&queue->lock); list_add_tail(&req.list_entry, &queue->pending_requests); spin_unlock(&queue->lock); wake_up(queue->worker_thread); // Notify worker thread to process the request. } ``` #### 4. Fuse Client vs Native Client (USRBIO) | 特性 | Fuse Client | Native Client (USRBIO) | |---------------------|--------------------------------------|---------------------------------------| | 上下文切换 | 多次进入退出内核 | 完全规避 | | 数据拷贝 | 至少两次 | 不发生 | | 并发控制 | 锁较多 | 几乎无锁 | | 性能表现(顺序写) | 较低 | 提升约20%-40% | 上述表格总结了两种模式的关键差异点。可以看出,在追求极致性能的应用场景中,采用 Native Client 方案无疑更具竞争力。 --- ###
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值