# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_memory/_device_memory_resource.pyx

from __future__ import annotations

import uuid
from dataclasses import dataclass

from cuda.core._device import Device
from cuda.core._memory._ipc import IPCAllocationHandle
from cuda.core._memory._memory_pool import _MemPool
from cuda.core._memory._peer_access_utils import PeerAccessibleBySetProxy


@dataclass
class DeviceMemoryResourceOptions:
    """Customizable :obj:`~_memory.DeviceMemoryResource` options.

    Attributes
    ----------
    ipc_enabled : bool, optional
        Specifies whether to create an IPC-enabled memory pool. When set to
        True, the memory pool and its allocations can be shared with other
        processes. (Default to False)

    max_size : int, optional
        Maximum pool size. When set to 0, defaults to a system-dependent value.
        (Default to 0)
    """
    ipc_enabled: bool = False
    max_size: int = 0

class DeviceMemoryResource(_MemPool):
    """
    A device memory resource managing a stream-ordered memory pool.

    Parameters
    ----------
    device_id : Device | int
        Device or Device ordinal for which a memory resource is constructed.

    options : DeviceMemoryResourceOptions
        Memory resource creation options.

        If set to `None`, the memory resource uses the driver's current
        stream-ordered memory pool for the specified `device_id`. If no memory
        pool is set as current, the driver's default memory pool for the device
        is used.

        If not set to `None`, a new memory pool is created, which is owned by
        the memory resource.

        When using an existing (current or default) memory pool, the returned
        device memory resource does not own the pool (`is_handle_owned` is
        `False`), and closing the resource has no effect.

    Notes
    -----
    To create an IPC-Enabled memory resource (MR) that is capable of sharing
    allocations between processes, specify ``ipc_enabled=True`` in the initializer
    option. Sharing an allocation is a two-step procedure that involves
    mapping a memory resource and then mapping buffers owned by that resource.
    These steps can be accomplished in several ways.

    An IPC-enabled memory resource can allocate memory buffers but cannot
    receive shared buffers. Mapping an MR to another process creates a "mapped
    memory resource" (MMR). An MMR cannot allocate memory buffers and can only
    receive shared buffers. MRs and MMRs are both of type
    :class:`DeviceMemoryResource` and can be distinguished via
    :attr:`DeviceMemoryResource.is_mapped`.

    An MR is shared via an allocation handle accessed through the
    :attr:`DeviceMemoryResource.allocation_handle` property. The allocation
    handle has a platform-specific interpretation; however, memory IPC is
    currently only supported for Linux, and in that case allocation handles
    are file descriptors. After sending an allocation handle to another
    process, it can be used to create an MMR by invoking
    :meth:`DeviceMemoryResource.from_allocation_handle`.

    Buffers can be shared as serializable descriptors accessed through the
    :attr:`Buffer.ipc_descriptor` property. In a receiving process, a shared
    buffer is created by invoking :meth:`Buffer.from_ipc_descriptor` with an
    MMR and buffer descriptor, where the MMR corresponds to the MR that
    created the described buffer.

    To help manage the association between memory resources and buffers, a
    registry is provided. Every MR has a unique identifier (UUID). MMRs can be
    registered by calling :meth:`DeviceMemoryResource.register` with the UUID
    of the corresponding MR. Registered MMRs can be looked up via
    :meth:`DeviceMemoryResource.from_registry`. When registering MMRs in this
    way, the use of buffer descriptors can be avoided. Instead, buffer objects
    can themselves be serialized and transferred directly. Serialization embeds
    the UUID, which is used to locate the correct MMR during reconstruction.

    IPC-enabled memory resources interoperate with the :mod:`multiprocessing`
    module to provide a simplified interface. This approach can avoid direct
    use of allocation handles, buffer descriptors, MMRs, and the registry. When
    using :mod:`multiprocessing` to spawn processes or send objects through
    communication channels such as :class:`multiprocessing.Queue`,
    :class:`multiprocessing.Pipe`, or :class:`multiprocessing.Connection`,
    :class:`Buffer` objects may be sent directly, and in such cases the process
    for creating MMRs and mapping buffers will be handled automatically.

    For greater efficiency when transferring many buffers, one may also send
    MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`,
    an MMR is created and registered in the receiving process. Subsequently,
    buffers may be serialized and transferred using ordinary :mod:`pickle`
    methods.  The reconstruction procedure uses the registry to find the
    associated MMR. Unpickling a :class:`Buffer` performs an IPC import from
    the embedded descriptor; only unpickle buffers received from trusted peers.

    Warning
    -------
    IPC descriptors and pickled buffers cross a trust boundary between
    cooperating same-host processes. A malicious peer can supply crafted
    descriptor fields. Use :meth:`Buffer.from_ipc_descriptor` only with
    descriptors from trusted peers, and do not unpickle buffers from
    untrusted sources.
    """

    def __cinit__(self, *args, **kwargs) -> None:
        ...

    def __init__(self, device_id: Device | int, options: DeviceMemoryResourceOptions | dict[str, object] | None=None) -> None:
        ...

    def __reduce__(self) -> tuple[object, ...]:
        ...

    @staticmethod
    def from_registry(uuid: uuid.UUID) -> DeviceMemoryResource:
        """
        Obtain a registered mapped memory resource.

        Raises
        ------
        RuntimeError
            If no mapped memory resource is found in the registry.
        """

    def register(self, uuid: uuid.UUID) -> DeviceMemoryResource:
        """
        Register a mapped memory resource.

        Returns
        -------
        The registered mapped memory resource. If one was previously registered
        with the given key, it is returned.
        """

    @classmethod
    def from_allocation_handle(cls, device_id: Device | int, alloc_handle: int | IPCAllocationHandle) -> DeviceMemoryResource:
        """Create a device memory resource from an allocation handle.

        Construct a new `DeviceMemoryResource` instance that imports a memory
        pool from a shareable handle. The memory pool is marked as owned, and
        the resource is associated with the specified `device_id`.

        Parameters
        ----------
        device_id : int | Device
            The ID of the device or a Device object for which the memory
            resource is created.

        alloc_handle : int | IPCAllocationHandle
            The shareable handle of the device memory resource to import. If an
            integer is supplied, it must represent a valid platform-specific
            handle. It is the caller's responsibility to close that handle.

        Returns
        -------
            A new device memory resource instance with the imported handle.
        """

    @property
    def allocation_handle(self) -> IPCAllocationHandle:
        """Shareable handle for this memory pool (requires IPC).

        The handle can be used to share the memory pool with other processes.
        The handle is cached in this `MemoryResource` and owned by it.
        """

    @property
    def device_id(self) -> int:
        """The associated device ordinal."""

    @property
    def peer_accessible_by(self) -> PeerAccessibleBySetProxy:
        """
        Get or set the devices that can access allocations from this memory
        pool. Access can be modified at any time and affects all allocations
        from this memory pool.

        Returns a set-like proxy of :obj:`~_device.Device` objects that manages
        peer access. Inputs are accepted as either :obj:`~_device.Device`
        objects or device-ordinal :class:`int` values.

        Examples
        --------
        >>> dmr = DeviceMemoryResource(0)
        >>> dmr.peer_accessible_by = {1}   # grant access to device 1
        >>> assert 1 in dmr.peer_accessible_by
        >>> dmr.peer_accessible_by.add(2)  # update access to include device 2
        >>> dmr.peer_accessible_by = []    # revoke peer access
        """

    @peer_accessible_by.setter
    def peer_accessible_by(self, devices) -> None:
        ...

    @property
    def is_device_accessible(self) -> bool:
        """Return True. This memory resource provides device-accessible buffers."""

    @property
    def is_host_accessible(self) -> bool:
        """Return False. This memory resource does not provide host-accessible buffers."""
__all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions']

def DMR_mempool_get_access(dmr: DeviceMemoryResource, device_id: int) -> str:
    """
    Probes peer access from the given device using cuMemPoolGetAccess.

    Parameters
    ----------
    device_id : int or Device
        The device to query access for.

    Returns
    -------
    str
        Access permissions: "rw" for read-write, "r" for read-only, "" for no access.
    """

def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]:
    ...