xoscar 0.9.0__cp312-cp312-macosx_10_13_x86_64.whl

xoscar/serialization/cuda.py

@@ -0,0 +1,111 @@
+# 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, Dict, List, Tuple
+
+import pandas as pd
+
+from ..utils import lazy_import
+from .core import Serializer, buffered
+
+cupy = lazy_import("cupy")
+cudf = lazy_import("cudf")
+
+
+class CupySerializer(Serializer):
+    @buffered
+    def serial(self, obj: Any, context: Dict):
+        if not (obj.flags["C_CONTIGUOUS"] or obj.flags["F_CONTIGUOUS"]):
+            obj = cupy.array(obj, copy=True)
+
+        header = obj.__cuda_array_interface__.copy()
+        header["strides"] = tuple(obj.strides)
+        header["lengths"] = [obj.nbytes]
+        buffer = cupy.ndarray(
+            shape=(obj.nbytes,), dtype=cupy.dtype("u1"), memptr=obj.data, strides=(1,)
+        )
+        return (header,), [buffer], True
+
+    def deserial(self, serialized: Tuple, context: Dict, subs: List):
+        (header,) = serialized
+        return cupy.ndarray(
+            shape=header["shape"],
+            dtype=header["typestr"],
+            memptr=cupy.asarray(subs[0]).data,
+            strides=header["strides"],
+        )
+
+
+class CudfSerializer(Serializer):
+    @staticmethod
+    def _get_ext_index_type(index_obj):
+        import cudf
+
+        multi_index_type = None
+        if isinstance(index_obj, pd.MultiIndex):
+            multi_index_type = "pandas"
+        elif isinstance(index_obj, cudf.MultiIndex):
+            multi_index_type = "cudf"
+
+        if multi_index_type is None:
+            return None
+        return {
+            "index_type": multi_index_type,
+            "names": list(index_obj.names),
+        }
+
+    @staticmethod
+    def _apply_index_type(obj, attr, header):
+        import cudf
+
+        multi_index_cls = (
+            pd.MultiIndex if header["index_type"] == "pandas" else cudf.MultiIndex
+        )
+        original_index = getattr(obj, attr)
+        if isinstance(original_index, (pd.MultiIndex, cudf.MultiIndex)):
+            return
+        new_index = multi_index_cls.from_tuples(original_index, names=header["names"])
+        setattr(obj, attr, new_index)
+
+    def serial(self, obj: Any, context: Dict):
+        header, buffers = obj.device_serialize()
+        if hasattr(obj, "columns"):
+            header["_ext_columns"] = self._get_ext_index_type(obj.columns)
+        if hasattr(obj, "index"):
+            header["_ext_index"] = self._get_ext_index_type(obj.index)
+        return (header,), buffers, True
+
+    def deserial(self, serialized: Tuple, context: Dict, buffers: List):
+        from cudf.core.abc import Serializable
+
+        (header,) = serialized
+        col_header = header.pop("_ext_columns", None)
+        index_header = header.pop("_ext_index", None)
+
+        result = Serializable.device_deserialize(header, buffers)
+
+        if col_header is not None:
+            self._apply_index_type(result, "columns", col_header)
+        if index_header is not None:
+            self._apply_index_type(result, "index", index_header)
+        return result
+
+
+if cupy is not None:
+    CupySerializer.register("cupy.ndarray")
+if cudf is not None:
+    CudfSerializer.register("cudf.DataFrame")
+    CudfSerializer.register("cudf.Series")
+    CudfSerializer.register("cudf.Index")

xoscar/serialization/exception.py

@@ -0,0 +1,48 @@
+# 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 pickle  # nosec  # pylint: disable=import_pickle
+
+from .core import Serializer, buffered, pickle_buffers, unpickle_buffers
+
+
+class UnpickleableError(Exception):
+    def __init__(self, raw_error: str | Exception):
+        if isinstance(raw_error, str):
+            super().__init__(raw_error)
+        else:
+            super().__init__(
+                f"Error cannot be pickled, "
+                f"error type: {type(raw_error)}, "
+                f"raw error:\n{raw_error}"
+            )
+
+
+class ExceptionSerializer(Serializer):
+    @buffered
+    def serial(self, obj: Exception, context: dict):
+        try:
+            buffers = pickle_buffers(obj)
+        except (TypeError, pickle.PicklingError):
+            buffers = pickle_buffers(UnpickleableError(obj))
+        return (), buffers, True
+
+    def deserial(self, serialized: tuple, context: dict, subs: list):
+        return unpickle_buffers(subs)
+
+
+ExceptionSerializer.register(Exception)

xoscar/serialization/mlx.py

@@ -0,0 +1,67 @@
+# 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.
+
+from typing import Any, List
+
+import numpy as np
+
+from ..utils import lazy_import
+from .core import Serializer, buffered
+
+mx = lazy_import("mlx.core")
+
+
+dtype_map = {
+    "b": np.int8,
+    "B": np.uint8,
+    "h": np.int16,
+    "H": np.uint16,
+    "i": np.int32,
+    "I": np.uint32,
+    "q": np.int64,
+    "Q": np.uint64,
+    "e": np.float16,
+    "f": np.float32,
+    "d": np.float64,
+}
+
+
+class MLXSerislizer(Serializer):
+    @buffered
+    def serial(self, obj: "mx.array", context: dict):  # type: ignore
+        ravel_obj = obj.reshape(-1).view(mx.uint8)
+        mv = memoryview(ravel_obj)
+        header = dict(
+            shape=obj.shape, format=mv.format, dtype=str(obj.dtype).rsplit(".", 1)[-1]
+        )
+        if not mv.c_contiguous:
+            # NOTE: we only consider c contiguous here,
+            # MLX has no way to create f contiguous arrays.
+            mv = memoryview(bytes(mv))
+        return (header,), [mv], True
+
+    def deserial(self, serialized: tuple, context: dict, subs: List[Any]):
+        header = serialized[0]
+        shape, format, dtype = header["shape"], header["format"], header["dtype"]
+        mv = memoryview(subs[0])
+        if mv.format != format:
+            dtype = dtype_map.get(format, np.uint8)
+            np_arr = np.frombuffer(mv, dtype=dtype).reshape(shape)  # parse
+            mv = memoryview(np_arr)  # recreate memoryview
+        ravel_array = mx.array(mv)
+        return ravel_array.view(getattr(mx, dtype)).reshape(shape)
+
+
+if mx is not None:
+    MLXSerislizer.register(mx.array)

xoscar/serialization/numpy.py

@@ -0,0 +1,82 @@
+# 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, Dict, List, Tuple
+
+import numpy as np
+
+from .core import Serializer, buffered, pickle_buffers, unpickle_buffers
+
+
+class NDArraySerializer(Serializer):
+    @buffered
+    def serial(self, obj: np.ndarray, context: Dict):
+        header: dict = {}
+        if obj.dtype.hasobject:
+            header["pickle"] = True
+            buffers = pickle_buffers(obj)
+            return (header,), buffers, True
+
+        order = "C"
+        if obj.flags.f_contiguous:
+            order = "F"
+        elif not obj.flags.c_contiguous:
+            obj = np.ascontiguousarray(obj)
+        try:
+            desc = np.lib.format.dtype_to_descr(obj.dtype)
+            dtype_new_order = None
+        except ValueError:
+            # for structured dtype, array[[field2, field1]] will create a view,
+            # and dtype_to_desc will fail due to the order
+            fields = obj.dtype.fields
+            new_fields = sorted(fields, key=lambda k: fields[k][1])
+            desc = np.lib.format.dtype_to_descr(obj.dtype[new_fields])
+            dtype_new_order = list(fields)
+        header.update(
+            dict(
+                pickle=False,
+                descr=desc,
+                dtype_new_order=dtype_new_order,
+                shape=list(obj.shape),
+                strides=list(obj.strides),
+                order=order,
+            )
+        )
+        return (header,), [memoryview(obj.ravel(order=order).view("uint8").data)], True  # type: ignore
+
+    def deserial(self, serialized: Tuple, context: Dict, subs: List[Any]):
+        header = serialized[0]
+        if header["pickle"]:
+            return unpickle_buffers(subs)
+
+        try:
+            dtype = np.lib.format.descr_to_dtype(header["descr"])
+        except AttributeError:  # pragma: no cover
+            # for older numpy versions, descr_to_dtype is not implemented
+            dtype = np.dtype(header["descr"])
+
+        dtype_new_order = header["dtype_new_order"]
+        if dtype_new_order:
+            dtype = dtype[dtype_new_order]
+        return np.ndarray(
+            shape=tuple(header["shape"]),
+            dtype=dtype,
+            buffer=subs[0],
+            strides=tuple(header["strides"]),
+            order=header["order"],
+        )
+
+
+NDArraySerializer.register(np.ndarray)

xoscar/serialization/pyfury.py

@@ -0,0 +1,37 @@
+import os
+import threading
+
+_fury = threading.local()
+_fury_not_installed = object()
+_register_class_list = set()
+
+
+def register_classes(*args):
+    instance = get_fury()
+    if instance is not None:
+        _register_class_list.update(args)
+        for c in _register_class_list:
+            instance.register_class(c)
+
+
+def get_fury():
+    if os.environ.get("USE_FURY") in ("1", "true", "True"):
+        instance = getattr(_fury, "instance", None)
+        if instance is _fury_not_installed:  # pragma: no cover
+            return None
+        if instance is not None:
+            return instance
+        else:
+            try:
+                import pyfury
+
+                _fury.instance = instance = pyfury.Fury(
+                    language=pyfury.Language.PYTHON, require_class_registration=False
+                )
+                for c in _register_class_list:  # pragma: no cover
+                    instance.register_class(c)
+                print("pyfury is enabled.")
+            except ImportError:  # pragma: no cover
+                print("pyfury is not installed.")
+                _fury.instance = _fury_not_installed
+            return instance

xoscar/serialization/scipy.py

@@ -0,0 +1,72 @@
+# 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, Dict, List, Tuple
+
+import numpy as np
+
+try:
+    import scipy.sparse as sps
+except ImportError:  # pragma: no cover
+    sps = None
+
+from .core import Serializer, buffered, deserialize, serialize
+
+
+class CsrMatrixSerializer(Serializer):
+    @buffered
+    def serial(self, obj: Any, context: Dict):
+        data_header, data_buffers = serialize(obj.data)
+        idx_header, idx_buffers = serialize(obj.indices)
+        indptr_header, indptr_buffers = serialize(obj.indptr)
+        header = (
+            data_header,  # data_header
+            len(data_buffers),  # data_buf_num
+            idx_header,  # idx_header
+            len(idx_buffers),  # idx_buf_num
+            indptr_header,  # indptr_header
+            obj.shape,  # shape
+        )
+        return header, data_buffers + idx_buffers + indptr_buffers, True
+
+    def deserial(self, serialized: Tuple, context: Dict, subs: List):
+        (
+            data_header,
+            data_buf_num,
+            idx_header,
+            idx_buf_num,
+            indptr_header,
+            shape,
+        ) = serialized
+        data_buffers = subs[:data_buf_num]
+        idx_buffers = subs[data_buf_num : data_buf_num + idx_buf_num]
+        indptr_buffers = subs[data_buf_num + idx_buf_num :]
+
+        data = deserialize(data_header, data_buffers)
+        indices = deserialize(idx_header, idx_buffers)
+        indptr = deserialize(indptr_header, indptr_buffers)
+        shape = tuple(shape)
+
+        empty_arr = np.zeros(0, dtype=data.dtype)
+
+        target_csr = sps.coo_matrix(
+            (empty_arr, (empty_arr,) * 2), dtype=data.dtype, shape=shape
+        ).tocsr()
+        target_csr.data, target_csr.indices, target_csr.indptr = data, indices, indptr
+        return target_csr
+
+
+if sps:  # pragma: no branch
+    CsrMatrixSerializer.register(sps.csr_matrix)

xoscar/serialization/torch.py

@@ -0,0 +1,180 @@
+# Copyright 2024 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 typing import Any, Dict, List, Tuple
+
+import numpy as np
+
+from ..utils import lazy_import
+from .core import Serializer, buffered
+
+# lazy import PyTorch to avoid enforced dependency
+torch = lazy_import("torch")
+cupy = lazy_import("cupy")
+
+
+def rmm_to_torch(buf):
+    cupy_arr = cupy.asarray(buf)  # zero-copy
+    torch_tensor = torch.utils.dlpack.from_dlpack(cupy_arr.toDlpack())  # zero-copy
+    return torch_tensor
+
+
+class TorchTensorSerializer(Serializer):
+    @buffered
+    def serial(self, obj: "torch.Tensor", context: Dict):  # type: ignore
+        # for cpu tensor, use memory viewpoint
+        if obj.device.type == "cpu":
+            # make sure tensor is contiguous
+            if not obj.is_contiguous():
+                obj = obj.contiguous()
+            # get memory viewpoint and collect header information
+            header = {
+                "shape": tuple(obj.shape),
+                "dtype": str(obj.dtype),
+                "device": obj.device.type,
+                "requires_grad": obj.requires_grad,
+                "strides": tuple(obj.stride()),
+            }
+            # Try zero-copy path: torch -> numpy
+            try:
+                np_array = obj.numpy()  # zero-copy if supported
+                buffer = memoryview(np_array)
+                header["format"] = "numpy"
+                return (header,), [buffer], True
+
+            except Exception:
+                # Fallback: copy to uint8 bytes
+                # This works for ANY dtype
+                byte_tensor = obj.view(torch.uint8).clone()  # explicit copy
+                np_array = byte_tensor.numpy()
+                buffer = memoryview(np_array)
+                header["format"] = "bytes"
+                return (header,), [buffer], True
+        elif obj.device.type == "cuda":
+            # for CUDA， use __cuda_array_interface__
+            if not (
+                obj.is_contiguous()
+                or obj.is_contiguous(memory_format=torch.channels_last)
+            ):
+                obj = obj.contiguous()
+
+            # get cuda array interface information
+            header = {
+                "shape": tuple(obj.shape),
+                "dtype": str(obj.dtype),
+                "device": obj.device.type,
+                "device_index": obj.device.index,
+                "requires_grad": obj.requires_grad,
+                "strides": tuple(obj.stride()),
+            }
+
+            # ---- Core idea: expose raw CUDA memory as a uint8 buffer (zero-copy) ----
+            # Get the underlying untyped storage that actually owns the CUDA memory
+            storage = obj.untyped_storage()
+
+            # Create a uint8 CUDA tensor that will act as a byte-view of the same memory
+            # This does NOT allocate new GPU memory; it only creates a new Tensor wrapper
+            buffer = torch.empty(
+                (storage.nbytes(),),
+                dtype=torch.uint8,
+                device=obj.device,
+            )
+
+            # Make the buffer tensor share the same storage (zero-copy)
+            buffer.set_(storage)  # type: ignore[attr-defined]
+
+            # Return: (metadata,), [raw CUDA buffer], mark as buffered
+            return (header,), [buffer], True
+        else:
+            # for unsupported device
+            raise NotImplementedError(f"Unsupported device type: {obj.device.type}")
+
+    def deserial(self, serialized: Tuple, context: Dict, subs: List[Any]):
+        header = serialized[0]
+        device = header["device"]
+        data_buffer = subs[0]
+
+        if device == "cpu":
+            fmt = header.get("format", "numpy")
+            dtype_name = header["dtype"].split(".")[-1]
+            dtype = getattr(torch, dtype_name)
+
+            if fmt == "numpy":
+                # zero-copy path
+                np_array = np.frombuffer(data_buffer, dtype=np.dtype(dtype_name))
+                tensor = torch.from_numpy(np_array).view(header["shape"])
+
+            else:
+                # bytes path (copy unavoidable)
+                byte_np = np.frombuffer(data_buffer, dtype=np.uint8)
+                byte_tensor = torch.from_numpy(byte_np)
+
+                tensor = torch.empty(0, dtype=dtype)
+                tensor = tensor.set_(
+                    byte_tensor.untyped_storage(),
+                    storage_offset=0,
+                    size=tuple(header["shape"]),
+                    stride=tuple(header["strides"]),
+                )
+
+            if header.get("requires_grad"):
+                tensor.requires_grad_(True)
+
+        elif device == "cuda":
+            # Unpack metadata
+            (header,) = serialized
+
+            # Raw CUDA buffer (uint8 tensor)
+            buffer = subs[0]
+            if not isinstance(buffer, torch.Tensor):
+                buffer = rmm_to_torch(buffer)
+            assert buffer.is_cuda, "buffer must be a CUDA tensor"
+            assert buffer.dtype == torch.uint8, "buffer must be uint8"
+
+            # Get the shared CUDA storage
+            storage = buffer.untyped_storage()
+
+            # Resolve original dtype
+            dtype_name = header["dtype"].split(".")[-1]
+            dtype = getattr(torch, dtype_name)
+
+            # Create an empty tensor wrapper with the correct dtype
+            # This does NOT allocate new GPU memory for data
+            tensor = torch.empty(0, device=buffer.device, dtype=dtype)
+
+            # Bind the tensor to the same storage with original shape/stride
+            # Pure zero-copy: only a new Tensor view, no data movement
+            tensor.set_(
+                storage,
+                storage_offset=0,
+                size=tuple(header["shape"]),
+                stride=tuple(header["strides"]),
+            )
+
+            # Restore requires_grad if needed
+            if header.get("requires_grad"):
+                tensor.requires_grad_(True)
+
+            return tensor
+        else:
+            raise NotImplementedError(f"Unsupported device type: {device}")
+
+        # recover requires_grad attributes
+        tensor.requires_grad = header["requires_grad"]
+        return tensor
+
+
+# only when torch is available, we register module
+if torch is not None:
+    TorchTensorSerializer.register("torch.Tensor")