chore: import upstream snapshot with attribution

python/tvm/rpc/proxy.py

@@ -0,0 +1,727 @@
+# 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: F841
+"""RPC proxy, allows both client/server to connect and match connection.
+
+In normal RPC, client directly connect to server's IP address.
+Sometimes this cannot be done when server do not have a static address.
+RPCProxy allows both client and server connect to the proxy server,
+the proxy server will forward the message between the client and server.
+"""
+
+# pylint: disable=unused-variable, unused-argument
+import asyncio
+import errno
+import logging
+import os
+import socket
+import struct
+import threading
+import time
+
+try:
+    import tornado
+    from tornado import gen, ioloop, websocket
+
+    from . import tornado_util
+except ImportError as error_msg:
+    raise ImportError(
+        f"RPCProxy module requires tornado package {error_msg}. Try 'pip install tornado'."
+    )
+
+from tvm.support.popen_pool import PopenWorker
+
+from . import _ffi_api, base
+from .base import TrackerCode
+from .server import _server_env
+
+
+class ForwardHandler:
+    """Forward handler to forward the message."""
+
+    def _init_handler(self):
+        """Initialize handler."""
+        self._init_message = b""
+        self._init_req_nbytes = 4
+        self._magic = None
+        self.timeout = None
+        self._rpc_key_length = None
+        self._done = False
+        self._proxy = ProxyServerHandler.current
+        assert self._proxy
+        self.rpc_key = None
+        self.match_key = None
+        self.forward_proxy = None
+        self.alloc_time = None
+
+    def __del__(self):
+        logging.info("Delete %s...", self.name())
+
+    def name(self):
+        """Name of this connection."""
+        return "RPCConnection"
+
+    def _init_step(self, message):
+        if self._magic is None:
+            assert len(message) == 4
+            self._magic = struct.unpack("<i", message)[0]
+            if self._magic != base.RPC_MAGIC:
+                logging.info("Invalid RPC magic from %s", self.name())
+                self.close()
+            self._init_req_nbytes = 4
+        elif self._rpc_key_length is None:
+            assert len(message) == 4
+            self._rpc_key_length = struct.unpack("<i", message)[0]
+            self._init_req_nbytes = self._rpc_key_length
+        elif self.rpc_key is None:
+            assert len(message) == self._rpc_key_length
+            self.rpc_key = message.decode("utf-8")
+            # match key is used to do the matching
+            self.match_key = self.rpc_key[7:].split()[0]
+            self.on_start()
+        else:
+            assert False
+
+    def on_start(self):
+        """Event when the initialization is completed"""
+        self._proxy.handler_ready(self)
+
+    def on_data(self, message):
+        """on data"""
+        assert isinstance(message, bytes)
+        if self.forward_proxy:
+            self.forward_proxy.send_data(message)
+        else:
+            while message and self._init_req_nbytes > len(self._init_message):
+                nbytes = self._init_req_nbytes - len(self._init_message)
+                self._init_message += message[:nbytes]
+                message = message[nbytes:]
+                if self._init_req_nbytes == len(self._init_message):
+                    temp = self._init_message
+                    self._init_req_nbytes = 0
+                    self._init_message = b""
+                    self._init_step(temp)
+            if message:
+                logging.info("Invalid RPC protocol, too many bytes %s", self.name())
+                self.close()
+
+    def on_error(self, err):
+        logging.info("%s: Error in RPC %s", self.name(), err)
+        self.close_pair()
+
+    def close_pair(self):
+        if self.forward_proxy:
+            self.forward_proxy.signal_close()
+            self.forward_proxy = None
+        self.close()
+
+    def on_close_event(self):
+        """on close event"""
+        assert not self._done
+        logging.info("RPCProxy:on_close_event %s ...", self.name())
+        if self.match_key:
+            key = self.match_key
+            if self._proxy._client_pool.get(key, None) == self:
+                self._proxy._client_pool.pop(key)
+            if self._proxy._server_pool.get(key, None) == self:
+                self._proxy._server_pool.pop(key)
+        self._done = True
+        self.forward_proxy = None
+
+
+class TCPHandler(tornado_util.TCPHandler, ForwardHandler):
+    """Event driven TCP handler."""
+
+    def __init__(self, sock, addr):
+        super().__init__(sock)
+        self._init_handler()
+        self.addr = addr
+
+    def name(self):
+        return f"TCPSocketProxy:{self.addr[0]!s}:{self.rpc_key}"
+
+    def send_data(self, message, binary=True):
+        self.write_message(message, True)
+
+    def on_message(self, message):
+        self.on_data(message)
+
+    def on_close(self):
+        logging.info("RPCProxy: on_close %s ...", self.name())
+        self._close_process = True
+
+        if self.forward_proxy:
+            self.forward_proxy.signal_close()
+            self.forward_proxy = None
+        self.on_close_event()
+
+
+class WebSocketHandler(websocket.WebSocketHandler, ForwardHandler):
+    """Handler for websockets."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._init_handler()
+
+    def name(self):
+        return f"WebSocketProxy:{self.rpc_key}"
+
+    def on_message(self, message):
+        self.on_data(message)
+
+    def data_received(self, _):
+        raise NotImplementedError()
+
+    def send_data(self, message):
+        try:
+            self.write_message(message, True)
+        except websocket.WebSocketClosedError as err:
+            self.on_error(err)
+
+    def on_close(self):
+        logging.info("RPCProxy: on_close %s ...", self.name())
+        if self.forward_proxy:
+            self.forward_proxy.signal_close()
+            self.forward_proxy = None
+        self.on_close_event()
+
+    def signal_close(self):
+        self.close()
+
+
+MIME_MAP = {
+    "js": "text/javascript",
+    "css": "text/css",
+    "wasm": "application/wasm",
+    "json": "application/json",
+}
+
+
+class RequestHandler(tornado.web.RequestHandler):
+    """Handles html request."""
+
+    def __init__(self, *args, **kwargs):
+        file_path = kwargs.pop("file_path")
+        self.format = file_path.split(".")[-1]
+
+        if file_path.endswith("html"):
+            self.page = open(file_path).read()
+            web_port = kwargs.pop("rpc_web_port", None)
+            if web_port:
+                self.page = self.page.replace(
+                    "ws://localhost:9190/ws", f"ws://localhost:{web_port}/ws"
+                )
+        else:
+            self.page = open(file_path, "rb").read()
+
+        super().__init__(*args, **kwargs)
+
+    def data_received(self, _):
+        pass
+
+    def get(self, *args, **kwargs):
+        if self.format in MIME_MAP:
+            self.set_header("Content-Type", MIME_MAP[self.format])
+        self.write(self.page)
+
+
+class ProxyServerHandler:
+    """Internal proxy server handler class."""
+
+    current = None
+
+    def __init__(
+        self,
+        sock,
+        listen_port,
+        web_port,
+        timeout_client,
+        timeout_server,
+        tracker_addr,
+        index_page=None,
+        resource_files=None,
+    ):
+        assert ProxyServerHandler.current is None
+        ProxyServerHandler.current = self
+        if web_port:
+            handlers = [(r"/ws", WebSocketHandler)]
+            if index_page:
+                handlers.append(
+                    (r"/", RequestHandler, {"file_path": index_page, "rpc_web_port": web_port})
+                )
+                logging.info("Serving RPC index html page at http://localhost:%d", web_port)
+            resource_files = resource_files if resource_files else []
+            for item in resource_files:
+                prefix, fname = item
+                if not prefix.endswith("/"):
+                    prefix += "/"
+                if not prefix.startswith("/"):
+                    prefix = "/" + prefix
+                basename = os.path.basename(fname)
+                pair = (rf"{prefix}{basename}", RequestHandler, {"file_path": fname})
+                handlers.append(pair)
+                logging.info(pair)
+            self.app = tornado.web.Application(handlers)
+            self.app.listen(web_port)
+
+        self.sock = sock
+        self.sock.setblocking(0)
+        self.loop = ioloop.IOLoop.current()
+
+        def event_handler(_, events):
+            self._on_event(events)
+
+        self.loop.add_handler(self.sock.fileno(), event_handler, self.loop.READ)
+        self._client_pool = {}
+        self._server_pool = {}
+        self.timeout_alloc = 5
+        self.timeout_client = timeout_client
+        self.timeout_server = timeout_server
+        # tracker information
+        self._listen_port = listen_port
+        self._tracker_addr = tracker_addr
+        self._tracker_conn = None
+        self._tracker_pending_puts = []
+        self._key_set = set()
+        self.update_tracker_period = 2
+        if tracker_addr:
+            logging.info("Tracker address:%s", str(tracker_addr))
+
+            def _callback():
+                self._update_tracker(True)
+
+            self.loop.call_later(self.update_tracker_period, _callback)
+        logging.info("RPCProxy: Websock port bind to %d", web_port)
+
+    def _on_event(self, _):
+        while True:
+            try:
+                conn, addr = self.sock.accept()
+                TCPHandler(conn, addr)
+            except OSError as err:
+                if err.args[0] in (errno.EAGAIN, errno.EWOULDBLOCK):
+                    break
+
+    def _pair_up(self, lhs, rhs):
+        lhs.forward_proxy = rhs
+        rhs.forward_proxy = lhs
+
+        lhs.send_data(struct.pack("<i", base.RPC_CODE_SUCCESS))
+        lhs.send_data(struct.pack("<i", len(rhs.rpc_key)))
+        lhs.send_data(rhs.rpc_key.encode("utf-8"))
+
+        rhs.send_data(struct.pack("<i", base.RPC_CODE_SUCCESS))
+        rhs.send_data(struct.pack("<i", len(lhs.rpc_key)))
+        rhs.send_data(lhs.rpc_key.encode("utf-8"))
+        logging.info("Pairup connect %s  and %s", lhs.name(), rhs.name())
+
+    def _regenerate_server_keys(self, keys):
+        """Regenerate keys for server pool"""
+        keyset = set(self._server_pool.keys())
+        new_keys = []
+        # re-generate the server match key, so old information is invalidated.
+        for key in tuple(keys):
+            rpc_key, _ = base.split_random_key(key)
+            handle = self._server_pool[key]
+            del self._server_pool[key]
+            new_key = base.random_key(rpc_key, keyset)
+            self._server_pool[new_key] = handle
+            keyset.add(new_key)
+            new_keys.append(new_key)
+        return new_keys
+
+    def _update_tracker(self, period_update=False):
+        """Update information on tracker."""
+        try:
+            if self._tracker_conn is None:
+                self._tracker_conn = socket.socket(
+                    base.get_addr_family(self._tracker_addr), socket.SOCK_STREAM
+                )
+                self._tracker_conn.connect(self._tracker_addr)
+                self._tracker_conn.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+                magic = struct.unpack("<i", base.recvall(self._tracker_conn, 4))[0]
+                if magic != base.RPC_TRACKER_MAGIC:
+                    self.loop.stop()
+                    raise RuntimeError(f"{self._tracker_addr} is not RPC Tracker")
+                # just connect to tracker, need to update all keys
+                self._tracker_pending_puts = self._server_pool.keys()
+
+            if self._tracker_conn and period_update:
+                # periodically update tracker information
+                # regenerate key if the key is not in tracker anymore
+                # and there is no in-coming connection after timeout_alloc
+                base.sendjson(self._tracker_conn, [TrackerCode.GET_PENDING_MATCHKEYS])
+                pending_keys = set(base.recvjson(self._tracker_conn))
+                update_keys = []
+                for k, v in self._server_pool.items():
+                    if k not in pending_keys:
+                        if v.alloc_time is None:
+                            v.alloc_time = time.time()
+                        elif time.time() - v.alloc_time > self.timeout_alloc:
+                            update_keys.append(k)
+                            v.alloc_time = None
+                if update_keys:
+                    logging.info(
+                        "RPCProxy: No incoming conn on %s, regenerate keys...", str(update_keys)
+                    )
+                    new_keys = self._regenerate_server_keys(update_keys)
+                    self._tracker_pending_puts += new_keys
+
+            need_update_info = False
+            # report new connections
+            for key in self._tracker_pending_puts:
+                rpc_key, _ = base.split_random_key(key)
+                base.sendjson(
+                    self._tracker_conn, [TrackerCode.PUT, rpc_key, (self._listen_port, key), None]
+                )
+                assert base.recvjson(self._tracker_conn) == TrackerCode.SUCCESS
+                if rpc_key not in self._key_set:
+                    self._key_set.add(rpc_key)
+                    need_update_info = True
+
+            if need_update_info:
+                keylist = "[" + ",".join(self._key_set) + "]"
+                cinfo = {"key": "server:proxy" + keylist, "addr": [None, self._listen_port]}
+                base.sendjson(self._tracker_conn, [TrackerCode.UPDATE_INFO, cinfo])
+                assert base.recvjson(self._tracker_conn) == TrackerCode.SUCCESS
+            self._tracker_pending_puts = []
+        except OSError as err:
+            logging.info(
+                "Lost tracker connection: %s, try reconnect in %g sec",
+                str(err),
+                self.update_tracker_period,
+            )
+            self._tracker_conn.close()
+            self._tracker_conn = None
+            self._regenerate_server_keys(self._server_pool.keys())
+
+        if period_update:
+
+            def _callback():
+                self._update_tracker(True)
+
+            self.loop.call_later(self.update_tracker_period, _callback)
+
+    def _handler_ready_tracker_mode(self, handler):
+        """tracker mode to handle handler ready."""
+        if handler.rpc_key.startswith("server:"):
+            key = base.random_key(handler.match_key, cmap=self._server_pool)
+            handler.match_key = key
+            self._server_pool[key] = handler
+            self._tracker_pending_puts.append(key)
+            self._update_tracker()
+        else:
+            if handler.match_key in self._server_pool:
+                self._pair_up(self._server_pool.pop(handler.match_key), handler)
+            else:
+                handler.send_data(struct.pack("<i", base.RPC_CODE_MISMATCH))
+                handler.signal_close()
+
+    def _handler_ready_proxy_mode(self, handler):
+        """Normal proxy mode when handler is ready."""
+        if handler.rpc_key.startswith("server:"):
+            pool_src, pool_dst = self._client_pool, self._server_pool
+            timeout = self.timeout_server
+        else:
+            pool_src, pool_dst = self._server_pool, self._client_pool
+            timeout = self.timeout_client
+
+        key = handler.match_key
+        if key in pool_src:
+            self._pair_up(pool_src.pop(key), handler)
+            return
+        if key not in pool_dst:
+            pool_dst[key] = handler
+
+            def cleanup():
+                """Cleanup client connection if timeout"""
+                if pool_dst.get(key, None) == handler:
+                    logging.info(
+                        "Timeout client connection %s, cannot find match key=%s",
+                        handler.name(),
+                        key,
+                    )
+                    pool_dst.pop(key)
+                    handler.send_data(struct.pack("<i", base.RPC_CODE_MISMATCH))
+                    handler.signal_close()
+
+            self.loop.call_later(timeout, cleanup)
+        else:
+            logging.info("Duplicate connection with same key=%s", key)
+            handler.send_data(struct.pack("<i", base.RPC_CODE_DUPLICATE))
+            handler.signal_close()
+
+    def handler_ready(self, handler):
+        """Report handler to be ready."""
+        logging.info("Handler ready %s", handler.name())
+        if self._tracker_addr:
+            self._handler_ready_tracker_mode(handler)
+        else:
+            self._handler_ready_proxy_mode(handler)
+
+    def run(self):
+        """Run the proxy server"""
+        ioloop.IOLoop.current().start()
+
+
+def _proxy_server(
+    listen_sock,
+    listen_port,
+    web_port,
+    timeout_client,
+    timeout_server,
+    tracker_addr,
+    index_page,
+    resource_files,
+):
+    asyncio.set_event_loop(asyncio.new_event_loop())
+    handler = ProxyServerHandler(
+        listen_sock,
+        listen_port,
+        web_port,
+        timeout_client,
+        timeout_server,
+        tracker_addr,
+        index_page,
+        resource_files,
+    )
+    handler.run()
+
+
+class PopenProxyServerState:
+    """Internal PopenProxy State for Popen"""
+
+    current = None
+
+    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,
+    ):
+        sock = socket.socket(base.get_addr_family((host, port)), socket.SOCK_STREAM)
+        self.port = None
+        for my_port in range(port, port_end):
+            try:
+                sock.bind((host, my_port))
+                self.port = my_port
+                break
+            except OSError as sock_err:
+                if sock_err.errno in [errno.EADDRINUSE]:
+                    continue
+                raise sock_err
+        if not self.port:
+            raise ValueError(f"cannot bind to any port in [{port}, {port_end})")
+        logging.info("RPCProxy: client port bind to %s:%d", host, self.port)
+        sock.listen(1)
+        self.thread = threading.Thread(
+            target=_proxy_server,
+            args=(
+                sock,
+                self.port,
+                web_port,
+                timeout_client,
+                timeout_server,
+                tracker_addr,
+                index_page,
+                resource_files,
+            ),
+        )
+        # start the server in a different thread
+        # so we can return the port directly
+        self.thread.start()
+
+
+def _popen_start_proxy_server(
+    host,
+    port=9091,
+    port_end=9199,
+    web_port=0,
+    timeout_client=600,
+    timeout_server=600,
+    tracker_addr=None,
+    index_page=None,
+    resource_files=None,
+):
+    # This is a function that will be sent to the
+    # Popen worker to run on a separate process.
+    # Create and start the server in a different thread
+    state = PopenProxyServerState(
+        host,
+        port,
+        port_end,
+        web_port,
+        timeout_client,
+        timeout_server,
+        tracker_addr,
+        index_page,
+        resource_files,
+    )
+    PopenProxyServerState.current = state
+    # returns the port so that the main can get the port number.
+    return state.port
+
+
+class Proxy:
+    """Start RPC proxy server 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
+
+    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):
+        try:
+            self.terminate()
+        except ImportError:
+            pass
+
+
+def websocket_proxy_server(url, key=""):
+    """Create a RPC server that uses an websocket that connects to a proxy.
+
+    Parameters
+    ----------
+    url : str
+        The url to be connected.
+
+    key : str
+        The key to identify the server.
+    """
+
+    def create_on_message(conn):
+        def _fsend(data):
+            data = bytes(data)
+            conn.write_message(data, binary=True)
+            return len(data)
+
+        on_message = _ffi_api.CreateEventDrivenServer(_fsend, "WebSocketProxyServer", "%toinit")
+        return on_message
+
+    @gen.coroutine
+    def _connect(key):
+        conn = yield websocket.websocket_connect(url)
+        on_message = create_on_message(conn)
+        temp = _server_env(None)
+        # Start connecton
+        conn.write_message(struct.pack("<i", base.RPC_MAGIC), binary=True)
+        key = "server:" + key
+        conn.write_message(struct.pack("<i", len(key)), binary=True)
+        conn.write_message(key.encode("utf-8"), binary=True)
+        msg = yield conn.read_message()
+        assert len(msg) >= 4
+        magic = struct.unpack("<i", msg[:4])[0]
+        if magic == base.RPC_CODE_DUPLICATE:
+            raise RuntimeError(f"key: {key} has already been used in proxy")
+        if magic == base.RPC_CODE_MISMATCH:
+            logging.info("RPCProxy do not have matching client key %s", key)
+        elif magic != base.RPC_CODE_SUCCESS:
+            raise RuntimeError(f"{url} is not RPC Proxy")
+        msg = msg[4:]
+
+        logging.info("Connection established with remote")
+
+        if msg:
+            on_message(bytearray(msg), 3)
+
+        while True:
+            try:
+                msg = yield conn.read_message()
+                if msg is None:
+                    break
+                on_message(bytearray(msg), 3)
+            except websocket.WebSocketClosedError as err:
+                break
+        logging.info("WebSocketProxyServer closed...")
+        temp.remove()
+        ioloop.IOLoop.current().stop()
+
+    ioloop.IOLoop.current().spawn_callback(_connect, key)
+    ioloop.IOLoop.current().start()