maxframe 0.1.0b5__cp39-cp39-macosx_11_0_arm64.whl

maxframe/protocol.py

@@ -0,0 +1,415 @@
+# 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 base64
+import enum
+import uuid
+from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+
+import pandas as pd
+
+from .core import OutputType, TileableGraph
+from .core.graph.entity import SerializableGraph
+from .lib import wrapped_pickle as pickle
+from .lib.tblib import pickling_support
+from .serialization import PickleContainer, RemoteException, pickle_buffers
+from .serialization.serializables import (
+    AnyField,
+    BoolField,
+    BytesField,
+    DictField,
+    EnumField,
+    FieldTypes,
+    Float64Field,
+    ListField,
+    ReferenceField,
+    Serializable,
+    SeriesField,
+    StringField,
+    TupleField,
+)
+
+pickling_support.install()
+
+BodyType = TypeVar("BodyType", bound="Serializable")
+
+
+class JsonSerializable(Serializable):
+    _ignore_non_existing_keys = True
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "JsonSerializable":
+        raise NotImplementedError
+
+    def to_json(self) -> dict:
+        raise NotImplementedError
+
+
+class ProtocolBody(Generic[BodyType], Serializable):
+    request_id: bytes = BytesField(
+        "request_id", default_factory=lambda: uuid.uuid4().bytes
+    )
+    body: BodyType = AnyField("body", default=None)
+
+
+class DagStatus(enum.Enum):
+    PREPARING = 0
+    RUNNING = 1
+    SUCCEEDED = 2
+    FAILED = 3
+    CANCELLING = 4
+    CANCELLED = 5
+
+
+class DimensionIndex(Serializable):
+    is_slice: bool = BoolField("is_slice", default=None)
+    is_int_index: bool = BoolField("is_int_index", default=None)
+    data: List = ListField("data", default=None)
+
+
+class ResultType(enum.Enum):
+    NULL = 0
+    ODPS_TABLE = 1
+    ODPS_VOLUME = 2
+
+
+class DataSerializeType(enum.Enum):
+    PICKLE = 0
+
+
+class VolumeDataMeta(Serializable):
+    output_type: OutputType = EnumField(
+        "output_type", OutputType, FieldTypes.int8, default=None
+    )
+    serial_type: DataSerializeType = EnumField(
+        "serial_type", DataSerializeType, FieldTypes.int8, default=None
+    )
+    shape: Tuple[int, ...] = TupleField("shape", FieldTypes.int64, default=None)
+    nsplits: Tuple[Tuple[int, ...], ...] = TupleField(
+        "nsplits", FieldTypes.tuple(FieldTypes.tuple(FieldTypes.int64)), default=None
+    )
+
+
+_result_type_to_info_cls: Dict[ResultType, Type["ResultInfo"]] = dict()
+
+
+class ResultInfo(JsonSerializable):
+    _result_type = ResultType.NULL
+
+    result_type: Optional[ResultType] = EnumField(
+        "result_type", ResultType, default=ResultType.NULL
+    )
+    slices: Optional[List[DimensionIndex]] = ListField(
+        "slices", FieldTypes.reference, default=None
+    )
+
+    @classmethod
+    def _cls_from_result_type(cls, result_type: ResultType):
+        if not _result_type_to_info_cls:
+            for dest_cls in globals().values():
+                if isinstance(dest_cls, type) and issubclass(dest_cls, ResultInfo):
+                    _result_type_to_info_cls[dest_cls._result_type] = dest_cls
+        return _result_type_to_info_cls[result_type]
+
+    @classmethod
+    def _json_to_kwargs(cls, serialized: dict) -> dict:
+        # todo retrieve slices from json once implemented
+        kw = serialized.copy()
+        kw["result_type"] = ResultType(kw["result_type"])
+        return kw
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "ResultInfo":
+        res_type = ResultType(serialized["result_type"])
+        res_cls = cls._cls_from_result_type(res_type)
+        return res_cls(**res_cls._json_to_kwargs(serialized))
+
+    def to_json(self) -> dict:
+        # todo convert slices to json once implemented
+        return {"result_type": self.result_type.value if self.result_type else None}
+
+
+ResultInfoType = TypeVar("ResultInfoType", bound=ResultInfo)
+
+
+class ODPSTableResultInfo(ResultInfo):
+    _result_type = ResultType.ODPS_TABLE
+
+    full_table_name: str = StringField("full_table_name", default=None)
+    partition_specs: Optional[List[str]] = ListField(
+        "partition_specs", FieldTypes.string, default=None
+    )
+
+    def __init__(self, result_type: ResultType = None, **kw):
+        result_type = result_type or ResultType.ODPS_TABLE
+        super().__init__(result_type=result_type, **kw)
+
+    def to_json(self) -> dict:
+        ret = super().to_json()
+        ret["full_table_name"] = self.full_table_name
+        if self.partition_specs:
+            ret["partition_specs"] = self.partition_specs
+        return ret
+
+
+class ODPSVolumeResultInfo(ResultInfo):
+    _result_type = ResultType.ODPS_VOLUME
+
+    volume_name: str = StringField("volume_name", default=None)
+    volume_path: str = StringField("volume_path", default=None)
+
+    def __init__(self, result_type: ResultType = None, **kw):
+        result_type = result_type or ResultType.ODPS_VOLUME
+        super().__init__(result_type=result_type, **kw)
+
+    def to_json(self) -> dict:
+        ret = super().to_json()
+        ret["volume_name"] = self.volume_name
+        ret["volume_path"] = self.volume_path
+        return ret
+
+
+class ErrorSource(enum.Enum):
+    PYTHON = 0
+
+
+class ErrorInfo(JsonSerializable):
+    error_messages: Optional[List[str]] = ListField("error_messages", FieldTypes.string)
+    error_tracebacks: Optional[List[List[str]]] = ListField(
+        "error_tracebacks", FieldTypes.list
+    )
+    raw_error_source: ErrorSource = EnumField(
+        "raw_error_source", ErrorSource, FieldTypes.int8
+    )
+    raw_error_data: Optional[Exception] = AnyField("raw_error_data")
+
+    @classmethod
+    def from_exception(cls, exc: Exception):
+        remote_exc = RemoteException.from_exception(exc)
+        messages, tracebacks = remote_exc.messages, remote_exc.tracebacks
+        return cls(messages, tracebacks, ErrorSource.PYTHON, exc)
+
+    def reraise(self):
+        if self.raw_error_source == ErrorSource.PYTHON:
+            raise self.raw_error_data
+        raise RemoteException(self.error_messages, self.error_tracebacks, [])
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "ErrorInfo":
+        kw = serialized.copy()
+        kw["raw_error_source"] = ErrorSource(serialized["raw_error_source"])
+        if kw.get("raw_error_data"):
+            bufs = [base64.b64decode(s) for s in kw["raw_error_data"]]
+            try:
+                kw["raw_error_data"] = pickle.loads(bufs[0], buffers=bufs[1:])
+            except:
+                kw["raw_error_data"] = None
+        return cls(**kw)
+
+    def to_json(self) -> dict:
+        ret = {
+            "error_messages": self.error_messages,
+            "error_tracebacks": self.error_tracebacks,
+            "raw_error_source": self.raw_error_source.value,
+        }
+        err_data_bufs = None
+        if isinstance(self.raw_error_data, (PickleContainer, RemoteException)):
+            err_data_bufs = self.raw_error_data.get_buffers()
+        elif isinstance(self.raw_error_data, BaseException):
+            err_data_bufs = pickle_buffers(self.raw_error_data)
+        if err_data_bufs:
+            ret["raw_error_data"] = [
+                base64.b64encode(s).decode() for s in err_data_bufs
+            ]
+        return ret
+
+
+class DagInfo(JsonSerializable):
+    session_id: str = StringField("session_id", default=None)
+    dag_id: str = StringField("dag_id", default=None)
+    status: DagStatus = EnumField("status", DagStatus, FieldTypes.int8, default=None)
+    progress: float = Float64Field("progress", default=None)
+    tileable_to_result_infos: Dict[str, ResultInfo] = DictField(
+        "tileable_to_result_infos",
+        FieldTypes.string,
+        FieldTypes.reference,
+        default_factory=dict,
+    )
+    error_info: Optional[ErrorInfo] = ReferenceField("error_info", default=None)
+    start_timestamp: Optional[float] = Float64Field("start_timestamp", default=None)
+    end_timestamp: Optional[float] = Float64Field("end_timestamp", default=None)
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "DagInfo":
+        kw = serialized.copy()
+        kw["status"] = DagStatus(kw["status"])
+        if kw.get("tileable_to_result_infos"):
+            kw["tileable_to_result_infos"] = {
+                k: ResultInfo.from_json(s)
+                for k, s in kw["tileable_to_result_infos"].items()
+            }
+        if kw.get("error_info"):
+            kw["error_info"] = ErrorInfo.from_json(kw["error_info"])
+        return DagInfo(**kw)
+
+    def to_json(self) -> dict:
+        ret = {
+            "session_id": self.session_id,
+            "dag_id": self.dag_id,
+            "status": self.status.value,
+            "progress": self.progress,
+            "start_timestamp": self.start_timestamp,
+            "end_timestamp": self.end_timestamp,
+        }
+        ret = {k: v for k, v in ret.items() if v is not None}
+        if self.tileable_to_result_infos:
+            ret["tileable_to_result_infos"] = {
+                k: v.to_json() for k, v in self.tileable_to_result_infos.items()
+            }
+        if self.error_info:
+            ret["error_info"] = self.error_info.to_json()
+        return ret
+
+
+class CreateSessionRequest(Serializable):
+    settings: Dict[str, Any] = DictField("settings", default=None)
+
+
+class SessionInfo(JsonSerializable):
+    session_id: str = StringField("session_id")
+    settings: Dict[str, Any] = DictField(
+        "settings", key_type=FieldTypes.string, default=None
+    )
+    start_timestamp: float = Float64Field("start_timestamp", default=None)
+    idle_timestamp: float = Float64Field("idle_timestamp", default=None)
+    dag_infos: Dict[str, Optional[DagInfo]] = DictField(
+        "dag_infos",
+        key_type=FieldTypes.string,
+        value_type=FieldTypes.reference,
+        default=None,
+    )
+    error_info: Optional[ErrorInfo] = ReferenceField("error_info", default=None)
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "SessionInfo":
+        kw = serialized.copy()
+        if kw.get("dag_infos"):
+            kw["dag_infos"] = {
+                k: DagInfo.from_json(v) for k, v in kw["dag_infos"].items()
+            }
+        if kw.get("error_info"):
+            kw["error_info"] = ErrorInfo.from_json(kw["error_info"])
+        return SessionInfo(**kw)
+
+    def to_json(self) -> dict:
+        ret = {
+            "session_id": self.session_id,
+            "settings": self.settings,
+            "start_timestamp": self.start_timestamp,
+            "idle_timestamp": self.idle_timestamp,
+        }
+        if self.dag_infos:
+            ret["dag_infos"] = {k: v.to_json() for k, v in self.dag_infos.items()}
+        if self.error_info:
+            ret["error_info"] = self.error_info.to_json()
+        return ret
+
+
+class ExecuteDagRequest(Serializable):
+    session_id: str = StringField("session_id")
+    dag: TileableGraph = ReferenceField(
+        "dag",
+        on_serialize=SerializableGraph.from_graph,
+        on_deserialize=lambda x: x.to_graph(),
+        default=None,
+    )
+    managed_input_infos: Dict[str, ResultInfo] = DictField(
+        "managed_input_infos",
+        key_type=FieldTypes.string,
+        value_type=FieldTypes.reference,
+        default=None,
+    )
+
+
+class SubDagInfo(Serializable):
+    subdag_id: str = StringField("subdag_id")
+    status: DagStatus = EnumField("status", DagStatus, FieldTypes.int8, default=None)
+    progress: float = Float64Field("progress", default=None)
+    error_info: Optional[ErrorInfo] = ReferenceField(
+        "error_info", reference_type=ErrorInfo, default=None
+    )
+    tileable_to_result_infos: Dict[str, ResultInfo] = DictField(
+        "tileable_to_result_infos",
+        FieldTypes.string,
+        FieldTypes.reference,
+        default_factory=dict,
+    )
+
+
+class ExecuteSubDagRequest(Serializable):
+    dag: TileableGraph = ReferenceField(
+        "dag",
+        on_serialize=SerializableGraph.from_graph,
+        on_deserialize=lambda x: x.to_graph(),
+        default=None,
+    )
+    settings: Dict[str, Any] = DictField("settings", default=None)
+
+
+class DecrefRequest(Serializable):
+    keys: List[str] = ListField("keys", FieldTypes.string, default=None)
+
+
+class DataFrameTableMeta(Serializable):
+    __slots__ = "_pd_column_names", "_pd_index_level_names"
+
+    table_name: Optional[str] = StringField("table_name", default=None)
+    type: OutputType = EnumField("type", OutputType, FieldTypes.int8, default=None)
+    table_column_names: List[str] = ListField(
+        "table_column_names", FieldTypes.string, default=None
+    )
+    table_index_column_names: List[str] = ListField(
+        "table_index_column_names", FieldTypes.string, default=None
+    )
+    pd_column_dtypes: pd.Series = SeriesField("pd_column_dtypes", default=None)
+    pd_column_level_names: List[Any] = ListField("pd_column_level_names", default=None)
+    pd_index_dtypes: pd.Series = SeriesField("pd_index_dtypes", default=None)
+
+    @property
+    def pd_column_names(self) -> List[Any]:
+        try:
+            return self._pd_column_names
+        except AttributeError:
+            self._pd_column_names = self.pd_column_dtypes.index.tolist()
+            return self._pd_column_names
+
+    @property
+    def pd_index_level_names(self) -> List[Any]:
+        try:
+            return self._pd_index_level_names
+        except AttributeError:
+            self._pd_index_level_names = self.pd_index_dtypes.index.tolist()
+            return self._pd_index_level_names
+
+    def __eq__(self, other: "Serializable") -> bool:
+        if not isinstance(other, type(self)):
+            return False
+        for k in self._FIELDS:
+            v = getattr(self, k, None)
+            is_same = v == getattr(other, k, None)
+            if callable(getattr(is_same, "all", None)):
+                is_same = is_same.all()
+            if not is_same:
+                return False
+        return True

maxframe/remote/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+# noinspection PyUnresolvedReferences
+from ..core import ExecutableTuple
+from .core import spawn
+from .run_script import run_script

maxframe/remote/core.py

@@ -0,0 +1,210 @@
+# 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 functools import partial
+
+from .. import opcodes
+from ..core import ENTITY_TYPE, ChunkData
+from ..core.operator import ObjectOperator, ObjectOperatorMixin
+from ..dataframe.core import DATAFRAME_TYPE, INDEX_TYPE, SERIES_TYPE
+from ..serialization.serializables import (
+    BoolField,
+    DictField,
+    FunctionField,
+    Int32Field,
+    ListField,
+)
+from ..tensor.core import TENSOR_TYPE
+from ..utils import build_fetch_tileable, find_objects, replace_objects
+
+
+class RemoteFunction(ObjectOperatorMixin, ObjectOperator):
+    _op_type_ = opcodes.REMOTE_FUNCATION
+    _op_module_ = "remote"
+
+    function = FunctionField("function")
+    function_args = ListField("function_args")
+    function_kwargs = DictField("function_kwargs")
+    retry_when_fail = BoolField("retry_when_fail")
+    resolve_tileable_input = BoolField("resolve_tileable_input", default=False)
+    n_output = Int32Field("n_output", default=None)
+
+    @property
+    def output_limit(self):
+        return self.n_output or 1
+
+    @property
+    def retryable(self) -> bool:
+        return self.retry_when_fail
+
+    @classmethod
+    def _no_prepare(cls, tileable):
+        return isinstance(
+            tileable, (TENSOR_TYPE, DATAFRAME_TYPE, SERIES_TYPE, INDEX_TYPE)
+        )
+
+    def _set_inputs(self, inputs):
+        raw_inputs = getattr(self, "_inputs", None)
+        super()._set_inputs(inputs)
+
+        function_inputs = iter(inp for inp in self._inputs)
+        mapping = {inp: new_inp for inp, new_inp in zip(inputs, self._inputs)}
+        if raw_inputs is not None:
+            for raw_inp in raw_inputs:
+                if self._no_prepare(raw_inp):
+                    if not isinstance(self._inputs[0], ChunkData):
+                        # not in tile, set_inputs from tileable
+                        mapping[raw_inp] = next(function_inputs)
+                    else:
+                        # in tile, set_inputs from chunk
+                        mapping[raw_inp] = build_fetch_tileable(raw_inp)
+                else:
+                    mapping[raw_inp] = next(function_inputs)
+        self.function_args = replace_objects(self.function_args, mapping)
+        self.function_kwargs = replace_objects(self.function_kwargs, mapping)
+
+    def __call__(self):
+        find_inputs = partial(find_objects, types=ENTITY_TYPE)
+        inputs = find_inputs(self.function_args) + find_inputs(self.function_kwargs)
+        if self.n_output is None:
+            return self.new_tileable(inputs)
+        else:
+            return self.new_tileables(
+                inputs, kws=[dict(i=i) for i in range(self.n_output)]
+            )
+
+
+def spawn(
+    func,
+    args=(),
+    kwargs=None,
+    retry_when_fail=False,
+    resolve_tileable_input=False,
+    n_output=None,
+    **kw,
+):
+    """
+    Spawn a function and return a MaxFrame Object which can be executed later.
+
+    Parameters
+    ----------
+    func : function
+        Function to spawn.
+    args: tuple
+       Args to pass to function
+    kwargs: dict
+       Kwargs to pass to function
+    retry_when_fail: bool, default False
+       If True, retry when function failed.
+    resolve_tileable_input: bool default False
+       If True, resolve tileable inputs as values.
+    n_output: int
+       Count of outputs for the function
+
+    Returns
+    -------
+    Object
+        MaxFrame Object.
+
+    Examples
+    --------
+    >>> import maxframe.remote as mr
+    >>> def inc(x):
+    >>>     return x + 1
+    >>>
+    >>> result = mr.spawn(inc, args=(0,))
+    >>> result
+    Object <op=RemoteFunction, key=e0b31261d70dd9b1e00da469666d72d9>
+    >>> result.execute().fetch()
+    1
+
+    List of spawned functions can be converted to :class:`maxframe.remote.ExecutableTuple`,
+    and `.execute()` can be called to run together.
+
+    >>> results = [mr.spawn(inc, args=(i,)) for i in range(10)]
+    >>> mr.ExecutableTuple(results).execute().fetch()
+    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+
+    MaxFrame Object returned by :meth:`maxframe.remote.spawn` can be treated
+    as arguments for other spawn functions.
+
+    >>> results = [mr.spawn(inc, args=(i,)) for i in range(10)]   # list of spawned functions
+    >>> def sum_all(xs):
+            return sum(xs)
+    >>> mr.spawn(sum_all, args=(results,)).execute().fetch()
+    55
+
+    inside a spawned function, new functions can be spawned.
+
+    >>> def driver():
+    >>>     results = [mr.spawn(inc, args=(i,)) for i in range(10)]
+    >>>     return mr.ExecutableTuple(results).execute().fetch()
+    >>>
+    >>> mr.spawn(driver).execute().fetch()
+    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+
+    MaxFrame tensor, DataFrame and so forth is available in spawned functions as well.
+
+    >>> import maxframe.tensor as mt
+    >>> def driver2():
+    >>>     t = mt.random.rand(10, 10)
+    >>>     return t.sum().to_numpy()
+    >>>
+    >>> mr.spawn(driver2).execute().fetch()
+    52.47844223908132
+
+    Argument of `n_output` can indicate that the spawned function will return multiple outputs.
+    This is important when some of the outputs may be passed to different functions.
+
+    >>> def triage(alist):
+    >>>     ret = [], []
+    >>>     for i in alist:
+    >>>         if i < 0.5:
+    >>>             ret[0].append(i)
+    >>>         else:
+    >>>             ret[1].append(i)
+    >>>     return ret
+    >>>
+    >>> def sum_all(xs):
+    >>>     return sum(xs)
+    >>>
+    >>> l = [0.4, 0.7, 0.2, 0.8]
+    >>> la, lb = mr.spawn(triage, args=(l,), n_output=2)
+    >>>
+    >>> sa = mr.spawn(sum_all, args=(la,))
+    >>> sb = mr.spawn(sum_all, args=(lb,))
+    >>> mr.ExecutableTuple([sa, sb]).execute().fetch()
+    >>> [0.6000000000000001, 1.5]
+    """
+    if not isinstance(args, tuple):
+        args = [args]
+    else:
+        args = list(args)
+    if kwargs is None:
+        kwargs = dict()
+    if not isinstance(kwargs, dict):
+        raise TypeError("kwargs has to be a dict")
+
+    op = RemoteFunction(
+        function=func,
+        function_args=args,
+        function_kwargs=kwargs,
+        retry_when_fail=retry_when_fail,
+        resolve_tileable_input=resolve_tileable_input,
+        n_output=n_output,
+        **kw,
+    )
+    if op.extra_params:
+        raise ValueError(f"Unexpected kw: {list(op.extra_params)[0]}")
+    return op()