maxframe 0.1.0b5__cp39-cp39-macosx_11_0_arm64.whl

maxframe-0.1.0b5.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: false
+Tag: cp39-cp39-macosx_11_0_arm64
+

maxframe-0.1.0b5.dist-info/top_level.txt

@@ -0,0 +1,3 @@
+build
+maxframe
+maxframe_client

maxframe_client/__init__.py

@@ -0,0 +1,17 @@
+# 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 .clients.framedriver import FrameDriverClient
+from .clients.spe import SPEClient
+from .session import MaxFrameRestSession

maxframe_client/clients/__init__.py

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

maxframe_client/clients/framedriver.py

@@ -0,0 +1,118 @@
+# 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
+
+import msgpack
+from tornado import httpclient
+
+from maxframe.core import TileableGraph
+from maxframe.protocol import (
+    CreateSessionRequest,
+    DagInfo,
+    DecrefRequest,
+    ExecuteDagRequest,
+    ProtocolBody,
+    ResultInfo,
+    SessionInfo,
+)
+from maxframe.typing_ import TimeoutType
+from maxframe.utils import (
+    format_timeout_params,
+    serialize_serializable,
+    wait_http_response,
+)
+
+
+class FrameDriverClient:
+    def __init__(self, endpoint: str):
+        self._endpoint = endpoint.rstrip("/")
+
+    async def create_session(self, settings: Dict[str, Any] = None) -> SessionInfo:
+        req_body: ProtocolBody[CreateSessionRequest] = ProtocolBody(
+            body=CreateSessionRequest(settings=settings)
+        )
+        req_url = f"{self._endpoint}/api/sessions"
+        resp = await httpclient.AsyncHTTPClient().fetch(
+            req_url, method="POST", body=serialize_serializable(req_body)
+        )
+        return SessionInfo.from_json(msgpack.loads(resp.body))
+
+    async def get_session(self, session_id: str) -> SessionInfo:
+        req_url = f"{self._endpoint}/api/sessions/{session_id}"
+        resp = await httpclient.AsyncHTTPClient().fetch(req_url, method="GET")
+        return SessionInfo.from_json(msgpack.loads(resp.body))
+
+    async def delete_session(self, session_id: str):
+        req_url = f"{self._endpoint}/api/sessions/{session_id}"
+        await httpclient.AsyncHTTPClient().fetch(req_url, method="DELETE")
+
+    async def submit_dag(
+        self,
+        session_id: str,
+        dag: TileableGraph,
+        managed_input_infos: Dict[str, ResultInfo] = None,
+    ) -> DagInfo:
+        req_url = f"{self._endpoint}/api/sessions/{session_id}/dags"
+        req_body = ExecuteDagRequest(session_id, dag, managed_input_infos)
+        resp = await httpclient.AsyncHTTPClient().fetch(
+            req_url,
+            method="POST",
+            body=serialize_serializable(ProtocolBody(body=req_body)),
+        )
+        return DagInfo.from_json(msgpack.loads(resp.body))
+
+    async def get_dag_info(self, session_id: str, dag_id: str) -> DagInfo:
+        req_url = f"{self._endpoint}/api/sessions/{session_id}/dags/{dag_id}"
+        resp = await httpclient.AsyncHTTPClient().fetch(req_url, method="GET")
+        return DagInfo.from_json(msgpack.loads(resp.body))
+
+    async def wait_dag(self, session_id: str, dag_id: str, timeout: TimeoutType = None):
+        query_part = format_timeout_params(timeout)
+        req_url = (
+            f"{self._endpoint}/api/sessions/{session_id}/dags/{dag_id}{query_part}"
+        )
+        try:
+            resp = await wait_http_response(
+                req_url, method="GET", request_timeout=timeout
+            )
+            info = DagInfo.from_json(msgpack.loads(resp.body))
+        except TimeoutError:
+            info = await self.get_dag_info(session_id, dag_id)
+        return info
+
+    async def cancel_dag(
+        self, session_id: str, dag_id: str, timeout: TimeoutType = None
+    ):
+        query_part = format_timeout_params(timeout)
+        req_url = (
+            f"{self._endpoint}/api/sessions/{session_id}/dags/{dag_id}{query_part}"
+        )
+        try:
+            resp = await wait_http_response(
+                req_url, method="DELETE", request_timeout=timeout
+            )
+            info = DagInfo.from_json(msgpack.loads(resp.body))
+        except TimeoutError:
+            info = await self.get_dag_info(session_id, dag_id)
+        return info
+
+    async def decref(self, session_id: str, tileable_keys: List[str]) -> None:
+        req_url = f"{self._endpoint}/api/sessions/{session_id}/lifecycle?decref"
+        req_body = DecrefRequest(tileable_keys)
+        await wait_http_response(
+            req_url,
+            method="POST",
+            body=serialize_serializable(ProtocolBody(body=req_body)),
+        )

maxframe_client/clients/spe.py

@@ -0,0 +1,104 @@
+# 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, Optional
+
+from tornado import httpclient
+
+from maxframe.core import TileableGraph
+from maxframe.protocol import ExecuteSubDagRequest, ProtocolBody, SubDagInfo
+from maxframe.typing_ import TimeoutType
+from maxframe.utils import (
+    deserialize_serializable,
+    format_timeout_params,
+    serialize_serializable,
+    wait_http_response,
+)
+
+
+class SPEClient:
+    def __init__(
+        self,
+        endpoint: str,
+        session_id: Optional[str] = None,
+        host: str = None,
+    ):
+        self._endpoint = endpoint.rstrip("/")
+        self._session_id = session_id
+        self._headers = {"Host": host}
+
+    @staticmethod
+    def _load_subdag_info(resp: httpclient.HTTPResponse) -> SubDagInfo:
+        res: ProtocolBody[SubDagInfo] = deserialize_serializable(resp.body)
+        return res.body
+
+    async def submit_subdag(
+        self, subdag: TileableGraph, settings: Dict[str, Any] = None
+    ) -> SubDagInfo:
+        req_url = f"{self._endpoint}/api/subdags"
+        req_body: ProtocolBody[ExecuteSubDagRequest] = ProtocolBody(
+            body=ExecuteSubDagRequest(dag=subdag, settings=settings),
+        )
+
+        if self._session_id is not None:
+            req_url += f"?session_id={self._session_id}"
+
+        resp = await httpclient.AsyncHTTPClient().fetch(
+            req_url,
+            method="POST",
+            headers=self._headers,
+            body=serialize_serializable(req_body),
+        )
+        return self._load_subdag_info(resp)
+
+    async def get_subdag_info(self, subdag_id: str) -> SubDagInfo:
+        req_url = f"{self._endpoint}/api/subdags/{subdag_id}?wait=0"
+        resp = await httpclient.AsyncHTTPClient().fetch(
+            req_url,
+            method="GET",
+            headers=self._headers,
+        )
+        return self._load_subdag_info(resp)
+
+    async def wait_subdag(
+        self, subdag_id: str, wait_timeout: TimeoutType = None
+    ) -> SubDagInfo:
+        req_url = f"{self._endpoint}/api/subdags/{subdag_id}"
+        params = format_timeout_params(wait_timeout)
+        try:
+            resp = await wait_http_response(
+                req_url + params,
+                method="GET",
+                headers=self._headers,
+                request_timeout=wait_timeout,
+            )
+            return self._load_subdag_info(resp)
+        except TimeoutError:
+            return await self.get_subdag_info(subdag_id)
+
+    async def cancel_subdag(
+        self, subdag_id: str, wait_timeout: TimeoutType = None
+    ) -> SubDagInfo:
+        req_url = f"{self._endpoint}/api/subdags/{subdag_id}"
+        params = format_timeout_params(wait_timeout)
+        try:
+            resp = await wait_http_response(
+                req_url + params,
+                method="DELETE",
+                headers=self._headers,
+                request_timeout=wait_timeout,
+            )
+            return self._load_subdag_info(resp)
+        except TimeoutError:
+            return await self.get_subdag_info(subdag_id)

maxframe_client/conftest.py

@@ -0,0 +1,15 @@
+# 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 maxframe.conftest import *  # noqa: F401

maxframe_client/fetcher.py

@@ -0,0 +1,264 @@
+# 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 json
+from abc import ABC, abstractmethod
+from numbers import Integral
+from typing import Any, Dict, List, Optional, Type, Union
+
+import pandas as pd
+import pyarrow as pa
+from odps import ODPS
+from odps.models import ExternalVolume, PartedVolume
+from odps.tunnel import TableTunnel
+from tornado import httpclient
+
+from maxframe.core import OBJECT_TYPE
+from maxframe.dataframe.core import DATAFRAME_TYPE
+from maxframe.lib import wrapped_pickle as pickle
+from maxframe.odpsio import HaloTableIO, arrow_to_pandas, build_dataframe_table_meta
+from maxframe.protocol import (
+    DataFrameTableMeta,
+    ODPSTableResultInfo,
+    ODPSVolumeResultInfo,
+    ResultInfo,
+    ResultType,
+)
+from maxframe.tensor.core import TENSOR_TYPE
+from maxframe.typing_ import PandasObjectTypes, TileableType
+from maxframe.utils import ToThreadMixin, deserialize_serializable
+
+_result_fetchers: Dict[ResultType, Type["ResultFetcher"]] = dict()
+
+
+def register_fetcher(fetcher_cls: Type["ResultFetcher"]):
+    _result_fetchers[fetcher_cls.result_type] = fetcher_cls
+    return fetcher_cls
+
+
+def get_fetcher_cls(result_type: ResultType) -> Type["ResultFetcher"]:
+    return _result_fetchers[result_type]
+
+
+class ResultFetcher(ABC):
+    result_type = None
+
+    def __init__(self, odps_entry: ODPS):
+        self._odps_entry = odps_entry
+
+    @abstractmethod
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ResultInfo,
+    ) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def fetch(
+        self,
+        tileable: TileableType,
+        info: ResultInfo,
+        indexes: List[Union[None, Integral, slice]],
+    ) -> Any:
+        raise NotImplementedError
+
+
+@register_fetcher
+class NullFetcher(ResultFetcher):
+    result_type = ResultType.NULL
+
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ResultInfo,
+    ) -> None:
+        return
+
+    async def fetch(
+        self,
+        tileable: TileableType,
+        info: ODPSTableResultInfo,
+        indexes: List[Union[None, Integral, slice]],
+    ) -> None:
+        return
+
+
+@register_fetcher
+class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
+    result_type = ResultType.ODPS_TABLE
+
+    def _get_table_comment(self, table_name: str) -> Optional[str]:
+        table = self._odps_entry.get_table(table_name)
+        return getattr(table, "comment", None)
+
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ODPSTableResultInfo,
+    ) -> None:
+        if isinstance(tileable, DATAFRAME_TYPE) and tileable.dtypes is None:
+            tb_comment = await self.to_thread(
+                self._get_table_comment, info.full_table_name
+            )
+            if tb_comment:  # pragma: no branch
+                comment_data = json.loads(tb_comment)
+
+                table_meta: DataFrameTableMeta = deserialize_serializable(
+                    base64.b64decode(comment_data["table_meta"])
+                )
+                tileable.refresh_from_table_meta(table_meta)
+
+        if tileable.shape and any(pd.isna(x) for x in tileable.shape):
+            part_specs = [None] if not info.partition_specs else info.partition_specs
+            tunnel = TableTunnel(self._odps_entry)
+            total_records = 0
+            for part_spec in part_specs:
+                session = tunnel.create_download_session(
+                    info.full_table_name, part_spec
+                )
+                total_records += session.count
+            new_shape_list = list(tileable.shape)
+            new_shape_list[-1] = total_records
+            tileable.params = {"shape": tuple(new_shape_list)}
+
+    def _read_single_source(
+        self,
+        table_meta: DataFrameTableMeta,
+        info: ODPSTableResultInfo,
+        indexes: List[Union[None, Integral, slice]],
+    ):
+        table_io = HaloTableIO(self._odps_entry)
+        read_kw = {}
+        row_step = None
+        if indexes:
+            if len(indexes) < 2:
+                indexes += [None]
+            row_sel, col_sel = indexes
+            if isinstance(row_sel, slice):
+                if row_sel.start or row_sel.stop:
+                    read_kw["start"] = row_sel.start
+                    read_kw["stop"] = row_sel.stop
+                    read_kw["reverse_range"] = (
+                        row_sel.step is not None and row_sel.step < 0
+                    )
+                    row_step = row_sel.step
+            elif isinstance(row_sel, int):
+                read_kw["start"] = row_sel
+                read_kw["stop"] = row_sel + 1
+                row_step = None
+            elif row_sel is not None:  # pragma: no cover
+                raise NotImplementedError(f"Does not support row index {row_sel!r}")
+
+            if isinstance(col_sel, (int, slice)):
+                data_cols = table_meta.table_column_names[col_sel]
+                if isinstance(col_sel, int):
+                    data_cols = [data_cols]
+                read_kw["columns"] = table_meta.table_index_column_names + data_cols
+            elif col_sel is not None:  # pragma: no cover
+                raise NotImplementedError(f"Does not support column index {row_sel!r}")
+
+        with table_io.open_reader(
+            info.full_table_name, info.partition_specs, **read_kw
+        ) as reader:
+            reader_count = reader.count
+            result = reader.read_all()
+
+        if not row_step:
+            return result
+
+        if row_step >= 0:
+            slice_start = 0
+            slice_stop = reader_count
+        else:
+            slice_start = reader_count - 1
+            slice_stop = None
+        return result[slice_start:slice_stop:row_step]
+
+    async def fetch(
+        self,
+        tileable: TileableType,
+        info: ODPSTableResultInfo,
+        indexes: List[Union[None, Integral, slice]],
+    ) -> PandasObjectTypes:
+        table_meta = build_dataframe_table_meta(tileable)
+        arrow_table: pa.Table = await self.to_thread(
+            self._read_single_source, table_meta, info, indexes
+        )
+        return arrow_to_pandas(arrow_table, table_meta)
+
+
+@register_fetcher
+class ODPSVolumeFetcher(ToThreadMixin, ResultFetcher):
+    result_type = ResultType.ODPS_VOLUME
+
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ODPSVolumeResultInfo,
+    ) -> None:
+        return
+
+    async def _read_parted_volume_data(
+        self, volume: PartedVolume, partition: str, file_name: str
+    ) -> bytes:
+        def sync_read():
+            with volume.open_reader(partition, file_name) as reader:
+                return reader.read()
+
+        return await self.to_thread(sync_read)
+
+    async def _read_external_volume_data(
+        self, volume: ExternalVolume, path: str, file_name: str
+    ) -> bytes:
+        signed_url = await self.to_thread(
+            volume.get_sign_url, path + "/" + file_name, "GET"
+        )
+        http_client = httpclient.AsyncHTTPClient()
+
+        resp = await http_client.fetch(signed_url)
+        if hasattr(resp, "status_code") and resp.code >= 400:
+            try:
+                import oss2.exceptions
+
+                oss_exc = oss2.exceptions.make_exception(resp.body)
+                raise oss_exc
+            except ImportError:
+                raise SystemError(resp.body)
+        return resp.body
+
+    async def _fetch_object(self, info: ODPSVolumeResultInfo) -> Any:
+        volume = await self.to_thread(self._odps_entry.get_volume, info.volume_name)
+        if isinstance(volume, PartedVolume):
+            byte_data = await self._read_parted_volume_data(
+                volume, info.volume_path, "data"
+            )
+        elif isinstance(volume, ExternalVolume):
+            byte_data = await self._read_external_volume_data(
+                volume, info.volume_path, "data"
+            )
+        else:
+            raise NotImplementedError(f"Volume type {type(volume)} not supported")
+        return pickle.loads(byte_data)
+
+    async def fetch(
+        self,
+        tileable: TileableType,
+        info: ODPSVolumeResultInfo,
+        indexes: List[Union[Integral, slice]],
+    ) -> Any:
+        if isinstance(tileable, (OBJECT_TYPE, TENSOR_TYPE)):
+            return await self._fetch_object(info)
+        raise NotImplementedError(f"Fetching {type(tileable)} not implemented")

maxframe_client/session/__init__.py

@@ -0,0 +1,22 @@
+# 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 .odps import MaxFrameRestSession
+
+
+def register_session_schemes(overwrite: bool = False):
+    from . import odps, task
+
+    odps.register_session_schemes(overwrite=overwrite)
+    task.register_session_schemes(overwrite=overwrite)

maxframe_client/session/consts.py

@@ -0,0 +1,36 @@
+# 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.
+
+# Restful Service
+RESTFUL_SESSION_INSECURE_SCHEME = "mf"
+RESTFUL_SESSION_SECURE_SCHEME = "mfs"
+
+# ODPS
+ODPS_SESSION_INSECURE_SCHEME = "mfo"
+ODPS_SESSION_SECURE_SCHEME = "mfos"
+
+# ODPS MaxFrameTask consts
+MAXFRAME_DEFAULT_PROTOCOL = "v1"
+# method names
+MAXFRAME_TASK_CREATE_SESSION_METHOD = "CREATE_SESSION"
+MAXFRAME_TASK_GET_SESSION_METHOD = "GET_SESSION"
+MAXFRAME_TASK_DELETE_SESSION_METHOD = "DELETE_SESSION"
+MAXFRAME_TASK_SUBMIT_DAG_METHOD = "SUBMIT_DAG"
+MAXFRAME_TASK_GET_DAG_INFO_METHOD = "GET_DAG_INFO"
+MAXFRAME_TASK_CANCEL_DAG_METHOD = "CANCEL_DAG"
+MAXFRAME_TASK_DECREF_METHOD = "DECREF"
+# serialization methods
+MAXFRAME_OUTPUT_MAXFRAME_FORMAT = "maxframe_v1"
+MAXFRAME_OUTPUT_JSON_FORMAT = "json"
+MAXFRAME_OUTPUT_MSGPACK_FORMAT = "msgpack"