maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/remote/core.py

@@ -0,0 +1,212 @@
+# Copyright 1999-2025 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 typing import List
+
+from .. import opcodes
+from ..core import ENTITY_TYPE, EntityData
+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 ..typing_ import TileableType
+from ..udf import BuiltinFunction
+from ..utils import 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 has_custom_code(self) -> bool:
+        return not isinstance(self.function, BuiltinFunction)
+
+    def check_inputs(self, inputs: List[TileableType]):
+        return
+
+    @classmethod
+    def _set_inputs(cls, op: "RemoteFunction", inputs: List[EntityData]):
+        raw_inputs = getattr(op, "_inputs", None)
+        super()._set_inputs(op, inputs)
+
+        function_inputs = iter(inp for inp in op._inputs)
+        mapping = {inp: new_inp for inp, new_inp in zip(inputs, op._inputs)}
+        if raw_inputs is not None:
+            for raw_inp in raw_inputs:
+                mapping[raw_inp] = next(function_inputs)
+        op.function_args = replace_objects(op.function_args, mapping)
+        op.function_kwargs = replace_objects(op.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=True,
+    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()

maxframe/remote/run_script.py

@@ -0,0 +1,124 @@
+# Copyright 1999-2025 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 os
+from typing import Any, BinaryIO, Dict, List, TextIO, Union
+
+from .. import opcodes
+from ..core import TILEABLE_TYPE, OutputType
+from ..core.operator import ObjectOperator, ObjectOperatorMixin
+from ..serialization.serializables import (
+    BoolField,
+    BytesField,
+    DictField,
+    Int32Field,
+    ListField,
+)
+from ..typing_ import SessionType, TileableType
+from ..utils import to_binary
+
+
+class RunScript(ObjectOperator, ObjectOperatorMixin):
+    _op_type_ = opcodes.RUN_SCRIPT
+
+    code: bytes = BytesField("code", default=None)
+    data: Dict[str, TileableType] = DictField("data", default=None)
+    retry_when_fail: bool = BoolField("retry_when_fail", default=None)
+    command_args: List[str] = ListField("command_args", default=None)
+    world_size: int = Int32Field("world_size", default=None)
+    rank: int = Int32Field("rank", default=None)
+
+    def __init__(self, command_args=None, **kw):
+        command_args = command_args or []
+        super().__init__(command_args=command_args, **kw)
+        if self.output_types is None:
+            self.output_types = [OutputType.object]
+
+    @property
+    def retryable(self):
+        return self.retry_when_fail
+
+    def has_custom_code(self) -> bool:
+        return True
+
+    def __call__(self, inputs):
+        return self.new_tileable(inputs)
+
+
+def _extract_inputs(data: Dict[str, TileableType] = None) -> List[TileableType]:
+    if data is not None and not isinstance(data, dict):
+        raise TypeError(
+            "`data` must be a dict whose key is variable name and value is data"
+        )
+
+    inputs = []
+    if data is not None:
+        for v in data.values():
+            if isinstance(v, TILEABLE_TYPE):
+                inputs.append(v)
+
+    return inputs
+
+
+def run_script(
+    script: Union[bytes, str, BinaryIO, TextIO],
+    data: Dict[str, TileableType] = None,
+    n_workers: int = 1,
+    command_argv: List[str] = None,
+    session: SessionType = None,
+    retry_when_fail: bool = True,
+    run_kwargs: Dict[str, Any] = None,
+):
+    """
+    Run script in MaxFrame cluster.
+
+    Parameters
+    ----------
+    script: str or file-like object
+        Script to run.
+    data: dict
+        Variable name to data.
+    n_workers: int
+        number of workers to run the script
+    command_argv: list
+        extra command args for script
+    session: MaxFrame session
+        if not provided, will use default one
+    retry_when_fail: bool, default False
+       If True, retry when function failed.
+    run_kwargs: dict
+        extra kwargs for session.run
+
+    Returns
+    -------
+    Object
+        MaxFrame Object.
+
+    """
+
+    if hasattr(script, "read"):
+        code = script.read()
+    else:
+        with open(os.path.abspath(script), "rb") as f:
+            code = f.read()
+
+    inputs = _extract_inputs(data)
+    op = RunScript(
+        data=data,
+        code=to_binary(code),
+        world_size=n_workers,
+        retry_when_fail=retry_when_fail,
+        command_args=command_argv,
+    )
+    return op(inputs).execute(session=session, **(run_kwargs or {}))

maxframe/serialization/__init__.py

@@ -0,0 +1,39 @@
+# Copyright 1999-2025 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 (
+    arrow,
+    blob,
+    exception,
+    maxframe_objects,
+    numpy,
+    pandas,
+    scipy,
+    serializables,
+)
+from .core import (
+    PickleContainer,
+    PickleHookOptions,
+    Serializer,
+    clear_type_cache,
+    deserialize,
+    load_member,
+    load_type,
+    pickle_buffers,
+    serialize,
+    serialize_with_spawn,
+)
+from .exception import RemoteException
+
+del arrow, blob, exception, maxframe_objects, numpy, pandas, scipy

maxframe/serialization/arrow.py

@@ -0,0 +1,107 @@
+# Copyright 1999-2025 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, Union
+
+from ..utils import arrow_type_from_str, extract_class_name
+from .core import Serializer, buffered
+
+try:
+    import pyarrow as pa
+
+    PA_ARRAY_TYPES = Union[pa.Array, pa.ChunkedArray]
+    PA_RECORD_TYPES = Union[pa.Table, pa.RecordBatch]
+except ImportError:  # pragma: no cover
+    pa = None
+    PA_ARRAY_TYPES = Any
+    PA_RECORD_TYPES = Any
+
+
+_TYPE_CHAR_ARROW_ARRAY = "A"
+_TYPE_CHAR_ARROW_CHUNKED_ARRAY = "C"
+_TYPE_CHAR_ARROW_REDUCED = "R"
+
+
+class ArrowDataTypeSerializer(Serializer):
+    def serial(self, obj: pa.DataType, context):
+        return [str(obj)], [], True
+
+    def deserial(self, serialized, context, subs):
+        return arrow_type_from_str(serialized[0])
+
+
+class ArrowArraySerializer(Serializer):
+    @buffered
+    def serial(self, obj: PA_ARRAY_TYPES, context: Dict):
+        if not isinstance(obj, (pa.Array, pa.ChunkedArray)):
+            raise NotImplementedError(f"Array type {type(obj)} not supported")
+
+        meth, extracted = obj.__reduce__()
+        meth_name = extract_class_name(meth)
+        return [_TYPE_CHAR_ARROW_REDUCED, meth_name, None], list(extracted), False
+
+    def deserial(self, serialized: List, context: Dict, subs: List):
+        array_type, data_type_str, sizes = serialized[:3]
+        if array_type == _TYPE_CHAR_ARROW_REDUCED:
+            if data_type_str == "pyarrow.lib#chunked_array":
+                return pa.chunked_array(*subs)
+            elif data_type_str == "pyarrow.lib#_restore_array":
+                return pa.lib._restore_array(*subs)
+            else:
+                raise NotImplementedError(f"Unknown array type: {array_type}")
+        if array_type == _TYPE_CHAR_ARROW_ARRAY:
+            data_type = arrow_type_from_str(data_type_str)
+            return pa.Array.from_buffers(data_type, sizes, subs)
+        elif array_type == _TYPE_CHAR_ARROW_CHUNKED_ARRAY:
+            data_type = arrow_type_from_str(data_type_str)
+            chunks = [
+                pa.Array.from_buffers(data_type, size, bufs)
+                for size, bufs in zip(sizes, subs)
+            ]
+            return pa.chunked_array(chunks)
+        else:  # pragma: no cover
+            raise NotImplementedError(f"Array type {array_type} not supported")
+
+
+class ArrowBatchSerializer(Serializer):
+    @buffered
+    def serial(self, obj: PA_RECORD_TYPES, context: Dict):
+        sink = pa.BufferOutputStream()
+        writer = pa.RecordBatchStreamWriter(sink, obj.schema)
+        if isinstance(obj, pa.Table):
+            batch_type = "T"
+            writer.write_table(obj)
+        else:
+            batch_type = "B"
+            writer.write_batch(obj)
+        writer.close()
+
+        buf = sink.getvalue()
+        buffers = [buf]
+        return [batch_type], buffers, True
+
+    def deserial(self, serialized: List, context: Dict, subs: List):
+        reader = pa.RecordBatchStreamReader(pa.BufferReader(subs[0]))
+        if serialized[0] == "T":  # Table
+            return reader.read_all()
+        else:  # RecordBatch
+            return reader.read_next_batch()
+
+
+if pa is not None:  # pragma: no branch
+    ArrowDataTypeSerializer.register(pa.DataType)
+    ArrowArraySerializer.register(pa.Array)
+    ArrowArraySerializer.register(pa.ChunkedArray)
+    ArrowBatchSerializer.register(pa.Table)
+    ArrowBatchSerializer.register(pa.RecordBatch)

maxframe/serialization/blob.py

@@ -0,0 +1,32 @@
+# Copyright 1999-2025 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 Dict
+
+from ..lib.dtypes_extension.blob import AbstractExternalBlob, SolidBlob
+from .core import Serializer
+
+
+class ExternalBlobSerializer(Serializer):
+    def serial(self, obj: AbstractExternalBlob, context: Dict):
+        _, vals = obj.__reduce__()
+        return [type(obj).__name__], list(vals), False
+
+    def deserial(self, serialized, context, subs):
+        cls_name = serialized[0]
+        cls = AbstractExternalBlob.get_cls_by_name(cls_name)
+        return cls(*subs)
+
+
+ExternalBlobSerializer.register(SolidBlob)

maxframe/serialization/core.pxd

@@ -0,0 +1,50 @@
+# Copyright 1999-2025 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 libc.stdint cimport int32_t, uint64_t
+
+
+cdef bint unpickle_allowed
+
+
+cdef class Serializer:
+    cdef int _serializer_id
+
+    cpdef bint is_public_data_exist(self, dict context, object key)
+    cpdef put_public_data(self, dict context, object key, object value)
+    cpdef get_public_data(self, dict context, object key)
+    cpdef serial(self, object obj, dict context)
+    cpdef deserial(self, list serialized, dict context, list subs)
+    cpdef on_deserial_error(
+        self,
+        list serialized,
+        dict context,
+        list subs_serialized,
+        int error_index,
+        object exc,
+    )
+
+
+cdef class Placeholder:
+    """
+    Placeholder object to reduce duplicated serialization
+
+    The object records object identifier and keeps callbacks
+    to replace itself in parent objects.
+    """
+    cdef public uint64_t id
+    cdef public list callbacks
+
+
+cdef Serializer get_deserializer(int32_t deserializer_id)

maxframe/serialization/core.pyi

@@ -0,0 +1,66 @@
+# Copyright 1999-2025 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 concurrent.futures import Executor
+from typing import Any, Callable, Dict, List, Type, TypeVar
+
+def buffered(func: Callable) -> Callable: ...
+def fast_id(obj: Any) -> int: ...
+
+LoadType = TypeVar("LoadType")
+
+def load_type(class_name: str, parent_class: LoadType) -> LoadType: ...
+def load_member(class_name: str, restrict_type: Type[LoadType]) -> LoadType: ...
+def reload_unpickle_flag(): ...
+
+class PickleContainer:
+    def __init__(self, buffers: List[bytes]): ...
+    def get(self) -> Any: ...
+    def get_buffers(self) -> List[bytes]: ...
+
+class Serializer:
+    serializer_id: int
+    def is_public_data_exist(self, context: Dict, key: Any) -> bool: ...
+    def put_public_data(self, context: Dict, key: Any, value: Any) -> None: ...
+    def get_public_data(self, context: Dict, key: Any) -> Any: ...
+    def serial(self, obj: Any, context: Dict): ...
+    def deserial(self, serialized: List, context: Dict, subs: List[Any]): ...
+    def on_deserial_error(
+        self,
+        serialized: List,
+        context: Dict,
+        subs_serialized: List,
+        error_index: int,
+        exc: BaseException,
+    ): ...
+    @classmethod
+    def register(cls, obj_type): ...
+    @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): ...
+async def serialize_with_spawn(
+    obj: Any,
+    context: Dict = None,
+    spawn_threshold: int = 100,
+    executor: Executor = None,
+): ...
+def deserialize(headers: List, buffers: List, context: Dict = None): ...