chore: import upstream snapshot with attribution

python/tvm/rpc/__init__.py

@@ -0,0 +1,37 @@
+# isort: skip_file
+# 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.
+"""Lightweight TVM RPC module.
+
+RPC enables connect to a remote server, upload and launch functions.
+This is useful to for cross-compile and remote testing,
+The compiler stack runs on local server, while we use RPC server
+to run on remote runtime which don't have a compiler available.
+
+The test program compiles the program on local server,
+upload and run remote RPC server, get the result back to verify correctness.
+
+TVM RPC server assumes that the user is trusted and needs to be
+used in a trusted network environment and encrypted channels.
+It allows writings of arbitrary files into the server and provide
+full remote code execution capabilities to anyone who can access this API.
+"""
+
+from .server import Server
+from .client import connect, connect_tracker
+from .client import RPCSession, LocalSession, PopenSession, TrackerSession
+from .minrpc import with_minrpc

python/tvm/rpc/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs for tvm.rpc"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("rpc", __name__)

python/tvm/rpc/base.py

@@ -0,0 +1,196 @@
+# 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.
+"""Base definitions for RPC."""
+# pylint: disable=invalid-name
+
+import errno
+import json
+import logging
+import random
+import socket
+import struct
+import time
+
+# Magic header for RPC data plane
+RPC_MAGIC = 0xFF271
+# magic header for RPC tracker(control plane)
+RPC_TRACKER_MAGIC = 0x2F271
+# sucess response
+RPC_CODE_SUCCESS = RPC_MAGIC + 0
+# duplicate key in proxy
+RPC_CODE_DUPLICATE = RPC_MAGIC + 1
+# cannot found matched key in server
+RPC_CODE_MISMATCH = RPC_MAGIC + 2
+
+logger = logging.getLogger("RPCServer")
+
+
+class TrackerCode:
+    """Enumeration code for the RPC tracker"""
+
+    FAIL = -1
+    SUCCESS = 0
+    PING = 1
+    STOP = 2
+    PUT = 3
+    REQUEST = 4
+    UPDATE_INFO = 5
+    SUMMARY = 6
+    GET_PENDING_MATCHKEYS = 7
+
+
+RPC_SESS_MASK = 128
+
+
+# Use "127.0.0.1" or "::1" if there is a need to force ip4 or ip6
+# connection for "localhost".
+def get_addr_family(addr):
+    res = socket.getaddrinfo(addr[0], addr[1], 0, 0, socket.IPPROTO_TCP)
+    return res[0][0]
+
+
+def recvall(sock, nbytes):
+    """Receive all nbytes from socket.
+
+    Parameters
+    ----------
+    sock: Socket
+       The socket
+
+    nbytes : int
+       Number of bytes to be received.
+    """
+    res = []
+    nread = 0
+    while nread < nbytes:
+        chunk = sock.recv(min(nbytes - nread, 1024))
+        if not chunk:
+            raise OSError("connection reset")
+        nread += len(chunk)
+        res.append(chunk)
+    return b"".join(res)
+
+
+def sendjson(sock, data):
+    """send a python value to remote via json
+
+    Parameters
+    ----------
+    sock : Socket
+        The socket
+
+    data : object
+        Python value to be sent.
+    """
+    data = json.dumps(data)
+    sock.sendall(struct.pack("<i", len(data)))
+    sock.sendall(data.encode("utf-8"))
+
+
+def recvjson(sock):
+    """receive python value from remote via json
+
+    Parameters
+    ----------
+    sock : Socket
+        The socket
+
+    Returns
+    -------
+    value : object
+        The value received.
+    """
+    size = struct.unpack("<i", recvall(sock, 4))[0]
+    data = json.loads((recvall(sock, size)).decode("utf-8"))
+    return data
+
+
+def random_key(prefix, delimiter=":", cmap=None):
+    """Generate a random key
+
+    Parameters
+    ----------
+    prefix : str
+        The string prefix
+
+    delimiter : str
+        The delimiter
+
+    cmap : dict
+        Conflict map
+
+    Returns
+    -------
+    key : str
+        The generated random key
+    """
+    while True:
+        key = f"{prefix}{delimiter}{random.random()}"
+        if not cmap or key not in cmap:
+            break
+    return key
+
+
+def split_random_key(key, delimiter=":"):
+    """Split a random key by delimiter into prefix and random part
+
+    Parameters
+    ----------
+    key : str
+        The generated random key
+
+    Returns
+    -------
+    prefix : str
+        The string prefix
+
+    random_part : str
+        The generated random
+    """
+    return key.rsplit(delimiter, 1)
+
+
+def connect_with_retry(addr, timeout=60, retry_period=5):
+    """Connect to a TPC address with retry
+
+    This function is only reliable to short period of server restart.
+
+    Parameters
+    ----------
+    addr : tuple
+        address tuple
+
+    timeout : float
+         Timeout during retry
+
+    retry_period : float
+         Number of seconds before we retry again.
+    """
+    tstart = time.time()
+    while True:
+        try:
+            sock = socket.socket(get_addr_family(addr), socket.SOCK_STREAM)
+            sock.connect(addr)
+            return sock
+        except OSError as sock_err:
+            if sock_err.args[0] not in (errno.ECONNREFUSED,):
+                raise sock_err
+            period = time.time() - tstart
+            if period > timeout:
+                raise RuntimeError(f"Failed to connect to server {addr!s}")
+            logger.warning(f"Cannot connect to tracker {addr!s}, retry in {retry_period:g} secs...")
+            time.sleep(retry_period)

python/tvm/rpc/client.py

@@ -0,0 +1,568 @@
+# 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.
+# pylint: disable=used-before-assignment,broad-exception-caught
+# ruff: noqa: F401
+"""RPC client tools"""
+
+import os
+import socket
+import stat
+import struct
+import time
+
+import tvm_ffi
+from tvm_ffi import DLDeviceType
+
+import tvm.runtime
+from tvm.support import utils
+
+from . import _ffi_api, base, server
+
+
+class RPCSession:
+    """RPC Client session module
+
+    Do not directly create the object, call connect
+    """
+
+    # pylint: disable=invalid-name
+    def __init__(self, sess):
+        self._sess = sess
+        self._tbl_index = _ffi_api.SessTableIndex(sess)
+        self._remote_funcs = {}
+
+    def system_lib(self):
+        """Get system-wide library module.
+
+        Returns
+        -------
+        module : runtime.Module
+            The system-wide library module.
+
+        See Also
+        --------
+        tvm.runtime.system_lib
+        """
+        return self.get_function("ffi.SystemLib")()
+
+    def get_function(self, name):
+        """Get function from the session.
+
+        Parameters
+        ----------
+        name : str
+            The name of the function
+
+        Returns
+        -------
+        f : Function
+            The result function.
+        """
+        return self._sess.get_function(name)
+
+    def device(self, dev_type, dev_id=0):
+        """Construct a remote device.
+
+        Parameters
+        ----------
+        dev_type: int or str
+
+        dev_id: int, optional
+
+        Returns
+        -------
+        dev: Device
+            The corresponding encoded remote device.
+        """
+        dev = tvm.runtime.device(dev_type, dev_id)
+        encode = (self._tbl_index + 1) * base.RPC_SESS_MASK
+        dev = tvm.runtime.device(dev.dlpack_device_type() + encode, dev.index)
+        dev._rpc_sess = self
+        return dev
+
+    def upload(self, data, target=None):
+        """Upload file to remote runtime temp folder
+
+        Parameters
+        ----------
+        data : str or bytearray
+            The file name or binary in local to upload.
+
+        target : str, optional
+            The path in remote
+        """
+        if isinstance(data, bytearray):
+            if not target:
+                raise ValueError("target must present when file is a bytearray")
+            blob = data
+        else:
+            blob = bytearray(open(data, "rb").read())
+            if not target:
+                target = os.path.basename(data)
+
+        if "upload" not in self._remote_funcs:
+            self._remote_funcs["upload"] = self.get_function("tvm.rpc.server.upload")
+        self._remote_funcs["upload"](target, blob)
+
+    def download(self, path):
+        """Download file from remote temp folder.
+
+        Parameters
+        ----------
+        path : str
+            The relative location to remote temp folder.
+
+        Returns
+        -------
+        blob : bytearray
+            The result blob from the file.
+        """
+        if "download" not in self._remote_funcs:
+            self._remote_funcs["download"] = self.get_function("tvm.rpc.server.download")
+        return self._remote_funcs["download"](path)
+
+    def remove(self, path):
+        """Remove file from remote temp folder.
+
+        Parameters
+        ----------
+        path: str
+            The relative location to remote temp folder.
+        """
+        if "remove" not in self._remote_funcs:
+            self._remote_funcs["remove"] = self.get_function("tvm.rpc.server.remove")
+        self._remote_funcs["remove"](path)
+
+    def listdir(self, path):
+        """ls files from remote temp folder.
+
+        Parameters
+        ----------
+        path: str
+            The relative location to remote temp folder.
+
+        Returns
+        -------
+        dirs: str
+            The files in the given directory with split token ','.
+        """
+        if "listdir" not in self._remote_funcs:
+            self._remote_funcs["listdir"] = self.get_function("tvm.rpc.server.listdir")
+        return self._remote_funcs["listdir"](path)
+
+    def load_module(self, path):
+        """Load a remote module, the file need to be uploaded first.
+
+        Parameters
+        ----------
+        path : str
+            The relative location to remote temp folder.
+
+        Returns
+        -------
+        m : Module
+            The remote module containing remote function.
+        """
+        return _ffi_api.LoadRemoteModule(self._sess, path)
+
+    def download_linked_module(self, path):
+        """Link a module in the remote and download it.
+
+        Parameters
+        ----------
+        path : str
+            The relative location to remote temp folder.
+
+        Returns
+        -------
+        blob : bytearray
+            The result blob from the file.
+
+        Note
+        ----
+        This function can be helpful when a linker
+        is not available on the local client.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            mod = build_module_with_cross_compilation()
+            # export the module as tar because a local linker is not available
+            mod.export_library("lib.tar")
+            remote.upload("lib.tar")
+            # invoke the linker on the remote to link the module as a library
+            # note that the library can only run on the same env as the remote
+            with open("lib.so", "wb") as file:
+                file.write(remote.download_linked_module("lib.tar"))
+        """
+        if "download_linked_module" not in self._remote_funcs:
+            self._remote_funcs["download_linked_module"] = self.get_function(
+                "tvm.rpc.server.download_linked_module"
+            )
+        return self._remote_funcs["download_linked_module"](path)
+
+    def cpu(self, dev_id=0):
+        """Construct CPU device."""
+        return self.device(DLDeviceType.kDLCPU, dev_id)
+
+    def cuda(self, dev_id=0):
+        """Construct CUDA GPU device."""
+        return self.device(DLDeviceType.kDLCUDA, dev_id)
+
+    def cl(self, dev_id=0):
+        """Construct OpenCL device."""
+        return self.device(DLDeviceType.kDLOpenCL, dev_id)
+
+    def vulkan(self, dev_id=0):
+        """Construct Vulkan device."""
+        return self.device(DLDeviceType.kDLVulkan, dev_id)
+
+    def metal(self, dev_id=0):
+        """Construct Metal device."""
+        return self.device(DLDeviceType.kDLMetal, dev_id)
+
+    def rocm(self, dev_id=0):
+        """Construct ROCm device."""
+        return self.device(DLDeviceType.kDLROCM, dev_id)
+
+    def ext_dev(self, dev_id=0):
+        """Construct extension device."""
+        return self.device(DLDeviceType.kDLExtDev, dev_id)
+
+    def hexagon(self, dev_id=0):
+        """Construct Hexagon device."""
+        return self.device(DLDeviceType.kDLHexagon, dev_id)
+
+    def webgpu(self, dev_id=0):
+        """Construct WebGPU device."""
+        return self.device(DLDeviceType.kDLWebGPU, dev_id)
+
+
+class LocalSession(RPCSession):
+    """RPCSession interface backed by local environment.
+
+    This class can be used to implement functions that
+    need to be ran both locally and remotely.
+    """
+
+    def __init__(self):
+        self._temp = server._server_env([])
+        RPCSession.__init__(self, _ffi_api.LocalSession())
+
+
+@tvm_ffi.register_global_func("rpc.PopenSession")
+def _popen_session(binary):
+    temp = utils.tempdir()
+
+    if isinstance(binary, bytes | bytearray):
+        path_exec = temp.relpath("server.minrpc")
+        with open(path_exec, "wb") as outfile:
+            outfile.write(binary)
+        os.chmod(path_exec, stat.S_IXUSR | stat.S_IRUSR)
+        path_exec = os.path.abspath(path_exec)
+    else:
+        path_exec = os.path.abspath(binary)
+        if not os.path.isfile(path_exec):
+            raise RuntimeError(f"{path_exec} does not exist.")
+        if not os.access(path_exec, os.X_OK):
+            raise RuntimeError(f"{path_exec} is not executable.")
+
+    sess = _ffi_api.CreatePipeClient(path_exec)
+    return sess
+
+
+class PopenSession(RPCSession):
+    """RPCSession interface backed by popen.
+
+    Parameters
+    ----------
+    binary : List[Union[str, bytes]]
+        The binary to be executed.
+    """
+
+    def __init__(self, binary):
+        RPCSession.__init__(self, _popen_session(binary))
+
+
+class TrackerSession:
+    """Tracker client session.
+
+    Parameters
+    ----------
+    addr : tuple
+        The address tuple
+    """
+
+    def __init__(self, addr):
+        self._addr = addr
+        self._sock = None
+        self._connect()
+
+    def __del__(self):
+        self.close()
+
+    def _connect(self):
+        self._sock = base.connect_with_retry(self._addr)
+        self._sock.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+        magic = struct.unpack("<i", base.recvall(self._sock, 4))[0]
+        if magic != base.RPC_TRACKER_MAGIC:
+            raise RuntimeError(f"{self._addr!s} is not RPC Tracker")
+
+    def close(self):
+        """Close the tracker connection."""
+        if self._sock:
+            self._sock.close()
+            self._sock = None
+
+    def summary(self):
+        """Get the summary dict of the tracker."""
+        base.sendjson(self._sock, [base.TrackerCode.SUMMARY])
+        value = base.recvjson(self._sock)
+        if value[0] != base.TrackerCode.SUCCESS:
+            raise RuntimeError(f"Invalid return value {value!s}")
+        return value[1]
+
+    def text_summary(self):
+        """Get a text summary of the tracker."""
+        data = self.summary()
+
+        total_ct = {}
+
+        res = ""
+        res += "Server List\n"
+        res += "------------------------------\n"
+        res += "server-address           key\n"
+        res += "------------------------------\n"
+        sorted_server = sorted(data["server_info"], key=lambda x: x["key"])
+        for item in sorted_server:
+            addr = item["addr"]
+            res += f"{':'.join(map(str, addr)):21s}    "
+            res += item["key"] + "\n"
+            key = item["key"].split(":")[1]  # 'server:rasp3b` -> 'rasp3b'
+            if key not in total_ct:
+                total_ct[key] = 0
+            total_ct[key] += 1
+        res += "------------------------------\n"
+        res += "\n"
+
+        # compute max length of device key
+        queue_info = data["queue_info"]
+        keys = list(queue_info.keys())
+        if keys:
+            keys.sort()
+            max_key_len = max([len(k) for k in keys])
+        else:
+            max_key_len = 0
+
+        res += "Queue Status\n"
+        title = f"{'key':<{max_key_len}s}   total  free  pending\n"
+        separate_line = "-" * len(title) + "\n"
+        res += separate_line + title + separate_line
+        for k in keys:
+            total = total_ct.get(k, 0)
+            free, pending = queue_info[k]["free"], queue_info[k]["pending"]
+            if total or pending:
+                res += f"{k:<{max_key_len}}   {total:<5d}  {free:<4d}  {pending:<7d}\n"
+        res += separate_line
+        return res
+
+    def request(
+        self,
+        key,
+        priority=1,
+        session_timeout=0,
+        max_retry=5,
+        session_constructor_args=None,
+    ):
+        """Request a new connection from the tracker.
+
+        Parameters
+        ----------
+        key : str
+            The type key of the device.
+
+        priority : int, optional
+            The priority of the request.
+
+        session_timeout : float, optional
+            The duration of the session, allows server to kill
+            the connection when duration is longer than this value.
+            When duration is zero, it means the request must always be kept alive.
+
+        max_retry : int, optional
+            Maximum number of times to retry before give up.
+
+        session_constructor_args : list, optional
+            List of additional arguments to passed as the remote session constructor.
+            The first element of the list is always a string specifying the name of
+            the session constructor, the following args are the positional args to that function.
+        """
+        last_err = None
+        for _ in range(max_retry):
+            try:
+                if self._sock is None:
+                    self._connect()
+                base.sendjson(self._sock, [base.TrackerCode.REQUEST, key, "", priority])
+                value = base.recvjson(self._sock)
+                if value[0] != base.TrackerCode.SUCCESS:
+                    raise RuntimeError(f"Invalid return value {value!s}")
+                url, port, matchkey = value[1]
+                return connect(
+                    url,
+                    port,
+                    matchkey,
+                    session_timeout,
+                    session_constructor_args=session_constructor_args,
+                )
+            except OSError as err:
+                self.close()
+                last_err = err
+            except RuntimeError as err:
+                last_err = err
+        raise RuntimeError(f"Cannot request {key} after {max_retry} retry, last_error:{last_err!s}")
+
+    def request_and_run(self, key, func, priority=1, session_timeout=0, max_retry=2):
+        """Request a resource from tracker and run the func.
+
+        This function safe-guard rare server node dropout during execution.
+        In such case, a new resource will be requested and func will be ran again.
+
+        Parameters
+        ----------
+        key : str
+            The type key of the device.
+
+        func : function of session -> value
+            A stateless function
+
+        priority : int, optional
+            The priority of the request.
+
+        session_timeout : float, optional
+            The duration of the session, allows server to kill
+            the connection when duration is longer than this value.
+            When duration is zero, it means the request must always be kept alive.
+
+        max_retry : int, optional
+            Maximum number of times to retry the function before give up.
+        """
+        last_err = None
+        for _ in range(max_retry):
+            try:
+                sess = self.request(key, priority=priority, session_timeout=session_timeout)
+                tstart = time.time()
+                return func(sess)
+            except RuntimeError as err:
+                duration = time.time() - tstart
+                # roughly estimate if the error is due to timeout termination
+                if session_timeout and duration >= session_timeout * 0.95:
+                    raise RuntimeError(f"Session timeout when running {func.__name__}")
+                last_err = err
+        raise RuntimeError(
+            f"Failed to run on {key} after {max_retry} retry, last_error:{last_err!s}"
+        )
+
+
+def connect(
+    url,
+    port,
+    key="",
+    session_timeout=0,
+    session_constructor_args=None,
+    enable_logging=False,
+):
+    """Connect to RPC Server
+
+    Parameters
+    ----------
+    url : str
+        The url of the host
+
+    port : int
+        The port to connect to
+
+    key : str, optional
+        Additional key to match server
+
+    session_timeout : float, optional
+        The duration of the session in seconds, allows server to kill
+        the connection when duration is longer than this value.
+        When duration is zero, it means the request must always be kept alive.
+
+    session_constructor_args: List
+        List of additional arguments to passed as the remote session constructor.
+        The first element of the list is always a string specifying the name of
+        the session constructor, the following args are the positional args to that function.
+
+    enable_logging: boolean
+        flag to enable/disable logging. Logging is disabled by default.
+
+    Returns
+    -------
+    sess : RPCSession
+        The connected session.
+
+    Examples
+    --------
+    Normal usage
+    .. code-block:: python
+
+        client = rpc.connect(server_url, server_port, server_key)
+
+    Session_constructor can be used to customize the session in the remote
+    The following code connects to a remote internal server via a proxy
+    by constructing another RPCClientSession on the proxy machine and use that
+    as the serving session of the proxy endpoint.
+
+    .. code-block:: python
+
+        client_via_proxy = rpc.connect(
+            proxy_server_url, proxy_server_port, proxy_server_key, enable_logging
+            session_constructor_args=[
+                "rpc.Connect", internal_url, internal_port, internal_key, internal_logging])
+
+    """
+    try:
+        if session_timeout:
+            key += f" -timeout={session_timeout}"
+        session_constructor_args = session_constructor_args if session_constructor_args else []
+        if not isinstance(session_constructor_args, list | tuple):
+            raise TypeError("Expect the session constructor to be a list or tuple")
+        sess = _ffi_api.Connect(url, port, key, enable_logging, *session_constructor_args)
+    except NameError:
+        raise RuntimeError("Please compile with USE_RPC=1")
+    return RPCSession(sess)
+
+
+def connect_tracker(url, port):
+    """Connect to a RPC tracker
+
+    Parameters
+    ----------
+    url : str
+        The url of the host
+
+    port : int
+        The port to connect to
+
+    Returns
+    -------
+    sess : TrackerSession
+        The connected tracker session.
+    """
+    return TrackerSession((url, port))

python/tvm/rpc/minrpc.py

@@ -0,0 +1,98 @@
+# 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.
+"""Utils to path."""
+
+import os
+
+import tvm_ffi
+
+from tvm import libinfo
+from tvm.support import cc
+
+
+def find_minrpc_server_libpath(server="posix_popen_server"):
+    """Get the path of minrpc server libary.
+
+    Parameters
+    ----------
+    server : str
+        The kind of built in minrpc server.
+
+    Returns
+    -------
+    path : str
+        The path to the min server library.
+    """
+    curr_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
+    source_dir = os.path.abspath(os.path.join(curr_dir, "..", "..", ".."))
+    minrpc_dir = os.path.join(source_dir, "src", "runtime", "rpc", "minrpc")
+    path = os.path.join(minrpc_dir, server, f"{server}.cc")
+
+    candidates = [path]
+    if not os.path.isfile(path):
+        raise RuntimeError(f"Cannot find minserver {server}, in candidates {candidates}")
+    return minrpc_dir, path
+
+
+def with_minrpc(compile_func, server="posix_popen_server"):
+    """Attach the compiler function with minrpc related options.
+
+    Parameters
+    ----------
+    compile_func : Union[str, Callable[[str, str, Optional[str]], None]]
+        The compilation function to decorate.
+
+    server : str
+        The server type.
+
+    Returns
+    -------
+    fcompile : function
+        The return compilation.
+    """
+    minrpc_dir, server_path = find_minrpc_server_libpath(server)
+    runtime_path = libinfo.find_libtvm_runtime()
+    tvm_ffi_path = tvm_ffi.libinfo.find_libtvm_ffi()
+
+    runtime_dir = os.path.abspath(os.path.dirname(runtime_path))
+    tvm_ffi_dir = os.path.abspath(os.path.dirname(tvm_ffi_path))
+    options = ["-std=c++17"]
+    # Make sure the rpath to the libtvm_runtime is set so we can do local tests.
+    # Note that however, this approach won't work on remote.
+    # Always recommend to link statically.
+    options += ["-Wl,-rpath=" + runtime_dir]
+    options += ["-Wl,-rpath=" + tvm_ffi_dir]
+    # Force the linker to retain the runtime shared lib dependency even if no
+    # symbol from it is directly referenced. The minrpc server binary only
+    # touches tvm_ffi APIs at link time but relies on global functions (e.g.
+    # ``rpc.CreateEventDrivenServer``) registered via static initializers
+    # inside libtvm_runtime; with the default ``--as-needed`` those
+    # registrations would never run.
+    options += ["-Wl,--no-as-needed"]
+    default_include_paths = [
+        libinfo.find_include_path(),
+        tvm_ffi.libinfo.find_include_path(),
+        tvm_ffi.libinfo.find_dlpack_include_path(),
+    ]
+    options += ["-I" + path for path in default_include_paths]
+    options += ["-I" + minrpc_dir]
+    fcompile = cc.cross_compiler(
+        compile_func, options=options, add_files=[server_path, runtime_path, tvm_ffi_path]
+    )
+    fcompile.__name__ = "with_minrpc"
+    fcompile.need_system_lib = True
+    return fcompile

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

python/tvm/rpc/server.py

@@ -0,0 +1,559 @@
+# 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: F401
+"""RPC server implementation.
+
+Note
+----
+Server is TCP based with the following protocol:
+- Initial handshake to the peer
+  - [RPC_MAGIC, keysize(int32), key-bytes]
+- The key is in format
+   - {server|client}:device-type[:random-key] [-timeout=timeout]
+"""
+
+# pylint: disable=invalid-name
+import ctypes
+import errno
+import logging
+import multiprocessing
+import os
+import select
+import socket
+import struct
+import sys
+import threading
+import time
+from pathlib import Path
+
+import tvm_ffi
+
+from tvm.runtime.module import load_module as _load_module
+from tvm.support import utils
+from tvm.support.popen_pool import PopenWorker
+
+# pylint: disable=unused-import
+from . import _ffi_api, base, testing
+from .base import TrackerCode
+
+logger = logging.getLogger("RPCServer")
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(
+    logging.Formatter(
+        fmt="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+    )
+)
+logger.addHandler(console_handler)
+logger.setLevel(logging.INFO)
+logger.propagate = False
+
+
+def _server_env(load_library, work_path=None):
+    """Server environment function return temp dir"""
+    if work_path:
+        temp = work_path
+    else:
+        temp = utils.tempdir()
+
+    # pylint: disable=unused-variable
+    @tvm_ffi.register_global_func("tvm.rpc.server.workpath", override=True)
+    def get_workpath(path):
+        return temp.relpath(path)
+
+    @tvm_ffi.register_global_func("tvm.rpc.server.load_module", override=True)
+    def load_module(file_name):
+        """Load module from remote side."""
+        path = temp.relpath(file_name)
+        m = _load_module(path)
+        logger.info("load_module %s", path)
+        return m
+
+    @tvm_ffi.register_global_func("tvm.rpc.server.download_linked_module", override=True)
+    def download_linked_module(file_name):
+        """Load module from remote side."""
+        # pylint: disable=import-outside-toplevel
+        path = temp.relpath(file_name)
+
+        if path.endswith(".o"):
+            # Extra dependencies during runtime.
+            from tvm.support import cc as _cc
+
+            _cc.create_shared(path + ".so", path)
+            path += ".so"
+        elif path.endswith(".tar"):
+            # Extra dependencies during runtime.
+            from tvm.support import cc as _cc
+            from tvm.support import tar as _tar
+
+            tar_temp = utils.tempdir(custom_path=path.replace(".tar", ""))
+            _tar.untar(path, tar_temp.temp_dir)
+            files = [tar_temp.relpath(x) for x in tar_temp.listdir()]
+            _cc.create_shared(path + ".so", files)
+            path += ".so"
+        elif path.endswith(".dylib") or path.endswith(".so"):
+            pass
+        else:
+            raise RuntimeError(f"Do not know how to link {file_name}")
+        logger.info("Send linked module %s to client", path)
+        return bytearray(open(path, "rb").read())
+
+    libs = []
+    load_library = load_library.split(":") if load_library else []
+    for file_name in load_library:
+        # Resolve the literal library file name against the current working
+        # directory, preserving the exact filename the caller requested.
+        resolved = tvm_ffi.libinfo._resolve_and_validate(
+            [Path.cwd() / file_name], cond=lambda p: p.is_file()
+        )
+        if resolved is None:
+            raise RuntimeError(f"Cannot find library {file_name}")
+        file_name = resolved
+        libs.append(ctypes.CDLL(file_name, ctypes.RTLD_GLOBAL))
+        logger.info("Load additional library %s", file_name)
+    temp.libs = libs
+    return temp
+
+
+def _serve_loop(sock, load_library, work_path):
+    _server_env(load_library, work_path)
+    _ffi_api.ServerLoop(sock.fileno())
+
+
+def _parse_server_opt(opts):
+    # parse client options
+    ret = {}
+    for kv in opts:
+        if kv.startswith("-timeout="):
+            ret["timeout"] = float(kv[9:])
+    return ret
+
+
+def _serving(sock, addr, opts, load_library):
+    logger.info(f"connected from {addr}")
+    work_path = utils.tempdir()
+    old_cwd = os.getcwd()
+    os.chdir(work_path.path)  # Avoiding file name conflict between sessions.
+    logger.info(f"start serving at {work_path.path}")
+
+    server_proc = multiprocessing.Process(target=_serve_loop, args=(sock, load_library, work_path))
+    server_proc.start()
+    server_proc.join(opts.get("timeout", None))  # Wait until finish or timeout.
+
+    if server_proc.is_alive():
+        logger.info("timeout in RPC session, kill..")
+        _ffi_api.ReturnException(
+            sock.fileno(),
+            f"RPCSessionTimeoutError: Your {opts['timeout']}s session has expired, "
+            f'try to increase the "session_timeout" value.',
+        )
+
+        try:
+            import psutil  # pylint: disable=import-outside-toplevel
+
+            # Terminate worker children firstly.
+            for child in psutil.Process(server_proc.pid).children(recursive=True):
+                child.terminate()
+        except ImportError:
+            # Don't dependent `psutil` hardly, because it isn't a pure Python
+            # package and maybe hard to be installed on some platforms.
+            pass
+        server_proc.terminate()
+    elif server_proc.exitcode != 0:
+        raise RuntimeError(
+            f"Child process {server_proc.pid} exited unsuccessfully "
+            f"with error code {server_proc.exitcode}"
+        )
+
+    logger.info(f"finish serving {addr}")
+    os.chdir(old_cwd)
+    work_path.remove()
+    sock.close()
+
+
+def _listen_loop(sock, port, rpc_key, tracker_addr, load_library, custom_addr):
+    """Listening loop of the server."""
+
+    def _accept_conn(listen_sock, tracker_conn, ping_period=2):
+        """Accept connection from the other places.
+
+        Parameters
+        ----------
+        listen_sock: Socket
+            The socket used by listening process.
+
+        tracker_conn : connection to tracker
+            Tracker connection
+
+        ping_period : float, optional
+            ping tracker every k seconds if no connection is accepted.
+        """
+        old_keyset = set()
+        # Report resource to tracker
+        if tracker_conn:
+            matchkey = base.random_key(rpc_key)
+            base.sendjson(tracker_conn, [TrackerCode.PUT, rpc_key, (port, matchkey), custom_addr])
+            assert base.recvjson(tracker_conn) == TrackerCode.SUCCESS
+        else:
+            matchkey = rpc_key
+
+        unmatch_period_count = 0
+        unmatch_timeout = 4
+        # Wait until we get a valid connection
+        while True:
+            if tracker_conn:
+                trigger = select.select([listen_sock], [], [], ping_period)
+                if listen_sock not in trigger[0]:
+                    base.sendjson(tracker_conn, [TrackerCode.GET_PENDING_MATCHKEYS])
+                    pending_keys = base.recvjson(tracker_conn)
+                    old_keyset.add(matchkey)
+                    # if match key not in pending key set
+                    # it means the key is acquired by a client but not used.
+                    if matchkey not in pending_keys:
+                        unmatch_period_count += 1
+                    else:
+                        unmatch_period_count = 0
+                    # regenerate match key if key is acquired but not used for a while
+                    if unmatch_period_count * ping_period > unmatch_timeout + ping_period:
+                        logger.info("no incoming connections, regenerate key ...")
+                        matchkey = base.random_key(rpc_key, cmap=old_keyset)
+                        base.sendjson(
+                            tracker_conn, [TrackerCode.PUT, rpc_key, (port, matchkey), custom_addr]
+                        )
+                        assert base.recvjson(tracker_conn) == TrackerCode.SUCCESS
+                        unmatch_period_count = 0
+                    continue
+            conn, addr = listen_sock.accept()
+            magic = struct.unpack("<i", base.recvall(conn, 4))[0]
+            if magic != base.RPC_MAGIC:
+                conn.close()
+                continue
+            keylen = struct.unpack("<i", base.recvall(conn, 4))[0]
+            key = (base.recvall(conn, keylen)).decode("utf-8")
+            arr = key.split()
+            expect_header = "client:" + matchkey
+            server_key = "server:" + rpc_key
+            if arr[0] != expect_header:
+                conn.sendall(struct.pack("<i", base.RPC_CODE_MISMATCH))
+                conn.close()
+                logger.warning("mismatch key from %s", addr)
+                continue
+            conn.sendall(struct.pack("<i", base.RPC_CODE_SUCCESS))
+            conn.sendall(struct.pack("<i", len(server_key)))
+            conn.sendall(server_key.encode("utf-8"))
+            return conn, addr, _parse_server_opt(arr[1:])
+
+    # Server logic
+    tracker_conn = None
+    while True:
+        try:
+            # step 1: setup tracker and report to tracker
+            if tracker_addr and tracker_conn is None:
+                tracker_conn = base.connect_with_retry(tracker_addr)
+                tracker_conn.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+                magic = struct.unpack("<i", base.recvall(tracker_conn, 4))[0]
+                if magic != base.RPC_TRACKER_MAGIC:
+                    raise RuntimeError(f"{tracker_addr!s} is not RPC Tracker")
+                # report status of current queue
+                cinfo = {"key": "server:" + rpc_key, "addr": (custom_addr, port)}
+                base.sendjson(tracker_conn, [TrackerCode.UPDATE_INFO, cinfo])
+                assert base.recvjson(tracker_conn) == TrackerCode.SUCCESS
+
+            # step 2: wait for in-coming connections
+            conn, addr, opts = _accept_conn(sock, tracker_conn)
+        except OSError:
+            # retry when tracker is dropped
+            if tracker_conn:
+                tracker_conn.close()
+                tracker_conn = None
+            continue
+        except RuntimeError as exc:
+            raise exc
+
+        # step 3: serving
+        _serving(conn, addr, opts, load_library)
+
+
+def _connect_proxy_loop(addr, key, load_library):
+    key = "server:" + key
+    retry_count = 0
+    max_retry = 5
+    retry_period = 5
+    while True:
+        try:
+            sock = socket.socket(base.get_addr_family(addr), socket.SOCK_STREAM)
+            sock.connect(addr)
+            sock.sendall(struct.pack("<i", base.RPC_MAGIC))
+            sock.sendall(struct.pack("<i", len(key)))
+            sock.sendall(key.encode("utf-8"))
+            magic = struct.unpack("<i", base.recvall(sock, 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:
+                logger.warning("RPCProxy do not have matching client key %s", key)
+            elif magic != base.RPC_CODE_SUCCESS:
+                raise RuntimeError(f"{addr!s} is not RPC Proxy")
+            keylen = struct.unpack("<i", base.recvall(sock, 4))[0]
+            remote_key = (base.recvall(sock, keylen)).decode("utf-8")
+
+            _serving(sock, addr, _parse_server_opt(remote_key.split()[1:]), load_library)
+            retry_count = 0
+        except OSError as err:
+            retry_count += 1
+            logger.warning("Error encountered %s, retry in %g sec", str(err), retry_period)
+            if retry_count > max_retry:
+                raise RuntimeError(f"Maximum retry error: last error: {err!s}")
+            time.sleep(retry_period)
+
+
+class PopenRPCServerState:
+    """Internal PopenRPCServer State"""
+
+    current = None
+
+    def __init__(
+        self,
+        host,
+        port=9091,
+        port_end=9199,
+        is_proxy=False,
+        tracker_addr=None,
+        key="",
+        load_library=None,
+        custom_addr=None,
+        silent=False,
+        reuse_addr=True,
+        timeout=None,
+    ):
+        # start update
+        self.host = host
+        self.port = port
+        self.libs = []
+        self.custom_addr = custom_addr
+
+        if silent:
+            logger.setLevel(logging.ERROR)
+
+        if not is_proxy:
+            sock = socket.socket(base.get_addr_family((host, port)), socket.SOCK_STREAM)
+
+            # Never set socket SO_REUSEADDR on Windows. The SO_REUSEADDR flag allow reusing the
+            # inactivate TIME_WATI state sockets on POSIX, but on Windows it will allow two or more
+            # activate sockets to bind on the same address and port if they all set SO_REUSEADDR,
+            # and result in indeterminate behavior.
+            if reuse_addr and sys.platform != "win32":
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            if timeout is not None:
+                sock.settimeout(timeout)
+            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})")
+            logger.info("bind to %s:%d", host, self.port)
+            sock.listen(1)
+            self.sock = sock
+            self.thread = threading.Thread(
+                target=_listen_loop,
+                args=(self.sock, self.port, key, tracker_addr, load_library, self.custom_addr),
+            )
+            self.thread.start()
+        else:
+            self.thread = threading.Thread(
+                target=_connect_proxy_loop, args=((host, port), key, load_library)
+            )
+            self.thread.start()
+
+
+def _popen_start_rpc_server(
+    host,
+    port=9091,
+    port_end=9199,
+    is_proxy=False,
+    tracker_addr=None,
+    key="",
+    load_library=None,
+    custom_addr=None,
+    silent=False,
+    no_fork=False,
+    server_init_callback=None,
+    reuse_addr=True,
+    timeout=None,
+):
+    if no_fork:
+        multiprocessing.set_start_method("spawn")
+
+    if server_init_callback:
+        server_init_callback()
+
+    # 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 = PopenRPCServerState(
+        host,
+        port,
+        port_end,
+        is_proxy,
+        tracker_addr,
+        key,
+        load_library,
+        custom_addr,
+        silent,
+        reuse_addr,
+        timeout,
+    )
+    PopenRPCServerState.current = state
+    # returns the port so that the main can get the port number.
+    return state.port
+
+
+class Server:
+    """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.
+
+    server_init_callback: Callable, optional
+        Additional initialization function when starting the server.
+
+    reuse_addr: bool, optional
+        Allows the kernel to reuse a local socket in TIME_WAIT state.
+
+    timeout: float, optional
+         set a timeout for all operations on the socket
+
+    Note
+    ----
+    TVM RPC server assumes that the user is trusted and needs to be
+    used in a trusted network environment and encrypted channels.
+    It allows writings of arbitrary files into the server and provide
+    full remote code execution capabilities to anyone who can access this API.
+
+    The RPC server only sees functions in the tvm namespace.
+    To bring additional custom functions to the server env, you can use server_init_callback.
+
+    .. code:: python
+
+        def server_init_callback():
+            import tvm
+            # must import mypackage here
+            import mypackage
+
+            tvm.register_global_func("function", mypackage.func)
+
+        server = rpc.Server(host, server_init_callback=server_init_callback)
+    """
+
+    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,
+        server_init_callback=None,
+        reuse_addr=True,
+        timeout=None,
+    ):
+        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,
+                server_init_callback,
+                reuse_addr,
+                timeout,
+            ],
+        )
+        # 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):
+        try:
+            self.terminate()
+        except ImportError:
+            pass

python/tvm/rpc/server_ios_launcher.py

@@ -0,0 +1,503 @@
+# 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: RUF012
+"""
+Python wrapper for running a RPC Server through iOS RPC
+on the iOS simulator using the simctl command line tool.
+"""
+
+# pylint: disable=invalid-name
+import json
+import os
+import subprocess
+import threading
+import time
+from enum import Enum
+from typing import AnyStr
+
+
+class OSName(Enum):
+    """The names of the operating systems available on the simulator."""
+
+    iOS = "iOS"
+    tvOS = "tvOS"
+    watchOS = "watchOS"
+
+
+class IOSDevice(Enum):
+    """The names of available iOS devices."""
+
+    iPhone = "iPhone"
+    iPod = "iPod"
+    iPad = "iPad"
+
+
+class RPCServerMode(Enum):
+    """Server modes available in the iOS RPC application."""
+
+    standalone = "standalone"
+    proxy = "proxy"
+    tracker = "tracker"
+
+
+def get_list_of_available_simulators() -> dict[AnyStr, list]:
+    """
+    List of simulators available on the system. Simulators are presented as a dictionary.
+    The dictionary key is the name of the operating system of the simulator.
+    The dictionary value is a list of all simulators with a given operating system.
+    """
+
+    with subprocess.Popen(
+        "xcrun simctl list devices available --json",
+        shell=True,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+    ) as proc:
+        out, _ = proc.communicate()
+        available_simulators = json.loads(out)["devices"]
+        available_simulators = {
+            key: value for key, value in available_simulators.items() if value != []
+        }
+    return available_simulators
+
+
+def grep_by_system(available_devices: dict[AnyStr, list], system_name: OSName) -> list[dict]:
+    """Search for simulators that use the target operating system."""
+
+    def find_index_of_substr(search_field: list[AnyStr], target: AnyStr) -> int:
+        for i, item in enumerate(search_field):
+            if target in item:
+                return i
+        raise ValueError("Search field doesn't content target")
+
+    keys = list(available_devices.keys())
+
+    return available_devices[keys[find_index_of_substr(keys, system_name.value)]]
+
+
+def grep_by_device(available_devices: list[dict], device_name: IOSDevice) -> list[dict]:
+    """Search for simulators that emulate a given device."""
+
+    return [item for item in available_devices if device_name.value in item["name"]]
+
+
+def get_device_uid(target_device: dict) -> AnyStr:
+    """Get a unique device ID."""
+
+    return target_device["udid"]
+
+
+def check_call_with_runtime_error(cmd: AnyStr, error_message: AnyStr) -> None:
+    """Calling the function `subprocess.check_call` and catching its possible thrown exception."""
+
+    try:
+        subprocess.check_call(cmd.split(" "))
+    except subprocess.CalledProcessError as called_process_error:
+        raise called_process_error from RuntimeError(error_message)
+
+
+def boot_device(udid: AnyStr) -> None:
+    """Boot the device by its unique ID."""
+
+    cmd = f"xcrun simctl boot {udid}"
+    error_message = f"Failed to boot device with unique id: {udid}"
+    check_call_with_runtime_error(cmd, error_message)
+    if not is_booted(udid):
+        raise RuntimeError(error_message)
+
+
+def shutdown_device(udid: AnyStr) -> None:
+    """Shutdown the device by its unique ID."""
+
+    cmd = f"xcrun simctl shutdown {udid}"
+    error_message = f"Failed to shut down device with unique id: {udid}"
+    check_call_with_runtime_error(cmd, error_message)
+    if not is_turned_off(udid):
+        raise RuntimeError(error_message)
+
+
+def deploy_bundle_to_simulator(udid: AnyStr, bundle_path: AnyStr) -> None:
+    """Deploy iOS RPC bundle <bundle_path> to simulator with its unique ID <udid>."""
+
+    check_call_with_runtime_error(
+        cmd=f"xcrun simctl install {udid} {bundle_path}",
+        error_message=f"Failed to deploy bundle <{bundle_path}> to device with unique id: {udid}",
+    )
+
+
+def delete_bundle_from_simulator(udid: AnyStr, bundle_id: AnyStr) -> None:
+    """Delete iOS RPC bundle <bundle_id> from simulator with its unique ID <udid>."""
+
+    check_call_with_runtime_error(
+        cmd=f"xcrun simctl uninstall {udid} {bundle_id}",
+        error_message=f"Failed to uninstall bundle <{bundle_id}> "
+        f"from device with unique id: {udid}",
+    )
+
+
+def launch_ios_rpc(
+    udid: AnyStr, bundle_id: AnyStr, host_url: AnyStr, host_port: int, key: AnyStr, mode: AnyStr
+):  # pylint: disable=too-many-arguments, consider-using-with
+    """
+    Launch iOS RPC application on simulator with No UI interconnection.
+
+    udid : str
+        Unique device ID.
+
+    bundle_id : str
+        iOS RPC bundle ID.
+
+    host_url : str
+        The tracker/proxy address.
+
+    host_port : int
+        The tracker/proxy port.
+
+    key : str
+        The key used to identify the device type in tracker.
+
+    mode : str
+        Server mode. See RPCServerMode.
+    """
+
+    cmd = (
+        f"xcrun simctl launch --console {udid} {bundle_id}"
+        f" --immediate_connect"
+        f" --host_url={host_url}"
+        f" --host_port={host_port}"
+        f" --key={key}"
+        f" --server_mode={mode}"
+        f" --verbose"
+    )
+    proc = subprocess.Popen(
+        cmd.split(" "),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        bufsize=1,
+        universal_newlines=True,
+    )
+    return proc
+
+
+def terminate_ios_rpc(udid: AnyStr, bundle_id: AnyStr) -> None:
+    """Terminate iOS RPC application."""
+
+    check_call_with_runtime_error(
+        cmd=f"xcrun simctl terminate {udid} {bundle_id}",
+        error_message=f"Failed to terminate bundle <{bundle_id}> "
+        f"from device with unique id: {udid}",
+    )
+
+
+def is_booted(udid: AnyStr) -> bool:
+    """Check that the device has booted."""
+
+    device = find_device(udid)
+    return device["state"] == "Booted"
+
+
+def is_turned_off(udid: AnyStr) -> bool:
+    """Check that the device has turned off."""
+
+    device = find_device(udid)
+    return device["state"] == "Shutdown"
+
+
+def check_booted_device(devices: list[dict]) -> dict:
+    """Check if there is already a booted device. If so, return this device."""
+
+    for device in devices:
+        if device["state"] == "Booted":
+            return device
+    return {}
+
+
+def find_device(udid: AnyStr) -> dict:
+    """Find device by its unique ID."""
+
+    return_value = {}
+    available_devices = get_list_of_available_simulators()
+    for devices in available_devices.values():
+        for device in devices:
+            if device["udid"] == udid:
+                return_value = device
+    return return_value
+
+
+class ServerIOSLauncher:
+    """
+    Python wrapper for launch iOS RPC to simulator.
+
+    mode : str
+        Server mode. See RPCServerMode.
+
+    host : str
+        The tracker/proxy address.
+
+    port : int
+        The tracker/proxy port.
+
+    key : str
+        The key used to identify the device type in tracker.
+    """
+
+    booted_devices = []
+    bundle_id = os.environ.get("BUNDLE_ID")
+    bundle_path = os.environ.get("BUNDLE_PATH")
+
+    class ConsoleMarkers(Enum):
+        """
+        Marker-messages that iOS RPC Server should print to the console output
+        when its states change (see apps/ios_rpc/tvmrpc/RPCServer.mm).
+
+        STOPPED : str
+            iOS RPC Server process was stopped
+
+        CALLSTACK : str
+            Call stack if RPC Server was stopped with an error.
+
+        CONNECTED : str
+            RPC Server reports that it successfully connected.
+
+        SERVER_IP : str
+            IP on which RPC Server started (for standalone mode).
+
+        SERVER_PORT : str
+            HOST on which RPC Server started (for standalone mode).
+        """
+
+        STOPPED = "PROCESS_STOPPED"
+        CALLSTACK = "First throw call stack"
+        CONNECTED = "[IOS-RPC] STATE: 2"
+        SERVER_IP = "[IOS-RPC] IP: "
+        SERVER_PORT = "[IOS-RPC] PORT: "
+
+    def __init__(self, mode, host, port, key):
+        if not ServerIOSLauncher.is_compatible_environment():
+            raise RuntimeError(
+                "Can't create ServerIOSLauncher instance."
+                " No environment variables set for iOS RPC Server."
+            )
+
+        self.host = host
+        self.port = port
+
+        self.external_booted_device = None
+        if not ServerIOSLauncher.booted_devices:
+            self._boot_or_find_booted_device()
+
+        self.udid = get_device_uid(
+            self.external_booted_device
+            if self.external_booted_device is not None
+            else ServerIOSLauncher.booted_devices[-1]
+        )
+
+        self.bundle_was_deployed = False
+        deploy_bundle_to_simulator(self.udid, self.bundle_path)
+        self.bundle_was_deployed = True
+
+        self.server_was_started = False
+        self.launch_process = launch_ios_rpc(self.udid, self.bundle_id, host, port, key, mode)
+        self._wait_launch_complete(
+            waiting_time=60,
+            hz=10,
+            should_print_host_and_port=mode == RPCServerMode.standalone.value,
+        )
+        self.server_was_started = True
+
+    def terminate(self):
+        """Terminate iOS RPC server."""
+
+        if self.bundle_was_deployed and self.server_was_started:
+            try:
+                terminate_ios_rpc(self.udid, self.bundle_id)
+                self.launch_process.terminate()
+                self.server_was_started = False
+            except RuntimeError as e:
+                print(e)
+        if self.bundle_was_deployed:
+            try:
+                delete_bundle_from_simulator(self.udid, self.bundle_id)
+                self.bundle_was_deployed = False
+            except RuntimeError as e:
+                print(e)
+
+    def __del__(self):
+        try:
+            self.terminate()
+        except ImportError:
+            pass
+
+    @staticmethod
+    def is_compatible_environment():
+        """Check that the current environment has the required variables."""
+
+        return bool(os.environ.get("BUNDLE_ID")) and bool(os.environ.get("BUNDLE_PATH"))
+
+    @staticmethod
+    def shutdown_booted_devices():
+        """Shutdown simulators that have been booted using this class."""
+
+        for device_meta in ServerIOSLauncher.booted_devices:
+            try:
+                shutdown_device(get_device_uid(device_meta))
+            except RuntimeError as e:
+                print(e)
+        ServerIOSLauncher.booted_devices = []
+
+    def _boot_or_find_booted_device(self):
+        """
+        Boot the required simulator if there is no suitable booted simulator
+        among the available simulators. If there is a suitable booted simulator,
+        then take it as a simulator to which the iOS RPC application will be deployed.
+        """
+
+        target_system = OSName.iOS
+        target_device_type = IOSDevice.iPhone
+        available_devices = get_list_of_available_simulators()
+        if not available_devices:
+            raise ValueError("No devices available in this environment")
+        target_devices = grep_by_system(available_devices, target_system)
+        if not target_devices:
+            raise ValueError(f"No available simulators for target system: {target_system.value}")
+        target_devices = grep_by_device(target_devices, target_device_type)
+        if not target_devices:
+            raise ValueError(
+                f"No available simulators for target device type: {target_device_type.value}"
+            )
+
+        maybe_booted = check_booted_device(target_devices)
+        if maybe_booted:
+            self.external_booted_device = maybe_booted
+        else:
+            take_latest_model = True
+            target_device = target_devices[-1 if take_latest_model else 0]
+            boot_device(get_device_uid(target_device))
+            ServerIOSLauncher.booted_devices.append(target_device)
+
+    def _wait_launch_complete(self, waiting_time, hz, should_print_host_and_port=False):
+        # pylint: disable=too-many-locals
+        """
+        Wait for the iOS RPC server to start.
+
+        waiting_time : int
+            The maximum waiting time during which it is necessary
+            to receive a message from RPC Server.
+
+        hz : int
+            The frequency of checking (in hertz) messages from RPC Server.
+            Checks for messages from the server will occur every 1 / hz second.
+
+        should_print_host_and_port : bool
+            A flag that indicates that RPC Server should print the host and port
+            on which it was started.
+            Used for standalone mode.
+        """
+
+        class Switch:
+            """A simple helper class for boolean switching."""
+
+            def __init__(self):
+                self._on = False
+
+            def toggle(self):
+                """Toggle flag."""
+                self._on = not self._on
+
+            @property
+            def on(self):
+                """Flag of this switch."""
+                return self._on
+
+        def watchdog():
+            for _ in range(waiting_time * hz):
+                time.sleep(1.0 / hz)
+                if switch_have_data.on:
+                    break
+            if not switch_have_data.on:
+                self.launch_process.terminate()
+                switch_process_was_terminated.toggle()
+
+        switch_have_data = Switch()
+        switch_process_was_terminated = Switch()
+        watchdog_thread = threading.Thread(target=watchdog)
+
+        host, port = None, None
+        watchdog_thread.start()
+        for line in self.launch_process.stdout:
+            if not switch_have_data.on:
+                switch_have_data.toggle()
+
+            found = str(line).find(ServerIOSLauncher.ConsoleMarkers.STOPPED.value)
+            if found != -1:
+                raise RuntimeError("[ERROR] Crash during RCP Server launch.. ")
+
+            found = str(line).find(ServerIOSLauncher.ConsoleMarkers.CALLSTACK.value)
+            if found != -1:
+                raise RuntimeError("[ERROR] Crash during RCP Server launch.. ")
+
+            found = str(line).find(ServerIOSLauncher.ConsoleMarkers.SERVER_IP.value)
+            if found != -1:
+                ip = str(line)[
+                    found + len(ServerIOSLauncher.ConsoleMarkers.SERVER_IP.value) :
+                ].rstrip("\n")
+                host = ip
+
+            found = str(line).find(ServerIOSLauncher.ConsoleMarkers.SERVER_PORT.value)
+            if found != -1:
+                port = str(line)[
+                    found + len(ServerIOSLauncher.ConsoleMarkers.SERVER_PORT.value) :
+                ].rstrip("\n")
+                port = int(port)
+
+            if str(line).find(ServerIOSLauncher.ConsoleMarkers.CONNECTED.value) != -1:
+                # rpc server reports that it successfully connected
+                break
+        watchdog_thread.join()
+
+        if switch_process_was_terminated.on:
+            raise TimeoutError("Can't get a response from the iOS Server.")
+        if should_print_host_and_port:
+            if host is None or port is None:
+                raise RuntimeError("No messages with actual host and port.")
+            self.port = port
+
+
+class ServerIOSContextManager:
+    """
+    Context manager for ServerIOSLauncher.
+    To work with ServerIOSLauncher, it is preferable to use this class
+    so that the terminate method is called in any case.
+    """
+
+    def __init__(self, mode, host, port, key):
+        self.__mode = mode
+        self.__host = host
+        self.__port = port
+        self.__key = key
+        self.__ios_rpc_server_launcher = None
+
+    def __enter__(self):
+        self.__ios_rpc_server_launcher = ServerIOSLauncher(
+            self.__mode, self.__host, self.__port, self.__key
+        )
+        return self.__ios_rpc_server_launcher
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.__ios_rpc_server_launcher is not None:
+            self.__ios_rpc_server_launcher.terminate()
+            self.__ios_rpc_server_launcher = None

python/tvm/rpc/testing.py

@@ -0,0 +1,76 @@
+# 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.
+
+# pylint: disable=invalid-name,unnecessary-comprehension
+"""Testing functions for the RPC server."""
+
+import numpy as np
+
+import tvm
+
+
+# RPC test functions to be registered for unit-tests purposes
+@tvm.register_global_func("rpc.test.addone")
+def _addone(x):
+    return x + 1
+
+
+@tvm.register_global_func("rpc.test.strcat")
+def _strcat(name, x):
+    return f"{name}:{x}"
+
+
+@tvm.register_global_func("rpc.test.except")
+def _remotethrow(name):
+    raise ValueError(f"{name}")
+
+
+@tvm.register_global_func("rpc.test.runtime_str_concat")
+def _strcat(x, y):
+    return x + y
+
+
+@tvm.register_global_func("rpc.test.remote_tensor_func")
+def _remote_tensor_func(y):
+    x = np.ones((3, 4))
+    np.testing.assert_equal(y.numpy(), x)
+
+
+@tvm.register_global_func("rpc.test.add_to_lhs")
+def _add_to_lhs(x):
+    return lambda y: x + y
+
+
+@tvm.register_global_func("rpc.test.remote_return_nd")
+def _my_module(name):
+    # Use closure to check the ref counter correctness
+    nd = tvm.runtime.tensor(np.zeros(10).astype("float32"))
+
+    if name == "get_arr":
+        return lambda: nd
+    if name == "ref_count":
+        # Imported lazily: rpc.server imports this module to register the
+        # rpc.test.* helpers, so a top-level ``import tvm.testing`` would drag
+        # tvm.testing (which imports pytest) into the plain ``import tvm`` path.
+        from tvm.testing import object_use_count
+
+        return lambda: object_use_count(nd)
+    if name == "get_elem":
+        return lambda idx: nd.numpy()[idx]
+    if name == "get_arr_elem":
+        return lambda arr, idx: arr.numpy()[idx]
+    raise RuntimeError("unknown name")

python/tvm/rpc/tornado_util.py

@@ -0,0 +1,129 @@
+# 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: F401
+"""Utilities used in tornado."""
+
+import errno
+import socket
+
+from tornado import ioloop
+
+
+class TCPHandler:
+    """TCP socket handler backed tornado event loop.
+
+    Parameters
+    ----------
+    sock : Socket
+        The TCP socket, will set it to non-blocking mode.
+    """
+
+    def __init__(self, sock):
+        self._sock = sock
+        self._ioloop = ioloop.IOLoop.current()
+        self._sock.setblocking(0)
+        self._pending_write = []
+        self._signal_close = False
+
+        def _event_handler(_, events):
+            self._event_handler(events)
+
+        self._ioloop.add_handler(
+            self._sock.fileno(), _event_handler, self._ioloop.READ | self._ioloop.ERROR
+        )
+
+    def signal_close(self):
+        """Signal the handler to close.
+
+        The handler will be closed after the existing
+        pending message are sent to the peer.
+        """
+        if not self._pending_write:
+            self.close()
+        else:
+            self._signal_close = True
+
+    def close(self):
+        """Close the socket"""
+        if self._sock is not None:
+            try:
+                self._ioloop.remove_handler(self._sock.fileno())
+                self._sock.close()
+            except OSError:
+                pass
+            self._sock = None
+            self.on_close()
+
+    def write_message(self, message, binary=True):
+        assert binary
+        if self._sock is None:
+            raise OSError("socket is already closed")
+        self._pending_write.append(message)
+        self._update_write()
+
+    def _event_handler(self, events):
+        """centeral event handler"""
+        if (events & self._ioloop.ERROR) or (events & self._ioloop.READ):
+            if self._update_read() and (events & self._ioloop.WRITE):
+                self._update_write()
+        elif events & self._ioloop.WRITE:
+            self._update_write()
+
+    def _update_write(self):
+        """Update the state on write"""
+        while self._pending_write:
+            try:
+                msg = self._pending_write[0]
+                if self._sock is None:
+                    return
+                nsend = self._sock.send(msg)
+                if nsend != len(msg):
+                    self._pending_write[0] = msg[nsend:]
+                else:
+                    self._pending_write.pop(0)
+            except OSError as err:
+                if err.args[0] in (errno.EAGAIN, errno.EWOULDBLOCK):
+                    break
+                self.on_error(err)
+
+        if self._pending_write:
+            self._ioloop.update_handler(
+                self._sock.fileno(), self._ioloop.READ | self._ioloop.ERROR | self._ioloop.WRITE
+            )
+        else:
+            if self._signal_close:
+                self.close()
+            else:
+                self._ioloop.update_handler(
+                    self._sock.fileno(), self._ioloop.READ | self._ioloop.ERROR
+                )
+
+    def _update_read(self):
+        """Update state when there is read event"""
+        try:
+            msg = bytes(self._sock.recv(4096))
+            if msg:
+                self.on_message(msg)
+                return True
+            # normal close, remote is closed
+            self.close()
+        except OSError as err:
+            if err.args[0] in (errno.EAGAIN, errno.EWOULDBLOCK):
+                pass
+            else:
+                self.on_error(err)
+        return False

python/tvm/rpc/tracker.py

@@ -0,0 +1,523 @@
+# 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.
+"""RPC Tracker, tracks and distributes the TVM RPC resources.
+
+This folder implements the tracker server logic.
+
+Note
+----
+Tracker is a TCP based rest api with the following protocol:
+- Initial handshake to the peer
+  - RPC_TRACKER_MAGIC
+- Normal message: [size(int32), json-data]
+- Each message is initiated by the client, and the tracker replies with a json.
+
+List of available APIs:
+
+- PING: check if tracker is alive
+  - input: [TrackerCode.PING]
+  - return: TrackerCode.SUCCESS
+- PUT: report resource to tracker
+  - input: [TrackerCode.PUT, [port, match-key]]
+  - return: TrackerCode.SUCCESS
+  - note: match-key is a randomly generated identify the resource during connection.
+- REQUEST: request a new resource from tracker
+  - input: [TrackerCode.REQUEST, [key, user, priority]]
+  - return: [TrackerCode.SUCCESS, [url, port, match-key]]
+"""
+# pylint: disable=invalid-name
+
+import asyncio
+import errno
+import heapq
+import json
+import logging
+import socket
+import struct
+import sys
+import threading
+
+from tvm.support.popen_pool import PopenWorker
+
+try:
+    from tornado import ioloop
+
+    from . import tornado_util
+except ImportError as error_msg:
+    raise ImportError(
+        f"RPCTracker module requires tornado package {error_msg}. Try 'pip install tornado'."
+    )
+
+from . import base
+from .base import RPC_TRACKER_MAGIC, TrackerCode
+
+logger = logging.getLogger("RPCTracker")
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(
+    logging.Formatter(
+        fmt="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+    )
+)
+logger.addHandler(console_handler)
+logger.setLevel(logging.INFO)
+logger.propagate = False
+
+# Maximum size in bytes for a single tracker message. Tracker frames carry
+# small JSON command tuples; 1 MiB is well above any legitimate payload and
+# bounds memory growth when a peer sends an oversized or malformed size
+# header on the wire.
+MAX_TRACKER_MSG_BYTES = 1 << 20
+
+
+class Scheduler:
+    """Abstract interface of scheduler."""
+
+    def put(self, value):
+        """Push a resource into the scheduler.
+
+        This function can trigger callbacks in the scheduler.
+
+        Parameters
+        ----------
+        value : object
+            The resource to be put in the scheduler.
+        """
+        raise NotImplementedError()
+
+    def request(self, user, priority, callback):
+        """Request a resource.
+
+        Parameters
+        ----------
+        user : str
+            The user who is requesting the resource.
+
+        priority : int
+            The job priority
+
+        callback : function: value->bool
+            Callback function to receive an resource when ready
+            returns True if the resource is consumed.
+        """
+        raise NotImplementedError()
+
+    def remove(self, value):
+        """Remove a resource in the scheduler
+
+        Parameters
+        ----------
+        value: object
+            The resource to remove
+        """
+
+    def summary(self):
+        """Get summary information of the scheduler."""
+        raise NotImplementedError()
+
+
+class PriorityScheduler(Scheduler):
+    """Priority based scheduler, FIFO based on request order"""
+
+    def __init__(self, key):
+        self._key = key
+        self._request_cnt = 0
+        self._lock = threading.Lock()
+        self._values = []
+        self._requests = []
+
+    def _schedule(self):
+        while self._requests and self._values:
+            value = self._values.pop(0)
+            item = heapq.heappop(self._requests)
+            callback = item[-1]
+            if callback(value[1:]):
+                value[0].pending_matchkeys.remove(value[-1])
+            else:
+                self._values.append(value)
+
+    def put(self, value):
+        self._values.append(value)
+        self._schedule()
+
+    def request(self, user, priority, callback):
+        with self._lock:
+            heapq.heappush(self._requests, (-priority, self._request_cnt, callback))
+            self._request_cnt += 1
+        self._schedule()
+
+    def remove(self, value):
+        if value in self._values:
+            self._values.remove(value)
+            self._schedule()
+
+    def summary(self):
+        """Get summary information of the scheduler."""
+        return {"free": len(self._values), "pending": len(self._requests)}
+
+
+class TCPEventHandler(tornado_util.TCPHandler):
+    """Base asynchronize message handler.
+
+    The tracker and client follows a simple message protocol.
+    The message is in form [nbytes(int32)] [json-str].
+    All the information is packed in json-str
+    """
+
+    def __init__(self, tracker, sock, addr):
+        super().__init__(sock)
+        self._data = bytearray()
+        self._tracker = tracker
+        self._msg_size = 0
+        self._addr = addr
+        self._init_req_nbytes = 4
+        self._info = {}
+        # list of pending match keys that has not been used.
+        self.pending_matchkeys = set()
+        self._tracker._connections.add(self)
+        self.put_values = []
+
+    def name(self):
+        """name of connection"""
+        return f"TCPSocket: {self._addr!s}"
+
+    def summary(self):
+        """Summary of this connection"""
+        return self._info
+
+    def _init_conn(self, message):
+        """Initialize the connection"""
+        if len(message) != 4:
+            logger.warning("Invalid connection from %s", self.name())
+            self.close()
+        magic = struct.unpack("<i", message)[0]
+        if magic != RPC_TRACKER_MAGIC:
+            logger.warning("Invalid magic from %s", self.name())
+            self.close()
+        self.write_message(struct.pack("<i", RPC_TRACKER_MAGIC), binary=True)
+        self._init_req_nbytes = 0
+
+    def on_message(self, message):
+        """Callback when a message is received.
+
+        Parameters
+        ----------
+        message : bytearray
+            The bytes received
+        """
+        assert isinstance(message, bytes)
+        if self._init_req_nbytes:
+            self._init_conn(message)
+            return
+
+        self._data += message
+
+        while True:
+            if self._msg_size == 0:
+                if len(self._data) >= 4:
+                    self._msg_size = struct.unpack("<i", self._data[:4])[0]
+                    if self._msg_size <= 0 or self._msg_size > MAX_TRACKER_MSG_BYTES:
+                        logger.warning(
+                            "Invalid msg_size %d from %s; closing connection",
+                            self._msg_size,
+                            self.name(),
+                        )
+                        self.close()
+                        return
+                    del self._data[:4]
+                else:
+                    return
+            if self._msg_size != 0 and len(self._data) >= self._msg_size:
+                msg = (bytes(self._data[: self._msg_size])).decode("utf-8")
+                del self._data[: self._msg_size]
+                self._msg_size = 0
+                try:
+                    self.call_handler(json.loads(msg))
+                except Exception:  # pylint: disable=broad-except
+                    logger.warning("Error handling message from %s", self.name(), exc_info=True)
+                    self.close()
+                    return
+            else:
+                return
+
+    def ret_value(self, data):
+        """return value to the output"""
+        data = json.dumps(data)
+        self.write_message(struct.pack("<i", len(data)), binary=True)
+        self.write_message(data.encode("utf-8"), binary=True)
+
+    def call_handler(self, args):
+        """Event handler when json request arrives."""
+        code = args[0]
+        if code == TrackerCode.PUT:
+            key = args[1]
+            port, matchkey = args[2]
+            self.pending_matchkeys.add(matchkey)
+            # got custom address (from rpc server)
+            if len(args) >= 4 and args[3] is not None:
+                value = (self, args[3], port, matchkey)
+            else:
+                value = (self, self._addr[0], port, matchkey)
+            self._tracker.put(key, value)
+            self.put_values.append(value)
+            self.ret_value(TrackerCode.SUCCESS)
+        elif code == TrackerCode.REQUEST:
+            key = args[1]
+            user = args[2]
+            priority = args[3]
+
+            def _cb(value):
+                # if the connection is already closed
+                if not self._sock:
+                    return False
+                try:
+                    self.ret_value([TrackerCode.SUCCESS, value])
+                except OSError:
+                    return False
+                return True
+
+            self._tracker.request(key, user, priority, _cb)
+        elif code == TrackerCode.PING:
+            self.ret_value(TrackerCode.SUCCESS)
+        elif code == TrackerCode.GET_PENDING_MATCHKEYS:
+            self.ret_value(list(self.pending_matchkeys))
+        elif code == TrackerCode.STOP:
+            # safe stop tracker
+            if self._tracker._stop_key == args[1]:
+                self.ret_value(TrackerCode.SUCCESS)
+                self._tracker.stop()
+            else:
+                self.ret_value(TrackerCode.FAIL)
+        elif code == TrackerCode.UPDATE_INFO:
+            info = args[1]
+            assert isinstance(info, dict)
+            if info["addr"][0] is None:
+                info["addr"][0] = self._addr[0]
+            self._info.update(info)
+            self.ret_value(TrackerCode.SUCCESS)
+        elif code == TrackerCode.SUMMARY:
+            status = self._tracker.summary()
+            self.ret_value([TrackerCode.SUCCESS, status])
+        else:
+            logger.warning("Unknown tracker code %d", code)
+            self.close()
+
+    def on_close(self):
+        self._tracker.close(self)
+
+    def on_error(self, err):
+        logger.warning("%s: Error in RPC Tracker: %s", self.name(), err)
+        self.close()
+
+
+class TrackerServerHandler:
+    """Tracker that tracks the resources."""
+
+    def __init__(self, sock, stop_key):
+        self._scheduler_map = {}
+        self._sock = sock
+        self._sock.setblocking(0)
+        self._ioloop = ioloop.IOLoop.current()
+        self._stop_key = stop_key
+        self._connections = set()
+
+        def _event_handler(_, events):
+            self._on_event(events)
+
+        self._ioloop.add_handler(self._sock.fileno(), _event_handler, self._ioloop.READ)
+
+    def _on_event(self, _):
+        while True:
+            try:
+                conn, addr = self._sock.accept()
+                TCPEventHandler(self, conn, addr)
+            except OSError as err:
+                if err.args[0] in (errno.EAGAIN, errno.EWOULDBLOCK):
+                    break
+
+    def create_scheduler(self, key):
+        """Create a new scheduler."""
+        return PriorityScheduler(key)
+
+    def put(self, key, value):
+        """Report a new resource to the tracker."""
+        if key not in self._scheduler_map:
+            self._scheduler_map[key] = self.create_scheduler(key)
+        self._scheduler_map[key].put(value)
+
+    def request(self, key, user, priority, callback):
+        """Request a new resource."""
+        if key not in self._scheduler_map:
+            self._scheduler_map[key] = self.create_scheduler(key)
+        self._scheduler_map[key].request(user, priority, callback)
+
+    def close(self, conn):
+        self._connections.remove(conn)
+        if "key" in conn._info:
+            for value in conn.put_values:
+                _, _, _, key = value
+                rpc_key, _ = base.split_random_key(key)
+                self._scheduler_map[rpc_key].remove(value)
+
+    def stop(self):
+        """Safely stop tracker."""
+        for conn in list(self._connections):
+            conn.close()
+        self._sock.close()
+        self._ioloop.stop()
+
+    def summary(self):
+        """Return a dict summarizing current status."""
+        qinfo = {}
+        for k, v in self._scheduler_map.items():
+            qinfo[k] = v.summary()
+        cinfo = []
+        # ignore client connections without key
+        for conn in self._connections:
+            res = conn.summary()
+            if res.get("key", "").startswith("server"):
+                cinfo.append(res)
+        return {"queue_info": qinfo, "server_info": cinfo}
+
+    def run(self):
+        """Run the tracker server"""
+        self._ioloop.start()
+
+
+def _tracker_server(listen_sock, stop_key):
+    asyncio.set_event_loop(asyncio.new_event_loop())
+    handler = TrackerServerHandler(listen_sock, stop_key)
+    handler.run()
+
+
+class PopenTrackerServerState:
+    """Internal PopenTrackerServer State"""
+
+    current = None
+
+    def __init__(self, host, port=9190, port_end=9199, silent=False, reuse_addr=True, timeout=None):
+        if silent:
+            logger.setLevel(logging.WARN)
+
+        sock = socket.socket(base.get_addr_family((host, port)), socket.SOCK_STREAM)
+
+        # Never set socket SO_REUSEADDR on Windows. The SO_REUSEADDR flag allow reusing the
+        # inactivate TIME_WATI state sockets on POSIX, but on Windows it will allow two or more
+        # activate sockets to bind on the same address and port if they all set SO_REUSEADDR,
+        # and result in indeterminate behavior.
+        if reuse_addr and sys.platform != "win32":
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        if timeout is not None:
+            sock.settimeout(timeout)
+        self.port = None
+        self.stop_key = base.random_key("tracker")
+        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})")
+        logger.info("bind to %s:%d", host, self.port)
+        sock.listen(1)
+        self.thread = threading.Thread(target=_tracker_server, args=(sock, self.stop_key))
+        self.thread.start()
+        self.host = host
+
+
+def _popen_start_tracker_server(
+    host, port=9190, port_end=9199, silent=False, reuse_addr=True, timeout=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 = PopenTrackerServerState(host, port, port_end, silent, reuse_addr, timeout)
+    PopenTrackerServerState.current = state
+    # returns the port so that the main can get the port number.
+    return (state.port, state.stop_key)
+
+
+class Tracker:
+    """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
+
+    reuse_addr: bool, optional
+        Allows the kernel to reuse a local socket in TIME_WAIT state.
+
+    timeout: float, optional
+         set a timeout for all operations on the socket
+
+    """
+
+    def __init__(
+        self, host="0.0.0.0", port=9190, port_end=9199, silent=False, reuse_addr=True, timeout=None
+    ):
+        if silent:
+            logger.setLevel(logging.WARN)
+        self.proc = PopenWorker()
+        # send the function
+        self.proc.send(
+            _popen_start_tracker_server, [host, port, port_end, silent, reuse_addr, timeout]
+        )
+        # 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