Ray Client 架构指南：深入解析 Ray 分布式运行时的 gRPC 客户端/服务器设计与实现

Ray Client 架构指南深入解析 Ray 分布式运行时的 gRPC 客户端/服务器设计与实现【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray导读本文以仓库内 python/ray/util/client/ARCHITECTURE.md 为骨架系统讲解 Ray Client 的整体架构它本质上是一对 gRPC 客户端与服务器服务器端运行ray.init()扮演普通的 Ray Driver并通过 gRPC 连接被远程客户端控制。读完本文你将掌握 Ray Client 的代码布局、gRPC 协议含数据通道与日志通道的设计动机、基于 pickle/cloudpickle 的自定义序列化协议如何弥合客户端桩对象与服务端真实对象以及 Ray Client 如何通过 client mode hook 与 Ray core 集成并了解其两套测试模式。总体架构一个受 gRPC 控制的远程 Driver核心模型Ray Client 是一对 gRPC 客户端与服务器服务器在远程节点上运行ray.init()行为与普通的 Ray Driver 完全一致只是它受 gRPC 连接控制。服务器负责所有簿记工作bookkeeping并为其接入的客户端保持对象的作用域scope。客户端用户代码运行的一端通过 gRPC 把请求发给服务器。代码布局上客户端代码位于ray/util/client服务器代码位于ray/util/client/server。查看服务器启动入口 python/ray/util/client/server/server.py 中的serve()可以看到服务器会同时注册三个 gRPC Servicertask_servicer RayletServicer(ray_connect_handler) data_servicer DataServicer(task_servicer) logs_servicer LogstreamServicer() ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server) ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server) ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server)依赖方向的刻意设计仓库约定ray/util/client避免直接 importray而服务器端本质上是另一个 Ray 应用因此允许直接依赖ray。这一分离有两个目的避免依赖环客户端不反向依赖 Ray core避免循环导入未来可拆分如果将来需要可以将任一端单独抽成独立仓库或独立安装包例如pip install ray_client。从源码看客户端代码确实大量通过from ray.util.client import ray即RayAPIStub而不是直接import ray工作而服务器端 server_pickler.py 则直接import ray、import ray.cloudpickle印证了这一约定。RayAPIStub以对象代替模块的 API 表面模块级全局变量ray类型为RayAPIStub在 python/ray/util/client/init.py 中定义它充当与ray包等价的 API 表面——ray命名空间中的函数都成了RayAPIStub对象上的方法。RayAPIStub是类而非模块这一微妙差异是后续 client mode hook 能够工作的前提见下文集成点章节因为类支持__getattr__可以动态地把未覆盖的 API 转发到客户端实现。客户端对象与服务端对象的对应关系根ray命名空间中的许多对象在客户端都有对应物大部分集中在 python/ray/util/client/common.py。例如ObjectRef - ClientObjectRef ActorID - ClientActorRef RemoteFunc - ClientRemoteFunc这条对应关系有实际的调试价值如果你在 bug 报告中看到的对象类型是ClientObjectRef那么它必然来自客户端代码路径由服务器返回后构造。在 common.py 中可以看到ClientObjectRef继承自raylet.ObjectRef内部持有Future延迟绑定 idClientActorRef继承自raylet.ActorID。两者的析构函数都会在对象失活时向服务器发送 release 请求这是客户端侧引用计数的体现见数据通道小节。ClientRemoteFunc、ClientActorClass、ClientActorHandle、ClientRemoteMethod则分别对应远程函数、Actor 类、Actor 句柄与 Actor 方法ray.remote在客户端模式下由remote_decorator把它们包装为ClientStub子类common.py。协议层gRPC 服务 pickle 数据编码Ray Client 存在两套协议需要区分清楚gRPC 服务API 表面用于实现远程、瘦客户端的各项功能数据编码协议函数与数据的编码方式。由于两端都是 Python这一层采用pickle实际是cloudpickle——客户端对象包括客户端桩对象通过 pickle 被透明地序列化并传输到服务器。简言之gRPC 服务定义了 API 形状pickle 协议定义了 API 之上的数据编码方式。gRPC 服务定义唯一的一份 gRPC 规范位于 src/ray/protobuf/ray_client.proto。proto 文件注释详尽每个字段都说明了用途。它定义了三个 serviceService用途RayletDriver一元unaryRPC 与传统请求/响应如Init、GetObject、PutObject、WaitObject、Schedule、Terminate、ClusterInfo、KV 操作等RayletDataStreamer双向流式数据通道Datapath承载全部请求/响应模式RayletLogStreamer双向流式日志通道Logstream一元 RPCget、put 与函数调用Client 最初以一组一元 RPC 起步功能刚好覆盖最常用的 API。作为 RPC API 的入门它们很适合用来理解协议工作方式Get / Put / Wait标准的对象存取与等待语义Schedule最有意思的一个——它隐含了一个前置的 Put把要执行的函数 put 上去然后再执行它。这些一元 RPC 至今仍保留但已处于可弃用ripe for deprecation状态核心问题在于它们不与持久连接绑定设想负载均衡器后面挂着多台 client-server任何客户端都可能通过一元 RPC 命中任意一台服务器的任意状态。由于我们需要持有 ray ObjectRef 等句柄、保证它们不因超出作用域而被丢弃这些句柄必须与已连接的客户端保持同步。一元 RPC 下x f.remote()可能打到服务器 A而后续的ray.get(x)却打到服务器 B——而 B 上根本没有对应的 ObjectRef。对应的消息结构ClientTask定义了RemoteExecType枚举FUNCTION、ACTOR、METHOD、STATIC_METHOD、NAMED_ACTOR并携带payload_id对已 put 函数的引用、args/kwargs已废弃改由序列化的data字段承载、client_id、namespace、任务选项options与baseline_options以及大对象分块的chunk_id/total_chunks。服务器端 server.py 的Schedule会根据task.type分派到_schedule_function、_schedule_actor、_schedule_method或_schedule_named_actor。数据通道双向流 ClientID 引用计数数据通道是客户端的双向流式连接Datapath它封装了与一元 RPC 相同的全部请求/响应模式见 ray_client.proto 中的DataRequest/DataResponse其oneof type覆盖 get、put、release、init、task、terminate、acknowledge 等全部消息。其设计要点ClientID 关联连接建立之初客户端用 UUID 生成ClientID与服务器关联。只要通道保持打开客户端就算在线。通过跟踪 ClientID服务器可以追踪它为某个特定客户端持有的全部资源并且一旦通道断开channel drops即可判定客户端已断连。引用计数客户端跟踪自己对各种 Ray client 对象的引用数当对象超出作用域时客户端可主动发送ReleaseRequest给服务器做乐观清理optimistic cleanup。其余情况引用计数全部在客户端完成服务器只需知道何时清理。服务器也可以在它认为安全时清理某客户端持有的全部引用。ReleaseRequest中ids是要释放的引用集合ray_client.proto。未来扩展文档指出将来可以增加显式的ClientDisconnection消息以区分客户端主动完成、永远不会回来与客户端正遭遇连接问题两种场景。在客户端实现中python/ray/util/client/dataclient.py 的DataClient维护一个专用线程ray_client_streaming_rpc运行双向流用req_idint32 递增计数器溢出回绕把异步请求与响应配对outstanding_requests记录未完成请求以便断线后重放ready_data存放阻塞式响应asyncio_waiting_data存放异步回调。服务器端 python/ray/util/client/server/dataservicer.py 的DataServicer同样处理这些请求并用OrderedResponseCache配合客户端的AcknowledgeRequest每收到 32 个响应发送一次 ACK见 dataclient.py做去重与缓存清理。大对象分块传输为了支持跨网络传输大对象Ray Client 将对象切成5 MiB 的块OBJECT_TRANSFER_CHUNK_SIZE见 common.pyput 与 task 的请求体在发送端被惰性分块dataclient.py 中的chunk_put/chunk_taskget 的响应在服务器端分块、客户端用ChunkCollector按序重组并支持从start_chunk_id断点续传。gRPC 单条消息上限被放宽到 2GiBGRPC_MAX_MESSAGE_SIZE超过 2GiB 的对象会触发用户警告建议改用 S3 等远程 URI 传输。此外common.py 还设置了 30 秒的 keepalive ping 与 600 秒的超时以适配 ELB 等负载均衡器的 60 秒空闲超时。日志通道独立于数据通道的双向流与数据通道类似还有一个关联的日志通道把日志回传给客户端它是独立的通道因为日志是附属品——即使因断连丢失一些日志通常也可以接受独立通道还让日志聚合器不必实现完整 API 即可接入它是双向流客户端发送LogSettingsRequest控制日志的开关enabled与级别loglevel服务器把产生的日志以LogDatamsglevelname流式回传ray_client.proto。协议注释约定level 0遵循 Python logging 的级别level -1表示 stdoutlevel -2表示 stderr。关于 CloudPickle自定义 Pickler 解决桩对象混用问题如本节开头所述pickle/cloudpickle 是把数据编码为可被 Python 执行的数据以走 gRPC 通用传输的方式。Ray Client 在 python/ray/util/client/client_pickler.py 与 python/ray/util/client/server/server_pickler.py 中提供了自己的 pickle/unpickle 子类目的是解决Client*桩对象混用的问题。用一个例子说明问题的由来假设RemoteFunc f()调用了另一个RemoteFunc g()。f()需要持有对g()的引用比如调用g.remote()于是f()在序列化时序列化数据里会包含一个RemoteFunc类的对象供 worker 端反序列化。在 Ray Client 中f()与g()都是ClientRemoteFunc行为同理。在早期版本里一个ClientRemoteFunc必须知道自己是在服务器端还是客户端并据此决定像普通 RemoteFunc 一样运行还是在客户端发起调用。这导致了一些棘手的 bug——尤其是把客户端桩对象传出去或更糟返回回来时设想f()调用g()而g()构造并返回了一个新的闭包h()。自定义 pickler 的解决方式只要客户端桩对象被序列化就用一个结构体写作时是一个元组PickleStub替代它反序列化时服务器填入对应的非桩对象。反之亦然——如果服务器要编码一个返回/响应中的ObjectRef就把元组放到线上客户端反序列化器再把它还原成ClientObjectRef。PickleStub是一个命名元组字段为type、client_id、ref_id、name、baseline_optionsclient_pickler.py。客户端ClientPickler.persistent_id对RayAPIStub、ClientObjectRef、ClientActorHandle、ClientRemoteFunc、ClientActorClass、ClientRemoteMethod分别生成对应 stub如Ray、Object、Actor、RemoteFunc、RemoteActor、RemoteMethod而服务器端ClientUnpickler.persistent_load按 stub 类型回填Object→self.server.object_refs[pid.client_id][pid.ref_id]Actor→self.server.actor_refs[pid.ref_id]RemoteFunc/RemoteActor→lookup_or_register_func/lookup_or_register_actorRemoteMethod→ 从 actor 句柄取方法server_pickler.py。反向路径上服务器ServerPickler.persistent_id遇到ray.ObjectRef/ray.actor.ActorHandle时把它们登记到对应客户端的对象/actor 表中并生成 stubServerUnpickler.persistent_load客户端侧再把Object/Actorstub 还原为ClientObjectRef/ClientActorHandleclient_pickler.py。这一设计带来的收益客户端侧尽可能只与桩对象打交道服务器侧永远看不到桩对象两端界限干净服务器侧出现了ClientObjectRef属于错误情形而非需要特殊处理的场景服务器侧可以像普通 Ray 一样工作、只处理普通 Ray 对象并在发送时把它们透明地编码成客户端对象因为永不相交never the twain shall meet建模与调试都容易得多。另外client_pickler.py 还处理了递归/自引用函数或 Actor 类在 put 自己过程中被编码时_ref为InProgressSentinel会生成RemoteFuncSelfReference/RemoteActorSelfReferencestub服务器端对应构造ClientReferenceFunction/ClientReferenceActorserver_pickler.py从而支持函数的参数里包含它自己这类递归闭包场景。与 Ray core 的集成点client mode hook为了提供与 Ray core 无缝的客户端体验需要包装一部分核心 Ray 函数例如ray.get()。Python 的动态特性在此帮了大忙如前所述RayAPIStub是类而非模块因此可以在 API 层实现__getattr__把调用重定向到任何想去的地方模块做不到这一点除非等 Python 3.6 废弃后使用 PEP 562 的__module_getattr__。如果raycore 本身是对象而非命名空间里的函数就不需要包装、直接替换实现即可但项目必须保持向后兼容。所有与 Ray core 的集成点都位于 python/ray/_private/client_mode_hook.pyclient_mode_hook装饰器用来包装 Ray core 函数。当client_mode_should_convert()根据环境变量返回True时装饰器生效把调用转发给ray/util/client对象getattr(ray, func.__name__)(*args, **kwargs)否则照常调用原函数client_mode_hook.py。模式开关is_client_mode_enabled默认关闭在ray.client(...).connect()或测试中开启RAY_CLIENT_MODE1环境变量可让整个进程默认开启 client mode主要用于测试见 client_mode_hook.py。disable_client_hook()上下文管理器线程本地地把 hook 状态置为 False服务器端所有真正执行 Ray 操作的地方都用它包住例如 server.py 中Init、Schedule、PutObject、WaitObject等实现都写在with disable_client_hook():内保证服务器端调用的是真实的 Ray 函数而不是客户端实现避免死循环。client_mode_convert_function/client_mode_convert_actor用于把预先注册的 RemoteFunction / ActorClass透明转换为ClientRemoteFunc/ClientActorClass典型场景是函数在库加载早期、尚未进入 client mode 时就被ray.remote装饰client_mode_hook.py。client_mode_wrap用于实现不属于主ray.*API、却需要服务器端执行的功能例如 Placement Group 的创建——客户端模式下会把函数包装成ray.remote(num_cpus0)任务去服务器端执行client_mode_hook.py。对应地python/ray/util/client/init.py 的_ClientContext.connect()会调用_explicitly_enable_client_mode()强制开启 client mode并在断连时由disconnect()恢复。连接生命周期与服务器部署虽然 ARCHITECTURE.md 以代码布局为重点但结合 python/ray/util/client/init.py 与 python/ray/util/client/server/server.py可以完整还原一条连接的生命周期启动服务器python -m ray.util.client.server --host host --port port --mode proxy|legacy|specific-server [--address ray集群地址]。默认 mode 为proxy即 Ray Client Proxy把请求转发给 Ray 集群legacy/specific-server则直接serve()起一个内嵌 ray.init 的服务器server.py。init_and_serve则是进程内启动并自动连接的便捷入口默认监听本地 50051 端口init.py。客户端连接ray.util.client.ray.connect(host:port, ...)创建Worker通过数据通道发送InitRequest内含 pickle 化的job_config与ray_init_kwargs服务器端RayletServicer.Init反序列化后调用ray_connect_handler执行ray.init()并校验客户端/服务器版本check_version_info可被ignore_version或RAY_IGNORE_VERSION_MISMATCH覆盖。使用与清理客户端通过Schedule提交任务、GetObject/PutObject存取对象连接关闭或通道断开时服务器通过 ClientID 追踪到的引用表object_refs、actor_refs、actor_owners被release_all清空server.py。断线重连DataClient在 RPC 错误后可恢复的情况下先 ping 服务器判断通道是否仍可用必要时重建 gRPC channel并把outstanding_requests中尚未确认的请求全部重放dataclient.py服务器端的ResponseCache/OrderedResponseCache则保证重放的请求不会在服务器上被执行两次幂等语义见 common.py。测试策略两种互补的模式测试 client 代码有两种主要方式方式一以客户端/服务器应用的视角测试明确知道自己正在测试一个客户端/服务器应用同时运行连接的两端调用已知是客户端侧的客户端观察服务器端的效果反之亦然。形如test_client*.py的测试集合采用该方式仓库中包括python/ray/tests/test_client.py核心客户端行为python/ray/tests/test_client_reconnect.py断线重连python/ray/tests/test_client_references.py引用计数与释放python/ray/tests/test_client_proxy.pyProxy 模式python/ray/tests/test_client_builder.pyray.client(...)builder APIpython/ray/tests/client_test_utils.py公共测试工具fixture。一般原则如果实现的是让 client 工作的特性应该放进test_client系列——因为你同时控制两端、可以直接测客户端代码。方式二把 Ray 的既有测试作为 fixture 跑在 client 模式上另一种方式是把 Ray 自己的测试当作 fixture像普通 Ray 一样测试 Ray 的 API。这是一个非常强大的模式某测试在此模式下通过意味着用户有信心一切工作方式与以往一致无论是否使用 client实现难度更高因为接入/关闭一个客户端/服务器对与单节点 ray 实例的 setup/shutdown 不同尤其是这些测试可能对 fixture 的搭建方式做了假设它是唯一能测试集成点的方式即验证 client mode hook 对 Ray core 函数的透明转发。因此如果修的是用户侧 API bug 或与 Ray core 的集成问题通常的做法是把一个既有单元测试改造/纳入 Ray Client 的测试集。小结从架构到调试回顾全文Ray Client 的架构可以归纳为三个层次API 层RayAPIStub以对象形式提供与ray包等价的 API 表面客户端桩对象ClientObjectRef、ClientActorRef、ClientRemoteFunc等在 common.py 中定义传输层gRPC 的三个 serviceray_client.proto分别承载一元 RPC、双向数据通道与日志通道ClientID 关联连接、ReleaseRequest 实现引用计数、分块与缓存机制支撑大对象与重连编码层client/server 两侧的自定义 picklerclient_pickler.py 与 server/server_pickler.py用PickleStub元组在客户端桩对象与服务端真实对象之间做无损翻译。对开发者而言理解这套分层最大的收益在于定位问题看到ClientObjectRef就知道它来自客户端路径看到PickleStub就知道正处于序列化转换边界而disable_client_hook()包裹的代码段则代表服务器端真实执行的 Ray 调用。结合 client_mode_hook.py 的转发逻辑与test_client*系列测试你可以快速判断一个 bug 是出在 API 转发、gRPC 传输还是序列化编码层。【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考