maxframe 0.1.0b5__cp311-cp311-macosx_10_9_universal2.whl → 1.0.0rc2__cp311-cp311-macosx_10_9_universal2.whl

maxframe_client/fetcher.py

@@ -12,11 +12,9 @@
 # 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
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import pandas as pd
 import pyarrow as pa
@@ -28,7 +26,7 @@ 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.odpsio import ODPSTableIO, arrow_to_pandas, build_dataframe_table_meta
 from maxframe.protocol import (
     DataFrameTableMeta,
     ODPSTableResultInfo,
@@ -38,7 +36,7 @@ from maxframe.protocol import (
 )
 from maxframe.tensor.core import TENSOR_TYPE
 from maxframe.typing_ import PandasObjectTypes, TileableType
-from maxframe.utils import ToThreadMixin, deserialize_serializable
+from maxframe.utils import ToThreadMixin
 
 _result_fetchers: Dict[ResultType, Type["ResultFetcher"]] = dict()
 
@@ -109,17 +107,12 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
         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 (
+            isinstance(tileable, DATAFRAME_TYPE)
+            and tileable.dtypes is None
+            and info.table_meta is not None
+        ):
+            tileable.refresh_from_table_meta(info.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
@@ -131,16 +124,39 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
                 )
                 total_records += session.count
             new_shape_list = list(tileable.shape)
-            new_shape_list[-1] = total_records
+            new_shape_list[0] = total_records
             tileable.params = {"shape": tuple(new_shape_list)}
 
+    @staticmethod
+    def _align_selection_with_shape(
+        row_sel: slice, shape: Tuple[Optional[int], ...]
+    ) -> dict:
+        size = shape[0]
+        if not row_sel.start and not row_sel.stop:
+            return {}
+        is_reversed = row_sel.step is not None and row_sel.step < 0
+        read_kw = {
+            "start": row_sel.start,
+            "stop": row_sel.stop,
+            "reverse_range": is_reversed,
+        }
+        if pd.isna(size):
+            return read_kw
+
+        if is_reversed and row_sel.start is not None:
+            read_kw["start"] = min(size - 1, row_sel.start)
+        if not is_reversed and row_sel.stop is not None:
+            read_kw["stop"] = min(size, row_sel.stop)
+        return read_kw
+
     def _read_single_source(
         self,
         table_meta: DataFrameTableMeta,
         info: ODPSTableResultInfo,
         indexes: List[Union[None, Integral, slice]],
+        shape: Tuple[Optional[int], ...],
     ):
-        table_io = HaloTableIO(self._odps_entry)
+        table_io = ODPSTableIO(self._odps_entry)
         read_kw = {}
         row_step = None
         if indexes:
@@ -148,13 +164,8 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
                 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
+                row_step = row_sel.step
+                read_kw = self._align_selection_with_shape(row_sel, shape)
             elif isinstance(row_sel, int):
                 read_kw["start"] = row_sel
                 read_kw["stop"] = row_sel + 1
@@ -173,8 +184,8 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
         with table_io.open_reader(
             info.full_table_name, info.partition_specs, **read_kw
         ) as reader:
-            reader_count = reader.count
             result = reader.read_all()
+            reader_count = result.num_rows
 
         if not row_step:
             return result
@@ -195,7 +206,7 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
     ) -> PandasObjectTypes:
         table_meta = build_dataframe_table_meta(tileable)
         arrow_table: pa.Table = await self.to_thread(
-            self._read_single_source, table_meta, info, indexes
+            self._read_single_source, table_meta, info, indexes, tileable.shape
         )
         return arrow_to_pandas(arrow_table, table_meta)
 

maxframe_client/session/odps.py

@@ -31,7 +31,12 @@ 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.odpsio import ODPSTableIO, pandas_to_arrow, pandas_to_odps_schema
 from maxframe.protocol import (
     DagInfo,
     DagStatus,
@@ -144,18 +149,23 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
         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, lifecycle=options.session.temp_table_lifecycle
+            table_name,
+            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]
@@ -178,7 +188,7 @@ 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
@@ -264,8 +274,10 @@ 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}
 
@@ -280,9 +292,37 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
                     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
@@ -344,7 +384,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):

maxframe_client/session/task.py

@@ -26,6 +26,7 @@ from odps.models import Instance, MaxFrameTask
 
 from maxframe.config import options
 from maxframe.core import TileableGraph
+from maxframe.errors import NoTaskServerResponseError, SessionAlreadyClosedError
 from maxframe.protocol import DagInfo, JsonSerializable, ResultInfo, SessionInfo
 from maxframe.utils import deserialize_serializable, serialize_serializable, to_str
 
@@ -82,6 +83,7 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
         self._running_cluster = running_cluster
         self._major_version = major_version
         self._output_format = output_format or MAXFRAME_OUTPUT_MSGPACK_FORMAT
+        self._deleted = False
 
         if nested_instance_id is None:
             self._nested = False
@@ -94,10 +96,18 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
         self, content: Union[bytes, str, dict], target_cls: Type[JsonSerializable]
     ):
         if isinstance(content, (str, bytes)):
+            if len(content) == 0:
+                content = "{}"
             json_data = json.loads(to_str(content))
         else:
             json_data = content
-        result_data = base64.b64decode(json_data["result"])
+        encoded_result = json_data.get("result")
+        if not encoded_result:
+            if self._deleted:
+                return None
+            else:
+                raise SessionAlreadyClosedError(self._instance.id)
+        result_data = base64.b64decode(encoded_result)
         if self._output_format == MAXFRAME_OUTPUT_MAXFRAME_FORMAT:
             return deserialize_serializable(result_data)
         elif self._output_format == MAXFRAME_OUTPUT_JSON_FORMAT:
@@ -178,6 +188,14 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             time.sleep(interval)
             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
+
     def get_session(self) -> SessionInfo:
         req_data = {"output_format": self._output_format}
         serialized = self._instance.put_task_info(
@@ -192,11 +210,8 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             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),
-            )
+            self._put_task_info(MAXFRAME_TASK_DELETE_SESSION_METHOD, req_data)
+        self._deleted = True
 
     def submit_dag(
         self,
@@ -211,9 +226,7 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             ).decode(),
             "output_format": self._output_format,
         }
-        res = self._instance.put_task_info(
-            self._task_name, MAXFRAME_TASK_SUBMIT_DAG_METHOD, json.dumps(req_data)
-        )
+        res = self._put_task_info(MAXFRAME_TASK_SUBMIT_DAG_METHOD, req_data)
         return self._deserial_task_info_result(res, DagInfo)
 
     def get_dag_info(self, dag_id: str) -> DagInfo:
@@ -222,9 +235,7 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             "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)
-        )
+        res = self._put_task_info(MAXFRAME_TASK_GET_DAG_INFO_METHOD, req_data)
         return self._deserial_task_info_result(res, DagInfo)
 
     def cancel_dag(self, dag_id: str) -> DagInfo:
@@ -233,23 +244,33 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             "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)
-        )
+        res = self._put_task_info(MAXFRAME_TASK_CANCEL_DAG_METHOD, 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)
-        )
+        self._put_task_info(MAXFRAME_TASK_DECREF_METHOD, req_data)
 
     def get_logview_address(self, dag_id=None, hours=None) -> Optional[str]:
+        """
+        Generate logview address
+
+        Parameters
+        ----------
+            dag_id: id of dag for which dag logview detail page to access
+            hours: hours of the logview address auth limit
+        Returns
+        -------
+            Logview address
+        """
         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
+        # notice: maxframe can't reuse subQuery else will conflict with mcqa when fetch resource data,
+        #         added dagId for maxframe so logview backend will return maxframe data format if
+        #         instance and dagId is provided.
+        dag_suffix = f"&dagId={dag_id}" if dag_id else ""
+        return self._instance.get_logview_address(hours) + dag_suffix
 
 
 class MaxFrameTaskSession(MaxFrameSession):

maxframe_client/tests/test_fetcher.py

@@ -17,19 +17,32 @@ import uuid
 import numpy as np
 import pandas as pd
 import pyarrow as pa
+import pytest
 from odps import ODPS
 
 import maxframe.dataframe as md
-from maxframe.odpsio import HaloTableIO
+from maxframe.config import options
+from maxframe.odpsio import ODPSTableIO
 from maxframe.protocol import ODPSTableResultInfo, ResultType
 from maxframe.tests.utils import tn
 
 from ..fetcher import ODPSTableFetcher
 
 
-async def test_table_fetcher():
+@pytest.fixture
+def switch_table_io(request):
+    old_use_common_table = options.use_common_table
+    try:
+        options.use_common_table = request.param
+        yield
+    finally:
+        options.use_common_table = old_use_common_table
+
+
+@pytest.mark.parametrize("switch_table_io", [False, True], indirect=True)
+async def test_table_fetcher(switch_table_io):
     odps_entry = ODPS.from_environments()
-    halo_table_io = HaloTableIO(odps_entry)
+    halo_table_io = ODPSTableIO(odps_entry)
     fetcher = ODPSTableFetcher(odps_entry)
 
     data = pd.DataFrame(
@@ -58,6 +71,11 @@ async def test_table_fetcher():
     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, [slice(None, 2000), 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

maxframe_client/tests/test_session.py

@@ -23,7 +23,10 @@ from odps import ODPS
 
 import maxframe.dataframe as md
 import maxframe.remote as mr
+from maxframe.config import options
+from maxframe.config.config import option_context
 from maxframe.core import ExecutableTuple, TileableGraph
+from maxframe.errors import NoTaskServerResponseError
 from maxframe.lib.aio import stop_isolation
 from maxframe.protocol import ResultInfo
 from maxframe.serialization import RemoteException
@@ -35,6 +38,7 @@ from maxframe_framedriver.app.tests.test_framedriver_webapp import (  # noqa: F4
 )
 
 from ..clients.framedriver import FrameDriverClient
+from ..session.odps import MaxFrameRestCaller
 
 pytestmark = pytest.mark.maxframe_engine(["MCSQL", "SPE"])
 
@@ -86,11 +90,25 @@ def test_simple_run_dataframe(start_mock_session):
         assert len(dag) == 2
         return await original_submit_dag(self, session_id, dag, managed_input_infos)
 
+    no_task_server_raised = False
+    original_get_dag_info = MaxFrameRestCaller.get_dag_info
+
+    async def patched_get_dag_info(self, dag_id: str):
+        nonlocal no_task_server_raised
+
+        if not no_task_server_raised:
+            no_task_server_raised = True
+            raise NoTaskServerResponseError
+        return await original_get_dag_info(self, dag_id)
+
     df["H"] = "extra_content"
 
     with mock.patch(
         "maxframe_client.clients.framedriver.FrameDriverClient.submit_dag",
         new=patched_submit_dag,
+    ), mock.patch(
+        "maxframe_client.session.odps.MaxFrameRestCaller.get_dag_info",
+        new=patched_get_dag_info,
     ):
         result = df.execute().fetch()
         assert len(result) == 1000
@@ -177,13 +195,32 @@ def test_run_dataframe_from_to_odps_table(start_mock_session):
         assert len(result_df) == 10
         assert len(result_df.columns) == 6
 
-        df = md.read_odps_table(table_obj, index_col="index").head(10).execute().fetch()
+        df = md.read_odps_table(table_obj, index_col="index").head(10).execute()
+        assert df.shape == (10, 5)
         assert len(df) == 10
         assert len(df.columns) == 5
     finally:
         odps_entry.delete_table(table_name, if_exists=True)
 
 
+def test_create_session_with_options(framedriver_app):  # noqa: F811
+    odps_entry = ODPS.from_environments()
+    framedriver_addr = f"mf://localhost:{framedriver_app.port}"
+    old_value = options.session.max_alive_seconds
+    session = None
+    try:
+        options.session.max_alive_seconds = 10
+        session = new_session(framedriver_addr, odps_entry=odps_entry)
+        session_id = session.session_id
+        session_conf = framedriver_app.session_manager.get_session_settings(session_id)
+        with option_context(session_conf) as session_options:
+            assert session_options.session.max_alive_seconds == 10
+    finally:
+        options.session.max_alive_seconds = old_value
+        if session is not None:
+            session.destroy()
+
+
 def test_run_and_fetch_series(start_mock_session):
     odps_entry = ODPS.from_environments()
 
@@ -244,7 +281,7 @@ def test_pivot_dataframe(start_mock_session):
     df = md.DataFrame(pd_df)
     pivot = df.pivot_table(values="D", index=["A", "B"], columns=["C"], aggfunc="sum")
     executed = pivot.execute()
-    assert pivot.shape == (2, 4)
+    assert pivot.shape == (4, 2)
     pd.testing.assert_index_equal(
         pivot.dtypes.index, pd.Index(["large", "small"], name="C")
     )
@@ -253,3 +290,13 @@ def test_pivot_dataframe(start_mock_session):
         values="D", index=["A", "B"], columns=["C"], aggfunc="sum"
     )
     pd.testing.assert_frame_equal(executed.to_pandas(), expected)
+
+
+def test_index_drop_duplicates(start_mock_session):
+    pd_idx = pd.Index(["lame", "cow", "lame", "beetle", "lame", "hippo"])
+    idx = md.Index(pd_idx)
+    executed = idx.drop_duplicates(keep="first").execute()
+    expected = pd_idx.drop_duplicates(keep="first")
+    pd.testing.assert_index_equal(
+        executed.to_pandas().sort_values(), expected.sort_values()
+    )

maxframe_client/clients/spe.py

@@ -1,104 +0,0 @@
-# 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)