tvm.rpc¶
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 don’t 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.
Classes:
|
Start RPC server on a separate process. |
|
RPC Client session module |
RPCSession interface backed by local environment. |
|
|
RPCSession interface backed by popen. |
|
Tracker client session. |
Functions:
|
Connect to RPC Server |
|
Connect to a RPC tracker |
|
Attach the compiler function with minrpc related options. |
- 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)¶
Start RPC server on a separate process.
This is a simple python implementation based on multi-processing. It is also possible to implement a similar C based server with TVM runtime which does not depend on the python.
- Parameters
host (str) – The host url of the server.
port (int) – The port to be bind to
port_end (int, optional) – The end port to search
is_proxy (bool, optional) – Whether the address specified is a proxy. If this is true, the host and port actually corresponds to the address of the proxy server.
tracker_addr (Tuple (str, int) , optional) – The address of RPC Tracker in tuple(host, ip) format. If is not None, the server will register itself to the tracker.
key (str, optional) – The key used to identify the device type in tracker.
load_library (str, optional) – List of additional libraries to be loaded during execution.
custom_addr (str, optional) – Custom IP Address to Report to RPC Tracker
silent (bool, optional) – Whether run this server in silent mode.
no_fork (bool, optional) – Whether forbid fork in multiprocessing.
server_init_callback (Callable, optional) – Additional initialization function when starting the server.
reuse_addr (bool, optional) – Allows the kernel to reuse a local socket in TIME_WAIT state.
timeout (float, optional) – set a timeout for all operations on the socket
Note
The RPC server only sees functions in the tvm namespace. To bring additional custom functions to the server env, you can use 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)
Methods:
Terminate the server process
- terminate()¶
Terminate the server process
- tvm.rpc.connect(url, port, key='', session_timeout=0, session_constructor_args=None, enable_logging=False)¶
Connect to RPC Server
- Parameters
url (str) – The url of the host
port (int) – The port to connect to
key (str, optional) – Additional key to match server
session_timeout (float, optional) – The duration of the session in seconds, allows server to kill the connection when duration is longer than this value. When duration is zero, it means the request must always be kept alive.
session_constructor_args (List) – List of additional arguments to passed as the remote session constructor. The first element of the list is always a string specifying the name of the session constructor, the following args are the positional args to that function.
enable_logging (boolean) – flag to enable/disable logging. Logging is disabled by default.
- Returns
sess – The connected session.
- Return type
Examples
Normal usage .. code-block:: python
client = rpc.connect(server_url, server_port, server_key)
Session_constructor can be used to customize the session in the remote The following code connects to a remote internal server via a proxy by constructing another RPCClientSession on the proxy machine and use that as the serving session of the proxy endpoint.
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)¶
Connect to a RPC tracker
- Parameters
- Returns
sess – The connected tracker session.
- Return type
- class tvm.rpc.RPCSession(sess)¶
RPC Client session module
Do not directly create the object, call connect
Methods:
Get system-wide library module.
get_function
(name)Get function from the session.
device
(dev_type[, dev_id])Construct a remote device.
upload
(data[, target])Upload file to remote runtime temp folder
download
(path)Download file from remote temp folder.
remove
(path)Remove file from remote temp folder.
load_module
(path)Load a remote module, the file need to be uploaded first.
download_linked_module
(path)Link a module in the remote and download it.
cpu
([dev_id])Construct CPU device.
cuda
([dev_id])Construct CUDA GPU device.
cl
([dev_id])Construct OpenCL device.
vulkan
([dev_id])Construct Vulkan device.
metal
([dev_id])Construct Metal device.
rocm
([dev_id])Construct ROCm device.
ext_dev
([dev_id])Construct extension device.
hexagon
([dev_id])Construct Hexagon device.
webgpu
([dev_id])Construct WebGPU device.
- system_lib()¶
Get system-wide library module.
- Returns
module – The system-wide library module.
- Return type
See also
- get_function(name)¶
Get function from the session.
- Parameters
name (str) – The name of the function
- Returns
f – The result function.
- Return type
Function
- device(dev_type, dev_id=0)¶
Construct a remote device.
- upload(data, target=None)¶
Upload file to remote runtime temp folder
- download(path)¶
Download file from remote temp folder.
- remove(path)¶
Remove file from remote temp folder.
- Parameters
path (str) – The relative location to remote temp folder.
- load_module(path)¶
Load a remote module, the file need to be uploaded first.
- download_linked_module(path)¶
Link a module in the remote and download it.
- Parameters
path (str) – The relative location to remote temp folder.
- Returns
blob – The result blob from the file.
- Return type
Note
This function can be helpful when a linker is not available on the local client.
Examples
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)¶
Construct CPU device.
- cuda(dev_id=0)¶
Construct CUDA GPU device.
- cl(dev_id=0)¶
Construct OpenCL device.
- vulkan(dev_id=0)¶
Construct Vulkan device.
- metal(dev_id=0)¶
Construct Metal device.
- rocm(dev_id=0)¶
Construct ROCm device.
- ext_dev(dev_id=0)¶
Construct extension device.
- hexagon(dev_id=0)¶
Construct Hexagon device.
- webgpu(dev_id=0)¶
Construct WebGPU device.
- class tvm.rpc.LocalSession¶
RPCSession interface backed by local environment.
This class can be used to implement functions that need to be ran both locally and remotely.
- class tvm.rpc.PopenSession(binary)¶
RPCSession interface backed by popen.
- class tvm.rpc.TrackerSession(addr)¶
Tracker client session.
- Parameters
addr (tuple) – The address tuple
Methods:
close
()Close the tracker connection.
summary
()Get the summary dict of the tracker.
Get a text summary of the tracker.
request
(key[, priority, session_timeout, ...])Request a new connection from the tracker.
request_and_run
(key, func[, priority, ...])Request a resource from tracker and run the func.
- close()¶
Close the tracker connection.
- summary()¶
Get the summary dict of the tracker.
- text_summary()¶
Get a text summary of the tracker.
- request(key, priority=1, session_timeout=0, max_retry=5, session_constructor_args=None)¶
Request a new connection from the tracker.
- Parameters
key (str) – The type key of the device.
priority (int, optional) – The priority of the request.
session_timeout (float, optional) – The duration of the session, allows server to kill the connection when duration is longer than this value. When duration is zero, it means the request must always be kept alive.
max_retry (int, optional) – Maximum number of times to retry before give up.
session_constructor_args (list, optional) – List of additional arguments to passed as the remote session constructor. The first element of the list is always a string specifying the name of the session constructor, the following args are the positional args to that function.
- request_and_run(key, func, priority=1, session_timeout=0, max_retry=2)¶
Request a resource from tracker and run the func.
This function safe-guard rare server node dropout during execution. In such case, a new resource will be requested and func will be ran again.
- Parameters
key (str) – The type key of the device.
func (function of session -> value) – A stateless function
priority (int, optional) – The priority of the request.
session_timeout (float, optional) – The duration of the session, allows server to kill the connection when duration is longer than this value. When duration is zero, it means the request must always be kept alive.
max_retry (int, optional) – Maximum number of times to retry the function before give up.
- tvm.rpc.with_minrpc(compile_func, server='posix_popen_server', runtime='libtvm')¶
Attach the compiler function with minrpc related options.