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

from collections.abc import MutableMapping
from .._deprecate.graph import DGLGraph

dtype_dict = ...
dtype_dict = ...
class NodeDataView(MutableMapping):
    """The data view class when G.nodes[...].data is called.

    See Also
    --------
    dgl.DGLGraph.nodes
    """
    __slots__ = ...
    def __init__(self, graph, nodes, graph_name) -> None:
        ...
    
    def __getitem__(self, key):
        ...
    
    def __setitem__(self, key, val): # -> None:
        ...
    
    def __delitem__(self, key): # -> None:
        ...
    
    def __len__(self): # -> int:
        ...
    
    def __iter__(self): # -> Iterator[Unknown]:
        ...
    
    def __repr__(self): # -> str:
        ...
    


class EdgeDataView(MutableMapping):
    """The data view class when G.edges[...].data is called.

    See Also
    --------
    dgl.DGLGraph.edges
    """
    __slots__ = ...
    def __init__(self, graph, edges, graph_name) -> None:
        ...
    
    def __getitem__(self, key):
        ...
    
    def __setitem__(self, key, val): # -> None:
        ...
    
    def __delitem__(self, key): # -> None:
        ...
    
    def __len__(self): # -> int:
        ...
    
    def __iter__(self): # -> Iterator[Unknown]:
        ...
    
    def __repr__(self): # -> str:
        ...
    


class Barrier:
    """ A barrier in the KVStore server used for one synchronization.

    All workers have to enter the barrier before any of them can proceed
    with any further computation.

    Parameters
    ----------
    num_workers: int
        The number of workers will enter the barrier.
    """
    def __init__(self, num_workers) -> None:
        ...
    
    def enter(self): # -> None:
        """ A worker enters the barrier.
        """
        ...
    
    def leave(self): # -> None:
        """ A worker notifies the server that it's going to leave the barrier.
        """
        ...
    
    def all_enter(self):
        """ Indicate that all workers have entered the barrier.
        """
        ...
    
    def all_leave(self):
        """ Indicate that all workers have left the barrier.
        """
        ...
    


class BarrierManager:
    """ The manager of barriers

    When a worker wants to enter a barrier, it creates the barrier if it doesn't
    exist. Otherwise, the worker will enter an existing barrier.

    The manager needs to know the number of workers in advance so that it can
    keep track of barriers and workers.

    Parameters
    ----------
    num_workers: int
        The number of workers that need to synchronize with barriers.
    """
    def __init__(self, num_workers) -> None:
        ...
    
    def enter(self, worker_id):
        """ A worker enters a barrier.

        Parameters
        ----------
        worker_id : int
            The worker that wants to enter a barrier.
        """
        ...
    
    def all_enter(self, worker_id, barrier_id):
        """ Indicate whether all workers have entered a specified barrier.
        """
        ...
    
    def leave(self, worker_id, barrier_id): # -> None:
        """ A worker leaves a barrier.

        This is useful for garbage collection of used barriers.
        """
        ...
    


def shared_mem_zero_initializer(shape, dtype, name):
    """Zero feature initializer in shared memory
    """
    ...

class InitializerManager:
    """Manage initializer.

    We need to convert built-in frame initializer to strings
    and send them to the graph store server through RPC.
    Through the conversion, we need to convert local built-in initializer
    to shared-memory initializer.
    """
    _fun2str = ...
    _str2fun = ...
    def serialize(self, init): # -> str:
        """Convert the initializer function to string.

        Parameters
        ----------
        init : callable
            the initializer function.

        Returns
        ------
        string
            The name of the built-in initializer function.
        """
        ...
    
    def deserialize(self, init): # -> (shape: Unknown, dtype: Unknown, name: Unknown) -> Unknown:
        """Convert the string to the initializer function.

        Parameters
        ----------
        init : string
            the name of the initializer function

        Returns
        -------
        callable
            The shared-memory initializer function.
        """
        ...
    


class SharedMemoryStoreServer:
    """The graph store server.

    The server loads graph structure and node embeddings and edge embeddings
    and store them in shared memory. The loaded graph can be identified by
    the graph name in the input argument.

    DGL graph accepts graph data of multiple formats:

    * NetworkX graph,
    * scipy matrix,
    * DGLGraph.

    If the input graph data is DGLGraph, the constructed DGLGraph only contains
    its graph index.

    Parameters
    ----------
    graph_data : graph data
        Data to initialize graph.
    graph_name : string
        Define the name of the graph, so the client can use the name to access the graph.
    multigraph : bool, optional
        Deprecated (Will be deleted in the future).
        Whether the graph would be a multigraph (default: True)
    num_workers : int
        The number of workers that will connect to the server.
    port : int
        The port that the server listens to.
    """
    def __init__(self, graph_data, graph_name, multigraph, num_workers, port) -> None:
        ...
    
    def __del__(self): # -> None:
        ...
    
    @property
    def ndata(self): # -> NodeDataView:
        """Return the data view of all the nodes.

        DGLGraph.ndata is an abbreviation of DGLGraph.nodes[:].data

        See Also
        --------
        dgl.DGLGraph.nodes
        """
        ...
    
    @property
    def edata(self): # -> EdgeDataView:
        """Return the data view of all the edges.

        DGLGraph.data is an abbreviation of DGLGraph.edges[:].data

        See Also
        --------
        dgl.DGLGraph.edges
        """
        ...
    
    def run(self): # -> None:
        """Run the graph store server.

        The server runs to process RPC requests from clients.
        """
        ...
    


class BaseGraphStore(DGLGraph):
    """The base class of the graph store.

    Shared-memory graph store and distributed graph store will be inherited from
    this base class. The graph stores only support large read-only graphs. Thus, many of
    DGLGraph APIs aren't supported.

    Specially, the graph store doesn't support the following methods:
        - ndata
        - edata
        - incidence_matrix
        - line_graph
        - reverse
    """
    def __init__(self, graph_data=..., multigraph=...) -> None:
        ...
    
    @property
    def ndata(self): # -> NoReturn:
        """Return the data view of all the nodes.

        DGLGraph.ndata is an abbreviation of DGLGraph.nodes[:].data
        """
        ...
    
    @property
    def edata(self): # -> NoReturn:
        """Return the data view of all the edges.

        DGLGraph.data is an abbreviation of DGLGraph.edges[:].data

        See Also
        --------
        dgl.DGLGraph.edges
        """
        ...
    
    def incidence_matrix(self, typestr, ctx=...): # -> NoReturn:
        """Return the incidence matrix representation of this graph.

        Parameters
        ----------
        typestr : str
            Can be either ``in``, ``out`` or ``both``
        ctx : context, optional (default=cpu)
            The context of returned incidence matrix.

        Returns
        -------
        SparseTensor
            The incidence matrix.
        """
        ...
    
    def line_graph(self, backtracking=..., shared=...): # -> NoReturn:
        """Return the line graph of this graph.

        See :func:`~dgl.transform.line_graph`.
        """
        ...
    
    def reverse(self, share_ndata=..., share_edata=...): # -> NoReturn:
        """Return the reverse of this graph.

        See :func:`~dgl.transform.reverse`.
        """
        ...
    


class SharedMemoryDGLGraph(BaseGraphStore):
    """Shared-memory DGLGraph.

    This is a client to access data in the shared-memory graph store that has loads
    the graph structure and node embeddings and edge embeddings to shared memory.
    It provides the DGLGraph interface.

    Parameters
    ----------
    graph_name : string
        Define the name of the graph.
    port : int
        The port that the server listens to.
    """
    def __init__(self, graph_name, port) -> None:
        ...
    
    def __del__(self): # -> None:
        ...
    
    @property
    def num_workers(self): # -> str | int | Any:
        """ The number of workers using the graph store.
        """
        ...
    
    @property
    def worker_id(self): # -> str | int | Any:
        """ The id of the current worker using the graph store.

        When a worker connects to a graph store, it is assigned with a worker id.
        This is useful for the graph store server to identify who is sending
        requests.

        The worker id is a unique number between 0 and num_workers.
        This is also useful for user's code. For example, user's code can
        use this number to decide how to assign GPUs to workers in multi-processing
        training.
        """
        ...
    
    def init_ndata(self, ndata_name, shape, dtype, ctx=...): # -> None:
        """Create node embedding.

        It first creates the node embedding in the server and maps it to the current process
        with shared memory.

        Parameters
        ----------
        ndata_name : string
            The name of node embedding
        shape : tuple
            The shape of the node embedding
        dtype : string
            The data type of the node embedding. The currently supported data types
            are "float32" and "int32".
        ctx : DGLContext
            The column context.
        """
        ...
    
    def init_edata(self, edata_name, shape, dtype, ctx=...): # -> None:
        """Create edge embedding.

        It first creates the edge embedding in the server and maps it to the current process
        with shared memory.

        Parameters
        ----------
        edata_name : string
            The name of edge embedding
        shape : tuple
            The shape of the edge embedding
        dtype : string
            The data type of the edge embedding. The currently supported data types
            are "float32" and "int32".
        ctx : DGLContext
            The column context.
        """
        ...
    
    def get_n_repr(self, u=...): # -> dict[Unknown, Unknown] | dict[Unknown, Unknown | FrameRef | LazyDict] | LazyDict:
        """Get node(s) representation.

        The returned feature tensor batches multiple node features on the first dimension.

        Parameters
        ----------
        u : node, container or tensor
            The node(s).

        Returns
        -------
        dict
            Representation dict from feature name to feature tensor.
        """
        ...
    
    def get_e_repr(self, edges=...): # -> dict[Unknown, Unknown] | dict[Unknown, Unknown | FrameRef | LazyDict] | LazyDict:
        """Get edge(s) representation.

        Parameters
        ----------
        edges : edges
            Edges can be a pair of endpoint nodes (u, v), or a
            tensor of edge ids. The default value is all the edges.

        Returns
        -------
        dict
            Representation dict
        """
        ...
    
    def set_n_repr(self, data, u=..., inplace=...): # -> None:
        """Set node(s) representation.

        `data` is a dictionary from the feature name to feature tensor. Each tensor
        is of shape (B, D1, D2, ...), where B is the number of nodes to be updated,
        and (D1, D2, ...) be the shape of the node representation tensor. The
        length of the given node ids must match B (i.e, len(u) == B).

        In the graph store, all updates are written inplace.

        Parameters
        ----------
        data : dict of tensor
            Node representation.
        u : node, container or tensor
            The node(s).
        inplace : bool
            The value is always True.
        """
        ...
    
    def set_e_repr(self, data, edges=..., inplace=...): # -> None:
        """Set edge(s) representation.

        `data` is a dictionary from the feature name to feature tensor. Each tensor
        is of shape (B, D1, D2, ...), where B is the number of edges to be updated,
        and (D1, D2, ...) be the shape of the edge representation tensor.

        In the graph store, all updates are written inplace.

        Parameters
        ----------
        data : tensor or dict of tensor
            Edge representation.
        edges : edges
            Edges can be a pair of endpoint nodes (u, v), or a
            tensor of edge ids. The default value is all the edges.
        inplace : bool
            The value is always True.
        """
        ...
    
    def apply_nodes(self, func=..., v=..., inplace=...): # -> None:
        """Apply the function on the nodes to update their features.

        If None is provided for ``func``, nothing will happen.

        In the graph store, all updates are written inplace.

        Parameters
        ----------
        func : callable or None, optional
            Apply function on the nodes. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        v : int, iterable of int, tensor, optional
            The node (ids) on which to apply ``func``. The default
            value is all the nodes.
        inplace : bool, optional
            The value is always True.
        """
        ...
    
    def apply_edges(self, func=..., edges=..., inplace=...): # -> None:
        """Apply the function on the edges to update their features.

        If None is provided for ``func``, nothing will happen.

        In the graph store, all updates are written inplace.

        Parameters
        ----------
        func : callable, optional
            Apply function on the edge. The function should be
            an :mod:`Edge UDF <dgl.udf>`.
        edges : valid edges type, optional
            Edges on which to apply ``func``. See :func:`send` for valid
            edges type. Default is all the edges.
        inplace: bool, optional
            The value is always True.
        """
        ...
    
    def group_apply_edges(self, group_by, func, edges=..., inplace=...): # -> None:
        """Group the edges by nodes and apply the function on the grouped edges to
         update their features.

        In the graph store, all updates are written inplace.

        Parameters
        ----------
        group_by : str
            Specify how to group edges. Expected to be either 'src' or 'dst'
        func : callable
            Apply function on the edge. The function should be
            an :mod:`Edge UDF <dgl.udf>`. The input of `Edge UDF` should
            be (bucket_size, degrees, *feature_shape), and
            return the dict with values of the same shapes.
        edges : valid edges type, optional
            Edges on which to group and apply ``func``. See :func:`send` for valid
            edges type. Default is all the edges.
        inplace: bool, optional
            The value is always True.
        """
        ...
    
    def recv(self, v=..., reduce_func=..., apply_node_func=..., inplace=...): # -> None:
        """Receive and reduce incoming messages and update the features of node(s) :math:`v`.

        Optionally, apply a function to update the node features after receive.

        In the graph store, all updates are written inplace.

        * `reduce_func` will be skipped for nodes with no incoming message.
        * If all ``v`` have no incoming message, this will downgrade to an :func:`apply_nodes`.
        * If some ``v`` have no incoming message, their new feature value will be calculated
          by the column initializer (see :func:`set_n_initializer`). The feature shapes and
          dtypes will be inferred.

        The node features will be updated by the result of the ``reduce_func``.

        Messages are consumed once received.

        The provided UDF maybe called multiple times so it is recommended to provide
        function with no side effect.

        Parameters
        ----------
        v : node, container or tensor, optional
            The node to be updated. Default is receiving all the nodes.
        reduce_func : callable, optional
            Reduce function on the node. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        apply_node_func : callable
            Apply function on the nodes. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        inplace: bool, optional
            The value is always True.
        """
        ...
    
    def send_and_recv(self, edges, message_func=..., reduce_func=..., apply_node_func=..., inplace=...): # -> None:
        """Send messages along edges and let destinations receive them.

        Optionally, apply a function to update the node features after receive.

        In the graph store, all updates are written inplace.

        This is a convenient combination for performing
        ``send(self, self.edges, message_func)`` and
        ``recv(self, dst, reduce_func, apply_node_func)``, where ``dst``
        are the destinations of the ``edges``.

        Parameters
        ----------
        edges : valid edges type
            Edges on which to apply ``func``. See :func:`send` for valid
            edges type.
        message_func : callable, optional
            Message function on the edges. The function should be
            an :mod:`Edge UDF <dgl.udf>`.
        reduce_func : callable, optional
            Reduce function on the node. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        apply_node_func : callable, optional
            Apply function on the nodes. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        inplace: bool, optional
            The value is always True.
        """
        ...
    
    def pull(self, v, message_func=..., reduce_func=..., apply_node_func=..., inplace=...): # -> None:
        """Pull messages from the node(s)' predecessors and then update their features.

        Optionally, apply a function to update the node features after receive.

        In the graph store, all updates are written inplace.

        * `reduce_func` will be skipped for nodes with no incoming message.
        * If all ``v`` have no incoming message, this will downgrade to an :func:`apply_nodes`.
        * If some ``v`` have no incoming message, their new feature value will be calculated
          by the column initializer (see :func:`set_n_initializer`). The feature shapes and
          dtypes will be inferred.

        Parameters
        ----------
        v : int, iterable of int, or tensor
            The node(s) to be updated.
        message_func : callable, optional
            Message function on the edges. The function should be
            an :mod:`Edge UDF <dgl.udf>`.
        reduce_func : callable, optional
            Reduce function on the node. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        apply_node_func : callable, optional
            Apply function on the nodes. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        inplace: bool, optional
            The value is always True.
        """
        ...
    
    def push(self, u, message_func=..., reduce_func=..., apply_node_func=..., inplace=...): # -> None:
        """Send message from the node(s) to their successors and update them.

        Optionally, apply a function to update the node features after receive.

        In the graph store, all updates are written inplace.

        Parameters
        ----------
        u : int, iterable of int, or tensor
            The node(s) to push messages out.
        message_func : callable, optional
            Message function on the edges. The function should be
            an :mod:`Edge UDF <dgl.udf>`.
        reduce_func : callable, optional
            Reduce function on the node. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        apply_node_func : callable, optional
            Apply function on the nodes. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        inplace: bool, optional
            The value is always True.
        """
        ...
    
    def update_all(self, message_func=..., reduce_func=..., apply_node_func=...): # -> None:
        """ Distribute the computation in update_all among all pre-defined workers.

        update_all requires that all workers invoke this method and will
        return only when all workers finish their own portion of computation.
        The number of workers are pre-defined. If one of them doesn't invoke the method,
        it won't return because some portion of computation isn't finished.

        Parameters
        ----------
        message_func : callable, optional
            Message function on the edges. The function should be
            an :mod:`Edge UDF <dgl.udf>`.
        reduce_func : callable, optional
            Reduce function on the node. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        apply_node_func : callable, optional
            Apply function on the nodes. The function should be
            a :mod:`Node UDF <dgl.udf>`.
        """
        ...
    
    def destroy(self): # -> None:
        """Destroy the graph store.

        This notifies the server that this client has terminated.
        """
        ...
    


def create_graph_store_server(graph_data, graph_name, store_type, num_workers, multigraph=..., port=...): # -> SharedMemoryStoreServer:
    """Create the graph store server.

    The server loads graph structure and node embeddings and edge embeddings.

    Currently, only shared-memory graph store server is supported, so `store_type`
    can only be "shared_mem".

    After the server runs, the graph store clients can access the graph data
    with the specified graph name.

    DGL graph accepts graph data of multiple formats:

    * NetworkX graph,
    * scipy matrix,
    * DGLGraph.

    If the input graph data is DGLGraph, the constructed DGLGraph only contains
    its graph index.

    Parameters
    ----------
    graph_data : graph data
        Data to initialize graph.
    graph_name : string
        Define the name of the graph.
    store_type : string
        The type of the graph store. The current option is "shared_mem".
    num_workers : int
        The number of workers that will connect to the server.
    multigraph : bool, optional
        Deprecated (Will be deleted in the future).
        Whether the graph would be a multigraph (default: True)
    port : int
        The port that the server listens to.

    Returns
    -------
    SharedMemoryStoreServer
        The graph store server
    """
    ...

def create_graph_from_store(graph_name, store_type, port=...): # -> SharedMemoryDGLGraph:
    """Create a client from the graph store.

    The client constructs the graph structure and node embeddings and edge embeddings
    that has been loaded by the graph store server.

    Currently, only shared-memory graph store server is supported, so `store_type`
    can only be "shared_memory".

    Parameters
    ----------
    graph_name : string
        Define the name of the graph.
    store_type : string
        The type of the graph store. The current option is "shared_mem".
    port : int
        The port that the server listens to.

    Returns
    -------
    SharedMemoryDGLGraph
        The shared-memory DGLGraph
    """
    ...

