maxframe 0.1.0b4__cp39-cp39-macosx_10_9_universal2.whl → 1.0.0__cp39-cp39-macosx_10_9_universal2.whl

maxframe_client/session/task.py

@@ -16,17 +16,23 @@ 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 import serializers
 from odps.errors import parse_instance_error
-from odps.models import Instance, Task
+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
 from maxframe.protocol import DagInfo, JsonSerializable, ResultInfo, SessionInfo
 from maxframe.utils import deserialize_serializable, serialize_serializable, to_str
 
@@ -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,
@@ -55,55 +62,6 @@ from .odps import MaxFrameServiceCaller, MaxFrameSession
 logger = logging.getLogger(__name__)
 
 
-class MaxFrameTask(Task):
-    __slots__ = ("_output_format", "_major_version", "_service_endpoint")
-    _root = "MaxFrame"
-    _anonymous_task_name = "AnonymousMaxFrameTask"
-
-    command = serializers.XMLNodeField("Command", default="CREATE_SESSION")
-
-    def __init__(self, **kwargs):
-        kwargs["name"] = kwargs.get("name") or self._anonymous_task_name
-        self._output_format = kwargs.pop(
-            "output_format", MAXFRAME_OUTPUT_MSGPACK_FORMAT
-        )
-        self._major_version = kwargs.pop("major_version", None)
-        self._service_endpoint = kwargs.pop("service_endpoint", None)
-        super().__init__(**kwargs)
-
-    def serial(self):
-        if self.properties is None:
-            self.properties = dict()
-
-        if odps_options.default_task_settings:
-            settings = odps_options.default_task_settings
-        else:
-            settings = dict()
-
-        if self._major_version is not None:
-            settings["odps.task.major.version"] = self._major_version
-
-        if "settings" in self.properties:
-            settings.update(json.loads(self.properties["settings"]))
-
-        # 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,
-            "odps.service.endpoint": self._service_endpoint,
-        }
-        if mf_version:
-            mf_opts["odps.maxframe.client_version"] = mf_version
-        settings.update(mf_opts)
-        self.properties["settings"] = json.dumps(settings)
-        return super().serial()
-
-
 class MaxFrameInstanceCaller(MaxFrameServiceCaller):
     _instance: Optional[Instance]
 
@@ -132,6 +90,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
@@ -140,14 +99,26 @@ 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]
     ):
         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:
@@ -159,13 +130,19 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
                 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)
+        mf_opts = {
+            "odps.maxframe.settings": json.dumps(self.get_settings_to_upload()),
+            "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 = MaxFrameTask(
-            name=self._task_name,
-            major_version=self._major_version,
-            output_format=self._output_format,
-            service_endpoint=self._odps_entry.endpoint,
-        )
+        task = self._create_maxframe_task()
         if not self._nested:
             self._task_name = task.name
             project = self._odps_entry.get_project(self._project)
@@ -210,11 +187,40 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             time.sleep(interval)
             interval = min(max_interval, interval * 2)
 
+    def _put_task_info(self, method_name: str, json_data: dict):
+        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
@@ -224,28 +230,28 @@ 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,
         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._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:
@@ -254,9 +260,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:
@@ -265,24 +269,39 @@ 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
+        # 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):
     schemes = [ODPS_SESSION_INSECURE_SCHEME, ODPS_SESSION_SECURE_SCHEME]
 
-    _instance: Instance
+    _caller: MaxFrameInstanceCaller
 
     @classmethod
     def _create_caller(
@@ -302,6 +321,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_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.io.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

@@ -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
@@ -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"])
 
@@ -82,15 +86,32 @@ 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
+
+    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
@@ -99,9 +120,12 @@ def test_simple_run_dataframe(start_mock_session):
     corner_top, corner_bottom = ExecutableTuple([df.iloc[:10], df.iloc[-10:]]).fetch()
     assert len(corner_top) == len(corner_bottom) == 10
 
-    # check ellipsis mark in DataFrame errors
+    # check ellipsis mark in DataFrame reprs
     df_str_repr = str(df)
     assert ".." in df_str_repr
+    # check ellipsis mark in Series reprs
+    series_str_repr = str(df.A.execute())
+    assert ".." in series_str_repr
 
     key = df.key
     assert odps_entry.exist_table(
@@ -109,13 +133,30 @@ 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)
+    retry_times = 10
+    while (
+        odps_entry.exist_table(
+            build_temp_table_name(start_mock_session, intermediate_key)
+        )
+        and retry_times > 0
+    ):
+        time.sleep(1)
+        retry_times -= 1
     assert not odps_entry.exist_table(
         build_temp_table_name(start_mock_session, intermediate_key)
     )
     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()
 
@@ -136,6 +177,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_session_empty_table")
+    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()
 
@@ -168,19 +228,38 @@ 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
         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()
 
@@ -207,7 +286,22 @@ def test_run_and_fetch_series(start_mock_session):
         )
 
 
-def test_run_remote_success(start_mock_session):
+def test_execute_with_tensor(oss_config, start_mock_session):
+    pd_df = pd.DataFrame(
+        {"angles": [0, 3, 4], "degrees": [360, 180, 360]},
+        index=["circle", "triangle", "rectangle"],
+    )
+    df = md.DataFrame(pd_df)
+
+    result = (df - [1, 2]).execute().fetch()
+    expected = pd_df - [1, 2]
+    # 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):
     def func(a, b):
         return a + b
 
@@ -218,7 +312,7 @@ def test_run_remote_success(start_mock_session):
     assert result == 21
 
 
-def test_run_remote_error(start_mock_session):
+def test_run_remote_error(oss_config, start_mock_session):
     def func():
         raise ValueError
 
@@ -226,3 +320,37 @@ def test_run_remote_error(start_mock_session):
 
     with pytest.raises((ValueError, RemoteException)):
         v.execute()
+
+
+def test_pivot_dataframe(start_mock_session):
+    pd_df = pd.DataFrame(
+        {
+            "A": "foo foo foo foo foo bar bar bar bar".split(),
+            "B": "one one one two two one one two two".split(),
+            "C": "small large large small small large small small large".split(),
+            "D": [1, 2, 2, 3, 3, 4, 5, 6, 7],
+            "E": [2, 4, 5, 5, 6, 6, 8, 9, 9],
+        }
+    )
+    df = md.DataFrame(pd_df)
+    pivot = df.pivot_table(values="D", index=["A", "B"], columns=["C"], aggfunc="sum")
+    executed = pivot.execute()
+    assert pivot.shape == (4, 2)
+    pd.testing.assert_index_equal(
+        pivot.dtypes.index, pd.Index(["large", "small"], name="C")
+    )
+
+    expected = pd_df.pivot_table(
+        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/core/entity/chunks.py

@@ -1,68 +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 ...serialization.serializables import BoolField, FieldTypes, TupleField
-from ...utils import tokenize
-from .core import Entity, EntityData
-
-
-class ChunkData(EntityData):
-    __slots__ = ()
-
-    is_broadcaster = BoolField("is_broadcaster", default=False)
-    # If the operator is a shuffle mapper, this flag indicates whether the current chunk is mapper chunk when
-    # the operator produce multiple chunks such as TensorUnique.
-    is_mapper = BoolField("is_mapper", default=None)
-    # optional fields
-    _index = TupleField("index", FieldTypes.uint32)
-
-    def __repr__(self):
-        if self.op.stage is None:
-            return (
-                f"{type(self).__name__} <op={type(self.op).__name__}, "
-                f"key={self.key}>"
-            )
-        else:
-            return (
-                f"{type(self).__name__} <op={type(self.op).__name__}, "
-                f"stage={self.op.stage.name}, key={self.key}>"
-            )
-
-    @property
-    def index(self):
-        return getattr(self, "_index", None)
-
-    @property
-    def device(self):
-        return self.op.device
-
-    def _update_key(self):
-        object.__setattr__(
-            self,
-            "_key",
-            tokenize(
-                type(self).__name__,
-                *(getattr(self, k, None) for k in self._keys_ if k != "_index"),
-            ),
-        )
-
-
-class Chunk(Entity):
-    _allow_data_type_ = (ChunkData,)
-
-    def __repr__(self):
-        return f"{type(self).__name__}({self._data.__repr__()})"
-
-
-CHUNK_TYPE = (Chunk, ChunkData)