maxframe 1.0.0rc3__cp39-cp39-macosx_10_9_universal2.whl → 1.0.0rc4__cp39-cp39-macosx_10_9_universal2.whl

maxframe_client/session/odps.py

@@ -18,12 +18,13 @@ 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 maxframe.config import options
 from maxframe.core import Entity, TileableGraph, build_fetch, enter_mode
@@ -65,6 +66,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 +79,43 @@ 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
+        mf_settings["session.default_schema"] = default_schema or "default"
+        return mf_settings
+
     @abc.abstractmethod
     def create_session(self) -> SessionInfo:
         raise NotImplementedError
@@ -86,7 +126,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 +183,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         self._tileable_to_infos = weakref.WeakKeyDictionary()
 
         self._caller = self._create_caller(odps_entry, address, **kwargs)
+        self._last_settings = None
 
     @classmethod
     def _create_caller(
@@ -149,13 +193,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 = 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 +208,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,
         )
@@ -217,10 +262,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 +290,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 +315,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 = 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 = new_settings
+        return update
+
     async def execute(self, *tileables, **kwargs) -> ExecutionInfo:
         tileables = [
             tileable.data if isinstance(tileable, Entity) else tileable
@@ -288,7 +352,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)
@@ -498,7 +565,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,7 +579,10 @@ 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)
 
@@ -551,7 +622,7 @@ class MaxFrameRestSession(MaxFrameSession):
 
     @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,7 +16,7 @@ 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
@@ -24,6 +24,12 @@ from odps import options as odps_options
 from odps.errors import parse_instance_error
 from odps.models import Instance, MaxFrameTask
 
+try:
+    from odps.errors import EmptyTaskInfoError
+except ImportError:  # pragma: no cover
+    # todo remove when pyodps>=0.12.0 is enforced
+    EmptyTaskInfoError = type("EmptyTaskInfoError", (Exception,), {})
+
 from maxframe.config import options
 from maxframe.core import TileableGraph
 from maxframe.errors import NoTaskServerResponseError, SessionAlreadyClosedError
@@ -36,6 +42,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 +99,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]
     ):
@@ -125,16 +136,8 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             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.settings": json.dumps(self.get_settings_to_upload()),
             "odps.maxframe.output_format": self._output_format,
         }
         if mf_version:
@@ -189,18 +192,39 @@ 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 TypeError:  # pragma: no cover
+                # todo remove when pyodps>=0.12.0 is enforced
+                resp_data = self._instance.put_task_info(
+                    self._task_name, method_name, json.dumps(json_data)
+                )
+                if resp_data:
+                    return resp_data
+                else:
+                    raise NoTaskServerResponseError(
+                        f"No response for request {method_name}. "
+                        f"Instance ID: {self._instance.id}"
+                    )
+            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 +241,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 +305,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 +325,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/tests/test_session.py

@@ -137,6 +137,15 @@ def test_simple_run_dataframe(start_mock_session):
     assert not odps_entry.exist_table(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):
     odps_entry = ODPS.from_environments()
 
@@ -189,7 +198,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
@@ -256,7 +265,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):