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()