maxframe 1.0.0rc3__cp310-cp310-macosx_10_9_universal2.whl → 1.1.0__cp310-cp310-macosx_10_9_universal2.whl

maxframe_client/session/odps.py

@@ -14,16 +14,19 @@
 
 import abc
 import asyncio
+import copy
 import logging
 import time
 import weakref
 from numbers import Integral
-from typing import Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
 from urllib.parse import urlparse
 
 import numpy as np
 import pandas as pd
 from odps import ODPS
+from odps import options as odps_options
+from odps.console import in_ipython_frontend
 
 from maxframe.config import options
 from maxframe.core import Entity, TileableGraph, build_fetch, enter_mode
@@ -65,6 +68,8 @@ from maxframe.utils import (
     ToThreadMixin,
     build_session_volume_name,
     build_temp_table_name,
+    str_to_bool,
+    sync_pyodps_options,
 )
 
 from ..clients.framedriver import FrameDriverClient
@@ -76,6 +81,45 @@ logger = logging.getLogger(__name__)
 
 
 class MaxFrameServiceCaller(metaclass=abc.ABCMeta):
+    def get_settings_to_upload(self) -> Dict[str, Any]:
+        sql_settings = (odps_options.sql.settings or {}).copy()
+        sql_settings.update(options.sql.settings or {})
+
+        quota_name = options.session.quota_name or getattr(
+            odps_options, "quota_name", None
+        )
+        lifecycle = options.session.table_lifecycle or odps_options.lifecycle
+        temp_lifecycle = (
+            options.session.temp_table_lifecycle or odps_options.temp_lifecycle
+        )
+
+        enable_schema = options.session.enable_schema
+        default_schema = options.session.default_schema
+        if hasattr(self, "_odps_entry"):
+            default_schema = default_schema or self._odps_entry.schema
+
+        # use flags in sql settings
+        if sql_settings.get("odps.default.schema"):
+            default_schema = sql_settings["odps.default.schema"]
+        if str_to_bool(
+            sql_settings.get("odps.namespace.schema") or "false"
+        ) or str_to_bool(
+            sql_settings.get("odps.sql.allow.namespace.schema") or "false"
+        ):
+            enable_schema = True
+
+        mf_settings = dict(options.to_dict(remote_only=True).items())
+        mf_settings["sql.settings"] = sql_settings
+        mf_settings["session.table_lifecycle"] = lifecycle
+        mf_settings["session.temp_table_lifecycle"] = temp_lifecycle
+        mf_settings["session.quota_name"] = quota_name
+        if enable_schema is not None:
+            mf_settings["session.enable_schema"] = enable_schema
+        if options.session.enable_high_availability is None:
+            mf_settings["session.enable_high_availability"] = not in_ipython_frontend()
+        mf_settings["session.default_schema"] = default_schema or "default"
+        return mf_settings
+
     @abc.abstractmethod
     def create_session(self) -> SessionInfo:
         raise NotImplementedError
@@ -86,7 +130,10 @@ class MaxFrameServiceCaller(metaclass=abc.ABCMeta):
 
     @abc.abstractmethod
     def submit_dag(
-        self, dag: TileableGraph, managed_input_infos: Dict[str, ResultInfo]
+        self,
+        dag: TileableGraph,
+        managed_input_infos: Dict[str, ResultInfo],
+        new_settings: Dict[str, Any] = None,
     ) -> DagInfo:
         raise NotImplementedError
 
@@ -140,6 +187,9 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         self._tileable_to_infos = weakref.WeakKeyDictionary()
 
         self._caller = self._create_caller(odps_entry, address, **kwargs)
+        self._last_settings = None
+        self._pull_interval = 1 if in_ipython_frontend() else 3
+        self._replace_internal_host = kwargs.get("replace_internal_host", True)
 
     @classmethod
     def _create_caller(
@@ -149,13 +199,14 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
     async def _init(self, _address: str):
         session_info = await self.ensure_async_call(self._caller.create_session)
+        self._last_settings = copy.deepcopy(self._caller.get_settings_to_upload())
         self._session_id = session_info.session_id
         await self._show_logview_address()
 
     def _upload_and_get_table_read_tileable(
         self, t: TileableType
     ) -> Optional[TileableType]:
-        schema, table_meta = pandas_to_odps_schema(t, unknown_as_string=True)
+        table_schema, table_meta = pandas_to_odps_schema(t, unknown_as_string=True)
         if self._odps_entry.exist_table(table_meta.table_name):
             self._odps_entry.delete_table(
                 table_meta.table_name, hints=options.sql.settings
@@ -163,7 +214,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         table_name = build_temp_table_name(self.session_id, t.key)
         table_obj = self._odps_entry.create_table(
             table_name,
-            schema,
+            table_schema,
             lifecycle=options.session.temp_table_lifecycle,
             hints=options.sql.settings,
         )
@@ -205,7 +256,12 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         self, t: TileableType
     ) -> Optional[TileableType]:
         vol_name = build_session_volume_name(self.session_id)
-        writer = ODPSVolumeWriter(self._odps_entry, vol_name, t.key)
+        writer = ODPSVolumeWriter(
+            self._odps_entry,
+            vol_name,
+            t.key,
+            replace_internal_host=self._replace_internal_host,
+        )
         io_handler = get_object_io_handler(t)
         io_handler().write_object(writer, t, t.op.data)
         return build_fetch(t).data
@@ -217,10 +273,11 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
             or t.inputs
         ):
             return None
-        if isinstance(t.op, PandasDataSourceOperator):
-            return self._upload_and_get_table_read_tileable(t)
-        else:
-            return self._upload_and_get_vol_read_tileable(t)
+        with sync_pyodps_options():
+            if isinstance(t.op, PandasDataSourceOperator):
+                return self._upload_and_get_table_read_tileable(t)
+            else:
+                return self._upload_and_get_vol_read_tileable(t)
 
     @enter_mode(kernel=True, build=True)
     def _scan_and_replace_local_sources(
@@ -244,7 +301,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
             for succ in successors:
                 graph.add_edge(replaced, succ)
-                succ.inputs = [replacements.get(t, t) for t in succ.inputs]
+                succ.op._set_inputs([replacements.get(t, t) for t in succ.inputs])
 
         graph.results = [replacements.get(t, t) for t in graph.results]
         return replacements
@@ -269,6 +326,24 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
                     infos[key] = ODPSTableResultInfo(full_table_name=t.op.table_name)
         return infos
 
+    def _get_diff_settings(self) -> Dict[str, Any]:
+        new_settings = self._caller.get_settings_to_upload()
+        if not self._last_settings:  # pragma: no cover
+            self._last_settings = copy.deepcopy(new_settings)
+            return new_settings
+
+        update = dict()
+        for k in new_settings.keys():
+            old_item = self._last_settings.get(k)
+            new_item = new_settings.get(k)
+            try:
+                if old_item != new_item:
+                    update[k] = new_item
+            except:  # noqa: E722  # nosec  # pylint: disable=bare-except
+                update[k] = new_item
+        self._last_settings = copy.deepcopy(new_settings)
+        return update
+
     async def execute(self, *tileables, **kwargs) -> ExecutionInfo:
         tileables = [
             tileable.data if isinstance(tileable, Entity) else tileable
@@ -288,7 +363,10 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
         replaced_infos = self._get_input_infos(list(source_replacements.values()))
         dag_info = await self.ensure_async_call(
-            self._caller.submit_dag, tileable_graph, replaced_infos
+            self._caller.submit_dag,
+            tileable_graph,
+            replaced_infos,
+            self._get_diff_settings(),
         )
 
         await self._show_logview_address(dag_info.dag_id)
@@ -312,18 +390,18 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         start_time = time.time()
         session_id = dag_info.session_id
         dag_id = dag_info.dag_id
-        wait_timeout = 10
         server_no_response_time = None
         with enter_mode(build=True, kernel=True):
             key_to_tileables = {t.key: t for t in tileables}
-
+            timeout_val = 0.1
             try:
                 while True:
                     elapsed_time = time.time() - start_time
+                    next_timeout_val = min(timeout_val * 2, self._pull_interval)
                     timeout_val = (
-                        min(self.timeout - elapsed_time, wait_timeout)
+                        min(self.timeout - elapsed_time, next_timeout_val)
                         if self.timeout
-                        else wait_timeout
+                        else next_timeout_val
                     )
                     if timeout_val <= 0:
                         raise TimeoutError("Running DAG timed out")
@@ -498,7 +576,8 @@ class MaxFrameRestCaller(MaxFrameServiceCaller):
     _client: FrameDriverClient
     _session_id: Optional[str]
 
-    def __init__(self, client: FrameDriverClient):
+    def __init__(self, odps_entry: ODPS, client: FrameDriverClient):
+        self._odps_entry = odps_entry
         self._client = client
         self._session_id = None
 
@@ -511,9 +590,14 @@ class MaxFrameRestCaller(MaxFrameServiceCaller):
         await self._client.delete_session(self._session_id)
 
     async def submit_dag(
-        self, dag: TileableGraph, managed_input_infos: Dict[str, ResultInfo]
+        self,
+        dag: TileableGraph,
+        managed_input_infos: Dict[str, ResultInfo] = None,
+        new_settings: Dict[str, Any] = None,
     ) -> DagInfo:
-        return await self._client.submit_dag(self._session_id, dag, managed_input_infos)
+        return await self._client.submit_dag(
+            self._session_id, dag, managed_input_infos, new_settings=new_settings
+        )
 
     async def get_dag_info(self, dag_id: str) -> DagInfo:
         return await self._client.get_dag_info(self._session_id, dag_id)
@@ -546,12 +630,12 @@ class MaxFrameRestSession(MaxFrameSession):
         real_endpoint = address.replace(f"{parsed_endpoint.scheme}://", f"{scheme}://")
 
         super().__init__(
-            real_endpoint, session_id, odps_entry=odps_entry, timeout=timeout
+            real_endpoint, session_id, odps_entry=odps_entry, timeout=timeout, **kwargs
         )
 
     @classmethod
     def _create_caller(cls, odps_entry: ODPS, address: str, **kwargs):
-        return MaxFrameRestCaller(FrameDriverClient(address))
+        return MaxFrameRestCaller(odps_entry, FrameDriverClient(address))
 
 
 def register_session_schemes(overwrite: bool = False):

maxframe_client/session/task.py

@@ -16,12 +16,12 @@ import base64
 import json
 import logging
 import time
-from typing import Dict, List, Optional, Type, Union
+from typing import Any, 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.errors import EmptyTaskInfoError, parse_instance_error
 from odps.models import Instance, MaxFrameTask
 
 from maxframe.config import options
@@ -36,6 +36,7 @@ except ImportError:
     mf_version = None
 
 from .consts import (
+    EMPTY_RESPONSE_RETRY_COUNT,
     MAXFRAME_DEFAULT_PROTOCOL,
     MAXFRAME_OUTPUT_JSON_FORMAT,
     MAXFRAME_OUTPUT_MAXFRAME_FORMAT,
@@ -92,6 +93,10 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             self._nested = True
             self._instance = odps_entry.get_instance(nested_instance_id)
 
+    @property
+    def instance(self):
+        return self._instance
+
     def _deserial_task_info_result(
         self, content: Union[bytes, str, dict], target_cls: Type[JsonSerializable]
     ):
@@ -120,21 +125,9 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             )
 
     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
-
+        task = MaxFrameTask(name=self._task_name, major_version=self._major_version)
         mf_opts = {
-            "odps.maxframe.settings": json.dumps(mf_settings),
+            "odps.maxframe.settings": json.dumps(self.get_settings_to_upload()),
             "odps.maxframe.output_format": self._output_format,
         }
         if mf_version:
@@ -189,18 +182,27 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             interval = min(max_interval, interval * 2)
 
     def _put_task_info(self, method_name: str, json_data: dict):
-        resp_data = self._instance.put_task_info(
-            self._task_name, method_name, json.dumps(json_data)
-        )
-        if not resp_data:
-            raise NoTaskServerResponseError(f"No response for request {method_name}")
-        return resp_data
+        for trial in range(EMPTY_RESPONSE_RETRY_COUNT):
+            try:
+                return self._instance.put_task_info(
+                    self._task_name,
+                    method_name,
+                    json.dumps(json_data),
+                    raise_empty=True,
+                )
+            except EmptyTaskInfoError as ex:
+                # retry when server returns HTTP 204, which is designed for retry
+                if ex.code != 204 or trial >= EMPTY_RESPONSE_RETRY_COUNT - 1:
+                    raise NoTaskServerResponseError(
+                        f"No response for request {method_name}. "
+                        f"Instance ID: {self._instance.id}. "
+                        f"Request ID: {ex.request_id}"
+                    ) from None
+                time.sleep(0.5)
 
     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)
-        )
+        serialized = self._put_task_info(MAXFRAME_TASK_GET_SESSION_METHOD, req_data)
         info: SessionInfo = self._deserial_task_info_result(serialized, SessionInfo)
         info.session_id = self._instance.id
         return info
@@ -217,13 +219,18 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
         self,
         dag: TileableGraph,
         managed_input_infos: Optional[Dict[str, ResultInfo]] = None,
+        new_settings: Dict[str, Any] = None,
     ) -> DagInfo:
+        new_settings_value = {
+            "odps.maxframe.settings": json.dumps(new_settings),
+        }
         req_data = {
             "protocol": MAXFRAME_DEFAULT_PROTOCOL,
             "dag": base64.b64encode(serialize_serializable(dag)).decode(),
             "managed_input_infos": base64.b64encode(
                 serialize_serializable(managed_input_infos)
             ).decode(),
+            "new_settings": json.dumps(new_settings_value),
             "output_format": self._output_format,
         }
         res = self._put_task_info(MAXFRAME_TASK_SUBMIT_DAG_METHOD, req_data)
@@ -276,7 +283,7 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
 class MaxFrameTaskSession(MaxFrameSession):
     schemes = [ODPS_SESSION_INSECURE_SCHEME, ODPS_SESSION_SECURE_SCHEME]
 
-    _instance: Instance
+    _caller: MaxFrameInstanceCaller
 
     @classmethod
     def _create_caller(
@@ -296,6 +303,15 @@ class MaxFrameTaskSession(MaxFrameSession):
             **kwargs,
         )
 
+    @property
+    def closed(self) -> bool:
+        if super().closed:
+            return True
+        if not self._caller or not self._caller.instance:
+            # session not initialized yet
+            return False
+        return self._caller.instance.is_terminated()
+
 
 def register_session_schemes(overwrite: bool = False):
     MaxFrameTaskSession.register_schemes(overwrite=overwrite)

maxframe_client/session/tests/test_task.py

@@ -53,10 +53,6 @@ def test_maxframe_instance_caller_creating_session():
         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"])

maxframe_client/tests/test_session.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import time
-from typing import Dict
+from typing import Any, Dict
 
 import mock
 import numpy as np
@@ -31,7 +31,7 @@ from maxframe.lib.aio import stop_isolation
 from maxframe.protocol import ResultInfo
 from maxframe.serialization import RemoteException
 from maxframe.session import new_session
-from maxframe.tests.utils import tn
+from maxframe.tests.utils import ensure_table_deleted, tn
 from maxframe.utils import build_temp_table_name
 from maxframe_framedriver.app.tests.test_framedriver_webapp import (  # noqa: F401
     framedriver_app,
@@ -86,9 +86,12 @@ def test_simple_run_dataframe(start_mock_session):
         session_id: str,
         dag: TileableGraph,
         managed_input_infos: Dict[str, ResultInfo] = None,
+        new_settings: Dict[str, Any] = None,
     ):
         assert len(dag) == 2
-        return await original_submit_dag(self, session_id, dag, managed_input_infos)
+        return await original_submit_dag(
+            self, session_id, dag, managed_input_infos, new_settings
+        )
 
     no_task_server_raised = False
     original_get_dag_info = MaxFrameRestCaller.get_dag_info
@@ -130,11 +133,19 @@ def test_simple_run_dataframe(start_mock_session):
     )
     assert odps_entry.exist_table(build_temp_table_name(start_mock_session, key))
     del df
-    time.sleep(5)
-    assert not odps_entry.exist_table(
-        build_temp_table_name(start_mock_session, intermediate_key)
+    ensure_table_deleted(
+        odps_entry, build_temp_table_name(start_mock_session, intermediate_key)
     )
-    assert not odps_entry.exist_table(build_temp_table_name(start_mock_session, key))
+    ensure_table_deleted(odps_entry, build_temp_table_name(start_mock_session, key))
+
+
+def test_run_and_fetch_slice(start_mock_session):
+    pd_df = pd.DataFrame(np.random.rand(1000, 5), columns=list("ABCDE"))
+    df = md.DataFrame(pd_df)
+    result = df.execute()
+
+    sliced = result.head(10).fetch()
+    assert len(sliced) == 10
 
 
 def test_run_empty_table(start_mock_session):
@@ -157,6 +168,25 @@ def test_run_empty_table(start_mock_session):
     empty_table.drop()
 
 
+def test_run_odps_query_without_schema(start_mock_session):
+    odps_entry = ODPS.from_environments()
+
+    table_name = tn("test_query_without_schema")
+    odps_entry.delete_table(table_name, if_exists=True)
+    test_table = odps_entry.create_table(table_name, "a double, b double", lifecycle=1)
+
+    with test_table.open_writer() as writer:
+        writer.write([123, 456])
+
+    df = md.read_odps_query(
+        f"select a, b, a + b as `special: name` from {table_name}", skip_schema=True
+    )
+    executed = df.execute().fetch()
+    assert len(executed.dtypes) == 3
+
+    test_table.drop()
+
+
 def test_run_dataframe_with_pd_source(start_mock_session):
     odps_entry = ODPS.from_environments()
 
@@ -189,7 +219,7 @@ def test_run_dataframe_from_to_odps_table(start_mock_session):
     table_name = build_temp_table_name(start_mock_session, "tmp_save")
     table_obj = odps_entry.get_table(table_name)
     try:
-        md.to_odps_table(md.DataFrame(pd_df), table_obj).execute().fetch()
+        md.to_odps_table(md.DataFrame(pd_df), table_obj, lifecycle=1).execute().fetch()
         with table_obj.open_reader() as reader:
             result_df = reader.to_pandas()
         assert len(result_df) == 10
@@ -237,9 +267,8 @@ def test_run_and_fetch_series(start_mock_session):
         pd.testing.assert_series_equal(pd_result, result)
 
         del s1
-        time.sleep(5)
-        assert not odps_entry.exist_table(
-            build_temp_table_name(start_mock_session, src_key)
+        ensure_table_deleted(
+            odps_entry, build_temp_table_name(start_mock_session, src_key)
         )
     finally:
         odps_entry.delete_table(
@@ -256,7 +285,10 @@ def test_execute_with_tensor(oss_config, start_mock_session):
 
     result = (df - [1, 2]).execute().fetch()
     expected = pd_df - [1, 2]
-    pd.testing.assert_frame_equal(result, expected)
+    # TODO: currently the record order in tensor reading from table is the index
+    # sorting order
+    expected.sort_index(axis=0, inplace=True)
+    pd.testing.assert_frame_equal(result, expected, check_like=True)
 
 
 def test_run_remote_success(oss_config, start_mock_session):