wehub-resource-sync
705 lines
23 KiB
Python
705 lines
23 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E712, F841
|
|
|
|
import gc
|
|
import multiprocessing
|
|
import os
|
|
import stat
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import tvm_ffi
|
|
|
|
pytest.importorskip("tornado") # tvm.rpc.proxy and tvm.rpc.tracker require tornado
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import rpc, te
|
|
from tvm.rpc.proxy import Proxy
|
|
from tvm.rpc.tracker import Tracker
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.support import cc, utils
|
|
from tvm.testing import env
|
|
|
|
if __name__ == "__main__":
|
|
# NOTE: must live here to avoid registering PackedFunc with libtvm_compiler.so twice.
|
|
tvm.testing.main()
|
|
|
|
|
|
# tkonolige: The issue as I understand it is this: multiprocessing's spawn
|
|
# method launches a new process and then imports the relevant modules. This
|
|
# means that all registered functions must exist at the top level scope. In
|
|
# this file they are, so all is well when we run this file directly.
|
|
# However, when run under pytest, the functions aren't registered on the
|
|
# server. I believe this is because pytest is also using multiprocessing to
|
|
# run individual functions. Somewhere along the way, the imports are being
|
|
# lost, so the server ends up not registering the functions.
|
|
pytestmark = pytest.mark.skipif(
|
|
# Windows does not support fork so we can enable Windows for testing
|
|
sys.platform.startswith("win") == False and multiprocessing.get_start_method() != "fork",
|
|
reason=(
|
|
"pytest + multiprocessing spawn method causes tvm.register_global_func to "
|
|
"not work on the rpc.Server."
|
|
),
|
|
)
|
|
|
|
# NOTE: When writing tests, wrap remote related checking in a sub-function
|
|
# to ensure all the remote resources destructs before the server terminates
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_bigendian_rpc():
|
|
"""Test big endian rpc when there is a PowerPC RPC server available"""
|
|
host = os.environ.get("TVM_POWERPC_TEST_HOST", None)
|
|
port = os.environ.get("TVM_POWERPC_TEST_PORT", 9090)
|
|
if host is None:
|
|
return
|
|
|
|
def verify_rpc(remote, target, shape, dtype):
|
|
A = te.placeholder(shape, dtype=dtype)
|
|
B = te.compute(A.shape, lambda i: A[i] + tvm.tirx.const(1, A.dtype))
|
|
f = tvm.compile(te.create_prim_func([A, B]), target=target)
|
|
|
|
dev = remote.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.randint(0, 256, size=shape).astype(A.dtype), device=dev)
|
|
b = tvm.runtime.tensor(np.zeros(shape).astype(A.dtype), device=dev)
|
|
temp = utils.tempdir()
|
|
path_dso = temp.relpath("dev_lib.o")
|
|
f.write_to_file(path_dso)
|
|
remote.upload(path_dso)
|
|
f = remote.load_module("dev_lib.o")
|
|
f(a, b)
|
|
tvm.testing.assert_allclose(a.numpy() + 1, b.numpy())
|
|
|
|
print("Test RPC connection to PowerPC...")
|
|
remote = rpc.connect(host, port)
|
|
target = {"kind": "llvm", "mtriple": "powerpc-linux-gnu"}
|
|
for dtype in ["float32", "float64", "int32", "int8"]:
|
|
verify_rpc(remote, target, (10,), dtype)
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_simple():
|
|
server = rpc.Server(key="x1")
|
|
client = rpc.connect("127.0.0.1", server.port, key="x1")
|
|
|
|
def check_remote():
|
|
f1 = client.get_function("rpc.test.addone")
|
|
assert f1(10) == 11
|
|
f3 = client.get_function("rpc.test.except")
|
|
|
|
with pytest.raises(RuntimeError):
|
|
f3("abc")
|
|
|
|
f2 = client.get_function("rpc.test.strcat")
|
|
assert f2("abc", 11) == "abc:11"
|
|
|
|
check_remote()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_runtime_string():
|
|
server = rpc.Server(key="x1")
|
|
client = rpc.connect("127.0.0.1", server.port, key="x1")
|
|
|
|
def check_remote():
|
|
func = client.get_function("rpc.test.runtime_str_concat")
|
|
x = tvm_ffi.core.String("abc")
|
|
y = tvm_ffi.core.String("def")
|
|
assert str(func(x, y)) == "abcdef"
|
|
|
|
check_remote()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_array():
|
|
server = rpc.Server()
|
|
remote = rpc.connect("127.0.0.1", server.port)
|
|
|
|
def check_remote():
|
|
x = np.ones((3, 4))
|
|
r_cpu = tvm.runtime.tensor(x, remote.cpu(0))
|
|
assert str(r_cpu.device).startswith("remote")
|
|
np.testing.assert_equal(r_cpu.numpy(), x)
|
|
fremote = remote.get_function("rpc.test.remote_tensor_func")
|
|
fremote(r_cpu)
|
|
|
|
check_remote()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_large_array():
|
|
# testcase of large array creation
|
|
server = rpc.Server()
|
|
remote = rpc.connect("127.0.0.1", server.port)
|
|
|
|
def check_remote():
|
|
dev = remote.cpu(0)
|
|
a_np = np.ones((5041, 720)).astype("float32")
|
|
b_np = np.ones((720, 192)).astype("float32")
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
b = tvm.runtime.tensor(b_np, dev)
|
|
np.testing.assert_equal(a.numpy(), a_np)
|
|
np.testing.assert_equal(b.numpy(), b_np)
|
|
|
|
check_remote()
|
|
|
|
|
|
@tvm.testing.skip_if_32bit(reason="skipping test for i386.")
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_echo():
|
|
def check(remote, local_session):
|
|
fecho = remote.get_function("testing.echo")
|
|
assert fecho(1, 2, 3) == 1
|
|
assert fecho(100, 2, 3) == 100
|
|
assert fecho("xyz") == "xyz"
|
|
assert bytes(fecho(bytearray(b"123"))) == b"123"
|
|
|
|
with pytest.raises(RuntimeError):
|
|
raise_err = remote.get_function("testing.test_raise_error")
|
|
raise_err("RuntimeError", "msg")
|
|
|
|
remote.cpu().sync()
|
|
# tests around system lib are not threadsafe by design
|
|
# and do not work well with multithread pytest
|
|
# skip local session as they are being tested elsewhere
|
|
if not local_session:
|
|
with pytest.raises(AttributeError):
|
|
f3 = remote.system_lib()["notexist"]
|
|
|
|
temp = rpc.server._server_env([])
|
|
server = rpc.Server()
|
|
client = rpc.connect("127.0.0.1", server.port)
|
|
check(rpc.LocalSession(), True)
|
|
|
|
check(client, False)
|
|
|
|
def check_minrpc():
|
|
if tvm.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None:
|
|
return
|
|
# Test minrpc server.
|
|
temp = utils.tempdir()
|
|
minrpc_exec = temp.relpath("minrpc")
|
|
tvm.rpc.with_minrpc(cc.create_executable)(minrpc_exec, [])
|
|
check(rpc.PopenSession(minrpc_exec), False)
|
|
# minrpc on the remote
|
|
server = rpc.Server()
|
|
client = rpc.connect(
|
|
"127.0.0.1",
|
|
server.port,
|
|
session_constructor_args=["rpc.PopenSession", open(minrpc_exec, "rb").read()],
|
|
)
|
|
check(client, False)
|
|
|
|
# skip for now until we upgrade to new FFI
|
|
# check_minrpc()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_file_exchange():
|
|
server = rpc.Server()
|
|
remote = rpc.connect("127.0.0.1", server.port)
|
|
|
|
def check_remote():
|
|
blob = bytearray(np.random.randint(0, 10, size=(10)))
|
|
remote.upload(blob, "dat.bin")
|
|
rev = remote.download("dat.bin")
|
|
assert rev == blob
|
|
|
|
check_remote()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_rpc_remote_module():
|
|
# graph
|
|
n = tvm.runtime.convert(102)
|
|
A = te.placeholder((n,), name="A")
|
|
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
|
|
mod = tvm.ir.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "myadd"))
|
|
|
|
server0 = rpc.Server(key="x0")
|
|
server1 = rpc.Server(key="x1")
|
|
|
|
client = rpc.connect(
|
|
"127.0.0.1",
|
|
server0.port,
|
|
key="x0",
|
|
session_constructor_args=["rpc.Connect", "127.0.0.1", server1.port, "x1", False],
|
|
)
|
|
|
|
def check_remote(remote):
|
|
temp = utils.tempdir()
|
|
dev = remote.cpu(0)
|
|
f = tvm.compile(mod, "llvm")
|
|
path_dso = temp.relpath("dev_lib.so")
|
|
f.export_library(path_dso)
|
|
remote.upload(path_dso)
|
|
f1 = remote.load_module("dev_lib.so")
|
|
a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
|
|
time_f = f1.time_evaluator(f1.entry_name, remote.cpu(0), number=10)
|
|
cost = time_f(a, b).mean
|
|
print(f"{cost:g} secs/op")
|
|
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
|
|
|
|
# Download the file from the remote
|
|
path_tar = temp.relpath("dev_lib.tar")
|
|
f.export_library(path_tar)
|
|
remote.upload(path_tar)
|
|
local_download_path = temp.relpath("dev_lib.download.so")
|
|
with open(local_download_path, "wb") as fo:
|
|
fo.write(remote.download_linked_module("dev_lib.tar"))
|
|
fupdated = tvm.runtime.load_module(local_download_path)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype), tvm.cpu(0))
|
|
b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), tvm.cpu(0))
|
|
fupdated(a, b)
|
|
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
|
|
|
|
def check_minrpc():
|
|
if tvm.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None:
|
|
return
|
|
# export to minrpc
|
|
temp = utils.tempdir()
|
|
# system lib prefix will trigger system lib build
|
|
f = tvm.compile(mod.with_attr("system_lib_prefix", ""), "llvm")
|
|
path_minrpc = temp.relpath("dev_lib.minrpc")
|
|
f.export_library(path_minrpc, fcompile=rpc.with_minrpc(cc.create_executable))
|
|
|
|
with pytest.raises(RuntimeError):
|
|
rpc.PopenSession("filenotexist")
|
|
|
|
# statrt the minrpc session.
|
|
remote = tvm.rpc.PopenSession(path_minrpc)
|
|
dev = remote.cpu(0)
|
|
f1 = remote.system_lib()
|
|
|
|
a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
|
|
time_f = f1.time_evaluator("myadd", remote.cpu(0), number=1)
|
|
cost = time_f(a, b).mean
|
|
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
|
|
|
|
# change to not executable
|
|
os.chmod(path_minrpc, stat.S_IRUSR)
|
|
with pytest.raises(RuntimeError):
|
|
rpc.PopenSession(path_minrpc)
|
|
|
|
def check_remote_link_cl(remote):
|
|
"""Test function to run remote code such as cl
|
|
|
|
This is not enabled because there is forking issue
|
|
of TVM runtime when server launches after OpenCL
|
|
runtime initializes. We leave it as an example
|
|
on how to do rpc when we want to do linking on remote.
|
|
"""
|
|
if not tvm.testing.device_enabled("opencl"):
|
|
print("Skip because opencl is not enabled")
|
|
return
|
|
temp = utils.tempdir()
|
|
dev = remote.cl(0)
|
|
|
|
s = tvm.s_tir.Schedule(mod)
|
|
|
|
x = s.get_loops(s.get_sblock("B"))
|
|
xo, xi = s.split(x, factors=[None, 32])
|
|
s.bind(xo, "blockIdx.x")
|
|
s.bind(xi, "threadIdx.x")
|
|
f = tvm.compile(s.mod, tvm.target.Target("opencl", host="llvm"))
|
|
path_tar = temp.relpath("myadd.tar")
|
|
f.export_library(path_tar)
|
|
remote.upload(path_tar)
|
|
fhost = remote.load_module("myadd.tar")
|
|
a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
|
|
fhost(a, b)
|
|
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
|
|
|
|
check_remote(rpc.LocalSession())
|
|
check_remote(client)
|
|
check_minrpc()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_return_func():
|
|
server = rpc.Server(key="x1")
|
|
client = rpc.connect("127.0.0.1", server.port, key="x1")
|
|
|
|
def check_remote():
|
|
f1 = client.get_function("rpc.test.add_to_lhs")
|
|
fadd = f1(10)
|
|
assert fadd(12) == 22
|
|
|
|
check_remote()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_session_constructor_args():
|
|
# start server
|
|
server0 = rpc.Server(key="x0")
|
|
server1 = rpc.Server(key="x1")
|
|
|
|
def check_multi_hop():
|
|
# use server0 as proxy to connect to server1
|
|
client = rpc.connect(
|
|
"127.0.0.1",
|
|
server0.port,
|
|
key="x0",
|
|
session_constructor_args=["rpc.Connect", "127.0.0.1", server1.port, "x1", False],
|
|
)
|
|
|
|
fecho = client.get_function("testing.echo")
|
|
assert fecho(1, 2, 3) == 1
|
|
assert fecho(100, 2, 3) == 100
|
|
assert fecho("xyz") == "xyz"
|
|
assert bytes(fecho(bytearray(b"123"))) == b"123"
|
|
|
|
nd = tvm.runtime.tensor([1, 2, 3], device=client.cpu(0))
|
|
assert nd.numpy()[1] == 2
|
|
|
|
def check_error_handling():
|
|
with pytest.raises(tvm.error.RPCError):
|
|
client = rpc.connect(
|
|
"127.0.0.1",
|
|
server0.port,
|
|
key="x0",
|
|
session_constructor_args=["rpc.NonExistingConstructor"],
|
|
)
|
|
|
|
check_multi_hop()
|
|
check_error_handling()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_return_tensor():
|
|
def run_arr_test():
|
|
server = rpc.Server(key="x1")
|
|
client = rpc.connect("127.0.0.1", server.port, key="x1")
|
|
m = client.get_function("rpc.test.remote_return_nd")
|
|
get_arr = m("get_arr")
|
|
get_elem = m("get_elem")
|
|
get_arr_elem = m("get_arr_elem")
|
|
|
|
arr = get_arr()
|
|
assert get_elem(0) == 0.0
|
|
assert get_arr_elem(arr, 0) == 0.0
|
|
|
|
del arr
|
|
gc.collect()
|
|
assert get_elem(0) == 0.0
|
|
|
|
run_arr_test()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_rpc_return_remote_object():
|
|
def check(client, is_local):
|
|
make_shape = client.get_function("ffi.Shape")
|
|
get_elem = client.get_function("rpc.testing.GetShapeElem")
|
|
get_size = client.get_function("rpc.testing.GetShapeSize")
|
|
shape = make_shape(2, 3)
|
|
assert get_elem(shape, 0) == 2
|
|
assert get_elem(shape, 1) == 3
|
|
assert get_size(shape) == 2
|
|
# Test free object by assigning to the same variable
|
|
shape = make_shape(0)
|
|
assert get_size(shape) == 1
|
|
assert get_elem(shape, 0) == 0
|
|
|
|
# start server
|
|
|
|
check(rpc.LocalSession(), True)
|
|
|
|
def check_remote():
|
|
server = rpc.Server(key="x1")
|
|
client = rpc.connect("127.0.0.1", server.port, key="x1")
|
|
check(client, False)
|
|
|
|
check_remote()
|
|
|
|
def check_minrpc():
|
|
if tvm.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None:
|
|
return
|
|
# Test minrpc server.
|
|
temp = utils.tempdir()
|
|
minrpc_exec = temp.relpath("minrpc")
|
|
tvm.rpc.with_minrpc(cc.create_executable)(minrpc_exec, [])
|
|
check(rpc.PopenSession(minrpc_exec), False)
|
|
# minrpc on the remote
|
|
server = rpc.Server()
|
|
client = rpc.connect(
|
|
"127.0.0.1",
|
|
server.port,
|
|
session_constructor_args=["rpc.PopenSession", open(minrpc_exec, "rb").read()],
|
|
)
|
|
check(client, False)
|
|
|
|
check_minrpc()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
def test_local_func():
|
|
client = rpc.LocalSession()
|
|
|
|
def check_remote():
|
|
f1 = client.get_function("rpc.test.add_to_lhs")
|
|
fadd = f1(10)
|
|
assert fadd(12) == 22
|
|
|
|
blob = bytearray(np.random.randint(0, 10, size=(10)))
|
|
client.upload(blob, "dat.bin")
|
|
rev = client.download("dat.bin")
|
|
assert rev == blob
|
|
|
|
check_remote()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
|
|
def test_rpc_tracker_register(device_key):
|
|
# test registration
|
|
tracker = Tracker(port=9000, port_end=10000)
|
|
server1 = rpc.Server(
|
|
host="127.0.0.1",
|
|
port=9000,
|
|
port_end=10000,
|
|
key=device_key,
|
|
tracker_addr=("127.0.0.1", tracker.port),
|
|
)
|
|
server2 = rpc.Server(
|
|
host="127.0.0.1",
|
|
port=9000,
|
|
port_end=10000,
|
|
key=device_key,
|
|
tracker_addr=("127.0.0.1", tracker.port),
|
|
custom_addr="test_addr", # this is a test address, which is unable to connect
|
|
)
|
|
time.sleep(1)
|
|
client = rpc.connect_tracker("127.0.0.1", tracker.port)
|
|
|
|
def exist_address(summary, key, host, port):
|
|
server_info = summary["server_info"]
|
|
for device in server_info:
|
|
if device["key"] == f"server:{key}":
|
|
addr = device["addr"]
|
|
if (host is None or host == addr[0]) and port == addr[1]:
|
|
return True
|
|
return False
|
|
|
|
summary = client.summary()
|
|
assert summary["queue_info"][device_key]["free"] == 2
|
|
assert exist_address(summary, device_key, "127.0.0.1", server1.port)
|
|
assert exist_address(summary, device_key, "test_addr", server2.port)
|
|
|
|
remote = client.request(device_key)
|
|
summary = client.summary()
|
|
assert summary["queue_info"][device_key]["free"] == 1
|
|
|
|
del remote
|
|
time.sleep(1)
|
|
|
|
summary = client.summary()
|
|
assert summary["queue_info"][device_key]["free"] == 2
|
|
|
|
server1.terminate()
|
|
time.sleep(1)
|
|
|
|
summary = client.summary()
|
|
assert summary["queue_info"][device_key]["free"] == 1
|
|
assert not exist_address(summary, device_key, "127.0.0.1", server1.port)
|
|
assert exist_address(summary, device_key, "test_addr", server2.port)
|
|
|
|
server2.terminate()
|
|
time.sleep(1)
|
|
|
|
summary = client.summary()
|
|
assert summary["queue_info"][device_key]["free"] == 0
|
|
assert not exist_address(summary, device_key, "test_addr", server2.port)
|
|
|
|
tracker.terminate()
|
|
|
|
|
|
def _target(host, port, device_key, timeout):
|
|
client = rpc.connect_tracker(host, port)
|
|
remote = client.request(device_key, session_timeout=timeout)
|
|
while True:
|
|
pass
|
|
remote.cpu()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
|
|
def test_rpc_tracker_request(device_key):
|
|
# test concurrent request
|
|
tracker = Tracker(port=9000, port_end=10000)
|
|
server = rpc.Server(
|
|
port=9000,
|
|
port_end=10000,
|
|
key=device_key,
|
|
tracker_addr=("127.0.0.1", tracker.port),
|
|
)
|
|
client = rpc.connect_tracker("127.0.0.1", tracker.port)
|
|
|
|
proc1 = multiprocessing.Process(target=_target, args=("127.0.0.1", tracker.port, device_key, 4))
|
|
proc2 = multiprocessing.Process(
|
|
target=_target, args=("127.0.0.1", tracker.port, device_key, 200)
|
|
)
|
|
proc1.start()
|
|
time.sleep(0.5)
|
|
proc2.start()
|
|
time.sleep(0.5)
|
|
|
|
summary = client.summary()
|
|
|
|
assert summary["queue_info"][device_key]["free"] == 0
|
|
assert summary["queue_info"][device_key]["pending"] == 1
|
|
|
|
proc1.terminate()
|
|
proc1.join()
|
|
time.sleep(0.5)
|
|
|
|
summary = client.summary()
|
|
assert summary["queue_info"][device_key]["free"] == 0
|
|
assert summary["queue_info"][device_key]["pending"] == 0
|
|
|
|
proc2.terminate()
|
|
proc2.join()
|
|
server.terminate()
|
|
tracker.terminate()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
|
|
def test_rpc_tracker_via_proxy(device_key):
|
|
"""
|
|
tracker
|
|
/ \
|
|
Host -- Proxy -- RPC server
|
|
"""
|
|
|
|
tracker_server = Tracker(port=9000, port_end=9100)
|
|
proxy_server = Proxy(
|
|
host=tracker_server.host,
|
|
port=8888,
|
|
port_end=8988,
|
|
tracker_addr=(tracker_server.host, tracker_server.port),
|
|
)
|
|
|
|
server1 = rpc.Server(
|
|
host=proxy_server.host,
|
|
port=proxy_server.port,
|
|
key=device_key,
|
|
tracker_addr=(tracker_server.host, tracker_server.port),
|
|
is_proxy=True,
|
|
)
|
|
server2 = rpc.Server(
|
|
host=proxy_server.host,
|
|
port=proxy_server.port,
|
|
key=device_key,
|
|
tracker_addr=(tracker_server.host, tracker_server.port),
|
|
is_proxy=True,
|
|
)
|
|
|
|
client = rpc.connect_tracker(tracker_server.host, tracker_server.port)
|
|
remote1 = client.request(device_key, session_timeout=30) # pylint: disable=unused-variable
|
|
remote2 = client.request(device_key, session_timeout=30) # pylint: disable=unused-variable
|
|
|
|
server2.terminate()
|
|
server1.terminate()
|
|
proxy_server.terminate()
|
|
tracker_server.terminate()
|
|
|
|
|
|
@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
|
|
@pytest.mark.parametrize("with_proxy", (True, False))
|
|
def test_rpc_session_timeout_error(with_proxy):
|
|
port = 9000
|
|
port_end = 10000
|
|
|
|
tracker = Tracker(port=port, port_end=port_end)
|
|
time.sleep(0.5)
|
|
tracker_addr = (tracker.host, tracker.port)
|
|
|
|
if with_proxy:
|
|
proxy = Proxy(host="0.0.0.0", port=port, port_end=port_end, tracker_addr=tracker_addr)
|
|
time.sleep(0.5)
|
|
server = rpc.Server(host=proxy.host, port=proxy.port, is_proxy=True, key="x1")
|
|
else:
|
|
server = rpc.Server(port=port, port_end=port_end, tracker_addr=tracker_addr, key="x1")
|
|
time.sleep(0.5)
|
|
|
|
rpc_sess = rpc.connect_tracker(*tracker_addr).request(key="x1", session_timeout=1)
|
|
|
|
with pytest.raises(tvm.error.RPCSessionTimeoutError):
|
|
f1 = rpc_sess.get_function("rpc.test.addone")
|
|
time.sleep(2)
|
|
f1(10)
|
|
|
|
server.terminate()
|
|
if with_proxy:
|
|
proxy.terminate()
|
|
tracker.terminate()
|
|
|
|
|
|
@pytest.mark.parametrize("call_with_unused_argument", [True, False])
|
|
def test_compiled_function_with_zero_arguments(call_with_unused_argument):
|
|
"""RPC functions do not require an argument
|
|
|
|
This is a regression test. When no arguments are provided, RPC
|
|
provides NULL as the `TVMFFIAny* args` argument to a PackedFunc.
|
|
However, previous implementations of `MakePackedAPI`
|
|
unconditionally asserted that the `args` pointer was non-null.
|
|
This assertion is now generated only when the function accepts
|
|
a non-zero number of arguments.
|
|
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def func_without_arg() -> T.int64:
|
|
return T.int64(42)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func_with_arg(unused: T.int64) -> T.int64:
|
|
return T.int64(42)
|
|
|
|
built = tvm.compile(Module, target="llvm")
|
|
|
|
server = tvm.rpc.Server(key="x1")
|
|
client = tvm.rpc.connect("127.0.0.1", server.port, key="x1")
|
|
|
|
libname = "libbuilt.so"
|
|
with tempfile.TemporaryDirectory(prefix="tvm_rpc_testing_") as temp_dir:
|
|
local_path = os.path.join(temp_dir, libname)
|
|
built.export_library(local_path)
|
|
client.upload(local_path)
|
|
|
|
remote_mod = client.load_module(libname)
|
|
|
|
if call_with_unused_argument:
|
|
res = remote_mod["func_with_arg"](0)
|
|
else:
|
|
res = remote_mod["func_without_arg"]()
|
|
|
|
assert res == 42
|