maxframe 0.1.0b5__cp39-cp39-macosx_11_0_arm64.whl

maxframe_client/session/task.py

@@ -0,0 +1,280 @@
+# 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
+import logging
+import time
+from typing import Dict, List, Optional, Type, Union
+
+import msgpack
+from odps import ODPS
+from odps import options as odps_options
+from odps.errors import parse_instance_error
+from odps.models import Instance, MaxFrameTask
+
+from maxframe.config import options
+from maxframe.core import TileableGraph
+from maxframe.protocol import DagInfo, JsonSerializable, ResultInfo, SessionInfo
+from maxframe.utils import deserialize_serializable, serialize_serializable, to_str
+
+try:
+    from maxframe import __version__ as mf_version
+except ImportError:
+    mf_version = None
+
+from .consts import (
+    MAXFRAME_DEFAULT_PROTOCOL,
+    MAXFRAME_OUTPUT_JSON_FORMAT,
+    MAXFRAME_OUTPUT_MAXFRAME_FORMAT,
+    MAXFRAME_OUTPUT_MSGPACK_FORMAT,
+    MAXFRAME_TASK_CANCEL_DAG_METHOD,
+    MAXFRAME_TASK_CREATE_SESSION_METHOD,
+    MAXFRAME_TASK_DECREF_METHOD,
+    MAXFRAME_TASK_DELETE_SESSION_METHOD,
+    MAXFRAME_TASK_GET_DAG_INFO_METHOD,
+    MAXFRAME_TASK_GET_SESSION_METHOD,
+    MAXFRAME_TASK_SUBMIT_DAG_METHOD,
+    ODPS_SESSION_INSECURE_SCHEME,
+    ODPS_SESSION_SECURE_SCHEME,
+)
+from .odps import MaxFrameServiceCaller, MaxFrameSession
+
+logger = logging.getLogger(__name__)
+
+
+class MaxFrameInstanceCaller(MaxFrameServiceCaller):
+    _instance: Optional[Instance]
+
+    def __init__(
+        self,
+        odps_entry: ODPS,
+        task_name: Optional[str] = None,
+        project: Optional[str] = None,
+        priority: Optional[int] = None,
+        running_cluster: Optional[str] = None,
+        nested_instance_id: Optional[str] = None,
+        major_version: Optional[str] = None,
+        output_format: Optional[str] = None,
+        **kwargs,
+    ):
+        if callable(odps_options.get_priority):
+            default_priority = odps_options.get_priority(odps_entry)
+        else:
+            default_priority = odps_options.priority
+        priority = priority if priority is not None else default_priority
+
+        self._odps_entry = odps_entry
+        self._task_name = task_name
+        self._project = project
+        self._priority = priority
+        self._running_cluster = running_cluster
+        self._major_version = major_version
+        self._output_format = output_format or MAXFRAME_OUTPUT_MSGPACK_FORMAT
+
+        if nested_instance_id is None:
+            self._nested = False
+            self._instance = None
+        else:
+            self._nested = True
+            self._instance = odps_entry.get_instance(nested_instance_id)
+
+    def _deserial_task_info_result(
+        self, content: Union[bytes, str, dict], target_cls: Type[JsonSerializable]
+    ):
+        if isinstance(content, (str, bytes)):
+            json_data = json.loads(to_str(content))
+        else:
+            json_data = content
+        result_data = base64.b64decode(json_data["result"])
+        if self._output_format == MAXFRAME_OUTPUT_MAXFRAME_FORMAT:
+            return deserialize_serializable(result_data)
+        elif self._output_format == MAXFRAME_OUTPUT_JSON_FORMAT:
+            return target_cls.from_json(json.loads(result_data))
+        elif self._output_format == MAXFRAME_OUTPUT_MSGPACK_FORMAT:
+            return target_cls.from_json(msgpack.loads(result_data))
+        else:
+            raise ValueError(
+                f"Serialization format {self._output_format} not supported"
+            )
+
+    def _create_maxframe_task(self) -> MaxFrameTask:
+        task = MaxFrameTask(
+            name=self._task_name,
+            major_version=self._major_version,
+            service_endpoint=self._odps_entry.endpoint,
+        )
+
+        # merge sql options
+        sql_settings = (odps_options.sql.settings or {}).copy()
+        sql_settings.update(options.sql.settings or {})
+
+        mf_settings = dict(options.to_dict(remote_only=True).items())
+        mf_settings["sql.settings"] = sql_settings
+
+        mf_opts = {
+            "odps.maxframe.settings": json.dumps(mf_settings),
+            "odps.maxframe.output_format": self._output_format,
+        }
+        if mf_version:
+            mf_opts["odps.maxframe.client_version"] = mf_version
+        task.update_settings(mf_opts)
+        return task
+
+    def create_session(self) -> SessionInfo:
+        task = self._create_maxframe_task()
+        if not self._nested:
+            self._task_name = task.name
+            project = self._odps_entry.get_project(self._project)
+            self._instance = project.instances.create(
+                task=task,
+                priority=self._priority,
+                running_cluster=self._running_cluster,
+            )
+            self._wait_instance_task_ready()
+            return self.get_session()
+        else:
+            result = self._instance.put_task_info(
+                self._task_name,
+                MAXFRAME_TASK_CREATE_SESSION_METHOD,
+                task.properties["settings"],
+            )
+            return self._deserial_task_info_result(result, SessionInfo)
+
+    def _parse_instance_result_error(self):
+        result_data = self._instance.get_task_result(self._task_name)
+        try:
+            info = self._deserial_task_info_result({"result": result_data}, SessionInfo)
+        except:
+            raise parse_instance_error(result_data)
+        info.error_info.reraise()
+
+    def _wait_instance_task_ready(
+        self, interval: float = 0.1, max_interval: float = 5.0, timeout: int = None
+    ):
+        check_time = time.time()
+        timeout = timeout or options.client.task_start_timeout
+        while True:
+            if self._instance.is_terminated(retry=True):
+                self._parse_instance_result_error()
+            status_json = json.loads(
+                self._instance.get_task_info(self._task_name, "status") or "{}"
+            )
+            if status_json.get("status") == "Running":
+                break
+            if time.time() - check_time > timeout:
+                raise TimeoutError("Check session startup time out")
+            time.sleep(interval)
+            interval = min(max_interval, interval * 2)
+
+    def get_session(self) -> SessionInfo:
+        req_data = {"output_format": self._output_format}
+        serialized = self._instance.put_task_info(
+            self._task_name, MAXFRAME_TASK_GET_SESSION_METHOD, json.dumps(req_data)
+        )
+        info: SessionInfo = self._deserial_task_info_result(serialized, SessionInfo)
+        info.session_id = self._instance.id
+        return info
+
+    def delete_session(self) -> None:
+        if not self._nested:
+            self._instance.stop()
+        else:
+            req_data = {"output_format": self._output_format}
+            self._instance.put_task_info(
+                self._task_name,
+                MAXFRAME_TASK_DELETE_SESSION_METHOD,
+                json.dumps(req_data),
+            )
+
+    def submit_dag(
+        self,
+        dag: TileableGraph,
+        managed_input_infos: Optional[Dict[str, ResultInfo]] = None,
+    ) -> DagInfo:
+        req_data = {
+            "protocol": MAXFRAME_DEFAULT_PROTOCOL,
+            "dag": base64.b64encode(serialize_serializable(dag)).decode(),
+            "managed_input_infos": base64.b64encode(
+                serialize_serializable(managed_input_infos)
+            ).decode(),
+            "output_format": self._output_format,
+        }
+        res = self._instance.put_task_info(
+            self._task_name, MAXFRAME_TASK_SUBMIT_DAG_METHOD, json.dumps(req_data)
+        )
+        return self._deserial_task_info_result(res, DagInfo)
+
+    def get_dag_info(self, dag_id: str) -> DagInfo:
+        req_data = {
+            "protocol": MAXFRAME_DEFAULT_PROTOCOL,
+            "dag_id": dag_id,
+            "output_format": self._output_format,
+        }
+        res = self._instance.put_task_info(
+            self._task_name, MAXFRAME_TASK_GET_DAG_INFO_METHOD, json.dumps(req_data)
+        )
+        return self._deserial_task_info_result(res, DagInfo)
+
+    def cancel_dag(self, dag_id: str) -> DagInfo:
+        req_data = {
+            "protocol": MAXFRAME_DEFAULT_PROTOCOL,
+            "dag_id": dag_id,
+            "output_format": self._output_format,
+        }
+        res = self._instance.put_task_info(
+            self._task_name, MAXFRAME_TASK_CANCEL_DAG_METHOD, json.dumps(req_data)
+        )
+        return self._deserial_task_info_result(res, DagInfo)
+
+    def decref(self, tileable_keys: List[str]) -> None:
+        req_data = {
+            "tileable_keys": ",".join(tileable_keys),
+        }
+        self._instance.put_task_info(
+            self._task_name, MAXFRAME_TASK_DECREF_METHOD, json.dumps(req_data)
+        )
+
+    def get_logview_address(self, dag_id=None, hours=None) -> Optional[str]:
+        hours = hours or options.session.logview_hours
+        subquery_suffix = f"&subQuery={dag_id}" if dag_id else ""
+        return self._instance.get_logview_address(hours) + subquery_suffix
+
+
+class MaxFrameTaskSession(MaxFrameSession):
+    schemes = [ODPS_SESSION_INSECURE_SCHEME, ODPS_SESSION_SECURE_SCHEME]
+
+    _instance: Instance
+
+    @classmethod
+    def _create_caller(
+        cls,
+        odps_entry: ODPS,
+        address: str,
+        priority: Optional[int] = None,
+        project: Optional[str] = None,
+        running_cluster: Optional[str] = None,
+        **kwargs,
+    ) -> MaxFrameServiceCaller:
+        return MaxFrameInstanceCaller(
+            odps_entry,
+            priority=priority,
+            running_cluster=running_cluster,
+            project=project,
+            **kwargs,
+        )
+
+
+def register_session_schemes(overwrite: bool = False):
+    MaxFrameTaskSession.register_schemes(overwrite=overwrite)

maxframe_client/session/tests/__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/session/tests/test_task.py

@@ -0,0 +1,85 @@
+# 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 json
+import os
+
+import mock
+from defusedxml import ElementTree
+from odps import ODPS
+from odps import options as odps_options
+
+from ...session.consts import MAXFRAME_OUTPUT_JSON_FORMAT
+from ...session.task import MaxFrameInstanceCaller, MaxFrameTask
+
+expected_file_dir = os.path.join(os.path.dirname(__file__), "expected-data")
+
+
+def test_maxframe_instance_caller_creating_session():
+    o = ODPS.from_environments()
+
+    def create_caller(**kwargs):
+        kw = dict(
+            odps_entry=o,
+            task_name="task_test",
+            major_version="test_version",
+            output_format=MAXFRAME_OUTPUT_JSON_FORMAT,
+            running_cluster="test_cluster",
+        )
+        kw.update(**kwargs)
+        return MaxFrameInstanceCaller(**kw)
+
+    def mock_create(self, task: MaxFrameTask, priority=None, running_cluster=None):
+        assert priority == 100
+        assert running_cluster == "test_cluster"
+        root = ElementTree.parse(
+            os.path.join(expected_file_dir, "create_session.xml")
+        ).getroot()
+        assert root.tag == "MaxFrame"
+        assert root.find("Name").text == "task_test"
+        assert root.find("Command").text == "CREATE_SESSION"
+        property_node = root.find("Config").find("Property")
+        assert property_node.find("Name").text == "settings"
+        setting_dict = json.loads(property_node.find("Value").text)
+        assert setting_dict["odps.task.major.version"] == "test_version"
+        assert (
+            setting_dict["odps.service.endpoint"]
+            == "http://100.69.248.78:8002/odps_dailyrunnew"
+        )
+
+        assert setting_dict["odps.maxframe.output_format"] == "json"
+        maxframe_setting_dict = json.loads(setting_dict["odps.maxframe.settings"])
+        assert maxframe_setting_dict["session.max_alive_seconds"] == 259200
+
+    with mock.patch.multiple(
+        target="maxframe_client.session.task.MaxFrameInstanceCaller",
+        _wait_instance_task_ready=mock.DEFAULT,
+        get_session=mock.DEFAULT,
+    ), mock.patch("odps.models.instances.BaseInstances.create", mock_create):
+        task_caller = create_caller(priority=100)
+        task_caller.create_session()
+
+        old_priority = odps_options.priority
+        old_get_priority = odps_options.get_priority
+        try:
+            task_caller = create_caller(priority=100)
+            odps_options.priority = 100
+            task_caller.create_session()
+
+            odps_options.priority = None
+            odps_options.get_priority = lambda _: 100
+            task_caller.create_session()
+        finally:
+            odps_options.priority = old_priority
+            odps_options.get_priority = old_get_priority

maxframe_client/tests/__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/tests/test_fetcher.py

@@ -0,0 +1,89 @@
+# 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 uuid
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+from odps import ODPS
+
+import maxframe.dataframe as md
+from maxframe.odpsio import HaloTableIO
+from maxframe.protocol import ODPSTableResultInfo, ResultType
+from maxframe.tests.utils import tn
+
+from ..fetcher import ODPSTableFetcher
+
+
+async def test_table_fetcher():
+    odps_entry = ODPS.from_environments()
+    halo_table_io = HaloTableIO(odps_entry)
+    fetcher = ODPSTableFetcher(odps_entry)
+
+    data = pd.DataFrame(
+        {
+            "_idx_0": np.arange(1000),
+            "a": np.random.rand(1000),
+            "b": np.random.randint(0, 10, 1000),
+            "c": np.random.choice(list("ABC"), 1000),
+        }
+    )
+
+    table_name = tn("mf_test_groupby_table_" + str(uuid.uuid4().hex))
+    odps_entry.delete_table(table_name, if_exists=True)
+
+    odps_entry.create_table(
+        table_name, "_idx_0 bigint, a double, b bigint, c string", lifecycle=1
+    )
+    with halo_table_io.open_writer(table_name) as writer:
+        writer.write(pa.Table.from_pandas(data, preserve_index=False))
+
+    raw_data = data[list("abc")]
+
+    tileable = md.read_pandas(data[list("abc")])
+    result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
+    fetched = await fetcher.fetch(tileable, result_info, [None, None])
+    assert len(fetched) == 1000
+    pd.testing.assert_frame_equal(raw_data, fetched)
+
+    result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
+    fetched = await fetcher.fetch(tileable, result_info, [2, None])
+    assert len(fetched) == 1
+    pd.testing.assert_frame_equal(raw_data.iloc[2:3], fetched)
+
+    result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
+    fetched = await fetcher.fetch(tileable, result_info, [slice(None, 10), None])
+    assert len(fetched) == 10
+    pd.testing.assert_frame_equal(raw_data.iloc[:10], fetched)
+
+    result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
+    fetched = await fetcher.fetch(tileable, result_info, [slice(None, 10, 3), None])
+    assert len(fetched) == 4
+    pd.testing.assert_frame_equal(raw_data.iloc[:10:3], fetched)
+
+    tileable = md.read_pandas(data[list("ab")])
+    result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
+    fetched = await fetcher.fetch(tileable, result_info, [slice(-5, None), slice(2)])
+    assert len(fetched) == 5
+    pd.testing.assert_frame_equal(raw_data.iloc[-5:, :2], fetched)
+
+    tileable = md.read_pandas(data[list("a")])
+    result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
+    fetched = await fetcher.fetch(tileable, result_info, [slice(-1, -6, -1), 0])
+    assert len(fetched) == 5
+    assert len(fetched.columns) == 1
+    pd.testing.assert_frame_equal(raw_data.iloc[-1:-6:-1, :1], fetched)
+
+    odps_entry.delete_table(table_name, if_exists=True)