maxframe 0.1.0b5__cp310-cp310-macosx_11_0_arm64.whl

maxframe/core/entity/executable.py

@@ -0,0 +1,337 @@
+# Copyright 1999-2024 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 atexit
+import concurrent.futures
+import queue
+import threading
+from typing import List
+from weakref import WeakKeyDictionary, ref
+
+from ...lib.aio import get_isolation
+from ...typing_ import SessionType, TileableType
+from ..mode import enter_mode
+
+
+class DecrefRunner:
+    def __init__(self):
+        self._decref_thread = None
+        self._queue = queue.Queue()
+
+    def start(self):
+        self._decref_thread = threading.Thread(
+            target=self._thread_body, name="DecrefThread"
+        )
+        self._decref_thread.daemon = True
+        self._decref_thread.start()
+
+    def _thread_body(self):
+        from ...session import SyncSession
+
+        while True:
+            key, session_ref, fut = self._queue.get()
+            if key is None:
+                break
+
+            session = session_ref()
+            if session is None:
+                fut.set_result(None)
+                continue
+            try:
+                s = SyncSession.from_isolated_session(session)
+                s.decref(key)
+                fut.set_result(None)
+            except (RuntimeError, ConnectionError, KeyError):
+                fut.set_result(None)
+            except (
+                Exception
+            ) as ex:  # pragma: no cover  # noqa: E722  # nosec  # pylint: disable=bare-except
+                fut.set_exception(ex)
+            finally:
+                del session
+
+    def stop(self):
+        if self._decref_thread:  # pragma: no branch
+            self._queue.put_nowait((None, None, None))
+            self._decref_thread.join(1)
+        self._decref_thread = None
+
+    def put(self, key: str, session_ref: ref):
+        if self._decref_thread is None:
+            self.start()
+
+        fut = concurrent.futures.Future()
+        self._queue.put_nowait((key, session_ref, fut))
+        return fut
+
+
+_decref_runner = DecrefRunner()
+atexit.register(_decref_runner.stop)
+
+
+class _TileableSession:
+    def __init__(self, tileable: TileableType, session: SessionType):
+        self._sess_id = id(session)
+        key = tileable.key
+
+        def cb(_, sess=ref(session)):
+            try:
+                cur_thread_ident = threading.current_thread().ident
+                decref_in_isolation = get_isolation().thread_ident == cur_thread_ident
+            except KeyError:
+                # isolation destroyed, no need to decref
+                return
+
+            fut = _decref_runner.put(key, sess)
+            if not decref_in_isolation:
+                # if decref in isolation, means that this tileable
+                # is not required for main thread, thus we do not need
+                # to wait for decref, otherwise, wait a bit
+                try:
+                    fut.result(0.5)
+                except concurrent.futures.TimeoutError:
+                    # ignore timeout
+                    pass
+
+        self.tileable = ref(tileable, cb)
+
+    def __eq__(self, other: "_TileableSession"):
+        return self._sess_id == other._sess_id
+
+
+class _TileableDataCleaner:
+    def __init__(self):
+        self._tileable_to_sessions = WeakKeyDictionary()
+
+    @enter_mode(build=True)
+    def register(self, tileable: TileableType, session: SessionType):
+        if tileable in self._tileable_to_sessions:
+            self._tileable_to_sessions[tileable].append(
+                _TileableSession(tileable, session)
+            )
+        else:
+            self._tileable_to_sessions[tileable] = [_TileableSession(tileable, session)]
+
+
+# we don't use __del__ to avoid potential Circular reference
+_cleaner = _TileableDataCleaner()
+
+
+def _get_session(executable: "_ExecutableMixin", session: SessionType = None):
+    from ...session import get_default_session
+
+    # if session is not specified, use default session
+    if session is None:
+        session = get_default_session()
+
+    return session
+
+
+class _ExecutableMixin:
+    __slots__ = ()
+    _executed_sessions: List[SessionType]
+
+    def execute(self, session: SessionType = None, **kw):
+        from ...session import execute
+
+        session = _get_session(self, session)
+        return execute(self, session=session, **kw)
+
+    def _check_session(self, session: SessionType, action: str):
+        if session is None:
+            if isinstance(self, tuple):
+                key = self[0].key
+            else:
+                key = self.key
+            raise ValueError(
+                f"Tileable object {key} must be executed first before {action}"
+            )
+
+    def _fetch(self, session: SessionType = None, **kw):
+        from ...session import fetch
+
+        session = _get_session(self, session)
+        self._check_session(session, "fetch")
+        return fetch(self, session=session, **kw)
+
+    def fetch(self, session: SessionType = None, **kw):
+        return self._fetch(session=session, **kw)
+
+    def fetch_log(
+        self,
+        session: SessionType = None,
+        offsets: List[int] = None,
+        sizes: List[int] = None,
+    ):
+        from ...session import fetch_log
+
+        session = _get_session(self, session)
+        self._check_session(session, "fetch_log")
+        return fetch_log(self, session=session, offsets=offsets, sizes=sizes)[0]
+
+    def _fetch_infos(self, fields=None, session=None, **kw):
+        from ...session import fetch_infos
+
+        session = _get_session(self, session)
+        self._check_session(session, "fetch_infos")
+        return fetch_infos(self, fields=fields, session=session, **kw)
+
+    def _attach_session(self, session: SessionType):
+        if session not in self._executed_sessions:
+            _cleaner.register(self, session)
+            self._executed_sessions.append(session)
+
+    def _detach_session(self, session: SessionType):
+        if session in self._executed_sessions:
+            sessions = _cleaner._tileable_to_sessions.get(self, [])
+            if sessions:
+                sessions.remove(_TileableSession(self, session))
+            if len(sessions) == 0:
+                del _cleaner._tileable_to_sessions[self]
+            self._executed_sessions.remove(session)
+
+
+class _ExecuteAndFetchMixin:
+    __slots__ = ()
+
+    def _execute_and_fetch(self, session: SessionType = None, **kw):
+        from ...session import ExecutionInfo, SyncSession, fetch
+
+        session = _get_session(self, session)
+        fetch_kwargs = kw.pop("fetch_kwargs", dict())
+        if session in self._executed_sessions:
+            # if has been executed, fetch directly.
+            return self.fetch(session=session, **fetch_kwargs)
+        ret = self.execute(session=session, **kw)
+        if isinstance(ret, ExecutionInfo):
+            # wait=False
+            aio_task = ret.aio_task
+
+            async def _wait():
+                await aio_task
+
+            def run():
+                asyncio.run_coroutine_threadsafe(_wait(), loop=ret.loop).result()
+                return fetch(self, session=session, **fetch_kwargs)
+
+            return SyncSession._execution_pool.submit(run)
+        else:
+            # wait=True
+            return self.fetch(session=session, **fetch_kwargs)
+
+
+class _ToObjectMixin(_ExecuteAndFetchMixin):
+    __slots__ = ()
+
+    def to_object(self, session: SessionType = None, **kw):
+        return self._execute_and_fetch(session=session, **kw)
+
+
+class ExecutableTuple(tuple, _ExecutableMixin, _ToObjectMixin):
+    def __init__(self, *args):
+        tuple.__init__(*args)
+
+        self._fields_to_idx = None
+        self._fields = None
+        self._raw_type = None
+
+        if len(args) == 1 and isinstance(args[0], tuple):
+            self._fields = getattr(args[0], "_fields", None)
+            if self._fields is not None:
+                self._raw_type = type(args[0])
+                self._fields_to_idx = {f: idx for idx, f in enumerate(self._fields)}
+
+        self._executed_sessions = []
+
+    def __getattr__(self, item):
+        if self._fields_to_idx is None or item not in self._fields_to_idx:
+            raise AttributeError(item)
+        return self[self._fields_to_idx[item]]
+
+    def __dir__(self):
+        result = list(super().__dir__()) + list(self._fields or [])
+        return sorted(result)
+
+    def __repr__(self):
+        if not self._fields:
+            return super().__repr__()
+        items = []
+        for k, v in zip(self._fields, self):
+            items.append(f"{k}={v!r}")
+        return "%s(%s)" % (self._raw_type.__name__, ", ".join(items))
+
+    def execute(self, session: SessionType = None, **kw):
+        from ...session import execute
+
+        if len(self) == 0:
+            return self
+
+        session = _get_session(self, session)
+        ret = execute(*self, session=session, **kw)
+
+        if session not in self._executed_sessions:
+            self._executed_sessions.append(session)
+
+        if kw.get("wait", True):
+            return self
+        else:
+            return ret
+
+    def _fetch(self, session: SessionType = None, **kw):
+        from ...session import fetch
+
+        session = _get_session(self, session)
+        self._check_session(session, "fetch")
+        return fetch(*self, session=session, **kw)
+
+    def _fetch_infos(self, fields=None, session=None, **kw):
+        from ...session import fetch_infos
+
+        session = _get_session(self, session)
+        self._check_session(session, "fetch_infos")
+        return fetch_infos(*self, fields=fields, session=session, **kw)
+
+    def fetch(self, session: SessionType = None, **kw):
+        if len(self) == 0:
+            return tuple()
+
+        session = _get_session(self, session)
+        ret = super().fetch(session=session, **kw)
+        if self._raw_type is not None:
+            ret = self._raw_type(*ret)
+        if len(self) == 1:
+            return (ret,)
+        return ret
+
+    def fetch_log(
+        self,
+        session: SessionType = None,
+        offsets: List[int] = None,
+        sizes: List[int] = None,
+    ):
+        from ...session import fetch_log
+
+        if len(self) == 0:
+            return []
+        session = self._get_session(session=session)
+        return fetch_log(*self, session=session, offsets=offsets, sizes=sizes)
+
+    def _get_session(self, session: SessionType = None):
+        if session is None:
+            for item in self:
+                session = _get_session(item, session)
+                if session is not None:
+                    return session
+        return session

maxframe/core/entity/fuse.py

@@ -0,0 +1,73 @@
+# Copyright 1999-2024 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 numpy as np
+
+from ...serialization.serializables import ReferenceField
+from .chunks import CHUNK_TYPE, Chunk, ChunkData
+
+
+class FuseChunkData(ChunkData):
+    __slots__ = ("_inited",)
+
+    _chunk = ReferenceField(
+        "chunk", CHUNK_TYPE, on_serialize=lambda x: x.data if hasattr(x, "data") else x
+    )
+
+    def __init__(self, *args, **kwargs):
+        self._inited = False
+        super().__init__(*args, **kwargs)
+        self._extra_params = {}
+        self._inited = True
+
+    @property
+    def chunk(self):
+        return self._chunk
+
+    @property
+    def composed(self):
+        # for compatibility, just return the topological ordering,
+        # once we apply optimization on the subgraph,
+        # `composed` is not needed any more and should be removed then.
+        assert getattr(self._op, "fuse_graph", None) is not None
+        fuse_graph = self._op.fuse_graph
+        return list(fuse_graph.topological_iter())
+
+    def __getattr__(self, attr):
+        if not self._inited:
+            return object.__getattribute__(self, attr)
+        if attr in self._extra_params:
+            return self._extra_params[attr]
+        try:
+            return getattr(self._chunk, attr)
+        except AttributeError:
+            return object.__getattribute__(self, attr)
+
+    def __setattr__(self, attr, value):
+        if attr == "params":
+            self._chunk.params = value
+        else:
+            super().__setattr__(attr, value)
+
+    @property
+    def nbytes(self):
+        return np.prod(self.shape) * self.dtype.itemsize
+
+
+class FuseChunk(Chunk):
+    __slots__ = ()
+    _allow_data_type_ = (FuseChunkData,)
+
+
+FUSE_CHUNK_TYPE = (FuseChunkData, FuseChunk)

maxframe/core/entity/objects.py

@@ -0,0 +1,100 @@
+# Copyright 1999-2024 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
+
+from ...serialization.serializables import FieldTypes, ListField
+from ...utils import skip_na_call
+from .chunks import Chunk, ChunkData
+from .core import Entity
+from .executable import _ToObjectMixin
+from .tileables import TileableData
+
+
+class ObjectChunkData(ChunkData):
+    # chunk whose data could be any serializable
+    __slots__ = ()
+    type_name = "Object"
+
+    def __init__(self, op=None, index=None, **kw):
+        super().__init__(_op=op, _index=index, **kw)
+
+    @property
+    def params(self) -> Dict[str, Any]:
+        # params return the properties which useful to rebuild a new chunk
+        return {
+            "index": self.index,
+        }
+
+    @params.setter
+    def params(self, new_params: Dict[str, Any]):
+        params = new_params.copy()
+        params.pop("index", None)  # index not needed to update
+        if params:  # pragma: no cover
+            raise TypeError(f"Unknown params: {list(params)}")
+
+    @classmethod
+    def get_params_from_data(cls, data: Any) -> Dict[str, Any]:
+        return dict()
+
+
+class ObjectChunk(Chunk):
+    __slots__ = ()
+    _allow_data_type_ = (ObjectChunkData,)
+    type_name = "Object"
+
+
+class ObjectData(TileableData, _ToObjectMixin):
+    __slots__ = ()
+    type_name = "Object"
+
+    # optional fields
+    _chunks = ListField(
+        "chunks",
+        FieldTypes.reference(ObjectChunkData),
+        on_serialize=skip_na_call(lambda x: [it.data for it in x]),
+        on_deserialize=skip_na_call(lambda x: [ObjectChunk(it) for it in x]),
+    )
+
+    def __init__(self, op=None, nsplits=None, **kw):
+        super().__init__(_op=op, _nsplits=nsplits, **kw)
+
+    def __repr__(self):
+        return f"Object <op={type(self.op).__name__}, key={self.key}>"
+
+    @property
+    def params(self):
+        # params return the properties which useful to rebuild a new tileable object
+        return dict()
+
+    @params.setter
+    def params(self, new_params: Dict[str, Any]):
+        params = new_params.copy()
+        if params:  # pragma: no cover
+            raise TypeError(f"Unknown params: {list(params)}")
+
+    def refresh_params(self):
+        # refresh params when chunks updated
+        # nothing needs to do for Object
+        pass
+
+
+class Object(Entity, _ToObjectMixin):
+    __slots__ = ()
+    _allow_data_type_ = (ObjectData,)
+    type_name = "Object"
+
+
+OBJECT_TYPE = (Object, ObjectData)
+OBJECT_CHUNK_TYPE = (ObjectChunk, ObjectChunkData)

maxframe/core/entity/output_types.py

@@ -0,0 +1,90 @@
+# Copyright 1999-2024 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 functools
+from enum import Enum
+
+from .fuse import FUSE_CHUNK_TYPE
+from .objects import OBJECT_TYPE
+
+
+class OutputType(Enum):
+    object = 1
+    tensor = 2
+    dataframe = 3
+    series = 4
+    index = 5
+    scalar = 6
+    categorical = 7
+    dataframe_groupby = 8
+    series_groupby = 9
+    df_or_series = 10
+
+    @classmethod
+    def serialize_list(cls, output_types):
+        return [ot.value for ot in output_types] if output_types is not None else None
+
+    @classmethod
+    def deserialize_list(cls, output_types):
+        return [cls(ot) for ot in output_types] if output_types is not None else None
+
+
+_OUTPUT_TYPE_TO_TILEABLE_TYPES = {OutputType.object: OBJECT_TYPE}
+_OUTPUT_TYPE_TO_FETCH_CLS = {}
+
+
+def register_output_types(output_type, tileable_types):
+    _OUTPUT_TYPE_TO_TILEABLE_TYPES[output_type] = tileable_types
+
+
+def register_fetch_class(output_type, fetch_cls, fetch_shuffle_cls):
+    _OUTPUT_TYPE_TO_FETCH_CLS[output_type] = (fetch_cls, fetch_shuffle_cls)
+
+
+def get_tileable_types(output_type):
+    return _OUTPUT_TYPE_TO_TILEABLE_TYPES[output_type]
+
+
+def get_fetch_class(output_type):
+    return _OUTPUT_TYPE_TO_FETCH_CLS[output_type]
+
+
+@functools.lru_cache(100)
+def _get_output_type_by_cls(cls):
+    for tp in OutputType.__members__.values():
+        try:
+            tileable_types = _OUTPUT_TYPE_TO_TILEABLE_TYPES[tp]
+            if issubclass(cls, tileable_types):
+                return tp
+        except KeyError:  # pragma: no cover
+            continue
+    raise TypeError("Output can only be tensor, dataframe or series")
+
+
+def get_output_types(*objs, unknown_as=None):
+    output_types = []
+    for obj in objs:
+        if obj is None:
+            continue
+        elif isinstance(obj, FUSE_CHUNK_TYPE):
+            obj = obj.chunk
+
+        try:
+            output_types.append(_get_output_type_by_cls(type(obj)))
+        except TypeError:
+            if unknown_as is not None:
+                output_types.append(unknown_as)
+            else:  # pragma: no cover
+                raise
+    return output_types