xoscar 0.7.15__cp39-cp39-manylinux2014_x86_64.manylinux_2_17_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)