xoscar 0.9.0__cp312-cp312-macosx_10_13_x86_64.whl

xoscar/batch.py

@@ -0,0 +1,256 @@
+# 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 __future__ import annotations
+
+import asyncio
+import inspect
+import textwrap
+from collections import namedtuple
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from .core import NO_LOCK_ATTRIBUTE_HINT
+
+
+def build_args_binder(func, remove_self: bool = True) -> Callable | None:
+    try:
+        spec = inspect.getfullargspec(func)
+    except TypeError:  # pragma: no cover
+        return None
+
+    sig_list = list(spec.args)
+    args_list = list(spec.args)
+    if remove_self:
+        args_list = args_list[1:]
+
+    if spec.varargs:
+        sig_list.append(f"*{spec.varargs}")
+        args_list.append(spec.varargs)
+    elif spec.kwonlyargs:
+        sig_list.append("*")
+
+    sig_list.extend(spec.kwonlyargs)
+    args_list.extend(spec.kwonlyargs)
+
+    if spec.varkw:
+        sig_list.append(f"**{spec.varkw}")
+        args_list.append(spec.varkw)
+
+    if getattr(func, "__name__").isidentifier():
+        ret_func_name = f"{func.__name__}_binder"
+        ret_type_name = f"_Args_{func.__name__}"
+    else:
+        ret_func_name = f"anon_{id(func)}_binder"
+        ret_type_name = f"_ArgsAnon_{id(func)}"
+
+    func_str = textwrap.dedent(
+        f"""
+    def {ret_func_name}({', '.join(sig_list)}):
+        return {ret_type_name}({', '.join(args_list)})
+    """
+    )
+
+    glob_vars = globals().copy()
+    glob_vars[ret_type_name] = namedtuple(ret_type_name, args_list)
+    loc_vars: dict[str, Any] = dict()
+    exec(func_str, glob_vars, loc_vars)
+    ext_func = loc_vars[ret_func_name]
+    ext_func.__defaults__ = spec.defaults
+    ext_func.__kwdefaults__ = spec.kwonlydefaults
+
+    return ext_func
+
+
+@dataclass
+class _DelayedArgument:
+    args: tuple
+    kwargs: dict
+
+
+class _ExtensibleCallable:
+    func: Callable
+    batch_func: Callable | None
+    is_async: bool
+    has_single_func: bool
+
+    def __call__(self, *args, **kwargs):
+        if self.is_async:
+            return self._async_call(*args, **kwargs)
+        else:
+            return self._sync_call(*args, **kwargs)
+
+    async def _async_call(self, *args, **kwargs):
+        try:
+            if self.has_single_func:
+                return await self.func(*args, **kwargs)
+        except NotImplementedError:
+            self.has_single_func = False
+
+        if self.batch_func is not None:
+            ret = await self.batch_func([args], [kwargs])
+            return None if ret is None else ret[0]
+        raise NotImplementedError
+
+    def _sync_call(self, *args, **kwargs):
+        try:
+            if self.has_single_func:
+                return self.func(*args, **kwargs)
+        except NotImplementedError:
+            self.has_single_func = False
+
+        if self.batch_func is not None:
+            return self.batch_func([args], [kwargs])[0]
+        raise NotImplementedError
+
+
+class _ExtensibleWrapper(_ExtensibleCallable):
+    def __init__(
+        self,
+        func: Callable,
+        batch_func: Callable | None = None,
+        bind_func: Callable | None = None,
+        is_async: bool = False,
+    ):
+        self.func = func
+        self.batch_func = batch_func
+        self.bind_func = bind_func
+        self.is_async = is_async
+        self.has_single_func = True
+
+    @staticmethod
+    def delay(*args, **kwargs):
+        return _DelayedArgument(args=args, kwargs=kwargs)
+
+    @staticmethod
+    def _gen_args_kwargs_list(delays):
+        args_list = [delay.args for delay in delays]
+        kwargs_list = [delay.kwargs for delay in delays]
+        return args_list, kwargs_list
+
+    async def _async_batch(self, args_list, kwargs_list):
+        # when there is only one call in batch, calling one-pass method
+        # will be more efficient
+        if len(args_list) == 0:
+            return []
+        elif len(args_list) == 1:
+            return [await self._async_call(*args_list[0], **kwargs_list[0])]
+        elif self.batch_func:
+            return await self.batch_func(args_list, kwargs_list)
+        else:
+            # this function has no batch implementation
+            # call it separately
+            tasks = [
+                asyncio.create_task(self.func(*args, **kwargs))
+                for args, kwargs in zip(args_list, kwargs_list)
+            ]
+            try:
+                return await asyncio.gather(*tasks)
+            except asyncio.CancelledError:
+                _ = [task.cancel() for task in tasks]
+                return await asyncio.gather(*tasks)
+
+    def _sync_batch(self, args_list, kwargs_list):
+        if len(args_list) == 0:
+            return []
+        elif self.batch_func:
+            return self.batch_func(args_list, kwargs_list)
+        else:
+            # this function has no batch implementation
+            # call it separately
+            return [
+                self.func(*args, **kwargs)
+                for args, kwargs in zip(args_list, kwargs_list)
+            ]
+
+    def batch(self, *delays):
+        args_list, kwargs_list = self._gen_args_kwargs_list(delays)
+        return self.call_with_lists(args_list, kwargs_list)
+
+    def call_with_lists(self, args_list, kwargs_list):
+        if self.is_async:
+            return self._async_batch(args_list, kwargs_list)
+        else:
+            return self._sync_batch(args_list, kwargs_list)
+
+    def bind(self, *args, **kwargs):
+        if self.bind_func is None:
+            raise TypeError(f"bind function not exist for method {self.func.__name__}")
+        return self.bind_func(*args, **kwargs)
+
+
+class _ExtensibleAccessor(_ExtensibleCallable):
+    func: Callable
+    batch_func: Callable | None
+
+    def __init__(self, func: Callable):
+        self.func = func
+        self.batch_func = None
+        self.bind_func = build_args_binder(func, remove_self=True)
+        self.is_async = asyncio.iscoroutinefunction(self.func)
+        self.has_single_func = True
+
+    def batch(self, func: Callable):
+        self.batch_func = func
+        return self
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            # calling from class
+            return self.func
+
+        func = self.func.__get__(instance, owner)
+        batch_func = (
+            self.batch_func.__get__(instance, owner)
+            if self.batch_func is not None
+            else None
+        )
+        bind_func = (
+            self.bind_func.__get__(instance, owner)
+            if self.bind_func is not None
+            else None
+        )
+
+        wrapper = _ExtensibleWrapper(
+            func, batch_func=batch_func, bind_func=bind_func, is_async=self.is_async
+        )
+
+        if (
+            getattr(self.func, NO_LOCK_ATTRIBUTE_HINT, None) is True
+            or getattr(self.batch_func, NO_LOCK_ATTRIBUTE_HINT, None) is True
+        ):
+            setattr(wrapper, NO_LOCK_ATTRIBUTE_HINT, True)
+        return wrapper
+
+
+def extensible(func: Callable):
+    """
+    `extensible` means this func could be functionality extended,
+    especially for batch operations.
+
+    Consider remote function calls, each function may have operations
+    like opening file, closing file, batching them can help to reduce the cost,
+    especially for remote function calls.
+
+    Parameters
+    ----------
+    func : callable
+        Function
+
+    Returns
+    -------
+    func
+    """
+    return _ExtensibleAccessor(func)

xoscar/collective/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+
+from .core import (
+    RankActor,
+    allgather,
+    allreduce,
+    alltoall,
+    broadcast,
+    gather,
+    init_process_group,
+    new_group,
+    reduce,
+    reduce_scatter,
+    scatter,
+)

xoscar/collective/backend/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2022-2025 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.

xoscar/collective/backend/nccl_backend.py

@@ -0,0 +1,160 @@
+# 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.
+
+# We need to extend cupy's inner class because an actor is a daemonic processes
+# which are not allowed to have children. However, the origin code in cupy
+# will create children processes.
+
+import queue
+import socket
+import threading
+from ctypes import sizeof
+
+from ...utils import lazy_import
+
+cupy = lazy_import("cupy")
+
+if cupy is not None:
+    import cupyx.distributed
+    from cupy.cuda import nccl
+    from cupyx.distributed import _klv_utils, _store, _store_actions
+
+    class ExceptionAwareThreading(threading.Thread):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self._exception = None
+            self.q = queue.Queue()
+
+        def run(self):
+            try:
+                super().run()
+                self.q.put(None)
+            except Exception as e:
+                self.q.put(e)
+
+        def join(self):
+            super().join()
+            if not self.q.empty():
+                exception = self.q.get()
+                if exception is not None:
+                    raise exception
+
+    class TCPStore:
+        # This is only used for initialization of nccl so we don't care
+        # too much about performance
+        def __init__(self, world_size):
+            self.storage = {}
+            self._thread = None
+            self._world_size = world_size
+            self._run = 1
+            # For implementing a barrier
+            self._lock = threading.Lock()
+            self._current_barrier = None
+
+        def __del__(self):
+            if not _store._exit_mode:
+                self.stop()
+
+        def _thread_request(self, c_socket):
+            with c_socket:
+                # Receive in KLV format
+                action_bytes = c_socket.recv(sizeof(_klv_utils.action_t))
+                if len(action_bytes) > 0:
+                    action_m = _klv_utils.action_t.from_buffer_copy(action_bytes)
+                    if action_m.length > 256:
+                        raise ValueError("Invalid length for message")
+                    value = bytearray(action_m.value)[: action_m.length]
+                    r = _store_actions.execute_action(action_m.action, value, self)
+                    if r is not None:
+                        c_socket.sendall(r.klv())
+
+        def _server_loop(self, host, port):
+            # This is for minimum info exchange during initialization
+            # a single connection allows to implement locking mechanics easily
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                s.bind((host, port))
+                s.listen()
+                s.settimeout(0.5)
+                while self._run == 1:
+                    try:
+                        c_socket, addr = s.accept()
+                    except socket.timeout:
+                        continue
+
+                    t = threading.Thread(
+                        target=self._thread_request, args=(c_socket,), daemon=True
+                    )
+                    t.start()
+
+        def run(self, host=_store._DEFAULT_HOST, port=_store._DEFAULT_PORT):
+            # Run the TCP store in a different process
+            t = ExceptionAwareThreading(target=self._server_loop, args=(host, port))
+            t.start()
+            self._thread = t
+
+        def stop(self):
+            if _store._exit_mode:
+                return  # Prevent shutdown errors
+            if self._thread is not None:
+                # acquire the lock
+                self._lock.acquire()
+                self._run = 0
+                self._lock.release()
+                self._thread.join()
+
+    class XoscarNCCLBackend(cupyx.distributed.NCCLBackend):
+        """Interface that uses NVIDIA's NCCL to perform communications.
+
+        Args:
+            n_devices (int): Total number of devices that will be used in the
+                distributed execution.
+            rank (int): Unique id of the GPU that the communicator is associated to
+                its value needs to be `0 <= rank < n_devices`.
+            host (str, optional): host address for the process rendezvous on
+                initialization. Defaults to `"127.0.0.1"`.
+            port (int, optional): port used for the process rendezvous on
+                initialization. Defaults to `13333`.
+            use_mpi(bool, optional): switch between MPI and use the included TCP
+                server for initialization & synchronization. Defaults to `False`.
+        """
+
+        def __init__(
+            self,
+            n_devices,
+            rank,
+            tcpstore,
+            host=_store._DEFAULT_HOST,
+            port=_store._DEFAULT_PORT,
+            use_mpi=False,
+        ):
+            self._tcpstore = tcpstore
+            super().__init__(n_devices, rank, host, port, use_mpi)
+
+        def _init_with_tcp_store(self, n_devices, rank, host, port):
+            nccl_id = None
+            if rank == 0:
+                self._tcpstore.run(host, port)
+                nccl_id = nccl.get_unique_id()
+                # get_unique_id return negative values due to cython issues
+                # with bytes && c strings. We shift them by 128 to
+                # make them positive and send them as bytes to the proxy store
+                shifted_nccl_id = bytes([b + 128 for b in nccl_id])
+                self._store_proxy["nccl_id"] = shifted_nccl_id
+                self._store_proxy.barrier()
+            else:
+                self._store_proxy.barrier()
+                nccl_id = self._store_proxy["nccl_id"]
+                nccl_id = tuple([int(b) - 128 for b in nccl_id])
+            self._comm = nccl.NcclCommunicator(n_devices, nccl_id, rank)

xoscar/collective/common.py

@@ -0,0 +1,102 @@
+# 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.
+from enum import IntEnum
+from typing import Dict, Type
+
+import numpy as np
+
+from ..utils import lazy_import
+from . import xoscar_pygloo as xp
+
+ReduceOpMappingGloo: Dict["CollectiveReduceOp", "xp.ReduceOp"] = {}
+AllReduceAlgorithmMappingGloo: Dict["AllReduceAlgorithm", "xp.AllreduceAlgorithm"] = {}
+
+
+def _register_reduce_op(reduce_op):
+    for op_type in reduce_op:
+        ReduceOpMappingGloo[op_type] = xp.ReduceOp(op_type)
+    return reduce_op
+
+
+def _register_allreduce_algo(algorithms):
+    for algo in algorithms:
+        AllReduceAlgorithmMappingGloo[algo] = xp.AllreduceAlgorithm(algo)
+    return algorithms
+
+
+@_register_reduce_op
+class CollectiveReduceOp(IntEnum):
+    SUM = 0
+    PRODUCT = 1
+    MIN = 2
+    MAX = 3
+    BAND = 4
+    BOR = 5
+    BXOR = 6
+    UNUSED = 7
+
+
+@_register_allreduce_algo
+class AllReduceAlgorithm(IntEnum):
+    UNSPECIFIED = 0
+    RING = 1
+    BCUBE = 2
+
+
+TypeMappingGloo: Dict[Type[np.dtype], "xp.GlooDataType_t"] = {
+    np.int8: xp.GlooDataType_t.glooInt8,  # type: ignore
+    np.uint8: xp.GlooDataType_t.glooUint8,  # type: ignore
+    np.int32: xp.GlooDataType_t.glooInt32,  # type: ignore
+    np.uint32: xp.GlooDataType_t.glooUint32,  # type: ignore
+    np.int64: xp.GlooDataType_t.glooInt64,  # type: ignore
+    np.uint64: xp.GlooDataType_t.glooUint64,  # type: ignore
+    np.float16: xp.GlooDataType_t.glooFloat16,  # type: ignore
+    np.float32: xp.GlooDataType_t.glooFloat32,  # type: ignore
+    np.float64: xp.GlooDataType_t.glooFloat64,  # type: ignore
+}
+cupy = lazy_import("cupy")
+if cupy is not None:
+    from cupy.cuda import nccl
+
+    TypeMappingNCCL: Dict[Type[np.dtype], int] = {
+        np.int8: nccl.NCCL_INT8,  # type: ignore
+        np.uint8: nccl.NCCL_UINT8,  # type: ignore
+        np.int32: nccl.NCCL_INT32,  # type: ignore
+        np.uint32: nccl.NCCL_UINT32,  # type: ignore
+        np.int64: nccl.NCCL_INT64,  # type: ignore
+        np.uint64: nccl.NCCL_UINT64,  # type: ignore
+        np.float16: nccl.NCCL_FLOAT16,  # type: ignore
+        np.float32: nccl.NCCL_FLOAT32,  # type: ignore
+        np.float64: nccl.NCCL_FLOAT64,  # type: ignore
+    }
+
+    ReduceOpMappingNCCL: Dict[CollectiveReduceOp, int] = {
+        CollectiveReduceOp.SUM: nccl.NCCL_SUM,
+        CollectiveReduceOp.PRODUCT: nccl.NCCL_PROD,
+        CollectiveReduceOp.MAX: nccl.NCCL_MAX,
+        CollectiveReduceOp.MIN: nccl.NCCL_MIN,
+    }
+
+    ReduceOpMappingNCCLStr: Dict[CollectiveReduceOp, str] = {
+        CollectiveReduceOp.SUM: "sum",
+        CollectiveReduceOp.PRODUCT: "prod",
+        CollectiveReduceOp.MAX: "max",
+        CollectiveReduceOp.MIN: "min",
+    }
+# Some static variables
+INVOKE_ERROR_MESSAGE = "Collective-related functions must be called in a process that is involved in collection communication."
+RANK_ADDRESS_ENV_KEY = "COLLECTIVE_RANK_ADDRESS"
+RENDEZVOUS_MASTER_IP_ENV_KEY = "COLLECTIVE_MASTER_IP"
+RENDEZVOUS_MASTER_PORT_ENV_KEY = "COLLECTIVE_MASTER_PORT"
+COLLECTIVE_DEVICE_ID_ENV_KEY = "COLLECTIVE_DEVICE_ID_FOR_AN_ACTOR"