Apache TVM RPC 模块完全指南远程设备连接、模块上传与资源调度【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址: https://gitcode.com/gh_mirrors/tvm7/tvmtvm.rpc是 Apache TVM 官方的轻量级远程过程调用RPC模块它解决了深度学习编译器本地编译、远端运行的核心诉求编译栈运行在本地服务器而 RPC 服务器则运行在没有编译器、只有运行时的远程设备如手机、嵌入式板卡、GPU 集群上。读完本文你将掌握 RPC 客户端/服务端/Tracker 的完整 API 用法、命令行启动方式、底层通信协议以及基于 Tracker 的资源调度模式能够独立搭建一套本地交叉编译 远程验证的测试与部署链路。本文内容以 docs/reference/api/python/rpc.rst 的tvm.rpc自动文档为核心骨架并结合 python/tvm/rpc/ 下的全部源码模块client.py、server.py、tracker.py、proxy.py、base.py、minrpc.py及 python/tvm/exec/ 下的命令行入口展开。一、RPC 模块的定位与典型应用场景tvm.rpc模块的模块级文档python/tvm/rpc/init.py给出了最权威的定义Lightweight TVM RPC module. RPC enables connect to a remote server, upload and launch functions. This is useful to for cross-compile and remote testing. The compiler stack runs on local server, while we use RPC server to run on remote runtime which dont have a compiler available. The test program compiles the program on local server, upload and run remote RPC server, get the result back to verify correctness.翻译成一句话RPC 让你连接到远程服务器上传并启动函数。它最典型的用途是交叉编译cross-compile与远程测试remote testing本地服务器运行完整的编译栈TVM 编译器通过 RPC 把编译产物如lib.so、lib.tar上传到远程 RPC 服务器在远端运行时上执行模块把计算结果传回本地做正确性验证。这套模式在 AutoTVM / AutoScheduler 的自动调优测量measure环节中大量使用——python/tvm/auto_scheduler/measure.py与python/tvm/autotvm/measure/measure_methods.py中均有通过 RPC Tracker 发起调优测量的代码路径。从源码结构看整个 RPC 体系由四类角色构成角色对应类/入口职责RPC Clientconnect()、RPCSession本地侧发起连接、上传模块、调用远端函数RPC ServerServer、python -m tvm.exec.rpc_server远端运行时接收客户端连接并执行任务RPC TrackerTracker、TrackerSession、python -m tvm.exec.rpc_tracker资源调度中心登记服务器并按 key 分发连接RPC Proxypython -m tvm.exec.rpc_proxy网关代理把 TCP 流量转发到 WebSocket 等特殊环境如浏览器端二、公开 API 总览tvm.rpc在包入口python/tvm/rpc/init.py导出了以下符号from .server import Server from .client import connect, connect_tracker from .client import RPCSession, LocalSession, PopenSession, TrackerSession from .minrpc import with_minrpcServer在独立进程中启动一个 RPC 服务器connect/connect_tracker分别建立到 RPC Server 与 RPC Tracker 的连接RPCSessionRPC 客户端会话对象由connect返回不要直接构造LocalSession/PopenSession本地环境 / 子进程方式模拟的会话用于同一段代码本地与远端都能跑的场景TrackerSessionTracker 客户端会话由connect_tracker返回with_minrpc把交叉编译函数包装为附带 miniRPC 服务器代码的编译函数。三、客户端会话connect 与 RPCSession3.1 建立连接connect()connect定义于 python/tvm/rpc/client.py函数签名如下connect(url, port, key, session_timeout0, session_constructor_argsNone, enable_loggingFalse)参数类型默认值含义urlstr必填RPC 服务器主机地址portint必填要连接的端口keystr附加的服务器匹配键与 Server 的key对应session_timeoutfloat0会话时长秒超时后服务器会杀掉连接0表示必须一直保持存活session_constructor_argsListNone传给远端会话构造函数的额外参数列表第一个元素始终是会话构造函数名字符串其后为位置参数enable_loggingboolFalse是否开启日志默认关闭两个值得注意的实现细节当设置了session_timeout时源码会自动在 key 上追加-timeout{session_timeout}后缀client.py#L531-L532服务端解析该后缀后按超时机制管理会话如果调用时NameError被抛出则说明本机编译 TVM 时未开启 RPC 支持错误信息为Please compile with USE_RPC1——即需要以USE_RPC1编译 TVM。基本用法import tvm from tvm import rpc client rpc.connect(server_url, server_port, server_key)session_constructor_args的一个典型场景是通过代理机器中转连接到内部服务器源码 docstring 中的示例client_via_proxy rpc.connect( proxy_server_url, proxy_server_port, proxy_server_key, enable_logging, session_constructor_args[ rpc.Connect, internal_url, internal_port, internal_key, internal_logging ], )3.2 会话对象RPCSessionRPCSessionclient.py#L34-L251封装了所有远端交互能力。官方文档明确说明不要直接创建该对象应通过connect()获取。方法签名说明get_functionget_function(name)从会话中按名字获取远端函数返回Functionsystem_libsystem_lib()获取系统级库模块等价于远端的tvm.runtime.system_libdevicedevice(dev_type, dev_id0)构造一个远端设备句柄dev_type可为 int 或 strcpu/cuda/cl/vulkan/metal/rocm/hexagon/webgpu/ext_devxxx(dev_id0)快捷构造对应类型CPU / CUDA / OpenCL / Vulkan / Metal / ROCm / Hexagon / WebGPU / 扩展设备的远端设备uploadupload(data, targetNone)上传文件到远端运行时的临时目录。data可为本地文件名str或bytearray当传bytearray时target必填传文件名时可省略target默认取文件名downloaddownload(path)从远端临时目录下载文件返回bytearrayremoveremove(path)删除远端临时目录中的文件listdirlistdir(path)列出远端临时目录下的文件返回以逗号,分隔的字符串load_moduleload_module(path)加载远端模块文件需先上传返回远端Moduledownload_linked_moduledownload_linked_module(path)在远端链接模块并下载结果。当本地缺少链接器时非常有用其中upload/download/remove/listdir/download_linked_module在底层分别调用了远端注册的tvm.rpc.server.upload、tvm.rpc.server.download、tvm.rpc.server.remove、tvm.rpc.server.listdir、tvm.rpc.server.download_linked_module函数并对调用结果做了会话级缓存见 client.py#L115-L215。device()的编码原理值得展开远端设备类型会叠加一个会话索引编码client.py#L89-L93dev nd.device(dev_type, dev_id) encode (self._tbl_index 1) * base.RPC_SESS_MASK # RPC_SESS_MASK 128 dev.device_type encode dev._rpc_sess self即通过RPC_SESS_MASK128对设备类型进行偏移编码使本地构造的Device能唯一对应到某个远端会话中的真实设备相关常量定义于 base.py#L58。3.3 远端链接模块download_linked_module 实战这是官方 docstring 中给出的完整示例适用于本地没有链接器、但远端有的情形mod build_module_with_cross_compilation() # 因为本地没有链接器先把模块导出为 tar mod.export_library(lib.tar) remote.upload(lib.tar) # 在远端调用链接器把模块链接成库 # 注意链接出的库只能在远端相同的环境下运行 with open(lib.so, wb) as file: file.write(remote.download_linked_module(lib.tar))服务端对.o、.tar、.dylib/.so三类文件分别做了处理.o直接调用create_shared生成.so.tar先解包再统一链接.dylib/.so直接原样下发server.py#L85-L111。四、Tracker 资源调度connect_tracker 与 TrackerSession当有多个 RPC 服务器时直接记住每个服务器的地址不现实。Tracker 作为资源调度中心负责按 key 登记服务器、按 key 分发连接。TrackerSessionclient.py#L300-L473由connect_tracker(url, port)创建。4.1 TrackerSession 方法方法签名说明requestrequest(key, priority1, session_timeout0, max_retry5, session_constructor_argsNone)向 Tracker 请求一台 key 匹配的服务器并返回已连接的RPCSessionrequest_and_runrequest_and_run(key, func, priority1, session_timeout0, max_retry2)请求资源并执行func(session)可防护执行过程中服务器节点掉线掉线时自动重新请求资源并重跑summarysummary()获取 Tracker 的摘要字典text_summarytext_summary()获取人类可读的文本摘要Server List Queue Statuscloseclose()关闭与 Tracker 的连接request参数说明key设备类型键例如rasp3b、gpu与 Server 启动时的--key对应priority请求优先级整数值越大越优先见PriorityScheduler中(-priority, ...)的堆排序实现 tracker.py#L151session_timeout会话时长秒0表示请求必须一直保持存活非零时服务器可在超时后杀连接max_retry失败前的最大重试次数session_constructor_args同connect额外传给远端会话构造函数的参数列表。request的实现细节它会与 Tracker 完成RPC_TRACKER_MAGIC握手随后发送[TrackerCode.REQUEST, key, , priority]JSON 消息收到[TrackerCode.SUCCESS, [url, port, matchkey]]后调用connect(url, port, matchkey, ...)建立真实会话client.py#L409-L424。request_and_run的容错逻辑在 client.py#L460-L473捕获TVMError后重试若设置了session_timeout且运行时长已达到session_timeout * 0.95则判定为会话超时并直接抛出RuntimeError。4.2 文本摘要输出示例text_summary()的输出格式如下来源client.py#L338-L380Server List ------------------------------ server-address key ------------------------------ 10.77.1.234:9091 server:rasp3b ------------------------------ Queue Status key total free pending ---- ----- ---- ------- rasp3b 1 1 0五、服务端Server 类与命令行启动5.1 Server 类Serverserver.py#L424-L547在独立进程中启动一个 RPC 服务器基于多进程 multiprocessing 的 Python 实现官方注释也提到可用不依赖 Python 的 C 实现。构造函数Server(host0.0.0.0, port9091, port_end9199, is_proxyFalse, tracker_addrNone, key, load_libraryNone, custom_addrNone, silentFalse, no_forkFalse, server_init_callbackNone, reuse_addrTrue, timeoutNone)参数类型默认值含义hoststr0.0.0.0服务器监听的主机地址portint9091绑定的起始端口port_endint9199端口搜索上界在[port, port_end)区间内寻找第一个可用端口is_proxyboolFalse该地址是否为代理为True时host/port实际是代理服务器地址tracker_addrTuple[str, int]NoneTracker 地址非None时服务器会向 Tracker 注册自己keystr用于在 Tracker 中标识设备类型的键load_librarystrNone执行期间额外加载的库列表冒号:分隔custom_addrstrNone上报给 Tracker 的自定义 IP 地址silentboolFalse静默模式no_forkboolFalse是否禁止 multiprocessing 的 fork 方式server_init_callbackCallableNone启动服务器时执行的额外初始化函数reuse_addrboolTrue允许内核复用TIME_WAIT状态的本地 socketWindows 上强制关闭见 server.py#L347-L352timeoutfloatNonesocket 所有操作的超时时间类还提供terminate()方法用于杀掉服务器进程。关键限制RPC 服务器只能看到tvm命名空间中的函数。要往服务器环境注入自定义函数必须使用server_init_callback官方 docstring 示例def server_init_callback(): import tvm # 必须在这里 import 自定义包 import mypackage tvm.register_func(function, mypackage.func) server rpc.Server(host, server_init_callbackserver_init_callback)启动服务器的 Python 方式server rpc.Server(host0.0.0.0, port9090, keyrasp3b, tracker_addr(10.77.1.234, 9190))5.2 命令行启动python -m tvm.exec.rpc_server更常用的是命令行方式入口在 python/tvm/exec/rpc_server.pypython -m tvm.exec.rpc_server --host 0.0.0.0 --port 9090 --port-end 9199 \ --tracker 10.77.1.234:9190 --key rasp3b全部命令行参数参数默认值说明--host0.0.0.0Tracker/RPC 绑定的主机 IP--port9090RPC 端口--port-end9199RPC 端口搜索上界--through-proxyFalse服务是否通过代理提供开启时 host/port 实际是代理地址--tracker无Tracker 地址格式host:port如10.77.1.234:9190指定后必须同时提供--key--key在 Tracker 中标识设备类型的键--load-library无额外加载的库--no-forkFalse默认 fork使用 spawn 模式避免 fork官方注释提示 Metal、OpenCL、ROCm 编译器可能因 fork 产生内部错误建议加此参数--custom-addr无上报给 Tracker 的自定义 IP--silentFalse静默模式命令行脚本在非静默模式下启动时会打印提示运行 ROCM/Metal 时 fork 会导致编译器内部错误建议使用--no-fork启动rpc_server.py#L97-L101。5.3 服务端会话生命周期_serving实现server.py#L137-L176揭示了服务端如何处理一次会话每个连接进入独立工作目录utils.tempdir()在multiprocessing.Process子进程中执行_serve_loop其内部调用 C 侧_ffi_api.ServerLoop(sock.fileno())若设置了timeout且子进程超时仍存活则先通过ReturnException告知客户端会话超时错误再终止子进程及其所有子进程。六、Tracker 服务端与状态查询6.1 Tracker 类Trackertracker.py#L437-L501在独立进程中启动 Tracker 服务器同样基于PopenWorker实现Tracker(host0.0.0.0, port9190, port_end9199, silentFalse, reuse_addrTrue, timeoutNone)参数含义与Server基本一致默认端口为9190。terminate()会先发送TrackerCode.STOP消息安全停止再兜底 kill 进程。6.2 命令行启动与查询启动 Trackerpython -m tvm.exec.rpc_tracker --host 0.0.0.0 --port 9190 --port-end 9199查询 Tracker 状态AutoTVM 测量代码中同样推荐使用见 python/tvm/autotvm/graph_tuner/base_graph_tuner.py#L397python -m tvm.exec.query_rpc_tracker --host0.0.0.0 --port91906.3 Tracker 内部实现Scheduler 与协议Tracker 的核心是TrackerServerHandlerPrioritySchedulertracker.py#L79-L162每个设备 key 对应一个独立的PriorityScheduler请求按(-priority, 请求计数)入堆实现高优先级优先、同优先级 FIFO服务器资源value与请求队列一旦同时非空_schedule立即撮合回调返回False表示资源未被消费则放回队列TrackerServerHandler.close()在连接断开时自动把该连接上报的资源从对应调度器中移除。Tracker 的连接处理基于tornado异步 I/OTCPEventHandler因此需要安装 tornado 依赖未安装时会提示pip install tornado见 tracker.py#L55-L61。七、代理模式与 Web RPCrpc_proxypython -m tvm.exec.rpc_proxypython/tvm/exec/rpc_proxy.py启动一个 RPC Web 代理其定位是允许重定向到基于 WebSocket 的 RPC 服务器浏览器。它同时充当静态资源服务器会查找 web/apps/browser/rpc_server.html、rpc_plugin.html以及 web/dist/ 下的tvmjs.bundle.js、tvmjs_runtime.wasi.js等前端资源文件rpc_proxy.py#L26-L63。服务端通过代理的场景在 Python API 侧对应两条路径服务端Server(..., is_proxyTrue, key...)此时host/port指向代理地址服务器作为代理的后端点运行_connect_proxy_loop见 server.py#L282-L313客户端通过connect的session_constructor_args指定rpc.Connect会话构造器穿透代理见 3.1 节示例。八、本地会话、子进程会话与 miniRPC这三类能力让同一段测试代码既能在本地跑、也能在远端跑成为可能。8.1 LocalSessionLocalSessionclient.py#L254-L263是由本地环境支撑的 RPCSession 接口官方 docstring 说明它可用于实现那些既需要本地运行、又需要远端运行的函数。它内部创建了本地服务器环境server._server_env([])并包装一个_ffi_api.LocalSession()。8.2 PopenSessionPopenSession(binary)client.py#L287-L297通过popen子进程方式执行指定的可执行文件作为会话后端。binary可为可执行文件路径str必须是存在的、可执行的文件二进制内容bytes/bytearray会被写入临时目录的server.minrpc文件并赋予执行权限S_IXUSR | S_IRUSR然后以管道客户端CreatePipeClient启动client.py#L266-L284。8.3 with_minrpcminiRPC 编译包装with_minrpc(compile_func, serverposix_popen_server, runtimelibtvm)minrpc.py#L47-L81把交叉编译函数包装为附带 miniRPC 服务器源码的编译函数常用于构建可独立运行的微型 RPC 可执行文件定位src/runtime/minrpc/{server}/{server}.cc服务器源文件用cc.cross_compiler附加-stdc17、rpath、include 路径等选项生成编译函数返回的fcompile设置了need_system_lib True。服务器类型对应 src/runtime/minrpc/ 目录下的实现如posix_popen_server。九、底层通信协议9.1 数据面RPC Server协议server.py#L17-L26 的模块文档说明了 TCP 协议- 与对端进行初始握手 - [RPC_MAGIC, keysize(int32), key-bytes] - key 的格式为 - {server|client}:device-type[:random-key] [-timeouttimeout]相关常量与编解码工具定义在 python/tvm/rpc/base.py常量值含义RPC_MAGIC0xFF271RPC 数据面魔数RPC_TRACKER_MAGIC0x2F271Tracker控制面魔数RPC_CODE_SUCCESSRPC_MAGIC 0握手成功响应RPC_CODE_DUPLICATERPC_MAGIC 1代理中 key 重复RPC_CODE_MISMATCHRPC_MAGIC 2服务器中找不到匹配 keyRPC_SESS_MASK128远端设备类型编码掩码服务端握手逻辑在_listen_loop._accept_connserver.py#L182-L249验证魔数后读取 key若前缀不是client: matchkey 则返回RPC_CODE_MISMATCH当使用 Tracker 时matchkey 是每次 PUT 时随机生成的base.random_keyTracker 会跟踪被客户端取得但尚未使用的 pending matchkeys若 key 长期未使用会重新生成unmatch_timeout为 4 个 ping 周期。9.2 控制面Tracker协议与 TrackerCodetracker.py#L23-L41 说明了 Tracker 的协议初始握手发送RPC_TRACKER_MAGIC普通消息[size(int32), json-data]所有消息由客户端发起Tracker 回以 json。TrackerCode枚举base.py#L44-L56代码值语义FAIL-1失败SUCCESS0成功PING1探测 Tracker 是否存活返回SUCCESSSTOP2携带 stop_key 安全停止 TrackerPUT3服务器上报资源输入[PUT, [port, match-key]]REQUEST4客户端请求资源输入[REQUEST, [key, user, priority]]返回[SUCCESS, [url, port, match-key]]UPDATE_INFO5更新连接信息服务器上报 key、addr 等SUMMARY6获取摘要GET_PENDING_MATCHKEYS7服务器轮询尚未被使用的 matchkey 集合消息封装使用sendjson/recvjsonbase.py#L89-L120即struct.pack(i, len(json)) UTF-8 编码的 JSON 字符串。十、测试支撑RPC 测试函数集python/tvm/rpc/testing.py 为单元测试注册了一批跨端可调用的示例函数可用于验证 RPC 链路是否正确注册名行为rpc.test.addone输入加 1rpc.test.strcat返回{name}:{x}rpc.test.except抛出ValueError测试远端异常传递rpc.test.runtime_str_concat字符串拼接rpc.test.remote_array_func校验传入数组是否等于np.ones((3, 4))rpc.test.add_to_lhs返回闭包lambda y: x yrpc.test.remote_return_nd通过闭包返回/查询tvm.nd.array引用计数等用于校验远端对象的引用计数正确性十一、端到端实战示例综合以上内容一个典型的本地编译 远端验证完整流程如下import tvm from tvm import rpc from tvm.contrib import utils # 1. 连接远端 RPC 服务器 remote rpc.connect(127.0.0.1, 9090, keylocal) # 2. 构造远端设备并上传编译产物 ctx remote.cpu(0) # 或用 remote.cuda(0) / remote.cl(0) 等 remote.upload(lib.so) # 上传到远端临时目录 # 3. 加载远端模块并调用 mod remote.load_module(lib.so) f mod.get_function(my_kernel) f(remote.nd.array(tvm.nd.array(data))) # 远端执行结果传回 # 4. 需要从远端取回结果时 blob remote.download(output.bin)若使用 Tracker 进行资源调度tracker rpc.connect_tracker(10.77.1.234, 9190) print(tracker.text_summary()) # 查看服务器与队列状态 remote tracker.request(rasp3b, session_timeout60) # 或执行一个可重试的完整任务 result tracker.request_and_run(rasp3b, lambda sess: sess.get_function(task)())需要说明的适用前提以上所有 API 与命令行工具都要求 TVM 以USE_RPC1编译且服务端侧远端设备只需运行时环境、无需完整编译器Tracker 服务器还需要 tornado 依赖。命令行的端口参数如--port 9090、Tracker 默认9190均以 python/tvm/exec/ 下入口脚本的实际默认值为准。【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址: https://gitcode.com/gh_mirrors/tvm7/tvm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考