EasyFL_with_PgFed/easyfl/registry/etcd_client.py

import random
import threading
from contextlib import contextmanager

import etcd3

from easyfl.registry import mock_etcd
from easyfl.registry.vclient import VirtualClient


class EtcdClient(object):
    """Etcd client to connect and communicate with etcd service.
    Etcd is the serves as the registry for remote training.
    Clients register themselves in etcd and server queries etcd to get client addresses.

    Args:
        name (str): The name of etcd.
        addrs (str): Etcd addresses, format: "<ip>:<port>,<ip>:<port>".
        base_dir (str): The prefix of all etcd requests, default to "backends".
        use_mock_etcd (bool): Whether use mocked etcd for testing.
    """

    ETCD_CLIENT_POOL_LOCK = threading.Lock()
    ETCD_CLIENT_POOL = {}
    ETCD_CLIENT_POOL_DESTROY = False

    class Event(object):
        def __init__(self, event, base_dir):
            self._event = event
            self._base_dir = base_dir

        def __getattr__(self, attr):
            return getattr(self._event, attr)

        @property
        def key(self):
            return EtcdClient.normalize_output_key(self._event.key, self._base_dir)

    def __init__(self, name, addrs, base_dir, use_mock_etcd=False):
        self._name = name
        self._base_dir = '/' + EtcdClient._normalize_input_key(base_dir)
        self._addrs = self._normalize_addr(addrs)
        if len(self._addrs) == 0:
            raise ValueError('Empty hosts EtcdClient')
        self._cur_addr_idx = random.randint(0, len(self._addrs) - 1)
        self._use_mock_etcd = use_mock_etcd

    def get_data(self, key):
        addr = self._get_next_addr()
        with EtcdClient.closing(self._name, addr, self._use_mock_etcd) as clnt:
            return clnt.get(self._generate_path(key))[0]

    def set_data(self, key, data):
        addr = self._get_next_addr()
        with EtcdClient.closing(self._name, addr, self._use_mock_etcd) as clnt:
            clnt.put(self._generate_path(key), data)

    def delete(self, key):
        addr = self._get_next_addr()
        with EtcdClient.closing(self._name, addr, self._use_mock_etcd) as clnt:
            return clnt.delete(self._generate_path(key))

    def delete_prefix(self, key):
        addr = self._get_next_addr()
        with EtcdClient.closing(self._name, addr, self._use_mock_etcd) as clnt:
            return clnt.delete_prefix(self._generate_path(key))

    def cas(self, key, old_data, new_data):
        addr = self._get_next_addr()
        with EtcdClient.closing(self._name, addr, self._use_mock_etcd) as clnt:
            etcd_path = self._generate_path(key)
            if old_data is None:
                return clnt.put_if_not_exists(etcd_path, new_data)
            return clnt.replace(etcd_path, old_data, new_data)

    def watch_key(self, key):
        addr = self._get_next_addr()
        with EtcdClient.closing(self._name, addr, self._use_mock_etcd) as clnt:
            notifier, cancel = clnt.watch(self._generate_path(key))

            def prefix_extractor(notifier, base_dir):
                while True:
                    try:
                        yield EtcdClient.Event(next(notifier), base_dir)
                    except StopIteration:
                        break

            return prefix_extractor(notifier, self._base_dir), cancel

    def get_prefix_kvs(self, prefix, ignore_prefix=False):
        addr = self._get_next_addr()
        kvs = []
        path = self._generate_path(prefix)
        with EtcdClient.closing(self._name, addr, self._use_mock_etcd) as clnt:
            for (data, key) in clnt.get_prefix(path, sort_order='ascend'):
                if ignore_prefix and key.key == path.encode():
                    continue
                nkey = EtcdClient.normalize_output_key(key.key, self._base_dir)
                kvs.append((nkey, data))
        return kvs

    def get_clients(self, prefix):
        """Retrieve client addresses from etcd using prefix.

        Args:
            prefix (str): the prefix of clients addresses; default is the docker image name "easyfl-client"

        Returns:
            list[:obj:`VirtualClient`]: A list of clients.
        """
        key_value_tuples = self.get_prefix_kvs(prefix)
        clients = []
        index = 0
        for (key_byte, value_byte) in key_value_tuples:
            key, value = key_byte.decode("utf-8"), value_byte.decode("utf-8")
            parts = key.split("/")
            if len(parts) <= 1:
                continue

            addr = parts[1]
            if not self._is_addr(addr):
                continue

            clients.append(VirtualClient(value, addr, index))
            index += 1
        return clients

    def _is_addr(self, address):
        return len(address.split(":")) > 1

    def _generate_path(self, key):
        return '/'.join([self._base_dir, self._normalize_input_key(key)])

    def _get_next_addr(self):
        return self._addrs[random.randint(0, len(self._addrs) - 1)]

    @staticmethod
    def _normalize_addr(addrs):
        naddrs = []
        for raw_addr in addrs.split(','):
            (host, port_str) = raw_addr.split(':')
            try:
                port = int(port_str)
                if port < 0 or port > 65535:
                    raise ValueError('port {} is out of range')
            except ValueError:
                raise ValueError('{} is not a valid port'.format(port_str))
            naddrs.append((host, port))
        return naddrs

    @staticmethod
    def _normalize_input_key(key):
        skip_cnt = 0
        while key[skip_cnt] == '.' or key[skip_cnt] == '/':
            skip_cnt += 1
        if skip_cnt > 0:
            return key[skip_cnt:]
        return key

    @staticmethod
    def normalize_output_key(key, base_dir):
        if isinstance(base_dir, str):
            assert key.startswith(base_dir.encode())
        else:
            assert key.startswith(base_dir)
        return key[len(base_dir) + 1:]

    @classmethod
    @contextmanager
    def closing(cls, name, addr, use_mock_etcd):
        clnt = None
        with cls.ETCD_CLIENT_POOL_LOCK:
            if (name in cls.ETCD_CLIENT_POOL and
                    len(cls.ETCD_CLIENT_POOL[name]) > 0):
                clnt = cls.ETCD_CLIENT_POOL[name][0]
                cls.ETCD_CLIENT_POOL[name] = cls.ETCD_CLIENT_POOL[name][1:]
        if clnt is None:
            try:
                if use_mock_etcd:
                    clnt = mock_etcd.MockEtcdClient(addr[0], addr[1])
                else:
                    clnt = etcd3.client(host=addr[0], port=addr[1])
            except Exception as e:
                clnt.close()
                raise e
        try:
            yield clnt
        except Exception as e:
            clnt.close()
            raise e
        else:
            with cls.ETCD_CLIENT_POOL_LOCK:
                if cls.ETCD_CLIENT_POOL_DESTROY:
                    clnt.close()
                else:
                    if name not in cls.ETCD_CLIENT_POOL:
                        cls.ETCD_CLIENT_POOL[name] = [clnt]
                    else:
                        cls.ETCD_CLIENT_POOL[name].append(clnt)

    @classmethod
    def destory_client_pool(cls):
        with cls.ETCD_CLIENT_POOL_LOCK:
            cls.ETCD_CLIENT_POOL_DESTROY = True
            for _, clnts in cls.ETCD_CLIENT_POOL.items():
                for clnt in clnts:
                    clnt.close()