xoscar 0.7.15__cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.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/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"