def test_popen_initializer():
initargs = [1, 2, 3]
proc = PopenWorker(initializer=initializer, initargs=initargs)
proc.send(after_initializer)
test_global_state_1, test_global_state_2, test_global_state_3 = proc.recv()
assert test_global_state_1 == initargs[0]
assert test_global_state_2 == initargs[1]
assert test_global_state_3 == initargs[2]
def test_popen_worker_reuses():
proc = PopenWorker(maximum_uses=None)
proc.send(os.getpid)
initial_pid = proc.recv()
proc.send(os.getpid)
assert proc.recv() == initial_pid
def test_popen_worker_recycles_with_initializer():
initargs = [1, 2, 3]
proc = PopenWorker(initializer=initializer,
initargs=initargs,
maximum_uses=3)
proc.send(os.getpid)
initial_pid = proc.recv()
proc.send(after_initializer)
assert list(proc.recv()) == initargs
proc.send(os.getpid)
assert proc.recv() == initial_pid
# The process should be recycled with this send.
proc.send(os.getpid)
assert proc.recv() != initial_pid
# But the initializer should've run this time as well.
proc.send(after_initializer)
assert list(proc.recv()) == initargs
def test_popen_worker():
proc = PopenWorker()
with pytest.raises(TimeoutError):
proc.send(identity_after, [1, 100], timeout=0.01)
proc.recv()
with pytest.raises(ChildProcessError):
proc.send(terminate_self)
proc.recv()
proc.send(identity_after, [2, 0])
assert proc.recv() == 2
proc.send(identity_after, [4, 0.0001])
assert proc.recv() == 4
def test_popen_worker_recycles():
proc = PopenWorker(maximum_uses=2)
proc.send(os.getpid)
initial_pid = proc.recv()
assert psutil.pid_exists(initial_pid)
proc.send(os.getpid)
assert proc.recv() == initial_pid
assert psutil.pid_exists(initial_pid)
proc.send(os.getpid)
assert proc.recv() != initial_pid
assert not psutil.pid_exists(initial_pid)
def test_popen_ffi():
proc = PopenWorker(register_ffi)
# call python function via ffi
initargs = [0]
proc.send(call_py_ffi, initargs)
assert proc.recv() == initargs[0]
# call cpp function via ffi
initargs = [1]
proc.send(call_cpp_ffi, initargs)
assert proc.recv() == initargs[0]
# call python function from cpp function via ffi
initargs = [2]
proc.send(call_cpp_py_ffi, initargs)
assert proc.recv() == initargs[0]
class Tracker(object):
"""Start RPC tracker on a separate process.
Python implementation based on PopenWorker.
Parameters
----------
host : str
The host url of the server.
port : int
The TCP port to be bind to
port_end : int, optional
The end TCP port to search
silent: bool, optional
Whether run in silent mode
"""
def __init__(self, host="0.0.0.0", port=9190, port_end=9199, silent=False):
if silent:
logger.setLevel(logging.WARN)
self.proc = PopenWorker()
# send the function
self.proc.send(
_popen_start_tracker_server,
[
host,
port,
port_end,
silent,
],
)
# receive the port
self.port, self.stop_key = self.proc.recv()
self.host = host
def _stop_tracker(self):
sock = socket.socket(base.get_addr_family((self.host, self.port)),
socket.SOCK_STREAM)
sock.connect(("127.0.0.1", self.port))
sock.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
magic = struct.unpack("<i", base.recvall(sock, 4))[0]
assert magic == base.RPC_TRACKER_MAGIC
base.sendjson(sock, [TrackerCode.STOP, self.stop_key])
assert base.recvjson(sock) == TrackerCode.SUCCESS
sock.close()
def terminate(self):
"""Terminate the server process"""
if self.proc:
if self.proc.is_alive():
self._stop_tracker()
self.proc.join(0.1)
if self.proc.is_alive():
logger.info("Terminating Tracker Server...")
self.proc.kill()
self.proc = None
def __del__(self):
try:
self.terminate()
except TypeError:
pass
class Server(object):
"""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.
"""
def __init__(
self,
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,
):
try:
if _ffi_api.ServerLoop is None:
raise RuntimeError("Please compile with USE_RPC=1")
except NameError:
raise RuntimeError("Please compile with USE_RPC=1")
self.proc = PopenWorker()
# send the function
self.proc.send(
_popen_start_rpc_server,
[
host,
port,
port_end,
is_proxy,
tracker_addr,
key,
load_library,
custom_addr,
silent,
no_fork,
],
)
# receive the port
self.port = self.proc.recv()
self.host = host
def terminate(self):
"""Terminate the server process"""
if self.proc:
self.proc.kill()
self.proc = None
def __del__(self):
self.terminate()
class Proxy(object):
"""Start RPC proxy server on a seperate process.
Python implementation based on PopenWorker.
Parameters
----------
host : str
The host url of the server.
port : int
The TCP port to be bind to
port_end : int, optional
The end TCP port to search
web_port : int, optional
The http/websocket port of the server.
timeout_client : float, optional
Timeout of client until it sees a matching connection.
timeout_server : float, optional
Timeout of server until it sees a matching connection.
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.
index_page : str, optional
Path to an index page that can be used to display at proxy index.
resource_files : str, optional
Path to local resources that can be included in the http request
"""
def __init__(
self,
host,
port=9091,
port_end=9199,
web_port=0,
timeout_client=600,
timeout_server=600,
tracker_addr=None,
index_page=None,
resource_files=None,
):
self.proc = PopenWorker()
# send the function
self.proc.send(
_popen_start_proxy_server,
[
host,
port,
port_end,
web_port,
timeout_client,
timeout_server,
tracker_addr,
index_page,
resource_files,
],
)
# receive the port
self.port = self.proc.recv()
self.host = host
def terminate(self):
"""Terminate the server process"""
if self.proc:
logging.info("Terminating Proxy Server...")
self.proc.kill()
self.proc = None
def __del__(self):
self.terminate()