xoscar 0.9.0__cp312-cp312-macosx_10_13_x86_64.whl

xoscar/collective/utils.py

@@ -0,0 +1,41 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import numpy as np
+
+from ..utils import lazy_import
+
+cupy = lazy_import("cupy")
+
+
+def convert_data_to_np_array(data):
+    if isinstance(data, np.ndarray):
+        return data
+    else:
+        return np.frombuffer(data, dtype="u1")
+
+
+def convert_data_to_cp_array(data):
+    if isinstance(data, cupy.ndarray):
+        return data
+    else:
+        return cupy.frombuffer(data, dtype="u1")
+
+
+def get_rank_address_via_env(env_key: str, err_message: str) -> str:
+    address = os.environ.get(env_key, None)
+    if address is None:
+        raise RuntimeError(err_message)
+    return address

xoscar/collective/xoscar_pygloo.pyi

@@ -0,0 +1,239 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+from ctypes import c_void_p
+from enum import IntEnum
+from typing import Callable, List, Optional
+
+import xoscar_pygloo
+
+class ReduceOp(IntEnum):
+    SUM = 0
+    PRODUCT = 1
+    MIN = 2
+    MAX = 3
+    BAND = 4
+    BOR = 5
+    BXOR = 6
+    UNUSED = 7
+
+class GlooDataType_t(IntEnum):
+    glooInt8 = 0
+    glooUint8 = 1
+    glooInt32 = 2
+    glooUint32 = 3
+    glooInt64 = 4
+    glooUint64 = 5
+    glooFloat16 = 6
+    glooFloat32 = 7
+    glooFloat64 = 8
+
+class AllreduceAlgorithm(IntEnum):
+    UNSPECIFIED = 0
+    RING = 1
+    BCUBE = 2
+
+def transport_tcp_available() -> bool: ...
+def transport_uv_available() -> bool: ...
+
+class Context:
+    rank: Optional[int] = None
+    size: Optional[int] = None
+    base: Optional[int] = 2
+    def getDevice(self) -> int: ...
+    def createUnboundBuffer(self, ptr: c_void_p, size: int): ...
+    def nextSlot(self, numToSkip: int) -> int: ...
+    def closeConnections(self) -> None: ...
+    def setTimeout(self, timeout: datetime.timedelta) -> None: ...
+    def getTimeout(self) -> datetime.timedelta: ...
+
+def allreduce(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    reduceop: Optional[ReduceOp] = ReduceOp.SUM,
+    algorithm: Optional[AllreduceAlgorithm] = AllreduceAlgorithm.RING,
+    tag: int = 0,
+) -> None: ...
+def allgather(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    tag: Optional[int] = 0,
+) -> None: ...
+def all_to_all(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    tag: Optional[int] = 0,
+) -> None: ...
+def allgatherv(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    tag: Optional[int] = 0,
+) -> None: ...
+def reduce(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    reduceop: Optional[ReduceOp] = ReduceOp.SUM,
+    root: Optional[int] = 0,
+    tag: Optional[int] = 0,
+) -> None: ...
+def scatter(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    root: Optional[int] = 0,
+    tag: Optional[int] = 0,
+) -> None: ...
+def gather(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    root: Optional[int] = 0,
+    tag: Optional[int] = 0,
+) -> None: ...
+def send(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    peer: Optional[int] = None,
+    tag: Optional[int] = 0,
+) -> None: ...
+def recv(
+    context: Optional[Context] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    peer: Optional[int] = None,
+    tag: Optional[int] = 0,
+) -> None: ...
+def broadcast(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    root: Optional[int] = 0,
+    tag: Optional[int] = 0,
+) -> None: ...
+def reduce_scatter(
+    context: Optional[Context] = None,
+    sendbuf: Optional[int] = None,
+    recvbuf: Optional[int] = None,
+    size: Optional[int] = None,
+    recvElems: Optional[List[int]] = None,
+    datatype: Optional[GlooDataType_t] = None,
+    reduceop: Optional[ReduceOp] = ReduceOp.SUM,
+) -> None: ...
+def barrier(context: Optional[Context] = None, tag: Optional[int] = 0) -> None: ...
+
+class rendezvous:
+    class Store:
+        def set(self, key: str, data: List[str]) -> None: ...
+        def get(self, key: str) -> str: ...
+
+    class TCPStoreOptions:
+        port: int
+        isServer: bool
+        numWorkers: Optional[int]
+        waitWorkers: bool
+        timeout: datetime.timedelta
+        multiTenant: bool
+
+    class TCPStore:
+        def __init__(
+            self,
+            host: str,
+            opts: rendezvous.TCPStoreOptions = rendezvous.TCPStoreOptions(),
+        ): ...
+        def set(self, key: str, value: bytes): ...
+        def get(self, key: str) -> bytes: ...
+        def wait(self, keys: List[str]): ...
+
+    class FileStore(Store):
+        def __init__(self, path: str) -> None: ...
+        def set(self, key: str, data: List[str]) -> None: ...
+        def get(self, key: str) -> str: ...
+
+    class HashStore(Store):
+        def __init__(self) -> None: ...
+        def set(self, key: str, data: List[str]) -> None: ...
+        def get(self, key: str) -> str: ...
+
+    class PrefixStore(Store):
+        def __init__(self, prefix: str, store: rendezvous.Store) -> None: ...
+        def set(self, key: str, data: List[str]) -> None: ...
+        def get(self, key: str) -> str: ...
+
+    class CustomStore(Store):
+        def __init__(self, real_store_py_object: object) -> None: ...
+        def delKeys(self, keys: List[str]) -> None: ...
+        def set(self, key: str, data: List[str]) -> None: ...
+        def get(self, key: str) -> str: ...
+
+    class Context(xoscar_pygloo.Context):
+        def connectFullMesh(
+            self, store: rendezvous.Store, dev: transport.Device
+        ) -> None: ...
+
+class transport:
+    class uv:
+        pass
+
+    class tcp:
+        class Device(transport.Device):
+            def __init__(self, attr: transport.tcp.attr) -> None: ...
+
+        class Context(xoscar_pygloo.Context):
+            def __init__(
+                self, device: transport.tcp.Device, rank: int, size: int
+            ) -> None: ...
+
+        class attr:
+            hostname: str
+            iface: str
+            ai_family: int
+            ai_socktype: int
+            ai_protocol: int
+            ai_addr: object
+            ai_addrlen: int
+            def __init__(self, string: Optional[str] = None) -> None: ...
+
+        def CreateDevice(self, src: transport.tcp.attr) -> transport.Device: ...
+
+    class Device:
+        def __str__(self) -> str: ...
+        def getPCIBusID(self) -> Callable[[], str]: ...
+        def getInterfaceSpeed(self) -> int: ...
+        def hasGPUDirect(self) -> bool: ...
+        def createContext(self, rank: int, size: int): ...

xoscar/constants.py

@@ -0,0 +1,23 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from pathlib import Path
+
+XOSCAR_TEMP_DIR = Path(os.getenv("XOSCAR_DIR", Path.home())) / ".xoscar"
+
+# unix socket.
+XOSCAR_UNIX_SOCKET_DIR = XOSCAR_TEMP_DIR / "socket"
+
+XOSCAR_CONNECT_TIMEOUT = 8

xoscar/context.pxd

@@ -0,0 +1,21 @@
+# Copyright 2022-2023 XProbe Inc.
+# derived from copyright 1999-2021 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+cdef class BaseActorContext:
+    cdef public str _address
+
+
+cpdef get_context()

xoscar/context.pyx

@@ -0,0 +1,368 @@
+# Copyright 2022-2023 XProbe Inc.
+# derived from copyright 1999-2021 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, List, Optional, Union
+from urllib.parse import urlparse
+
+from ._utils cimport new_actor_id, new_random_id
+from .core cimport ActorRef, BufferRef, FileObjectRef
+
+
+cdef dict _backend_context_cls = dict()
+
+cdef object _context = None
+
+
+cdef class BaseActorContext:
+    # allocate strategy is for Indigen backend only
+    support_allocate_strategy = False
+
+    """
+    Base class for actor context. Every backend need to implement
+    actor context for their own.
+    """
+
+    def __init__(self, address: str = None):
+        self._address = address
+
+    async def create_actor(
+        self,
+        object actor_cls,
+        *args,
+        object uid=None,
+        object address=None,
+        **kwargs,
+    ):
+        """
+        Stub method for creating an actor in current context.
+
+        Parameters
+        ----------
+        actor_cls : Actor
+            Actor class
+        args : tuple
+            args to be passed into actor_cls.__init__
+        uid : identifier
+            Actor identifier
+        address : str
+            Address to locate the actor
+        kwargs : dict
+            kwargs to be passed into actor_cls.__init__
+
+        Returns
+        -------
+        ActorRef
+
+        """
+        raise NotImplementedError
+
+    async def has_actor(self, ActorRef actor_ref):
+        """
+        Check if actor exists in current context
+
+        Parameters
+        ----------
+        actor_ref : ActorRef
+            Reference to an actor
+
+        Returns
+        -------
+        bool
+        """
+        raise NotImplementedError
+
+    async def destroy_actor(self, ActorRef actor_ref):
+        """
+        Destroy an actor by its reference
+
+        Parameters
+        ----------
+        actor_ref : ActorRef
+            Reference to an actor
+
+        Returns
+        -------
+        bool
+        """
+        raise NotImplementedError
+
+    async def kill_actor(self, ActorRef actor_ref):
+        """
+        Force to kill an actor, take care this is a dangerous operation,
+        it may lead to the result that other actors are killed as well.
+        Hence, unless you are knowing what you are doing and know how
+        to recover possible effected actors, DO NOT USE this method!
+
+        Parameters
+        ----------
+        actor_ref : ActorRef
+            Reference to an actor
+
+        Returns
+        -------
+        bool
+        """
+
+    async def send(
+        self,
+        ActorRef actor_ref,
+        object message,
+        bint wait_response=True,
+        object profiling_context=None,
+    ):
+        """
+        Send a message to given actor by its reference
+
+        Parameters
+        ----------
+        actor_ref : ActorRef
+            Reference to an actor
+        message : object
+            Message to send to an actor, need to comply to Actor.__on_receive__
+        wait_response : bool
+            Whether to wait for responses from the actor.
+        profiling_context: ProfilingContext
+            The profiling context.
+
+        Returns
+        -------
+        object
+        """
+        raise NotImplementedError
+
+    async def actor_ref(self, *args, **kwargs):
+        """
+        Create a reference to an actor
+
+        Returns
+        -------
+        ActorRef
+        """
+        raise NotImplementedError
+
+    async def wait_actor_pool_recovered(self, str address, str main_address = None):
+        """
+        Wait until an actor pool is recovered
+
+        Parameters
+        ----------
+        address
+            address of the actor pool
+        main_address
+            address of the main pool
+        """
+        raise NotImplementedError
+
+    async def get_pool_config(self, str address):
+        """
+        Get config of actor pool with given address
+
+        Parameters
+        ----------
+        address
+            address of the actor pool
+
+        Returns
+        -------
+
+        """
+        raise NotImplementedError
+
+    def buffer_ref(self, str address, object buf) -> BufferRef:
+        """
+        Create a reference to a buffer
+
+        Parameters
+        ----------
+        address
+            address of the actor pool
+        buf
+            buffer object
+
+        Returns
+        -------
+        BufferRef
+        """
+        return BufferRef.create(buf, address, new_random_id(32))
+
+    def file_object_ref(self, str address, object file_object) -> FileObjectRef:
+        """
+        Create a reference to an aio file object
+
+        Parameters
+        ----------
+        address
+            address of the actor pool
+        file_object
+            aio file object
+
+        Returns
+        -------
+        FileObjectRef
+        """
+        return FileObjectRef.create(file_object, address, new_random_id(32))
+
+    async def copy_to_buffers(self, local_buffers: List, remote_buffer_refs: List[BufferRef], block_size: Optional[int] = None):
+        """
+        Copy local buffers to remote buffers.
+        Parameters
+        ----------
+        local_buffers
+            Local buffers.
+        remote_buffer_refs
+            Remote buffer refs
+        block_size
+            Transfer block size when non-ucx
+        """
+        raise NotImplementedError
+
+    async def copy_to_fileobjs(self, local_fileobjs: list, remote_fileobj_refs: List[FileObjectRef], block_size: Optional[int] = None):
+        """
+        Copy local file objs to remote file objs.
+        Parameters
+        ----------
+        local_fileobjs
+            Local file objs.
+        remote_fileobj_refs
+            Remote file object refs
+        block_size
+            Transfer block size when non-ucx
+        """
+        raise NotImplementedError
+
+
+cdef class ClientActorContext(BaseActorContext):
+    """
+    Default actor context. This context will keep references to other contexts
+    given their protocol scheme (i.e., `ray://xxx`).
+    """
+    cdef dict _backend_contexts
+
+    def __init__(self, address: str = None):
+        BaseActorContext.__init__(self, address)
+        self._backend_contexts = dict()
+
+    cdef inline object _get_backend_context(self, object address):
+        if address is None:
+            raise ValueError('address has to be provided')
+        if '://' not in address:
+            scheme = None
+        else:
+            scheme = urlparse(address).scheme or None
+        try:
+            return self._backend_contexts[scheme]
+        except KeyError:
+            context = self._backend_contexts[scheme] = \
+                _backend_context_cls[scheme](address)
+            return context
+
+    def create_actor(
+        self,
+        object actor_cls,
+        *args,
+        object uid=None,
+        object address=None,
+        **kwargs,
+    ):
+        context = self._get_backend_context(address)
+        uid = uid or new_actor_id()
+        return context.create_actor(actor_cls, *args, uid=uid, address=address, **kwargs)
+
+    def has_actor(self, ActorRef actor_ref):
+        context = self._get_backend_context(actor_ref.address)
+        return context.has_actor(actor_ref)
+
+    def destroy_actor(self, ActorRef actor_ref):
+        context = self._get_backend_context(actor_ref.address)
+        return context.destroy_actor(actor_ref)
+
+    def kill_actor(self, ActorRef actor_ref):
+        context = self._get_backend_context(actor_ref.address)
+        return context.kill_actor(actor_ref)
+
+    def actor_ref(self, *args, **kwargs):
+        from ._utils import create_actor_ref
+
+        actor_ref = create_actor_ref(*args, **kwargs)
+        context = self._get_backend_context(actor_ref.address)
+        return context.actor_ref(actor_ref)
+
+    def send(
+        self,
+        ActorRef actor_ref,
+        object message,
+        bint wait_response=True,
+        object profiling_context=None
+    ):
+        context = self._get_backend_context(actor_ref.address)
+        return context.send(
+            actor_ref,
+            message,
+            wait_response=wait_response,
+            profiling_context=profiling_context,
+        )
+
+    def wait_actor_pool_recovered(self, str address, str main_address = None):
+        context = self._get_backend_context(address)
+        return context.wait_actor_pool_recovered(address, main_address)
+
+    def get_pool_config(self, str address):
+        context = self._get_backend_context(address)
+        return context.get_pool_config(address)
+
+    def buffer_ref(self, str address, buf: Any) -> BufferRef:
+        context = self._get_backend_context(address)
+        return context.buffer_ref(address, buf)
+
+    def file_object_ref(self, str address, object file_object) -> FileObjectRef:
+        context = self._get_backend_context(address)
+        return context.file_object_ref(address, file_object)
+
+    def copy_to(self, local_buffers_or_fileobjs: list, remote_refs: List[Union[BufferRef, FileObjectRef]], block_size: Optional[int] = None):
+        if len(local_buffers_or_fileobjs) == 0 or len(remote_refs) == 0:
+            raise ValueError("Nothing to transfer since the length of `local_buffers_or_fileobjs` or `remote_refs` is 0.")
+        assert (
+            len({ref.address for ref in remote_refs}) == 1
+        ), "remote_refs for `copy_to` can support only 1 destination"
+        assert len(local_buffers_or_fileobjs) == len(remote_refs), (
+            f"Buffers or fileobjs from local and remote must have same size, "
+            f"local: {len(local_buffers_or_fileobjs)}, remote: {len(remote_refs)}"
+        )
+        if block_size is not None:
+            assert (
+                block_size > 0
+            ), f"`block_size` option must be greater than 0, current value: {block_size}."
+        remote_ref = remote_refs[0]
+        address = remote_ref.address
+        context = self._get_backend_context(address)
+        if isinstance(remote_ref, BufferRef):
+            return context.copy_to_buffers(local_buffers_or_fileobjs, remote_refs, block_size)
+        else:
+            return context.copy_to_fileobjs(local_buffers_or_fileobjs, remote_refs, block_size)
+
+
+def register_backend_context(scheme, cls):
+    assert issubclass(cls, BaseActorContext)
+    _backend_context_cls[scheme] = cls
+
+
+cpdef get_context():
+    """
+    Get an actor context. If not in an actor environment,
+    ClientActorContext will be used
+    """
+    global _context
+    if _context is None:
+        _context = ClientActorContext()
+    return _context