wehub-resource-sync
560 lines
19 KiB
Python
560 lines
19 KiB
Python
# 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
|