maxframe 1.3.1__cp39-cp39-win_amd64.whl → 2.0.0b1__cp39-cp39-win_amd64.whl

maxframe/core/entity/objects.py

@@ -16,6 +16,7 @@ from typing import Any, Dict, Type
 
 from ...serialization import load_type
 from ...serialization.serializables import StringField
+from ...utils import extract_class_name
 from .core import Entity
 from .executable import _ToObjectMixin
 from .tileables import TileableData
@@ -36,7 +37,7 @@ class ObjectData(TileableData, _ToObjectMixin):
         if getattr(cls, "_entity_class", None) is not None:
             return cls._entity_class
         assert cls.__qualname__[-4:] == "Data"
-        target_class_name = cls.__module__ + "#" + cls.__qualname__[:-4]
+        target_class_name = extract_class_name(cls)[:-4]
         cls._entity_class = load_type(target_class_name, Object)
         return cls._entity_class
 
@@ -57,15 +58,18 @@ class ObjectData(TileableData, _ToObjectMixin):
         if isinstance(obj_cls, type):
             if isinstance(obj_cls, type) and issubclass(obj_cls, Object):
                 obj_cls = obj_cls.get_data_class()
-            kw["object_class"] = obj_cls.__module__ + "#" + obj_cls.__qualname__
+            kw["object_class"] = extract_class_name(obj_cls)
 
         super().__init__(_op=op, _nsplits=nsplits, **kw)
         if self.object_class is None and type(self) is not ObjectData:
             cls = type(self)
-            self.object_class = cls.__module__ + "#" + cls.__qualname__
+            self.object_class = extract_class_name(cls)
 
     def __repr__(self):
-        return f"Object <op={type(self.op).__name__}, key={self.key}>"
+        return (
+            f"{type(self).__name__[:-4]} "
+            f"<op={type(self.op).__name__}, key={self.key}>"
+        )
 
     @property
     def params(self):
@@ -75,6 +79,7 @@ class ObjectData(TileableData, _ToObjectMixin):
     @params.setter
     def params(self, new_params: Dict[str, Any]):
         params = new_params.copy()
+        params.pop("object_class", None)
         if params:  # pragma: no cover
             raise TypeError(f"Unknown params: {list(params)}")
 
@@ -89,11 +94,20 @@ class Object(Entity, _ToObjectMixin):
     _allow_data_type_ = (ObjectData,)
     type_name = "Object"
 
+    def __new__(cls, data=None, **kw):
+        if (
+            cls is Object
+            and isinstance(data, ObjectData)
+            and type(data) is not ObjectData
+        ):
+            return data.get_entity_class()(data, **kw)
+        return super().__new__(cls)
+
     @classmethod
     def get_data_class(cls) -> Type[ObjectData]:
         if getattr(cls, "_data_class", None) is not None:
             return cls._data_class
-        target_class_name = cls.__module__ + "#" + cls.__qualname__ + "Data"
+        target_class_name = extract_class_name(cls) + "Data"
         cls._data_class = load_type(target_class_name, ObjectData)
         return cls._data_class
 

maxframe/core/entity/output_types.py

@@ -39,6 +39,17 @@ class OutputType(Enum):
         return [cls(ot) for ot in output_types] if output_types is not None else None
 
 
+df_output_types = {
+    OutputType.dataframe,
+    OutputType.dataframe_groupby,
+    OutputType.series,
+    OutputType.series_groupby,
+    OutputType.df_or_series,
+    OutputType.index,
+    OutputType.categorical,
+}
+
+
 _OUTPUT_TYPE_TO_TILEABLE_TYPES = {OutputType.object: OBJECT_TYPE}
 _OUTPUT_TYPE_TO_FETCH_CLS = {}
 

maxframe/core/entity/tests/test_objects.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ....utils import extract_class_name
 from ..objects import Object, ObjectData
 
 
@@ -26,18 +27,16 @@ class TestSubObject(Object):
 def test_object_init():
     assert TestSubObjectData.get_entity_class() is TestSubObject
 
-    obj = ObjectData(
-        object_class=TestSubObjectData.__module__ + "#" + TestSubObjectData.__name__
-    )
-    assert isinstance(obj, TestSubObjectData)
+    obj_data = ObjectData(object_class=extract_class_name(TestSubObjectData))
+    assert isinstance(obj_data, TestSubObjectData)
+    obj = Object(obj_data)
+    assert isinstance(obj, TestSubObject)
 
-    obj = ObjectData(object_class=TestSubObjectData)
-    assert isinstance(obj, TestSubObjectData)
+    obj_data = ObjectData(object_class=TestSubObjectData)
+    assert isinstance(obj_data, TestSubObjectData)
 
-    obj = ObjectData(
-        object_class=TestSubObject.__module__ + "#" + TestSubObject.__name__
-    )
-    assert isinstance(obj, TestSubObjectData)
+    obj_data = ObjectData(object_class=extract_class_name(TestSubObject))
+    assert isinstance(obj_data, TestSubObjectData)
 
-    obj = ObjectData(object_class=TestSubObject)
-    assert isinstance(obj, TestSubObjectData)
+    obj_data = ObjectData(object_class=TestSubObject)
+    assert isinstance(obj_data, TestSubObjectData)

maxframe/core/entity/tileables.py

@@ -172,6 +172,8 @@ class TileableData(EntityData, _ExecutableMixin):
     def __init__(self: TileableType, *args, **kwargs):
         if kwargs.get("chunks") is not None:
             self._chunks = kwargs.pop("chunks")
+        if "nsplits" in kwargs:
+            kwargs["_nsplits"] = kwargs.pop("nsplits")
         if kwargs.get("_nsplits", None) is not None:
             kwargs["_nsplits"] = tuple(tuple(s) for s in kwargs["_nsplits"])
 
@@ -198,7 +200,7 @@ class TileableData(EntityData, _ExecutableMixin):
 
     @property
     def chunks(self) -> list:
-        return getattr(self, "_chunks", None)
+        return getattr(self, "_chunks", None) or self._extra_params.get("_chunks")
 
     @property
     def nsplits(self):

maxframe/core/entity/utils.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
+
+from ...typing_ import TileableType
 from ...utils import calc_nsplits, has_unknown_shape
 
 
@@ -22,3 +25,15 @@ def refresh_tileable_shape(tileable):
         shape = tuple(sum(ns) for ns in nsplits)
         tileable._nsplits = nsplits
         tileable._shape = shape
+
+
+def fill_chunk_slices(tileable: TileableType) -> None:
+    if tileable is None or not tileable.nsplits or not tileable.chunks:
+        # errors may occurred
+        return
+    acc_splits = tuple((0,) + tuple(np.cumsum(sp)) for sp in tileable.nsplits)
+    for chunk in tileable.chunks:
+        chunk.slices = list(
+            (acc_splits[i][chunk.index[i]], acc_splits[i][chunk.index[i] + 1])
+            for i in range(tileable.ndim)
+        )

maxframe/core/graph/__init__.py

@@ -14,4 +14,9 @@
 
 from .builder import TileableGraphBuilder
 from .core import DAG, DirectedGraph, GraphContainsCycleError
-from .entity import EntityGraph, GraphSerializer, TileableGraph
+from .entity import EntityGraph, GraphSerializer, SerializableGraph, TileableGraph
+
+try:
+    from ._internal import SubDAG
+except ImportError:  # pragma: no cover
+    pass

maxframe/core/graph/builder/base.py

@@ -41,6 +41,10 @@ class AbstractGraphBuilder(ABC):
     ):  # pylint: disable=no-self-use
         return node not in visited
 
+    @classmethod
+    def _get_node_outputs(cls, node: EntityType):
+        return node.op.outputs
+
     def _add_nodes(
         self,
         graph: EntityGraph,
@@ -70,7 +74,7 @@ class AbstractGraphBuilder(ABC):
                     graph.add_node(c)
                 if not graph.has_successor(c, node):
                     graph.add_edge(c, node)
-                for out in c.op.outputs:
+                for out in self._get_node_outputs(c):
                     if self._if_add_node(out, visited):
                         nodes.append(out)
 

maxframe/core/graph/core.pyx

@@ -309,13 +309,19 @@ cdef class DirectedGraph:
         return val
 
     def _extract_operators(self, node):
+        if getattr(node, "opfunc", None) is not None:
+            return [node.opfunc]
         return [node.op]
 
     def to_dot(
         self,
         graph_attrs=None,
         node_attrs=None,
-        trunc_key=5, result_chunk_keys=None, show_columns=False):
+        trunc_key=5,
+        result_chunk_keys=None,
+        show_columns=False,
+        with_op_id=True,
+    ):
 
         sio = StringIO()
         sio.write('digraph {\n')
@@ -344,20 +350,25 @@ cdef class DirectedGraph:
 
         for node in self.iter_nodes():
             for op in self._extract_operators(node):
-                op_name = type(op).__name__
+                op_name = op.op_name
                 if op.stage is not None:
                     op_name = f'{op_name}:{op.stage.name}'
                 if op.key in visited:
                     continue
+
+                op_key_id = op.key[:trunc_key]
+                if with_op_id:
+                    op_key_id += f"_{id(op)}"
+
                 for input_chunk in (op.inputs or []):
                     if input_chunk.key not in visited:
                         sio.write(f'"Chunk:{self._gen_chunk_key(input_chunk, trunc_key)}" {chunk_style}\n')
                         visited.add(input_chunk.key)
                     if op.key not in visited:
-                        sio.write(f'"{op_name}:{op.key[:trunc_key]}_{id(op)}" {operator_style}\n')
+                        sio.write(f'"{op_name}:{op_key_id}" {operator_style}\n')
                         visited.add(op.key)
                     sio.write(f'"Chunk:{self._gen_chunk_key(input_chunk, trunc_key)}" -> '
-                              f'"{op_name}:{op.key[:trunc_key]}_{id(op)}"\n')
+                              f'"{op_name}:{op_key_id}"\n')
 
                 for output_chunk in (op.outputs or []):
                     if output_chunk.key not in visited:
@@ -367,9 +378,9 @@ cdef class DirectedGraph:
                         sio.write(f'"Chunk:{self._gen_chunk_key(output_chunk, trunc_key)}" {tmp_chunk_style}\n')
                         visited.add(output_chunk.key)
                     if op.key not in visited:
-                        sio.write(f'"{op_name}:{op.key[:trunc_key]}_{id(op)}" {operator_style}\n')
+                        sio.write(f'"{op_name}:{op_key_id}" {operator_style}\n')
                         visited.add(op.key)
-                    sio.write(f'"{op_name}:{op.key[:trunc_key]}_{id(op)}" -> '
+                    sio.write(f'"{op_name}:{op_key_id}" -> '
                               f'"Chunk:{self._gen_chunk_key(output_chunk, trunc_key)}"')
                     if show_columns:
                         sio.write(f' [ label={get_col_names(output_chunk)} ]')

maxframe/core/graph/entity.py

@@ -16,10 +16,17 @@ from abc import ABCMeta, abstractmethod
 from typing import Dict, Iterable, List
 
 from ...core import Tileable
-from ...serialization.core import buffered
-from ...serialization.serializables import BoolField, DictField, ListField, Serializable
+from ...serialization.core import buffered, load_type
+from ...serialization.serializables import (
+    BoolField,
+    DictField,
+    ListField,
+    Serializable,
+    StringField,
+)
 from ...serialization.serializables.core import SerializableSerializer
-from ...utils import tokenize
+from ...utils import extract_class_name, tokenize
+from ..operator import Fetch
 from .core import DAG
 
 
@@ -110,11 +117,10 @@ class SerializableGraph(Serializable):
     _predecessors = DictField("predecessors")
     _successors = DictField("successors")
     _results = ListField("results")
+    _graph_cls = StringField("graph_cls")
 
     @classmethod
     def from_graph(cls, graph: EntityGraph) -> "SerializableGraph":
-        from ..operator import Fetch
-
         return SerializableGraph(
             _is_chunk=False,
             _fetch_nodes=[chunk for chunk in graph if isinstance(chunk.op, Fetch)],
@@ -122,10 +128,16 @@ class SerializableGraph(Serializable):
             _predecessors=graph._predecessors,
             _successors=graph._successors,
             _results=graph.results,
+            _graph_cls=extract_class_name(type(graph)),
         )
 
     def to_graph(self) -> EntityGraph:
-        graph = TileableGraph(self._results)
+        graph_cls = (
+            load_type(self._graph_cls, EntityGraph)
+            if hasattr(self, "_graph_cls") and self._graph_cls
+            else TileableGraph
+        )
+        graph = graph_cls(self._results)
         graph._nodes.update(self._nodes)
         graph._predecessors.update(self._predecessors)
         graph._successors.update(self._successors)

maxframe/core/operator/__init__.py

@@ -20,8 +20,13 @@ from .base import (
     OperatorStage,
     VirtualOperator,
 )
-from .core import TileableOperatorMixin, estimate_size, execute
-from .fetch import Fetch, FetchMixin, FetchShuffle, ShuffleFetchType
-from .objects import MergeDictOperator, ObjectFetch, ObjectOperator, ObjectOperatorMixin
+from .core import (
+    TileableOperatorMixin,
+    estimate_size,
+    estimate_tileable_execution_size,
+    execute,
+)
+from .fetch import Fetch, FetchMixin, FetchShuffle
+from .objects import ObjectFetch, ObjectOperator, ObjectOperatorMixin
 from .shuffle import MapReduceOperator, ShuffleProxy
 from .utils import add_fetch_builder, build_fetch

maxframe/core/operator/base.py

@@ -16,7 +16,7 @@ import weakref
 from copy import deepcopy
 from enum import Enum
 from functools import lru_cache, partial
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, MutableMapping, Optional, Tuple, Type, Union
 
 from ...serialization.core import Placeholder
 from ...serialization.serializables import (
@@ -88,6 +88,8 @@ class SchedulingHint(Serializable):
     priority = Int32Field("priority", default=None)
     expect_engine = StringField("expect_engine", default=None)
     expect_resources = DictField("expect_resources", FieldTypes.string, default=None)
+    # id of gang scheduling for machine learning trainings
+    gang_scheduling_id = StringField("gang_scheduling_id", default=None)
 
     @classproperty
     @lru_cache(1)
@@ -127,7 +129,7 @@ def _install_scheduling_hint_properties(cls: Type["Operator"]):
 class OperatorLogicKeyGeneratorMixin:
     """
     This generator will generate an unique and deterministic key for operator compute logic. It should be same
-    for different run if the compute logic doesn't change. This id will be used in subtask speculative
+    for different run if the compute logic doesn't change. This id will be used in substep speculative
     execution and hbo scheduling and so on.
     """
 
@@ -264,6 +266,10 @@ class Operator(Base, OperatorLogicKeyGeneratorMixin, metaclass=OperatorMetaclass
             else:
                 extras[k] = kwargs.pop(k)
 
+    @property
+    def op_name(self) -> str:
+        return type(self).__name__
+
     def __repr__(self):
         if self.stage is None:
             return f"{type(self).__name__} <key={self.key}>"
@@ -280,12 +286,13 @@ class Operator(Base, OperatorLogicKeyGeneratorMixin, metaclass=OperatorMetaclass
     def _get_inputs_data(cls, inputs):
         return [cls._get_entity_data(inp) for inp in inputs]
 
-    def _set_inputs(self, inputs):
+    @classmethod
+    def _set_inputs(cls, op: "Operator", inputs: List[EntityData]):
         if inputs is not None:
-            inputs = self._get_inputs_data(inputs)
-        if hasattr(self, "check_inputs"):
-            self.check_inputs(inputs)
-        setattr(self, "_inputs", inputs)
+            inputs = cls._get_inputs_data(inputs)
+        if hasattr(op, "check_inputs"):
+            op.check_inputs(inputs)
+        setattr(op, "_inputs", inputs)
 
     def replace_input(self, index: int, replaced_input: ENTITY_TYPE):
         """
@@ -299,7 +306,7 @@ class Operator(Base, OperatorLogicKeyGeneratorMixin, metaclass=OperatorMetaclass
             The replaced input object.
         """
         self.inputs[index] = replaced_input
-        self._set_inputs(self.inputs)
+        self._set_inputs(self, self.inputs)
 
     @property
     def inputs(self) -> List[Union[ENTITY_TYPE]]:
@@ -310,7 +317,7 @@ class Operator(Base, OperatorLogicKeyGeneratorMixin, metaclass=OperatorMetaclass
 
     @inputs.setter
     def inputs(self, vals):
-        self._set_inputs(vals)
+        self._set_inputs(self, vals)
 
     @property
     def output_limit(self) -> int:
@@ -358,6 +365,9 @@ class Operator(Base, OperatorLogicKeyGeneratorMixin, metaclass=OperatorMetaclass
     def is_gpu(self) -> bool:
         return self.gpu
 
+    def has_custom_code(self) -> bool:
+        return False
+
     @property
     def retryable(self) -> bool:
         return True
@@ -386,6 +396,18 @@ class Operator(Base, OperatorLogicKeyGeneratorMixin, metaclass=OperatorMetaclass
             new_op.extra_params = deepcopy(extra_params)
         return new_op
 
+    @classmethod
+    def estimate_size(
+        cls, ctx: MutableMapping[str, Union[int, float]], op: "Operator"
+    ) -> None:
+        try:
+            total_input_size = sum(ctx[inp.key] for inp in (op.inputs or ()))
+            for out in op.outputs:
+                ctx[out.key] = total_input_size // len(op.outputs)
+        except KeyError:
+            for out in op.outputs:
+                ctx[out.key] = float("inf")
+
     def on_output_modify(self, new_output):
         # when `create_view` is True, if the output is modified,
         # the modification should be set back to the input.
@@ -445,6 +467,7 @@ class HasInput(Operator):
     def input(self):
         return self._input
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self._input = self._inputs[0]
+    @classmethod
+    def _set_inputs(cls, op: "Operator", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op._input = op._inputs[0]

maxframe/core/operator/core.py

@@ -12,11 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 import sys
-from typing import Any, Callable, Dict, List, Type
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from .. import TileableGraph
+
 try:
     from numpy.core._exceptions import UFuncTypeError
 except ImportError:  # pragma: no cover
@@ -117,7 +122,7 @@ class TileableOperatorMixin:
         if output_limit is None:
             output_limit = getattr(self, "output_limit")
 
-        self._set_inputs(inputs)
+        self._set_inputs(self, inputs)
         if self.gpu is None:
             self.gpu = self._check_if_gpu(self._inputs)
         if getattr(self, "_key", None) is None:
@@ -182,20 +187,13 @@ class TileableOperatorMixin:
 
         return self.new_tileables(inputs, kws=kws, **kw)[0]
 
-    @classmethod
-    def concat_tileable_chunks(cls, tileable: TileableType):
-        raise NotImplementedError
-
-    @classmethod
-    def create_tileable_from_chunks(
-        cls, chunks: List[ChunkType], inputs: List[TileableType] = None, **kw
-    ) -> TileableType:
-        raise NotImplementedError
-
-    def get_fetch_op_cls(self, obj: ChunkType):
+    def get_fetch_op_cls(self, obj: Union[ChunkType, OutputType]):
         from .shuffle import ShuffleProxy
 
-        output_types = get_output_types(obj, unknown_as=OutputType.object)
+        if isinstance(obj, OutputType):
+            output_types = [obj or OutputType.object]
+        else:
+            output_types = get_output_types(obj, unknown_as=OutputType.object)
         fetch_cls, fetch_shuffle_cls = get_fetch_class(output_types[0])
         if isinstance(self, ShuffleProxy):
             cls = fetch_shuffle_cls
@@ -274,3 +272,28 @@ def estimate_size(results: Dict[str, Any], op: OperatorType):
                 _op_type_to_size_estimator[type(op)] = size_estimator
                 return size_estimator(results, op)
         raise KeyError(f"No handler found for op: {op} to estimate size")
+
+
+def estimate_tileable_execution_size(
+    tileable_graph: "TileableGraph",
+    fetch_sizes: Optional[Dict[str, int]] = None,
+) -> Union[int, float]:
+    ctx = dict()
+    ctx.update(fetch_sizes)
+    ref_counts = defaultdict(lambda: 0)
+    max_size = 0
+
+    for tileable in tileable_graph:
+        for key in set(inp.key for inp in tileable.inputs or ()):
+            ref_counts[key] += 1
+
+    for tileable in tileable_graph.topological_iter():
+        estimate_size(ctx, tileable.op)
+        max_size = max(max_size, sum(ctx.values()))
+        if math.isinf(max_size):
+            return float("inf")
+        for inp in set(inp.key for inp in tileable.inputs or ()):
+            ref_counts[inp] -= 1
+            if ref_counts[inp] <= 0:
+                ctx.pop(inp, None)
+    return max_size

maxframe/core/operator/fetch.py

@@ -12,16 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import enum
-
 from ... import opcodes
-from ...serialization.serializables import (
-    FieldTypes,
-    Int32Field,
-    ListField,
-    ReferenceField,
-    StringField,
-)
+from ...serialization.serializables import Int32Field, StringField
 from .base import Operator
 from .core import TileableOperatorMixin
 
@@ -42,12 +34,7 @@ class FetchMixin(TileableOperatorMixin):
 class FetchShuffle(Operator):
     _op_type_ = opcodes.FETCH_SHUFFLE
 
-    source_keys = ListField("source_keys", FieldTypes.string)
-    n_mappers = Int32Field("n_mappers")
-    n_reducers = Int32Field("n_reducers")
-    shuffle_fetch_type = ReferenceField("shuffle_fetch_type")
-
-
-class ShuffleFetchType(enum.Enum):
-    FETCH_BY_KEY = 0
-    FETCH_BY_INDEX = 1
+    shuffle_key = StringField("shuffle_key", default=None)
+    reducer_id = Int32Field("reducer_id", default=None)
+    n_mappers = Int32Field("n_mappers", default=None)
+    n_reducers = Int32Field("n_reducers", default=None)

maxframe/core/operator/objects.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...serialization.serializables import BoolField
 from ..entity import OutputType, register_fetch_class
 from .base import Operator
 from .core import TileableOperatorMixin
@@ -42,22 +41,3 @@ class ObjectFetch(FetchMixin, ObjectOperatorMixin, Fetch):
 
 
 register_fetch_class(OutputType.object, ObjectFetch, None)
-
-
-class MergeDictOperator(ObjectOperator, ObjectOperatorMixin):
-    _merge = BoolField("merge")
-
-    def __init__(self, merge=None, **kw):
-        super().__init__(_merge=merge, **kw)
-
-    @property
-    def merge(self):
-        return self._merge
-
-    @classmethod
-    def concat_tileable_chunks(cls, tileable):
-        assert not tileable.is_coarse()
-
-        op = cls(merge=True)
-        chunk = cls(merge=True).new_chunk(tileable.chunks)
-        return op.new_tileable([tileable], chunks=[chunk], nsplits=((1,),))