TCPRandomWalk.py

import importlib
import json
import logging
import lzma
import pickle
import time
import weakref
from collections import deque
from ctypes import c_int, c_long
from multiprocessing.sharedctypes import Value
from queue import Empty
import traceback
from multiprocessing import Lock

import torch
import zmq
import faulthandler
faulthandler.enable()

import torch.multiprocessing as mp

from decentralizepy.communication.Communication import Communication

HELLO = b"HELLO"
BYE = b"BYE"


class TCPRandomWalkBase(Communication):
    """
    TCPRandomWalkBase that copies only the encrypt and decrypt functions from TCP.py
    This dummy is needed as the sharing interfaces can call encrpt and decrypt,
    so we need the functions both in TCPRandomWalk and TCPRandomWalkInternal

    """

    def addr(self, rank, machine_id):
        """
        Returns TCP address of the process.

        Parameters
        ----------
        rank : int
            Local rank of the process
        machine_id : int
            Machine id of the process

        Returns
        -------
        str
            Full address of the process using TCP

        """
        machine_addr = self.ip_addrs[str(machine_id)]
        port = rank + self.offset
        return "tcp://{}:{}".format(machine_addr, port)
    
    def __init__(
        self,
        rank,
        machine_id,
        mapping,
        total_procs,
        addresses_filepath,
        compress=False,
        offset=2000,
        compression_package=None,
        compression_class=None,
    ):
        """
        Constructor

        Parameters
        ----------
        rank : int
            Local rank of the process
        machine_id : int
            Machine id of the process
        mapping : decentralizepy.mappings.Mapping
            uid, rank, machine_id invertible mapping
        total_procs : int
            Total number of processes
        addresses_filepath : str
            JSON file with machine_id -> ip mapping
        compression_package : str
            Import path of a module that implements the compression.Compression.Compression class
        compression_class : str
            Name of the compression class inside the compression package

        """
        super().__init__(rank, machine_id, mapping, total_procs)
        self.addresses_filepath = addresses_filepath
        self.compress = compress
        self.offset = 10000 + offset
        self.compression_package = compression_package
        self.compression_class = compression_class


        with open(addresses_filepath) as addrs:
            self.ip_addrs = json.load(addrs)

        self.identity = str(self.uid).encode()

        if compression_package and compression_class:
            compressor_module = importlib.import_module(compression_package)
            compressor_class = getattr(compressor_module, compression_class)
            self.compressor = compressor_class()
            logging.info(f"Using the {compressor_class} to compress the data")
        else:
            assert not self.compress

        self.total_data = 0 # dummy values
        self.total_meta = 0

    def encrypt(self, data):
        """
        Encode data as python pickle.

        Parameters
        ----------
        data : dict
            Data dict to send

        Returns
        -------
        byte
            Encoded data

        """
        logging.debug("in encrypt")
        if self.compress:
            logging.debug("in encrypt: compress")
            if "indices" in data:
                data["indices"] = self.compressor.compress(data["indices"])

            assert "params" in data
            data["params"] = self.compressor.compress_float(data["params"])
            data_len = len(pickle.dumps(data["params"]))
            output = pickle.dumps(data)

            # the compressed meta data gets only a few bytes smaller after pickling
            self.total_meta.value += (len(output) - data_len)
            self.total_data.value += data_len
        else:
            logging.debug("in encrypt: else")
            output = pickle.dumps(data)
            # centralized testing uses its own instance
            if type(data) == dict:
                assert "params" in data
                data_len = len(pickle.dumps(data["params"]))
                self.total_meta.value += (len(output) - data_len)
                self.total_data.value += data_len
        return output

    def decrypt(self, sender, data):
        """
        Decode received pickle data.

        Parameters
        ----------
        sender : byte
            sender of the data
        data : byte
            Data received

        Returns
        -------
        tuple
            (sender: int, data: dict)

        """
        logging.debug("in decrypt")
        sender = int(sender.decode())
        if self.compress:
            logging.debug("in decrypt:comp")
            data = pickle.loads(data)
            if "indices" in data:
                data["indices"] = self.compressor.decompress(data["indices"])
            if "params" in data:
                data["params"] = self.compressor.decompress_float(data["params"])
        else:
            logging.debug("in decrypt:else")
            data = pickle.loads(data)
        return sender, data


class TCPRandomWalk(TCPRandomWalkBase):
    """
    Wrapper for TCPRandomWalkInternal, mostly a copy of Communication
    Connect

    """

    def __init__(
        self,
        rank,
        machine_id,
        mapping,
        total_procs,
        addresses_filepath,
        compress=False,
        offset=2000,
        compression_package=None,
        compression_class=None,
        sampler="equi",
    ):
        """
        Constructor

        Parameters
        ----------
        rank : int
            Local rank of the process
        machine_id : int
            Machine id of the process
        mapping : decentralizepy.mappings.Mapping
            uid, rank, machine_id invertible mapping
        total_procs : int
            Total number of processes

        """
        super().__init__(rank, machine_id, mapping, total_procs, addresses_filepath, compress, offset, compression_package, compression_class)
        self.sampler = sampler
        self.send_queue = mp.Queue(1000)
        self.recv_queue = mp.Queue(1000)
        # Since we are only adding these do not need to share the same lock, intermediary results may be inconsistent
        self.lock = mp.Lock()
        self.total_data = mp.Value(c_long, 0, lock = self.lock)
        self.total_meta = mp.Value(c_long, 0, lock = self.lock)
        self.total_bytes = mp.Value(c_long, 0, lock = self.lock)
        self.flag_running = mp.Value(c_int, 0, lock=False)

    def connect_neighbors(self, neighbors):
        """
        Spawns TCPRandomWalkInternal. It will connect to the neighbours.
        This function should only be called once.

        Parameters
        ----------
        neighbors : list(int)
            List of neighbors

        """
        self.flag_running.value = 1
        self.ctx = mp.start_processes(
            lambda *args: TCPRandomWalkInternal(*(args[1:])),
            args=(
                self.rank,
                self.machine_id,
                self.mapping,
                self.total_procs,
                self.addresses_filepath,
                self.send_queue,
                self.recv_queue,
                self.flag_running,
                neighbors,
                self.total_data,
                self.total_meta,
                self.total_bytes,
                self.compress,
                self.offset,
                self.compression_package,
                self.compression_class,
                self.sampler,
            ),
            start_method="fork",
            join=False,
        )

    def receive(self, block=True, timeout=None):
        """
        Returns a received message. It blocks if no message has been received.

        Returns
        ----------
        dict
            Received and decrypted data

        """
        logging.debug("Receive in TCPRandomWalk")
        try:
            return self.recv_queue.get(block=block, timeout=None)  # already decrypted
            logging.debug("Receive in TCPRandomWalk; post get")
        except Empty:
            return None, None

    def send(self, uid, data, encrypt=True):
        """
        Send a message to a process.

        Parameters
        ----------
        uid : int
            Neighbor's unique ID
        data : dict
            Message as a Python dictionary

        """
        if uid is not None:
            logging.debug("Send to %i in TCPRandomWalk", uid)
        self.send_queue.put((uid, data, encrypt))

    def disconnect_neighbors(self):
        """
        Disconnects all neighbors.

        """
        print("disconnect_neighbors")
        self.flag_running.value = 0
        time.sleep(4)
        self.send_queue.close()  # this crashes
        self.recv_queue.close()

        #del self.lock
        self.send_queue.join_thread()
        self.recv_queue.join_thread()
        self.ctx.join()
        print(f"disconnect_neighbors: joined {self.uid}")


class TCPRandomWalkInternal(TCPRandomWalkBase):
    def __init__(
        self,
        rank,
        machine_id,
        mapping,
        total_procs,
        addresses_filepath,
        send_queue: mp.Queue,
        recv_queue: mp.Queue,
        flag_running,
        neighbors,
        total_data,
        total_meta,
        total_bytes,
        compress=False,
        offset=2000,
        compression_package=None,
        compression_class=None,
        sampler="equi",
    ):
        """
        Constructor

        Parameters
        ----------
        rank : int
            Local rank of the process
        machine_id : int
            Machine id of the process
        mapping : decentralizepy.mappings.Mapping
            uid, rank, machine_id invertible mapping
        total_procs : int
            Total number of processes
        addresses_filepath : str
            JSON file with machine_id -> ip mapping

        """
        super().__init__(rank, machine_id, mapping, total_procs, addresses_filepath, compress, offset,
                         compression_package, compression_class)
        print("post super")
        try:
            self.sampler = sampler
            self.context = zmq.Context()
            self.router = self.context.socket(zmq.ROUTER)
            self.router.setsockopt(zmq.IDENTITY, self.identity)
            self.router.bind(self.addr(rank, machine_id))
            self.sent_disconnections = False
            self.compress = compress
            self.total_data = total_data
            self.total_meta = total_meta
            self.total_bytes = total_bytes

            self.send_queue = send_queue
            self.recv_queue = recv_queue
            self.flag_running = flag_running
            self.neighbors = neighbors

            self.random_generator = torch.Generator()
            self.rw_seed = (
                self.random_generator.seed()
            )  # new random seed from the random device
            logging.info("Machine %i has random seed %i for RW", self.uid, self.rw_seed)
            self.rw_messages_stat = []
            self.rw_double_count_stat = []
            if self.sampler == "equi":
                logging.info("rw_samper is rw_sampler_equi")
                self.rw_sampler = weakref.WeakMethod(
                    self.rw_sampler_equi
                )  # self.rw_sampler_equi_check_history
            elif self.sampler == "equi_check_history":
                logging.info("rw_samper is rw_sampler_equi_check_history")
                self.rw_sampler = weakref.WeakMethod(self.rw_sampler_equi_check_history)
            else:
                logging.info("rw_samper is rw_sampler_equi (default)")
                self.rw_sampler = weakref.WeakMethod(
                    self.rw_sampler_equi
                )  # self.rw_sampler_equi_check_history
            self.peer_sockets = dict()
            self.barrier = set()
            self.connect_neighbors(self.neighbors)


            self.comm_loop()
        except BaseException as e:
            error_message = traceback.format_exc()
            print(error_message)
            print("GOT EXCEPTION")
            logging.debug("GOT EXCEPTION")
            logging.debug(error_message)

        while not self.recv_queue.empty():
            print(f"{self.uid}: clear rcv")
            _ = self.recv_queue.get_nowait()
        self.recv_queue.close()
        self.recv_queue.join_thread()
        print(f"{self.uid}: joined recv")
        while not self.send_queue.empty():
            print(f"{self.uid}: clear snd")
            _ = self.send_queue.get_nowait()
        print(f"{self.uid}: joined send")
        self.send_queue.close()
        self.send_queue.join_thread()
        print("end")

    def __del__(self):
        """
        Destroys zmq context

        """
        self.context.destroy(linger=0)
        logging.info(f"TCPRandomWalkInternal for {self.uid} is destroyed")
        # exit(1)
        # uncomment during debugging
        #self.recv_queue.close()
        #self.send_queue.close()

    def connect_neighbors(self, neighbors):
        """
        Connects all neighbors. Sends HELLO. Waits for HELLO.
        Caches any data received while waiting for HELLOs.

        Parameters
        ----------
        neighbors : list(int)
            List of neighbors

        Raises
        ------
        RuntimeError
            If received BYE while waiting for HELLO

        """
        logging.info("Sending connection request to neighbors")
        for uid in neighbors:
            logging.debug("Connecting to my neighbour: {}".format(uid))
            id = str(uid).encode()
            req = self.context.socket(zmq.DEALER)
            req.setsockopt(zmq.IDENTITY, self.identity)
            req.connect(self.addr(*self.mapping.get_machine_and_rank(uid)))
            self.peer_sockets[id] = req
            req.send(HELLO)

        num_neighbors = len(neighbors)
        while len(self.barrier) < num_neighbors:
            sender, recv = self.router.recv_multipart()

            if recv == HELLO:
                logging.debug("Received {} from {}".format(HELLO, sender))
                self.barrier.add(sender)
            elif recv == BYE:
                logging.debug("Received {} from {}".format(BYE, sender))
                raise RuntimeError(
                    "A neighbour wants to disconnect before training started!"
                )
            else:
                logging.debug(
                    "Received message from {} @ connect_neighbors".format(sender)
                )

                self.recv_queue.put(self.decrypt(sender, recv))

        logging.info("Connected to all neighbors")
        self.initialized = True

    def comm_loop(self):
        # TODO: May want separate loops for send and receive
        sleep_time = 0.001
        while True:
            successes = 1
            flushed = True
            try:
                sender, recv = self.router.recv_multipart(flags=zmq.NOBLOCK)
                logging.debug("comm_loop received from %i", int(sender.decode()))
                self.receive(sender, recv)
                successes += 5
                flushed = False
            except zmq.ZMQError:
                # logging.debug("zmq error due to empty recv_multipart")
                sleep_time *= 1.1

            try:
                uid, data, compress = self.send_queue.get_nowait()
                # send may block if send queue is full
                if uid is not None:
                    logging.debug("comm_loop will send to %i", uid)
                else:
                    logging.debug("comm_loop will send rw")
                self.send(uid, data, compress)
                successes += 5
                flushed = False
            except Empty:
                sleep_time *= 1.1
                # logging.debug("empty send queue %i", self.send_queue.empty())

            sleep_time /= successes
            sleep_time = max(0.00001, sleep_time)
            sleep_time = min(0.005, sleep_time)
            # logging.debug("comm loop sleeping for %f", sleep_time)
            time.sleep(sleep_time)

            if self.flag_running.value == 0 and flushed:
                break
        self.disconnect_neighbors()
        print(f"DISCONNECTED {self.uid}")

    def receive(self, sender, recv):
        """
        Returns ONE message received.

        Returns
        ----------
        dict
            Received and decrypted data

        Raises
        ------
        RuntimeError
            If received HELLO

        """
        if recv == HELLO:
            logging.debug("Received {} from {}".format(HELLO, sender))
            # TODO: how to properly shut down here
            self.flag_running.value = False
            self.recv_queue.close()
            self.send_queue.close()
            raise RuntimeError(
                "A neighbour wants to connect when everyone is connected!"
            )
        elif recv == BYE:
            logging.debug("Received {} from {}".format(BYE, sender))
            self.barrier.remove(sender)
        else:
            # TODO: here process receive of rw messages
            logging.debug("Received message from {}".format(sender))
            src, data = self.decrypt(sender, recv)
            # ned new data object such that forwarding does not change the internal fields
            # We add our selves to visited so, it could trigger "RW message was already once received"
            new_data = data.copy()
            if data.get("rw", False):
                new_data["visited"] = new_data["visited"].copy()
                src = new_data["visited"][
                    0
                ]  # the original src is the neighbour that sent it to us.
                logging.debug(
                    "Received message from {} was rw {}".format(sender, data["visited"])
                )
                if data["fuel"] > 0:
                    new_neighbor = self.rw_sampler()(data)
                    if new_neighbor == None:
                        logging.info(
                            "dropped rw message due to no new neigbor being available: %s",
                            str(data["visited"]),
                        )
                    else:
                        data["fuel"] -= 1
                        assert data["fuel"] + 1 == new_data["fuel"]
                        visited = data["visited"]
                        visited.append(self.uid)
                        data["visited"] = visited
                        logging.debug(
                            "Forward rw {} to {}".format(data["visited"], new_neighbor)
                        )
                        self.send_queue.put((new_neighbor, data, True))
                else:
                    # the message has no more fuel so it is dropped
                    logging.info("dropped rw message with fuel %i ", data["fuel"])
            self.recv_queue.put((src, new_data))
            logging.debug(
                "Received message from {} after putting into the queue".format(sender)
            )

    def send(self, uid, data, encrypt = True):
        """
        Send a message to a process.

        Parameters
        ----------
        uid : int
            Neighbor's unique ID. If it is none then it means we sample the neighbour!
        data : dict
            Message as a Python dictionary

        """
        assert self.initialized == True
        if encrypt:
            rw = data.get("rw", False)
            logging.debug("send: rw? {}".format(rw))
            to_send = self.encrypt(data)
        else:
            rw = False
            logging.debug("send: rw? {}".format(rw))
            to_send = data
        data_size = len(to_send)
        self.total_bytes.value += data_size
        if uid is None and rw:  # a rw message
            if self.flag_running.value == 1:  # Do not send rw if we are shutting down
                uid = self.rw_sampler()(data)
                logging.debug("send: rw to {}".format(uid))
            else:
                logging.debug("send: rw dropped due to being in shutdown")
                return
        id = str(uid).encode()
        self.peer_sockets[id].send(to_send)
        logging.debug("{} sent the message to {}.".format(self.uid, uid))
        logging.info("Sent this round: {}".format(data_size))

    def disconnect_neighbors(self):
        """
        Disconnects all neighbors.

        """
        assert self.initialized == True
        if not self.sent_disconnections:
            logging.info("Disconnecting neighbors")
            for sock in self.peer_sockets.values():
                sock.send(BYE)
            self.sent_disconnections = True
            while len(self.barrier):
                sender, recv = self.router.recv_multipart()
                if recv == BYE:
                    logging.debug("Received {} from {}".format(BYE, sender))
                    self.barrier.remove(sender)
                else:
                    # this can happen now due to async
                    logging.info("Received unexpected message from {}".format(sender))
                    sender, data = self.decrypt(sender, recv)
                    if data.get("rw", False):
                        logging.info("Message was rw {}".format(data["visited"]))
                    else:
                        logging.info("Message was normal")
                    # raise RuntimeError(
                    #    "Received a message when expecting BYE from {}".format(sender)
                    # )
        for sock in self.peer_sockets.values():
            sock.close()
        self.router.close()

    def rw_sampler_equi(self, message):

        index = torch.randint(
            0, len(self.neighbors), size=(1,), generator=self.random_generator
        ).item()
        logging.debug(
            "rw_sampler_equi selected index {} of {} {}".format(
                index, self.neighbors, type(self.neighbors)
            )
        )
        return list(self.neighbors)[index]

    def rw_sampler_equi_check_history(self, message):
        if message is None:  # RW starts from here
            index = torch.randint(
                0, len(self.neighbors), size=(1,), generator=self.random_generator
            ).item()
            logging.debug(
                "rw_sampler_equi_check_history selected index {} of {} {}".format(
                    index, self.neighbors, type(self.neighbors)
                )
            )
            return list(self.neighbors)[index]
        else:
            visited = set(message["visited"])
            neighbors = self.neighbors  # is already a set
            possible_neigbors = neighbors.difference(visited)
            if len(possible_neigbors) == 0:
                return None
            else:
                index = torch.randint(
                    0,
                    len(possible_neigbors),
                    size=(1,),
                    generator=self.random_generator,
                ).item()
                return list(possible_neigbors)[index]

    def rw_sampler_mh(self, message):
        # Metro hastings version of the sampler
        # Samples the neighbour based on the MH weights, allows self loops (will just decrease fuel and reroll!)
        pass