xoscar 0.9.0__cp312-cp312-macosx_10_13_x86_64.whl

xoscar/profiling.py

@@ -0,0 +1,260 @@
+# 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.
+
+import asyncio
+import copy
+import heapq
+import json
+import logging
+import operator
+import os
+from collections import Counter
+from collections.abc import Mapping
+
+from .backends.message import SendMessage, TellMessage
+
+logger = logging.getLogger(__name__)
+
+XOSCAR_ENABLE_PROFILING = int(os.environ.get("XOSCAR_ENABLE_PROFILING", 0))
+
+
+class _ProfilingOptionDescriptor:
+    def __init__(self, _type, default):
+        self._name = None
+        self._type = _type
+        self._default = default
+
+    def __get__(self, obj, cls):
+        if obj is None:
+            return self
+        v = obj._options.get(self._name)
+        if v is None:
+            v = os.environ.get(f"XOSCAR_PROFILING_{self._name.upper()}", self._default)
+        if v is not None:
+            v = self._type(v)
+        # Cache the value.
+        obj.__dict__[self._name] = v
+        return v
+
+    def set_name(self, name: str):
+        self._name = name
+
+
+class _ProfilingOptionsMeta(type):
+    def __init__(cls, name, bases, classdict):
+        super(_ProfilingOptionsMeta, cls).__init__(name, bases, classdict)
+        for k, v in classdict.items():
+            if isinstance(v, _ProfilingOptionDescriptor):
+                v.set_name(k)
+
+
+class _ProfilingOptions(metaclass=_ProfilingOptionsMeta):
+    debug_interval_seconds = _ProfilingOptionDescriptor(float, default=None)
+    slow_calls_duration_threshold = _ProfilingOptionDescriptor(int, default=1)
+
+    def __init__(self, options):
+        if isinstance(options, Mapping):
+            invalid_keys = options.keys() - type(self).__dict__.keys()
+            if invalid_keys:
+                raise ValueError(f"Invalid profiling options: {invalid_keys}")
+            self._options = options
+        elif options in (True, False, None):
+            self._options = {}
+        else:
+            raise ValueError(f"Invalid profiling options: {options}")
+
+
+class DummyOperator:
+    @staticmethod
+    def set(key, value):
+        pass
+
+    @staticmethod
+    def inc(key, value):
+        pass
+
+    @staticmethod
+    def nest(key):
+        return DummyOperator
+
+    @staticmethod
+    def values():
+        return []
+
+    @staticmethod
+    def empty():
+        return True
+
+
+class ProfilingDataOperator:
+    __slots__ = ("_target",)
+
+    def __init__(self, target):
+        self._target = target
+
+    def set(self, key, value):
+        self._target[key] = value
+
+    def inc(self, key, value):
+        old = self._target.get(key, 0)
+        self._target[key] = old + value
+
+    def nest(self, key):
+        v = self._target.setdefault(key, {})
+        if not isinstance(v, dict):
+            raise TypeError(
+                f"The value type of key {key} is {type(v)}, but a dict is expected."
+            )
+        return ProfilingDataOperator(v)
+
+    def values(self):
+        return self._target.values()
+
+    def empty(self):
+        return len(self._target) == 0
+
+
+class _CallStats:
+    _call_counter: Counter
+    _slow_calls: list
+
+    def __init__(self, options: _ProfilingOptions):
+        self._options = options
+        self._call_counter = Counter()
+        self._slow_calls = []
+
+    def collect(self, message, duration: float):
+        key = (message.actor_ref.uid, message.content[0])
+        self._call_counter[key] += 1
+        if duration < self._options.slow_calls_duration_threshold:
+            return
+        slow_call_key = (
+            duration,
+            message.actor_ref.uid,
+            message.actor_ref.address,
+            message.content,
+        )
+        try:
+            if len(self._slow_calls) < 10:
+                heapq.heappush(self._slow_calls, slow_call_key)
+            else:
+                heapq.heapreplace(self._slow_calls, slow_call_key)
+        except TypeError:
+            pass
+
+    def to_dict(self) -> dict:
+        most_calls = {}
+        for name_tuple, count in self._call_counter.most_common(10):
+            uid, method_name = name_tuple
+            most_calls[f"{uid.decode('utf-8')}.{method_name}"] = count
+        slow_calls = {}
+        for duration, uid, address, content in sorted(
+            self._slow_calls, key=operator.itemgetter(0), reverse=True
+        ):
+            method_name, _batch, args, kwargs = content
+            slow_calls[
+                f"[{address}]{uid.decode('utf-8')}.{method_name}(args={args}, kwargs={kwargs})"
+            ] = duration
+        return {"most_calls": most_calls, "slow_calls": slow_calls}
+
+
+class _ProfilingData:
+    _instance = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def set_instance(cls, inst: "_ProfilingData"):
+        cls._instance = inst
+
+    def __init__(self):
+        self._data = {}
+        self._call_stats = {}
+        self._debug_task = {}
+
+    def init(self, task_id: str, options=None):
+        options = _ProfilingOptions(options)
+        logger.info(
+            "Init profiling data for task %s with debug interval seconds %s.",
+            task_id,
+            options.debug_interval_seconds,
+        )
+        self._data[task_id] = {
+            "general": {},
+            "serialization": {},
+            "most_calls": {},
+            "slow_calls": {},
+        }
+        self._call_stats[task_id] = _CallStats(options)
+
+        async def _debug_profiling_log():
+            while True:
+                try:
+                    r = self._data.get(task_id, None)
+                    if r is None:
+                        logger.info("Profiling debug log break.")
+                        break
+                    r = copy.copy(r)  # shadow copy is enough.
+                    r.update(self._call_stats.get(task_id).to_dict())
+                    logger.warning("Profiling debug:\n%s", json.dumps(r, indent=4))
+                except Exception:
+                    logger.exception("Profiling debug log failed.")
+                await asyncio.sleep(options.debug_interval_seconds)
+
+        if options.debug_interval_seconds is not None:
+            self._debug_task[task_id] = task = asyncio.create_task(
+                _debug_profiling_log()
+            )
+            task.add_done_callback(lambda _: self._debug_task.pop(task_id, None))
+
+    def pop(self, task_id: str):
+        logger.info("Pop profiling data of task %s.", task_id)
+        debug_task = self._debug_task.pop(task_id, None)
+        if debug_task is not None:
+            debug_task.cancel()
+        r = self._data.pop(task_id, None)
+        if r is not None:
+            r.update(self._call_stats.pop(task_id).to_dict())
+        return r
+
+    def collect_actor_call(self, message, duration: float):
+        if self._call_stats:
+            message_type = type(message)
+            if message_type is SendMessage or message_type is TellMessage:
+                for stats in self._call_stats.values():
+                    stats.collect(message, duration)
+
+    def __getitem__(self, item):
+        key = item if isinstance(item, tuple) else (item,)
+        v = None
+        d = self._data
+        for k in key:
+            v = d.get(k, None)
+            if v is None:
+                break
+            else:
+                d = v
+        return DummyOperator if v is None else ProfilingDataOperator(v)
+
+
+def get_profiling_data() -> _ProfilingData:
+    return _ProfilingData.get_instance()
+
+
+ProfilingData = get_profiling_data()

xoscar/serialization/__init__.py

@@ -0,0 +1,20 @@
+# 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 . import cuda, exception, mlx, numpy, scipy, torch
+from .aio import AioDeserializer, AioSerializer
+from .core import Serializer, deserialize, serialize, serialize_with_spawn
+
+del cuda, numpy, scipy, mlx, exception, torch

xoscar/serialization/aio.py

@@ -0,0 +1,141 @@
+# 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.
+
+import asyncio
+import struct
+from io import BytesIO
+from typing import Any, BinaryIO, Union
+
+import cloudpickle
+
+from ..utils import lazy_import
+from .core import deserialize, serialize_with_spawn
+
+rmm = lazy_import("rmm")
+
+DEFAULT_SERIALIZATION_VERSION = 1
+DEFAULT_SPAWN_THRESHOLD = 100
+BUFFER_SIZES_NAME = "buf_sizes"
+
+
+class AioSerializer:
+    def __init__(self, obj: Any, compress=0):
+        self._obj = obj
+        self._compress = compress
+
+    async def _get_buffers(self):
+        headers, buffers = await serialize_with_spawn(
+            self._obj, spawn_threshold=DEFAULT_SPAWN_THRESHOLD
+        )
+
+        def _is_cuda_buffer(buf: Union["rmm.DeviceBuffer", BinaryIO]):  # type: ignore
+            return hasattr(buf, "__cuda_array_interface__")
+
+        is_cuda_buffers = [_is_cuda_buffer(buf) for buf in buffers]
+        headers[0]["is_cuda_buffers"] = list(is_cuda_buffers)
+
+        # add buffer lengths into headers
+        headers[0][BUFFER_SIZES_NAME] = [
+            buf.nbytes if hasattr(buf, "nbytes") else len(buf) for buf in buffers
+        ]
+        header = cloudpickle.dumps(headers)
+
+        # gen header buffer
+        header_bio = BytesIO()
+        # write version first
+        header_bio.write(struct.pack("B", DEFAULT_SERIALIZATION_VERSION))
+        # write header length
+        header_bio.write(struct.pack("<Q", len(header)))
+        # write compression
+        header_bio.write(struct.pack("<H", self._compress))
+
+        out_buffers = [header_bio.getbuffer(), header]
+        out_buffers.extend(buffers)
+
+        return out_buffers
+
+    async def run(self):
+        return await self._get_buffers()
+
+
+MALFORMED_MSG = "Received malformed data, please check Xoscar version on both side."
+
+
+def get_header_length(header_bytes: bytes):
+    version = struct.unpack("B", header_bytes[:1])[0]
+    # now we only have default version
+    if version != DEFAULT_SERIALIZATION_VERSION:
+        # when version not matched,
+        # we will immediately abort the connection
+        # EOFError will be captured by channel
+        raise EOFError(MALFORMED_MSG)
+    # header length
+    header_length = struct.unpack("<Q", header_bytes[1:9])[0]
+    # compress
+    _ = struct.unpack("<H", header_bytes[9:])[0]
+    return header_length
+
+
+class AioDeserializer:
+    def __init__(self, file):
+        self._file = file
+
+    def _readexactly(self, n: int):
+        # asyncio StreamReader may not guarantee to read n bytes
+        # for it we need to call `readexactly` instead
+        read = (
+            self._file.readexactly
+            if hasattr(self._file, "readexactly")
+            else self._file.read
+        )
+        return read(n)
+
+    async def _get_obj_header_bytes(self):
+        try:
+            header_bytes = bytes(await self._readexactly(11))
+        except ConnectionResetError:
+            raise EOFError("Server may be closed")
+        if len(header_bytes) == 0:
+            raise EOFError("Received empty bytes")
+        header_length = get_header_length(header_bytes)
+        return await self._readexactly(header_length)
+
+    async def _get_obj(self):
+        header = cloudpickle.loads(await self._get_obj_header_bytes())
+        # get buffer size
+        buffer_sizes = header[0].pop(BUFFER_SIZES_NAME)
+        # get buffers
+        buffers = [await self._readexactly(size) for size in buffer_sizes]
+        # get num of objs
+        num_objs = header[0].get("_N", 0)
+
+        if num_objs <= DEFAULT_SPAWN_THRESHOLD:
+            return deserialize(header, buffers)
+        else:
+            return await asyncio.to_thread(deserialize, header, buffers)
+
+    async def run(self):
+        return await self._get_obj()
+
+    async def get_size(self):
+        # extract header
+        header_bytes = await self._get_obj_header_bytes()
+        header = cloudpickle.loads(header_bytes)
+        # get buffer size
+        buffer_sizes = header[0].pop(BUFFER_SIZES_NAME)
+        return 11 + len(header_bytes) + sum(buffer_sizes)
+
+    async def get_header(self):
+        return cloudpickle.loads(await self._get_obj_header_bytes())

xoscar/serialization/core.pxd

@@ -0,0 +1,28 @@
+# Copyright 2022-2023 XProbe Inc.
+# derived from copyright 1999-2022 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 Serializer:
+    cdef int _serializer_id
+
+    cpdef serial(self, object obj, dict context)
+    cpdef deserial(self, tuple serialized, dict context, list subs)
+    cpdef on_deserial_error(
+        self,
+        tuple serialized,
+        dict context,
+        list subs_serialized,
+        int error_index,
+        object exc,
+    )

xoscar/serialization/core.pyi

@@ -0,0 +1,57 @@
+# 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
+
+from concurrent.futures import Executor
+from typing import Any, Callable
+
+def buffered(func: Callable) -> Callable: ...
+def fast_id(obj: Any) -> int: ...
+
+class Serializer:
+    serializer_id: int
+    def serial(self, obj: Any, context: dict): ...
+    def deserial(self, serialized: tuple, context: dict, subs: list[Any]): ...
+    def on_deserial_error(
+        self,
+        serialized: tuple,
+        context: dict,
+        subs_serialized: list,
+        error_index: int,
+        exc: BaseException,
+    ): ...
+    @classmethod
+    def register(cls, obj_type, name: str | None = None): ...
+    @classmethod
+    def unregister(cls, obj_type): ...
+
+class Placeholder:
+    id: int
+    callbacks: list[Callable]
+    def __init__(self, id_: int): ...
+    def __hash__(self): ...
+    def __eq__(self, other): ...
+
+def serialize(obj: Any, context: dict | None = None): ...
+async def serialize_with_spawn(
+    obj: Any,
+    context: dict | None = None,
+    spawn_threshold: int = 100,
+    executor: Executor | None = None,
+): ...
+def deserialize(headers: list, buffers: list, context: dict | None = None): ...
+def pickle_buffers(obj: Any) -> list: ...
+def unpickle_buffers(buffers: list) -> Any: ...