def test_rpc_tracker_register():
# test registration
tracker = Tracker('localhost', port=9000, port_end=10000)
device_key = 'test_device'
server = rpc.Server('localhost',
port=9000,
port_end=10000,
key=device_key,
tracker_addr=(tracker.host, tracker.port))
time.sleep(1)
client = rpc.connect_tracker(tracker.host, tracker.port)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 1
remote = client.request(device_key)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 0
del remote
time.sleep(1)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 1
server.terminate()
time.sleep(1)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 0
tracker.terminate()
def test_rpc_tracker_register():
# test registration
tracker = Tracker(port=9000, port_end=10000)
device_key = "test_device"
server = rpc.Server(
port=9000,
port_end=10000,
key=device_key,
tracker_addr=("127.0.0.1", tracker.port),
)
time.sleep(1)
client = rpc.connect_tracker("127.0.0.1", tracker.port)
summary = client.summary()
assert summary["queue_info"][device_key]["free"] == 1
remote = client.request(device_key)
summary = client.summary()
assert summary["queue_info"][device_key]["free"] == 0
del remote
time.sleep(1)
summary = client.summary()
assert summary["queue_info"][device_key]["free"] == 1
server.terminate()
time.sleep(1)
summary = client.summary()
assert summary["queue_info"][device_key]["free"] == 0
tracker.terminate()
def __init__(self,
priority=1,
n_parallel=1,
timeout=10,
number=3,
repeat=1,
min_repeat_ms=0,
cooldown_interval=0.0,
enable_cpu_cache_flush=False):
ctx = tvm.context("cuda", 0)
if ctx.exist:
cuda_arch = "sm_" + "".join(ctx.compute_version.split('.'))
set_cuda_target_arch(cuda_arch)
host = '0.0.0.0'
self.tracker = Tracker(host, port=9000, port_end=10000, silent=True)
device_key = '$local$device$%d' % self.tracker.port
self.server = Server(host,
port=self.tracker.port,
port_end=10000,
key=device_key,
use_popen=True,
silent=True,
tracker_addr=(self.tracker.host,
self.tracker.port))
self.runner = RPCRunner(device_key, host, self.tracker.port, priority,
n_parallel, timeout, number, repeat,
min_repeat_ms, cooldown_interval,
enable_cpu_cache_flush)
# Wait for the processes to start
time.sleep(0.5)
def test_rpc_tracker_register():
# test registration
tracker = Tracker('localhost', port=9000, port_end=10000)
device_key = 'test_device'
server = rpc.Server('localhost', port=9000, port_end=10000,
key=device_key,
tracker_addr=(tracker.host, tracker.port))
time.sleep(1)
client = rpc.connect_tracker(tracker.host, tracker.port)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 1
remote = client.request(device_key)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 0
del remote
time.sleep(1)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 1
server.terminate()
time.sleep(1)
summary = client.summary()
assert summary['queue_info'][device_key]['free'] == 0
tracker.terminate()
class LocalRPCMeasureContext:
""" A context wrapper for running RPCRunner locally.
This will launch a local RPC Tracker and local RPC Server.
TODO(FrozenGene): Add cpu cache flush to this RPC context.
Parameters
----------
priority : int = 1
The priority of this run request, larger is more prior.
n_parallel : int = 1
The number of tasks run in parallel.
timeout : int = 10
The timeout limit (in second) for each run.
This is used in a wrapper of the multiprocessing.Process.join().
number : int = 3
The number of times to run the generated code for taking average.
We call these runs as one `repeat` of measurement.
repeat : int = 1
The number of times to repeat the measurement.
In total, the generated code will be run (1 + number x repeat) times,
where the first "1" is warm up and will be discarded.
The returned result contains `repeat` costs,
each of which is an average of `number` costs.
min_repeat_ms : int = 0
The minimum duration of one `repeat` in milliseconds.
By default, one `repeat` contains `number` runs. If this parameter is set,
the parameters `number` will be dynamically adjusted to meet the
minimum duration requirement of one `repeat`.
i.e., When the run time of one `repeat` falls below this time, the `number` parameter
will be automatically increased.
cooldown_interval : float = 0.0
The cool down interval between two measurements.
"""
def __init__(self, priority=1, n_parallel=1, timeout=10, number=3, repeat=1,
min_repeat_ms=0, cooldown_interval=0.0):
ctx = tvm.context("cuda", 0)
if ctx.exist:
cuda_arch = "sm_" + "".join(ctx.compute_version.split('.'))
set_cuda_target_arch(cuda_arch)
host = '0.0.0.0'
self.tracker = Tracker(host, port=9000, port_end=10000, silent=True)
device_key = '$local$device$%d' % self.tracker.port
self.server = Server(host, port=self.tracker.port, port_end=10000,
key=device_key, use_popen=True, silent=True,
tracker_addr=(self.tracker.host, self.tracker.port))
self.runner = RPCRunner(device_key, host, self.tracker.port, priority,
n_parallel, timeout, number, repeat,
min_repeat_ms, cooldown_interval)
# Wait for the processes to start
time.sleep(0.5)
def __del__(self):
# Close the tracker and server before exit
self.tracker.terminate()
self.server.terminate()
def test_rpc_tracker_request():
# test concurrent request
tracker = Tracker("localhost", port=9000, port_end=10000)
device_key = "test_device"
server = rpc.Server(
"localhost",
port=9000,
port_end=10000,
key=device_key,
tracker_addr=(tracker.host, tracker.port),
)
client = rpc.connect_tracker(tracker.host, tracker.port)
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()
proc1 = multiprocessing.Process(target=target,
args=(tracker.host, tracker.port,
device_key, 4))
proc2 = multiprocessing.Process(target=target,
args=(tracker.host, 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()
def __init__(
self,
priority=1,
n_parallel=1,
timeout=10,
number=3,
repeat=1,
min_repeat_ms=0,
cooldown_interval=0.0,
enable_cpu_cache_flush=False,
):
# pylint: disable=import-outside-toplevel
from tvm.rpc.tracker import Tracker
from tvm.rpc.server import Server
dev = tvm.device("cuda", 0)
if dev.exist:
cuda_arch = "sm_" + "".join(dev.compute_version.split("."))
set_cuda_target_arch(cuda_arch)
host = "0.0.0.0"
self.tracker = Tracker(host, port=9000, port_end=10000, silent=True)
device_key = "$local$device$%d" % self.tracker.port
self.server = Server(
host,
port=self.tracker.port,
port_end=10000,
key=device_key,
use_popen=True,
silent=True,
tracker_addr=(self.tracker.host, self.tracker.port),
)
self.runner = RPCRunner(
device_key,
host,
self.tracker.port,
priority,
n_parallel,
timeout,
number,
repeat,
min_repeat_ms,
cooldown_interval,
enable_cpu_cache_flush,
)
# Wait for the processes to start
time.sleep(0.5)
def test_rpc_tracker_request():
# test concurrent request
tracker = Tracker('localhost', port=9000, port_end=10000)
device_key = 'test_device'
server = rpc.Server('localhost', port=9000, port_end=10000,
key=device_key,
tracker_addr=(tracker.host, tracker.port))
client = rpc.connect_tracker(tracker.host, tracker.port)
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()
proc1 = multiprocessing.Process(target=target,
args=(tracker.host, tracker.port, device_key, 4))
proc2 = multiprocessing.Process(target=target,
args=(tracker.host, 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()
def __init__(
self,
priority=1,
n_parallel=1,
timeout=10,
number=3,
repeat=1,
min_repeat_ms=0,
cooldown_interval=0.0,
enable_cpu_cache_flush=False,
device=0,
):
# pylint: disable=import-outside-toplevel
from tvm.rpc.tracker import Tracker
from tvm.rpc.server import Server
self.tracker = Tracker(port=9000, port_end=10000, silent=True)
device_key = "$local$device$%d" % self.tracker.port
self.server = Server(
port=self.tracker.port,
port_end=10000,
key=device_key,
silent=True,
tracker_addr=("127.0.0.1", self.tracker.port),
)
self.runner = RPCRunner(
device_key,
"127.0.0.1",
self.tracker.port,
priority,
n_parallel,
timeout,
number,
repeat,
min_repeat_ms,
cooldown_interval,
enable_cpu_cache_flush,
device,
)
# Wait for the processes to start
time.sleep(0.5)
def test_rpc_tracker_request():
# test concurrent request
tracker = Tracker(port=9000, port_end=10000)
device_key = "test_device"
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()
def test_rpc_tracker_via_proxy():
"""
tracker
/ \
Host -- Proxy -- RPC server
"""
device_key = "test_device"
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()
def __init__(
self,
tracker_key: str = "key",
silent: bool = False,
no_fork: bool = False,
) -> None:
self.tracker = Tracker(
silent=silent,
port=9190,
port_end=12345,
)
self.server = Server(
host="0.0.0.0",
is_proxy=False,
tracker_addr=(self.tracker.host, self.tracker.port),
key=tracker_key,
silent=silent,
no_fork=no_fork,
port=9190,
port_end=12345,
)
self.tracker_host = self.tracker.host
self.tracker_port = self.tracker.port
self.tracker_key = tracker_key
class LocalRPCMeasureContext:
"""A context wrapper for running RPCRunner locally.
This will launch a local RPC Tracker and local RPC Server.
Parameters
----------
priority : int = 1
The priority of this run request, larger is more prior.
n_parallel : int = 1
The number of tasks run in parallel.
timeout : int = 10
The timeout limit (in second) for each run.
This is used in a wrapper of the multiprocessing.Process.join().
number : int = 3
The number of times to run the generated code for taking average.
We call these runs as one `repeat` of measurement.
repeat : int = 1
The number of times to repeat the measurement.
In total, the generated code will be run (1 + number x repeat) times,
where the first "1" is warm up and will be discarded.
The returned result contains `repeat` costs,
each of which is an average of `number` costs.
min_repeat_ms : int = 0
The minimum duration of one `repeat` in milliseconds.
By default, one `repeat` contains `number` runs. If this parameter is set,
the parameters `number` will be dynamically adjusted to meet the
minimum duration requirement of one `repeat`.
i.e., When the run time of one `repeat` falls below this time, the `number` parameter
will be automatically increased.
cooldown_interval : float = 0.0
The cool down interval between two measurements.
enable_cpu_cache_flush: bool = False
Whether to flush cache on CPU between repeated measurements.
Flushing cache can make the measured latency of one operator closer to
its actual latency during end-to-end inference.
To make this option effective, the argument `number` should also be set to 1.
This is only has effect on CPU task.
"""
def __init__(
self,
priority=1,
n_parallel=1,
timeout=10,
number=3,
repeat=1,
min_repeat_ms=0,
cooldown_interval=0.0,
enable_cpu_cache_flush=False,
):
# pylint: disable=import-outside-toplevel
from tvm.rpc.tracker import Tracker
from tvm.rpc.server import Server
dev = tvm.device("cuda", 0)
if dev.exist:
cuda_arch = "sm_" + "".join(dev.compute_version.split("."))
set_cuda_target_arch(cuda_arch)
host = "0.0.0.0"
self.tracker = Tracker(host, port=9000, port_end=10000, silent=True)
device_key = "$local$device$%d" % self.tracker.port
self.server = Server(
host,
port=self.tracker.port,
port_end=10000,
key=device_key,
use_popen=True,
silent=True,
tracker_addr=(self.tracker.host, self.tracker.port),
)
self.runner = RPCRunner(
device_key,
host,
self.tracker.port,
priority,
n_parallel,
timeout,
number,
repeat,
min_repeat_ms,
cooldown_interval,
enable_cpu_cache_flush,
)
# Wait for the processes to start
time.sleep(0.5)
def __del__(self):
# Close the tracker and server before exit
self.tracker.terminate()
self.server.terminate()
time.sleep(0.5)
def test_rpc_tracker_register():
# test registration
tracker = Tracker(port=9000, port_end=10000)
device_key = "test_device"
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"] == "server:%s" % 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()
