import tvm import logging import numpy as np import time import multiprocessing from tvm.contrib import rpc def rpc_proxy_check(): """This is a simple test function for RPC Proxy It is not included as nosetests, because: - It depends on tornado - It relies on the fact that Proxy starts before client and server connects, which is often the case but not always User can directly run this script to verify correctness. """ try: from tvm.contrib import rpc_proxy web_port = 8888 prox = rpc_proxy.Proxy("localhost", web_port=web_port) def check(): if not tvm.module.enabled("rpc"): return @tvm.register_func("rpc.test2.addone") def addone(x): return x + 1 @tvm.register_func("rpc.test2.strcat") def addone(name, x): return "%s:%d" % (name, x) server = multiprocessing.Process( target=rpc_proxy.websocket_proxy_server, args=("ws://localhost:%d/ws" % web_port,"x1")) # Need to make sure that the connection start after proxy comes up time.sleep(0.1) server.deamon = True server.start() client = rpc.connect(prox.host, prox.port, key="x1") f1 = client.get_function("rpc.test2.addone") assert f1(10) == 11 f2 = client.get_function("rpc.test2.strcat") assert f2("abc", 11) == "abc:11" check() except ImportError: print("Skipping because tornado is not avaliable...") if __name__ == "__main__": logging.basicConfig(level=logging.INFO) rpc_proxy_check()