text: # 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 client tools"""
  import os
  import stat
  import socket
  import struct
  import time
  
  import tvm._ffi
  from tvm.contrib import util
  from tvm._ffi.base import TVMError
  from tvm.runtime import ndarray as nd
  
  from . import base
  from . import server
  from . import _ffi_api
  
  
  class RPCSession(object):
      """RPC Client session module
  
      Do not directly create the obhect, 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("runtime.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 context(self, dev_type, dev_id=0):
          """Construct a remote context.
  
          Parameters
          ----------
          dev_type: int or str
  
          dev_id: int, optional
  
          Returns
          -------
          ctx: TVMContext
              The corresponding encoded remote context.
          """
          ctx = nd.context(dev_type, dev_id)
          encode = (self._tbl_index + 1) * base.RPC_SESS_MASK
          ctx.device_type += encode
          ctx._rpc_sess = self
          return ctx
  
      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 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 cpu(self, dev_id=0):
          """Construct CPU device."""
          return self.context(1, dev_id)
  
      def gpu(self, dev_id=0):
          """Construct GPU device."""
          return self.context(2, dev_id)
  
      def cl(self, dev_id=0):
          """Construct OpenCL device."""
          return self.context(4, dev_id)
  
      def vulkan(self, dev_id=0):
          """Construct Vulkan device."""
          return self.context(7, dev_id)
  
      def metal(self, dev_id=0):
          """Construct Metal device."""
          return self.context(8, dev_id)
  
      def opengl(self, dev_id=0):
          """Construct OpenGL device."""
          return self.context(11, dev_id)
  
      def ext_dev(self, dev_id=0):
          """Construct extension device."""
          return self.context(12, 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_func("rpc.PopenSession")
  def _popen_session(binary):
      temp = util.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(object):
      """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("%s is not RPC Tracker" % str(self._addr))
  
      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("Invalid return value %s" % str(value))
          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\tkey\n"
          res += "----------------------------\n"
          for item in data["server_info"]:
              addr = item["addr"]
              res += addr[0] + ":" + str(addr[1]) + "\t"
              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 = ("%%-%ds" % max_key_len + "   total  free  pending\n") % 'key'
          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 += ("%%-%ds" % max_key_len + "   %-5d  %-4d  %-7d\n") % \
                         (k, total, free, pending)
          res += separate_line
          return res
  
      def request(self, key, priority=1, session_timeout=0, max_retry=5):
          """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.
          """
          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("Invalid return value %s" % str(value))
                  url, port, matchkey = value[1]
                  return connect(url, port, matchkey, session_timeout)
              except socket.error as err:
                  self.close()
                  last_err = err
              except TVMError as err:
                  last_err = err
          raise RuntimeError(
              "Cannot request %s after %d retry, last_error:%s" % (
                  key, max_retry, str(last_err)))
  
      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 TVMError 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(
                          "Session timeout when running %s" % func.__name__)
                  last_err = err
          raise RuntimeError(
              "Failed to run on %s after %d retry, last_error:%s" % (
                  key, max_retry, str(last_err)))
  
  
  def connect(url, port, key="", session_timeout=0, session_constructor_args=None):
      """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, 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.
  
      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,
              session_constructor_args=[
                  "rpc.Connect", internal_url, internal_port, internal_key])
  
      """
      try:
          if session_timeout:
              key += " -timeout=%s" % str(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, *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))
  