maxframe 0.1.0b5__cp311-cp311-macosx_11_0_arm64.whl

maxframe/remote/run_script.py

@@ -0,0 +1,121 @@
+# 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 os
+from typing import Any, BinaryIO, Dict, List, TextIO, Union
+
+from .. import opcodes
+from ..core import TILEABLE_TYPE, OutputType
+from ..core.operator import MergeDictOperator
+from ..serialization.serializables import (
+    BoolField,
+    BytesField,
+    DictField,
+    Int32Field,
+    ListField,
+)
+from ..typing_ import SessionType, TileableType
+from ..utils import to_binary
+
+
+class RunScript(MergeDictOperator):
+    _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 __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 = False,
+    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,26 @@
+# 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 . import arrow, exception, maxframe_objects, numpy, pandas, scipy, serializables
+from .core import (
+    PickleContainer,
+    Serializer,
+    deserialize,
+    pickle_buffers,
+    serialize,
+    serialize_with_spawn,
+)
+from .exception import RemoteException
+
+del arrow, numpy, scipy, pandas, maxframe_objects, exception

maxframe/serialization/arrow.py

@@ -0,0 +1,95 @@
+# 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, List, Union
+
+from ..utils import arrow_type_from_str
+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"
+
+
+class ArrowArraySerializer(Serializer):
+    @buffered
+    def serial(self, obj: PA_ARRAY_TYPES, context: Dict):
+        data_type = str(obj.type)
+        if isinstance(obj, pa.Array):
+            array_type = _TYPE_CHAR_ARROW_ARRAY
+            buffers = obj.buffers()
+            sizes = len(obj)
+        elif isinstance(obj, pa.ChunkedArray):
+            array_type = _TYPE_CHAR_ARROW_CHUNKED_ARRAY
+            buffers = [c.buffers() for c in obj.chunks]
+            sizes = [len(c) for c in obj.chunks]
+        else:  # pragma: no cover
+            raise NotImplementedError(f"Array type {type(obj)} not supported")
+        return [array_type, data_type, sizes], buffers, True
+
+    def deserial(self, serialized: List, context: Dict, subs: List):
+        array_type, data_type_str, sizes = serialized[:3]
+        data_type = arrow_type_from_str(data_type_str)
+        if array_type == _TYPE_CHAR_ARROW_ARRAY:
+            return pa.Array.from_buffers(data_type, sizes, subs)
+        elif array_type == _TYPE_CHAR_ARROW_CHUNKED_ARRAY:
+            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
+    ArrowArraySerializer.register(pa.Array)
+    ArrowArraySerializer.register(pa.ChunkedArray)
+    ArrowBatchSerializer.register(pa.Table)
+    ArrowBatchSerializer.register(pa.RecordBatch)

maxframe/serialization/core.pxd

@@ -0,0 +1,44 @@
+# 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 libc.stdint cimport int32_t, uint64_t
+
+
+cdef class Serializer:
+    cdef int _serializer_id
+
+    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,61 @@
+# 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 concurrent.futures import Executor
+from typing import Any, Callable, Dict, List, 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: ...
+
+class PickleContainer:
+    def __init__(self, buffers: List[bytes]): ...
+    def get(self) -> Any: ...
+    def get_buffers(self) -> List[bytes]: ...
+
+class Serializer:
+    serializer_id: int
+    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): ...