maxframe 0.1.0b4__cp311-cp311-win32.whl → 1.0.0__cp311-cp311-win32.whl

maxframe_client/session/odps.py

@@ -14,28 +14,44 @@
 
 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, enter_mode
+from maxframe.core import Entity, TileableGraph, build_fetch, enter_mode
+from maxframe.core.operator import Fetch
 from maxframe.dataframe import read_odps_table
 from maxframe.dataframe.core import DATAFRAME_TYPE, SERIES_TYPE
 from maxframe.dataframe.datasource import PandasDataSourceOperator
 from maxframe.dataframe.datasource.read_odps_table import DataFrameReadODPSTable
-from maxframe.odpsio import HaloTableIO, pandas_to_arrow, pandas_to_odps_schema
+from maxframe.errors import (
+    MaxFrameError,
+    NoTaskServerResponseError,
+    SessionAlreadyClosedError,
+)
+from maxframe.io.objects import get_object_io_handler
+from maxframe.io.odpsio import (
+    ODPSTableIO,
+    ODPSVolumeWriter,
+    pandas_to_arrow,
+    pandas_to_odps_schema,
+)
 from maxframe.protocol import (
     DagInfo,
     DagStatus,
     ODPSTableResultInfo,
+    ODPSVolumeResultInfo,
     ResultInfo,
     SessionInfo,
 )
@@ -46,8 +62,15 @@ from maxframe.session import (
     Profiling,
     Progress,
 )
+from maxframe.tensor.datasource import ArrayDataSource
 from maxframe.typing_ import TileableType
-from maxframe.utils import ToThreadMixin, build_temp_table_name
+from maxframe.utils import (
+    ToThreadMixin,
+    build_session_volume_name,
+    build_temp_table_name,
+    str_to_bool,
+    sync_pyodps_options,
+)
 
 from ..clients.framedriver import FrameDriverClient
 from ..fetcher import get_fetcher_cls
@@ -58,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
@@ -68,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
 
@@ -84,6 +149,9 @@ class MaxFrameServiceCaller(metaclass=abc.ABCMeta):
     def decref(self, tileable_keys: List[str]) -> None:
         raise NotImplementedError
 
+    def get_logview_address(self, dag_id=None, hours=None) -> Optional[str]:
+        return None
+
 
 class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     _odps_entry: Optional[ODPS]
@@ -119,6 +187,8 @@ 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
 
     @classmethod
     def _create_caller(
@@ -128,28 +198,32 @@ 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_read_tileable(self, t: TileableType) -> Optional[TileableType]:
-        if (
-            not isinstance(t.op, PandasDataSourceOperator)
-            or t.op.get_data() is None
-            or t.inputs
-        ):
-            return None
-
-        schema, table_meta = pandas_to_odps_schema(t, unknown_as_string=True)
+    def _upload_and_get_table_read_tileable(
+        self, t: TileableType
+    ) -> Optional[TileableType]:
+        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)
+            self._odps_entry.delete_table(
+                table_meta.table_name, hints=options.sql.settings
+            )
         table_name = build_temp_table_name(self.session_id, t.key)
-        table_obj = self._odps_entry.create_table(table_name, schema)
+        table_obj = self._odps_entry.create_table(
+            table_name,
+            table_schema,
+            lifecycle=options.session.temp_table_lifecycle,
+            hints=options.sql.settings,
+        )
 
         data = t.op.get_data()
         batch_size = options.session.upload_batch_size
 
         if len(data):
-            halo_client = HaloTableIO(self._odps_entry)
-            with halo_client.open_writer(table_obj.full_table_name) as writer:
+            table_client = ODPSTableIO(self._odps_entry)
+            with table_client.open_writer(table_obj.full_table_name) as writer:
                 for batch_start in range(0, len(data), batch_size):
                     if isinstance(data, pd.Index):
                         batch = data[batch_start : batch_start + batch_size]
@@ -172,13 +246,35 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
                 read_tileable.name = t.name
         else:  # INDEX_TYPE
             if list(read_tileable.names) != list(t.names):
-                read_tileable.names = t.names
+                read_tileable.rename(t.names, inplace=True)
         read_tileable._key = t.key
         read_tileable.params = t.params
         return read_tileable.data
 
+    def _upload_and_get_vol_read_tileable(
+        self, t: TileableType
+    ) -> Optional[TileableType]:
+        vol_name = build_session_volume_name(self.session_id)
+        writer = ODPSVolumeWriter(self._odps_entry, vol_name, t.key)
+        io_handler = get_object_io_handler(t)
+        io_handler().write_object(writer, t, t.op.data)
+        return build_fetch(t).data
+
+    def _upload_and_get_read_tileable(self, t: TileableType) -> Optional[TileableType]:
+        if (
+            not isinstance(t.op, (ArrayDataSource, PandasDataSourceOperator))
+            or t.op.get_data() is None
+            or t.inputs
+        ):
+            return None
+        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_pandas_sources(
+    def _scan_and_replace_local_sources(
         self, graph: TileableGraph
     ) -> Dict[TileableType, TileableType]:
         """Replaces Pandas data sources with temp table sources in the graph"""
@@ -199,7 +295,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
@@ -207,16 +303,41 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     @enter_mode(kernel=True, build=True)
     def _get_input_infos(self, tileables: List[TileableType]) -> Dict[str, ResultInfo]:
         """Generate ResultInfo structs from generated temp tables"""
+        vol_name = build_session_volume_name(self.session_id)
+
         infos = dict()
         for t in tileables:
             key = t.key
-            if not isinstance(t.op, DataFrameReadODPSTable):
-                if not isinstance(t.inputs[0].op, DataFrameReadODPSTable):
-                    continue
-                t = t.inputs[0]
-            infos[key] = ODPSTableResultInfo(full_table_name=t.op.table_name)
+            if isinstance(t.op, DataFrameReadODPSTable):
+                infos[key] = ODPSTableResultInfo(full_table_name=t.op.table_name)
+            else:
+                if isinstance(t.op, Fetch):
+                    infos[key] = ODPSVolumeResultInfo(
+                        volume_name=vol_name, volume_path=t.key
+                    )
+                elif t.inputs and isinstance(t.inputs[0].op, DataFrameReadODPSTable):
+                    t = t.inputs[0]
+                    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
@@ -226,7 +347,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         tileable_graph, to_execute_tileables = gen_submit_tileable_graph(
             self, tileables, tileable_to_copied
         )
-        source_replacements = self._scan_and_replace_pandas_sources(tileable_graph)
+        source_replacements = self._scan_and_replace_local_sources(tileable_graph)
 
         # we need to manage uploaded data sources with refcounting mechanism
         # as nodes in tileable_graph are copied, we need to use original nodes
@@ -236,9 +357,14 @@ 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)
+
         progress = Progress()
         profiling = Profiling()
         aio_task = asyncio.create_task(
@@ -256,25 +382,55 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         self, dag_info: DagInfo, tileables: List, progress: Progress
     ):
         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")
 
-                    dag_info: DagInfo = await self.ensure_async_call(
-                        self._caller.get_dag_info, dag_id
-                    )
+                    try:
+                        dag_info: DagInfo = await self.ensure_async_call(
+                            self._caller.get_dag_info, dag_id
+                        )
+                        server_no_response_time = None
+                    except (NoTaskServerResponseError, SessionAlreadyClosedError) as ex:
+                        # when we receive SessionAlreadyClosedError after NoTaskServerResponseError
+                        #  is received, it is possible that task server is restarted and
+                        #  SessionAlreadyClosedError might be flaky. Otherwise, the error
+                        #  should be raised.
+                        if (
+                            isinstance(ex, SessionAlreadyClosedError)
+                            and not server_no_response_time
+                        ):
+                            raise
+                        server_no_response_time = server_no_response_time or time.time()
+                        if (
+                            time.time() - server_no_response_time
+                            > options.client.task_restart_timeout
+                        ):
+                            raise MaxFrameError(
+                                "Failed to get valid response from service. "
+                                f"Session {self._session_id}."
+                            ) from None
+                        await asyncio.sleep(timeout_val)
+                        continue
+
+                    if dag_info is None:
+                        raise SystemError(
+                            f"Cannot find DAG with ID {dag_id} in session {session_id}"
+                        )
                     progress.value = dag_info.progress
                     if dag_info.status != DagStatus.RUNNING:
                         break
@@ -294,6 +450,8 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
             for key, result_info in dag_info.tileable_to_result_infos.items():
                 t = key_to_tileables[key]
+                fetcher = get_fetcher_cls(result_info.result_type)(self._odps_entry)
+                await fetcher.update_tileable_meta(t, result_info)
                 self._tileable_to_infos[t] = result_info
 
     def _get_data_tileable_and_indexes(
@@ -334,7 +492,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
                 data_tileable, indexes = self._get_data_tileable_and_indexes(tileable)
                 info = self._tileable_to_infos[data_tileable]
                 fetcher = get_fetcher_cls(info.result_type)(self._odps_entry)
-                results.append(await fetcher.fetch(tileable, info, indexes))
+                results.append(await fetcher.fetch(data_tileable, info, indexes))
         return results
 
     async def decref(self, *tileable_keys):
@@ -388,12 +546,32 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     async def get_mutable_tensor(self, name: str):
         raise NotImplementedError
 
+    async def get_logview_address(self, hours=None) -> Optional[str]:
+        return await self.get_dag_logview_address(None, hours)
+
+    async def get_dag_logview_address(self, dag_id=None, hours=None) -> Optional[str]:
+        return await self.ensure_async_call(
+            self._caller.get_logview_address, dag_id, hours
+        )
+
+    async def _show_logview_address(self, dag_id=None, hours=None):
+        identity = f"Session ID: {self._session_id}"
+        if dag_id:
+            identity += f", DAG ID: {dag_id}"
+
+        logview_addr = await self.get_dag_logview_address(dag_id, hours)
+        if logview_addr:
+            logger.info("%s, Logview: %s", identity, logview_addr)
+        else:
+            logger.info("%s, Logview address does not exist", identity)
+
 
 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
 
@@ -406,9 +584,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)
@@ -446,7 +629,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):