Class TrackerSession

python/tvm/rpc/client.py:302–478

Tracker client session. Parameters ---------- addr : tuple The address tuple

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