tvm.rpc
轻量级 TVM RPC 模块。
RPC 允许连接到远程服务器,上传和启动函数。这对于交叉编译和远程测试非常有用。编译器堆栈在本地服务器上运行,而我们使用 RPC 服务器在没有可用编译器的远程运行时上运行。
测试程序在本地服务器上编译程序,上传并在远程 RPC 服务器上运行,取回结果以验证正确性。
TVM RPC 服务器假定用户是可信任的,需要在可信任的网络环境和加密通道中使用。它允许将任意文件写入服务器,并为任何可以访问此 API 的人提供完整的远程代码执行能力。
类
|
在单独的进程中启动 RPC 服务器。 |
|
RPC 客户端会话模块 |
由本地环境支持的 RPCSession 接口。 |
|
|
由 popen 支持的 RPCSession 接口。 |
|
Tracker 客户端会话。 |
函数
|
连接到 RPC 服务器 |
|
连接到 RPC tracker |
|
将编译器函数附加到 minrpc 相关选项。 |
- class tvm.rpc.Server(host='0.0.0.0', port=9091, port_end=9199, is_proxy=False, tracker_addr=None, key='', load_library=None, custom_addr=None, silent=False, no_fork=False, server_init_callback=None, reuse_addr=True, timeout=None)
在单独的进程中启动 RPC 服务器。
这是一个基于多处理的简单 Python 实现。也可以使用不依赖 Python 的 TVM 运行时来实现类似的基于 C 的服务器。
- 参数:
host (str) – 服务器的主机 url。
port (int) – 要绑定的端口
port_end (int, optional) – 要搜索的结束端口
is_proxy (bool, optional) – 指定的地址是否为代理。如果为 true,则主机和端口实际上对应于代理服务器的地址。
tracker_addr (Tuple (str, int) , optional) – RPC Tracker 的地址,格式为 tuple(host, ip)。如果不是 None,服务器将自己注册到 tracker。
key (str, optional) – 用于在 tracker 中标识设备类型的 key。
load_library (str, optional) – 执行期间要加载的其他库的列表。
custom_addr (str, optional) – 要报告给 RPC Tracker 的自定义 IP 地址
silent (bool, optional) – 是否以静默模式运行此服务器。
no_fork (bool, optional) – 是否禁止在多处理中 fork。
server_init_callback (Callable, optional) – 启动服务器时的其他初始化函数。
reuse_addr (bool, optional) – 允许内核重用 TIME_WAIT 状态下的本地套接字。
timeout (float, optional) – 设置套接字上所有操作的超时时间
注意
TVM RPC 服务器假定用户是可信任的,需要在可信任的网络环境和加密通道中使用。它允许将任意文件写入服务器,并为任何可以访问此 API 的人提供完整的远程代码执行能力。
RPC 服务器仅查看 tvm 命名空间中的函数。要将其他自定义函数引入服务器环境,可以使用 server_init_callback。
def server_init_callback(): import tvm # must import mypackage here import mypackage tvm.register_func("function", mypackage.func) server = rpc.Server(host, server_init_callback=server_init_callback)
方法
终止服务器进程
- terminate()
终止服务器进程
- tvm.rpc.connect(url, port, key='', session_timeout=0, session_constructor_args=None, enable_logging=False)
连接到 RPC 服务器
- 参数:
url (str) – 主机的 url
port (int) – 要连接的端口
key (str, optional) – 用于匹配服务器的其他 key
session_timeout (float, optional) – 会话的持续时间(秒),允许服务器在持续时间超过此值时终止连接。当持续时间为零时,表示请求必须始终保持活动状态。
session_constructor_args (List) – 作为远程会话构造函数的附加参数列表。列表的第一个元素始终是一个字符串,指定会话构造函数的名称,后面的参数是该函数的位置参数。
enable_logging (boolean) – 用于启用/禁用日志记录的标志。默认情况下禁用日志记录。
- 返回:
sess – 连接的会话。
- 返回类型:
示例
正常用法 .. code-block:: python
client = rpc.connect(server_url, server_port, server_key)
Session_constructor 可用于自定义远程会话。以下代码通过在代理计算机上构建另一个 RPCClientSession 并将其用作代理端点的服务会话,从而通过代理连接到远程内部服务器。
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])
- tvm.rpc.connect_tracker(url, port)
连接到 RPC tracker
- 参数:
- 返回:
sess – 连接的 tracker 会话。
- 返回类型:
- class tvm.rpc.RPCSession(sess)
RPC 客户端会话模块
不要直接创建对象,调用 connect
方法
获取系统范围的库模块。
get_function
(name)从会话中获取函数。
device
(dev_type[, dev_id])构造远程设备。
upload
(data[, target])将文件上传到远程运行时临时文件夹
download
(path)从远程临时文件夹下载文件。
remove
(path)从远程临时文件夹中删除文件。
listdir
(path)ls 来自远程临时文件夹的文件。
load_module
(path)加载远程模块,需要先上传文件。
download_linked_module
(path)链接远程模块并下载它。
cpu
([dev_id])构造 CPU 设备。
cuda
([dev_id])构造 CUDA GPU 设备。
cl
([dev_id])构造 OpenCL 设备。
vulkan
([dev_id])构造 Vulkan 设备。
metal
([dev_id])构造 Metal 设备。
rocm
([dev_id])构造 ROCm 设备。
ext_dev
([dev_id])构造扩展设备。
hexagon
([dev_id])构造 Hexagon 设备。
webgpu
([dev_id])构造 WebGPU 设备。
- system_lib()
获取系统范围的库模块。
- 返回:
module – 系统范围的库模块。
- 返回类型:
runtime.Module
另请参阅
tvm.runtime.system_lib
- device(dev_type, dev_id=0)
构造远程设备。
- upload(data, target=None)
将文件上传到远程运行时临时文件夹
- download(path)
从远程临时文件夹下载文件。
- listdir(path)
ls 来自远程临时文件夹的文件。
- load_module(path)
加载远程模块,需要先上传文件。
- 参数:
path (str) – 相对于远程临时文件夹的位置。
- 返回:
m – 包含远程函数的远程模块。
- 返回类型:
Module
- download_linked_module(path)
链接远程模块并下载它。
注意
当本地客户端上没有链接器时,此函数可能很有用。
示例
mod = build_module_with_cross_compilation() # export the module as tar because a local linker is not available mod.export_library("lib.tar") remote.upload("lib.tar") # invoke the linker on the remote to link the module as a library # note that the library can only run on the same env as the remote with open("lib.so", "wb") as file: file.write(remote.download_linked_module("lib.tar"))
- cpu(dev_id=0)
构造 CPU 设备。
- cuda(dev_id=0)
构造 CUDA GPU 设备。
- cl(dev_id=0)
构造 OpenCL 设备。
- vulkan(dev_id=0)
构造 Vulkan 设备。
- metal(dev_id=0)
构造 Metal 设备。
- rocm(dev_id=0)
构造 ROCm 设备。
- ext_dev(dev_id=0)
构造扩展设备。
- hexagon(dev_id=0)
构造 Hexagon 设备。
- webgpu(dev_id=0)
构造 WebGPU 设备。
- class tvm.rpc.LocalSession
由本地环境支持的 RPCSession 接口。
此类可用于实现需要在本地和远程运行的函数。
- class tvm.rpc.PopenSession(binary)
由 popen 支持的 RPCSession 接口。
- class tvm.rpc.TrackerSession(addr)
Tracker 客户端会话。
- 参数:
addr (tuple) – 地址元组
方法
close
()关闭 tracker 连接。
summary
()获取 tracker 的摘要字典。
获取 tracker 的文本摘要。
request
(key[, priority, session_timeout, ...])从 tracker 请求新连接。
request_and_run
(key, func[, priority, ...])从 tracker 请求资源并运行 func。
- close()
关闭 tracker 连接。
- summary()
获取 tracker 的摘要字典。
- text_summary()
获取 tracker 的文本摘要。
- request(key, priority=1, session_timeout=0, max_retry=5, session_constructor_args=None)
从 tracker 请求新连接。
- request_and_run(key, func, priority=1, session_timeout=0, max_retry=2)
从 tracker 请求资源并运行 func。
此函数可以安全地防止执行期间罕见的服务器节点掉线。在这种情况下,将请求新资源并再次运行 func。