wehub-resource-sync
524 lines
17 KiB
Python
524 lines
17 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.
|
|
"""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
|