"""
This type stub file was generated by pyright.
"""

import os

if os.name != 'nt':
    ...
def read_ip_config(filename): # -> dict[Unknown, Unknown]:
    """Read network configuration information of kvstore from file.

    The format of configuration file should be:

        [ip] [base_port] [server_count]

        172.31.40.143 30050 2
        172.31.36.140 30050 2
        172.31.47.147 30050 2
        172.31.30.180 30050 2

    Note that, DGL KVStore supports multiple servers that can share data with each other
    on the same machine via shared-tensor. So the server_count should be >= 1.

    Parameters
    ----------
    filename : str
        name of configuration file.

    Returns
    -------
    dict
        server namebook. e.g.,

        [server_id]:[machine_id, ip, port, group_count]

          {0:[0, '172.31.40.143', 30050, 2],
           1:[0, '172.31.40.143', 30051, 2],
           2:[1, '172.31.36.140', 30050, 2],
           3:[1, '172.31.36.140', 30051, 2],
           4:[2, '172.31.47.147', 30050, 2],
           5:[2, '172.31.47.147', 30051, 2],
           6:[3, '172.31.30.180', 30050, 2],
           7:[3, '172.31.30.180', 30051, 2]}
    """
    ...

def get_type_str(dtype): # -> Literal['float16', 'float32', 'float64', 'uint8', 'int8', 'int16', 'int32', 'int64']:
    """Get data type string
    """
    ...

class KVServer:
    """KVServer is a lightweight key-value store service for DGL distributed training.

    In practice, developers can use KVServer to hold large-scale graph features or
    graph embeddings across machines in a distributed setting. Also, user can re-wriite _push_handler()
    and _pull_handler() API to support flexibale algorithms.

    DGL kvstore supports multiple-servers on single-machine. That means we can lunach many servers on the same machine and all of
    these servers will share the same shared-memory tensor for load-balance.

    Note that, DO NOT use KVServer in multiple threads on Python because this behavior is not defined.

    For now, KVServer can only run in CPU. We will support GPU KVServer in the future.

    Parameters
    ----------
    server_id : int
        KVServer's ID (start from 0).
    server_namebook: dict
        IP address namebook of KVServer, where key is the KVServer's ID
        (start from 0) and value is the server's machine_id, IP address and port, e.g.,

          {0:'[0, 172.31.40.143, 30050],
           1:'[0, 172.31.40.143, 30051],
           2:'[1, 172.31.36.140, 30050],
           3:'[1, 172.31.36.140, 30051],
           4:'[2, 172.31.47.147, 30050],
           5:'[2, 172.31.47.147, 30051],
           6:'[3, 172.31.30.180, 30050],
           7:'[3, 172.31.30.180, 30051]}

    num_client : int
        Total number of client nodes.
    queue_size : int
        Sise (bytes) of kvstore message queue buffer (~20 GB on default).
        Note that the 20 GB is just an upper-bound number and DGL will not allocate 20GB memory.
    net_type : str
        networking type, e.g., 'socket' (default) or 'mpi' (do not support yet).
    """
    def __init__(self, server_id, server_namebook, num_client, queue_size=..., net_type=...) -> None:
        ...
    
    def __del__(self): # -> None:
        """Finalize KVServer
        """
        ...
    
    def set_global2local(self, name, global2local=...): # -> None:
        """Set data mapping of global ID to local ID.

        Parameters
        ----------
        name : str
            data name
        global2local : list or tensor (mx.ndarray or torch.tensor)
            A data mapping of global ID to local ID. KVStore will use global ID by default
            if the global2local is not been set.

            Note that, if the global2local is None KVServer will read shared-tensor.
        """
        ...
    
    def set_partition_book(self, name, partition_book=...): # -> None:
        """Partition book contains the data mapping of global ID to machine ID.

        Parameters
        ----------
        name : str
            data name
        partition_book : list or tensor (mx.ndarray or torch.tensor)
            Mapping global ID to target machine ID.

        Note that, if the partition_book is None KVClient will read shared-tensor by name.
        """
        ...
    
    def init_data(self, name, data_tensor=...): # -> None:
        """Initialize data tensor on KVServe.

        Parameters
        ----------
        name : str
            data name
        data_tensor : tensor (mx.ndarray or torch.tensor)
            data tensor

            Note that, if the data_tensor is None KVServer will read shared-tensor.
        """
        ...
    
    def get_id(self):
        """Get current server id

        Return
        ------
        int
            KVServer ID
        """
        ...
    
    def get_addr(self):
        """Get current server IP address and port

        Return
        ------
        str
            IP address and port
        """
        ...
    
    def get_machine_id(self):
        """Get local machine ID

        Return
        -------
        int
            machine ID
        """
        ...
    
    def get_group_count(self):
        """Get count of server inside a machine

        Return
        ------
        int
            count of server
        """
        ...
    
    def get_message_count(self): # -> int:
        """Get total message count on current KVServer

        Return
        ------
        int
            count of message
        """
        ...
    
    def print(self): # -> None:
        """Print server information (Used by debug)
        """
        ...
    
    def start(self): # -> None:
        """Start service of KVServer.

        The start() api performs the following things:

          1. Get connected with all client nodes.
          2. Recv client address information.
          3. assign client ID to each client node.
          4. send shared-tensor information to each client node.
          5. Service loop for listening requests from client nodes.

        """
        ...
    


class KVClient:
    """KVClient is used to push/pull tensors to/from KVServer. If the server node and client node are on the
    same machine, they can commuincate with each other using local shared-memory tensor, instead of TCP/IP connections.

    Note that, DO NOT use KVClient in multiple threads on Python because this behavior is not defined.

    For now, KVClient can only run in CPU, and we will support GPU KVClient in the future.

    Parameters
    ----------
    server_namebook: dict
        IP address namebook of KVServer, where key is the KVServer's ID
        (start from 0) and value is the server's machine_id, IP address and port, and group_count, e.g.,

          {0:'[0, 172.31.40.143, 30050, 2],
           1:'[0, 172.31.40.143, 30051, 2],
           2:'[1, 172.31.36.140, 30050, 2],
           3:'[1, 172.31.36.140, 30051, 2],
           4:'[2, 172.31.47.147, 30050, 2],
           5:'[2, 172.31.47.147, 30051, 2],
           6:'[3, 172.31.30.180, 30050, 2],
           7:'[3, 172.31.30.180, 30051, 2]}

    queue_size : int
        Sise (bytes) of kvstore message queue buffer (~20 GB on default).
    net_type : str
        networking type, e.g., 'socket' (default) or 'mpi'.
    """
    def __init__(self, server_namebook, queue_size=..., net_type=...) -> None:
        ...
    
    def __del__(self): # -> None:
        """Finalize KVClient
        """
        ...
    
    def connect(self): # -> None:
        """Connect to all the KVServer nodes

        The connect() api performs the following things:

          1. Get connected with all server nodes.
          2. Send client address information to server.
          3. Recv client ID from server.
          4. Recv shared-tensor information from server.

        """
        ...
    
    def init_data(self, name, shape, dtype, target_name): # -> None:
        """Send message to kvserver to initialize new data and
        get corresponded shared-tensor (e.g., partition_book, g2l) on kvclient.

        The new data will be initialized to zeros.

        Note that, this API must be invoked after the conenct() API.

        Parameters
        ----------
        name : str
            data name
        shape : list or tuple of int
            data shape
        dtype : dtype
            data type
        target_name : str
            target name is used to find existing partition_book and g2l mapping.
        """
        ...
    
    def print(self): # -> None:
        """Print client information (Used by debug)
        """
        ...
    
    def get_id(self): # -> int:
        """Get current client id

        Return
        ------
        int
            KVClient ID
        """
        ...
    
    def get_addr(self): # -> Any | str:
        """Get current client IP address

        Return
        ------
        str
            IP address
        """
        ...
    
    def get_machine_id(self): # -> int:
        """Get local machine ID

        Return
        -------
        int
            machine ID
        """
        ...
    
    def get_data_name_list(self): # -> list[Unknown]:
        """Get all the data name

        Return
        ------
        list of str
            name list
        """
        ...
    
    def get_data_meta(self, name): # -> tuple[Unknown, Unknown, Unknown]:
        """Get meta data (data_type, data_shape, partition_book) of the target shared-tensor

        Parameter
        ---------
        name : str
            data name

        Return
        ------
        tuple
            (data_type, data_shape, partition_book)
        """
        ...
    
    def push(self, name, id_tensor, data_tensor): # -> None:
        """Push data to KVServer.

        Note that push() is an async operation that will return immediately after calling.

        Parameters
        ----------
        name : str
            data name
        id_tensor : tensor (mx.ndarray or torch.tensor)
            a vector storing the global data ID
        data_tensor : tensor (mx.ndarray or torch.tensor)
            a tensor with the same row size of data ID
        """
        ...
    
    def pull(self, name, id_tensor):
        """Pull message from KVServer.

        Parameters
        ----------
        name : str
            data name
        id_tensor : tensor (mx.ndarray or torch.tensor)
            a vector storing the ID list

        Returns
        -------
        tensor
            a data tensor with the same row size of id_tensor.
        """
        ...
    
    def barrier(self): # -> None:
        """Barrier for all client nodes

        This API will be blocked untill all the clients call this API.
        """
        ...
    
    def shut_down(self): # -> None:
        """Shut down all KVServer nodes.

        We usually invoke this API by just one client (e.g., client_0).
        """
        ...
    


