maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/extension.py

@@ -0,0 +1,131 @@
+# 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 abc
+import itertools
+from typing import TYPE_CHECKING, Iterable, Optional, Tuple, Type
+
+if TYPE_CHECKING:
+    from .codegen import DAGCodeGenerator
+
+    try:
+        import maxframe_framedriver.runners
+    except ImportError:
+        pass
+
+
+class MaxFrameExtension(metaclass=abc.ABCMeta):
+    @classmethod
+    @abc.abstractmethod
+    def get_codegen(cls) -> Type["DAGCodeGenerator"]:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def get_runner(cls) -> Type["maxframe_framedriver.runners.SubDagRunner"]:
+        raise NotImplementedError
+
+    @classmethod
+    async def destroy_session(cls, session_id: str) -> None:
+        """
+        This interface will be called once a session in the framedriver is ended.
+
+        Parameters
+        ----------
+        session_id : str
+            The session id.
+        """
+        pass
+
+    @classmethod
+    async def create_session(cls, session_id: str) -> None:
+        """
+        create the session
+        Parameters
+        ----------
+        session_id : str
+            The session id.
+        """
+        pass
+
+    @classmethod
+    async def reload_session(cls, session_id: str) -> None:
+        """
+        Reload the session state when the session is recovered from failover.
+
+        Parameters
+        ----------
+        session_id : str
+            The session id.
+        """
+        pass
+
+    @classmethod
+    def init_service_extension(cls) -> None:
+        """
+        Init the services of the extension before the app is actually run and will be
+        called only once.
+        """
+        pass
+
+    @classmethod
+    def destroy_service_extension(cls) -> None:
+        """
+        Destroy the services of the extension before the app is actually stopped and
+        will be called only once.
+        """
+        pass
+
+    @classmethod
+    def get_logview_collector_class(cls) -> Optional[Type]:
+        """
+        Get the logview collector class for this engine.
+
+        Returns
+        -------
+        Optional[Type]
+            The collector class, or None if this engine doesn't provide logview support.
+        """
+        pass
+
+    @classmethod
+    def get_logview_data_field_name(cls) -> Optional[str]:
+        """
+        Get the field name to store engine-specific logview data in LogviewSessionDetail.
+
+        Returns
+        -------
+        Optional[str]
+            Field name (e.g., 'dpes'), or None if no logview support.
+        """
+        pass
+
+
+def iter_extensions() -> Iterable[Tuple[str, MaxFrameExtension]]:
+    try:
+        from importlib.metadata import entry_points
+
+        eps = entry_points(group="maxframe.extension")
+    except (ImportError, TypeError):
+        from importlib_metadata import entry_points
+
+        eps = entry_points(group="maxframe.extension")
+
+    if callable(getattr(eps, "values", None)):
+        _it = itertools.chain(*eps.values())
+    else:
+        _it = eps
+
+    for entry_point in _it:
+        yield entry_point.name, entry_point.load()

maxframe/io/__init__.py

@@ -0,0 +1,13 @@
+# 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.

maxframe/io/objects/__init__.py

@@ -0,0 +1,24 @@
+# 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 .core import (
+    AbstractObjectIOHandler,
+    get_object_io_handler,
+    register_object_io_handler,
+)
+
+# isort: off
+from . import tensor
+
+del tensor

maxframe/io/objects/core.py

@@ -0,0 +1,156 @@
+# 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 abc import ABCMeta, abstractmethod
+from typing import Any, Dict, Type, Union
+
+import msgpack
+import numpy as np
+
+from ...core import Entity, EntityData, OutputType, get_output_types
+from ...lib import wrapped_pickle as pickle
+from ...typing_ import SlicesType, TileableType
+from ...utils import deserialize_serializable, serialize_serializable
+from ..odpsio.volumeio import ODPSVolumeReader, ODPSVolumeWriter
+
+_MetaType = Dict[str, Any]
+
+_META_FILE_NAME = ".meta"
+_META_PICKLED_KEYS_KEY = ".pickled_keys"
+_META_SERIALIZED_VALUE_KEY = ".serialized_values"
+
+
+_io_handler_dispatcher = dict()
+
+
+def register_object_io_handler(out_type: OutputType):
+    def wrapper(handler_cls):
+        _io_handler_dispatcher[out_type] = handler_cls
+        return handler_cls
+
+    return wrapper
+
+
+def get_object_io_handler(
+    tileable_data_type: Union[Entity, EntityData, Type[EntityData], OutputType]
+) -> Type["AbstractObjectIOHandler"]:
+    output_type = tileable_data_type
+    if not isinstance(tileable_data_type, OutputType):
+        if hasattr(tileable_data_type, "output_type"):
+            output_type = tileable_data_type.output_type
+        else:
+            [output_type] = get_output_types(tileable_data_type)
+    return _io_handler_dispatcher[output_type]
+
+
+class AbstractObjectIOHandler(metaclass=ABCMeta):
+    def _prepare_meta_for_serial(self, tileable: TileableType) -> Dict[str, Any]:
+        return tileable.params.copy()
+
+    def read_object_meta(
+        self, reader: ODPSVolumeReader, tileable: TileableType
+    ) -> Dict[str, Any]:
+        file_bin = reader.read_file(_META_FILE_NAME)
+        meta_obj = msgpack.loads(file_bin)
+        pickled_keys = meta_obj.pop(_META_PICKLED_KEYS_KEY, None) or []
+        serialized_values = meta_obj.pop(_META_SERIALIZED_VALUE_KEY, None) or {}
+        for k in pickled_keys:
+            try:
+                meta_obj[k] = deserialize_serializable(serialized_values[k])
+            except:
+                meta_obj[k] = pickle.loads(meta_obj[k])
+        return meta_obj
+
+    @abstractmethod
+    def read_object_body(
+        self,
+        reader: ODPSVolumeReader,
+        params: Dict[str, Any],
+        extra_params: Dict[str, Any] = None,
+        slices: SlicesType = None,
+    ) -> Any:
+        raise NotImplementedError
+
+    def read_object(
+        self,
+        reader: ODPSVolumeReader,
+        tileable: TileableType,
+        extra_params: Dict[str, Any] = None,
+        slices: SlicesType = None,
+    ) -> Any:
+        obj_meta = self.read_object_meta(reader, tileable)
+        return self.read_object_body(reader, obj_meta, extra_params, slices)
+
+    @abstractmethod
+    def write_object_body(
+        self, writer: ODPSVolumeWriter, params: Dict[str, Any], value: Any
+    ):
+        raise NotImplementedError
+
+    @classmethod
+    def _transform_serializable_meta(cls, meta: Dict[str, Any]) -> Dict[str, Any]:
+        to_pack = meta.copy()
+        pickled_keys = []
+        serialized_values = {}
+        for k, v in meta.items():
+            if not isinstance(v, (str, bytes, int, float, bool, type(None))):
+                to_pack[k] = pickle.dumps(v)
+                serialized_values[k] = serialize_serializable(v)
+                pickled_keys.append(k)
+        to_pack[_META_PICKLED_KEYS_KEY] = pickled_keys
+        to_pack[_META_SERIALIZED_VALUE_KEY] = serialized_values
+        return to_pack
+
+    def write_object_meta(
+        self,
+        writer: ODPSVolumeWriter,
+        tileable: TileableType,
+        extra_meta: Dict[str, Any] = None,
+    ):
+        meta_obj = self._prepare_meta_for_serial(tileable)
+        if extra_meta:
+            meta_obj.update(extra_meta)
+
+        def encode_np_num(obj):
+            if isinstance(obj, np.generic) and obj.shape == ():
+                return obj.item()
+            return obj
+
+        meta_obj = self._transform_serializable_meta(meta_obj)
+        writer.write_file(
+            _META_FILE_NAME, msgpack.dumps(meta_obj, default=encode_np_num)
+        )
+
+    def write_object(
+        self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
+    ):
+        self.write_object_meta(writer, tileable)
+        self.write_object_body(writer, tileable.params, value)
+
+
+@register_object_io_handler(OutputType.object)
+class ObjectIOHandler(AbstractObjectIOHandler):
+    def read_object_body(
+        self,
+        reader: ODPSVolumeReader,
+        params: Dict[str, Any],
+        extra_params: Dict[str, Any] = None,
+        slices: SlicesType = None,
+    ) -> Any:
+        return pickle.loads(reader.read_file("data"))
+
+    def write_object_body(
+        self, writer: ODPSVolumeWriter, params: Dict[str, Any], value: Any
+    ):
+        writer.write_file("data", pickle.dumps(value))

maxframe/io/objects/tensor.py

@@ -0,0 +1,133 @@
+# 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 itertools
+import struct
+from io import BytesIO
+from typing import Any, Dict
+
+import msgpack
+import numpy as np
+
+from ...config import options
+from ...core import OutputType
+from ...lib import wrapped_pickle as pickle
+from ...lib.version import parse as parse_version
+from ...typing_ import SlicesType, TileableType
+from ..odpsio import ODPSVolumeReader, ODPSVolumeWriter
+from .core import AbstractObjectIOHandler, register_object_io_handler
+
+try:
+    from maxframe import __version__ as mf_version
+except ImportError:
+    mf_version = None
+
+
+@register_object_io_handler(OutputType.scalar)
+@register_object_io_handler(OutputType.tensor)
+class TensorIOHandler(AbstractObjectIOHandler):
+    def _prepare_meta_for_serial(self, tileable: TileableType) -> Dict[str, Any]:
+        meta = super()._prepare_meta_for_serial(tileable)
+        meta["nsplits"] = (
+            getattr(tileable, "nsplits", None) or (np.nan,) * tileable.ndim
+        )
+        return meta
+
+    def write_object_meta(
+        self,
+        writer: ODPSVolumeWriter,
+        tileable: TileableType,
+        extra_meta: Dict[str, Any] = None,
+    ):
+        extra_meta = extra_meta or dict()
+        extra_meta["nsplits"] = (
+            getattr(tileable, "nsplits", None) or ((np.nan,),) * tileable.ndim
+        )
+        super().write_object_meta(writer, tileable, extra_meta=extra_meta)
+
+    @staticmethod
+    def _get_chunk_file_name(params: Dict[str, Any]) -> str:
+        # fixme remove this when all v1.0 clients removed
+        mf_ver = options.session.client_version or mf_version
+        if mf_ver:
+            client_major = parse_version(
+                options.session.client_version or mf_version
+            ).major
+        else:
+            client_major = None
+
+        if client_major == 1:
+            # returns v1.0 file name as we do not support tiled tensors
+            return "0,0.dat"
+
+        ndim = len(params.get("shape") or (0,))
+        data_index = params.get("index") or (0,) * ndim
+        return ",".join(str(ix) for ix in data_index) + ".dat"
+
+    def _read_single_chunk(
+        self, reader: ODPSVolumeReader, params: Dict[str, Any]
+    ) -> Any:
+        file_name = self._get_chunk_file_name(params)
+        body = reader.read_file(file_name)
+        bio = BytesIO(body)
+        (header_len,) = struct.unpack("<I", bio.read(4))
+        header_data = msgpack.loads(bio.read(header_len))
+
+        pickled = bio.read(header_data[0])
+        bufs = [bio.read(size) for size in header_data[1:]]
+        return pickle.loads(pickled, buffers=bufs)
+
+    def read_object_body(
+        self,
+        reader: ODPSVolumeReader,
+        params: Dict[str, Any],
+        extra_params: Dict[str, Any] = None,
+        slices: SlicesType = None,
+    ) -> Any:
+        if "index" in params:
+            return self._read_single_chunk(reader, params)
+
+        tileable_params = params.copy()
+        tileable_params.update(extra_params or {})
+
+        # todo implements slices argument for head and tail fetching
+        assert slices is None
+
+        chunk_shape = tuple(len(x) for x in params["nsplits"])
+        block_array = np.empty(shape=chunk_shape, dtype="O")
+        for idx in itertools.product(*(range(sp) for sp in chunk_shape)):
+            chunk_params = tileable_params.copy()
+            chunk_params.pop("nsplits")
+            chunk_params["index"] = idx
+            block_array[idx] = self._read_single_chunk(reader, chunk_params)
+
+        return np.block(block_array.tolist())
+
+    def write_object_body(
+        self, writer: ODPSVolumeWriter, params: Dict[str, Any], value: Any
+    ):
+        def data_gen():
+            bufs = []
+            pickled = pickle.dumps(value, buffer_callback=bufs.append)
+            header_data = msgpack.dumps(
+                [len(pickled)] + [len(buf.raw()) for buf in bufs]
+            )
+            yield struct.pack("<I", len(header_data))
+            yield header_data
+            yield pickled
+            for buf in bufs:
+                yield buf
+
+        file_name = self._get_chunk_file_name(params)
+        writer.write_file(file_name, data_gen())

maxframe/io/objects/tests/__init__.py

@@ -0,0 +1,13 @@
+# 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.

maxframe/io/objects/tests/test_object_io.py

@@ -0,0 +1,85 @@
+# 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 numpy as np
+import pytest
+from odps import ODPS
+
+from ....core import OutputType
+from ....core.operator import ObjectOperatorMixin, Operator
+from ....tensor.datasource import ArrayDataSource, scalar
+from ....tests.utils import create_test_volume, get_test_unique_name, tn
+from ...odpsio import ODPSVolumeReader, ODPSVolumeWriter
+from ..core import get_object_io_handler
+
+
+class TestObjectOp(Operator, ObjectOperatorMixin):
+    def __call__(self):
+        self._output_types = [OutputType.object]
+        return self.new_tileable([])
+
+
+@pytest.fixture(scope="module")
+def create_volume(oss_config):
+    with create_test_volume(
+        tn("test_object_io_vol_" + get_test_unique_name(5)), oss_config
+    ) as test_vol_name:
+        yield test_vol_name
+
+
+@pytest.mark.parametrize(
+    "obj_creator, data",
+    [
+        (lambda: TestObjectOp()(), "abcdefg"),
+        (lambda: scalar(123), 123),
+    ],
+)
+def test_object_and_scalar_io(create_volume, obj_creator, data):
+    obj = obj_creator()
+
+    odps_entry = ODPS.from_environments()
+
+    reader = ODPSVolumeReader(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
+    writer = ODPSVolumeWriter(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
+
+    handler = get_object_io_handler(obj)()
+    handler.write_object(writer, obj, data)
+    assert data == handler.read_object(reader, obj)
+
+
+def test_tensor_object_io(create_volume):
+    data = np.array([[4, 9, 2], [3, 5, 7], [8, 1, 6]])
+    obj = ArrayDataSource(data, dtype=data.dtype)(data.shape)
+
+    odps_entry = ODPS.from_environments()
+
+    reader = ODPSVolumeReader(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
+    writer = ODPSVolumeWriter(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
+
+    # test write and read full object
+    handler = get_object_io_handler(obj)()
+    handler.write_object(writer, obj, data)
+    np.testing.assert_equal(data, handler.read_object(reader, obj))
+
+    # test read single chunk
+    params = {"index": (0, 0)}
+    np.testing.assert_equal(data, handler.read_object_body(reader, params))

maxframe/io/odpsio/__init__.py

@@ -0,0 +1,24 @@
+# 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 .arrow import arrow_to_pandas, pandas_to_arrow
+from .schema import (
+    arrow_schema_to_odps_schema,
+    build_dataframe_table_meta,
+    cast_df_with_possible_nans,
+    odps_schema_to_pandas_dtypes,
+    pandas_to_odps_schema,
+)
+from .tableio import HaloTableIO, ODPSTableIO, TunnelTableIO
+from .volumeio import ODPSVolumeReader, ODPSVolumeWriter