Optimize tileable graph construction (#2583)

Author: wjsi

mars/_utils.pyx

@@ -18,7 +18,6 @@ import pickle
 import pkgutil
 import types
 import uuid
-from collections import deque
 from datetime import date, datetime, timedelta, tzinfo
 from enum import Enum
 from functools import lru_cache, partial
@@ -142,7 +141,7 @@ cdef class TypeDispatcher:
 cdef inline build_canonical_bytes(tuple args, kwargs):
     if kwargs:
         args = args + (kwargs,)
-    return str([tokenize_handler(arg) for arg in args]).encode('utf-8')
+    return pickle.dumps(tokenize_handler(args))
 
 
 def tokenize(*args, **kwargs):
@@ -155,26 +154,31 @@ def tokenize_int(*args, **kwargs):
 
 cdef class Tokenizer(TypeDispatcher):
     def __call__(self, object obj, *args, **kwargs):
-        if hasattr(obj, '__mars_tokenize__') and not isinstance(obj, type):
-            return super().__call__(obj.__mars_tokenize__(), *args, **kwargs)
-        if callable(obj):
-            if PDTick is not None and not isinstance(obj, PDTick):
-                return tokenize_function(obj)
-
         try:
             return super().__call__(obj, *args, **kwargs)
         except KeyError:
+            if hasattr(obj, '__mars_tokenize__') and not isinstance(obj, type):
+                if len(args) == 0 and len(kwargs) == 0:
+                    return obj.__mars_tokenize__()
+                else:
+                    return super().__call__(obj.__mars_tokenize__(), *args, **kwargs)
+            if callable(obj):
+                if PDTick is not None and not isinstance(obj, PDTick):
+                    return tokenize_function(obj)
+
             try:
                 return cloudpickle.dumps(obj)
             except:
                 raise TypeError(f'Cannot generate token for {obj}, type: {type(obj)}') from None
 
 
 cdef inline list iterative_tokenize(object ob):
-    dq = deque(ob)
-    h_list = []
-    while dq:
-        x = dq.pop()
+    cdef list dq = [ob]
+    cdef int dq_pos = 0
+    cdef list h_list = []
+    while dq_pos < len(dq):
+        x = dq[dq_pos]
+        dq_pos += 1
         if isinstance(x, (list, tuple)):
             dq.extend(x)
         elif isinstance(x, set):
@@ -188,7 +192,6 @@ cdef inline list iterative_tokenize(object ob):
 
 cdef inline tuple tokenize_numpy(ob):
     cdef int offset
-    cdef str data
 
     if not ob.shape:
         return str(ob), ob.dtype
@@ -288,6 +291,11 @@ cdef list tokenize_sqlalchemy_selectable(ob):
     return iterative_tokenize([str(ob)])
 
 
+cdef list tokenize_enum(ob):
+    cls = type(ob)
+    return iterative_tokenize([id(cls), cls.__name__, ob.name])
+
+
 @lru_cache(500)
 def tokenize_function(ob):
     if isinstance(ob, partial):
@@ -342,8 +350,8 @@ tokenize_handler.register(np.ndarray, tokenize_numpy)
 tokenize_handler.register(dict, lambda ob: iterative_tokenize(sorted(ob.items())))
 tokenize_handler.register(set, lambda ob: iterative_tokenize(sorted(ob)))
 tokenize_handler.register(np.random.RandomState, lambda ob: iterative_tokenize(ob.get_state()))
-tokenize_handler.register(Enum, lambda ob: iterative_tokenize((type(ob), ob.name)))
 tokenize_handler.register(memoryview, lambda ob: mmh3_hash_from_buffer(ob))
+tokenize_handler.register(Enum, tokenize_enum)
 tokenize_handler.register(pd.Index, tokenize_pandas_index)
 tokenize_handler.register(pd.Series, tokenize_pandas_series)
 tokenize_handler.register(pd.DataFrame, tokenize_pandas_dataframe)

mars/core/base.py

@@ -46,17 +46,31 @@ def _keys_(self):
             setattr(cls, member, slots)
             return slots
 
+    @property
+    def _copy_tags_(self):
+        cls = type(self)
+        member = f"__copy_tags_{cls.__name__}"
+        try:
+            return getattr(cls, member)
+        except AttributeError:
+            slots = sorted(
+                f.attr_name
+                for k, f in self._FIELDS.items()
+                if k not in self._no_copy_attrs_
+            )
+            setattr(cls, member, slots)
+            return slots
+
     @property
     def _values_(self):
-        return [
-            getattr(self, k, None) for k in self._keys_ if k not in self._no_copy_attrs_
-        ]
+        return [self._FIELD_VALUES.get(k) for k in self._copy_tags_]
 
     def __mars_tokenize__(self):
-        if hasattr(self, "_key"):
+        try:
+            return self._key
+        except AttributeError:  # pragma: no cover
+            self._update_key()
             return self._key
-        else:
-            return (type(self), *self._values_)
 
     def _obj_set(self, k, v):
         object.__setattr__(self, k, v)

mars/core/entity/tileables.py

@@ -25,7 +25,7 @@
 from ...typing import OperandType, TileableType, ChunkType
 from ...utils import on_serialize_shape, on_deserialize_shape, on_serialize_nsplits
 from ..base import Base
-from ..mode import enter_mode, is_build_mode
+from ..mode import enter_mode
 from .chunks import Chunk
 from .core import EntityData, Entity
 from .executable import _ExecutableMixin
@@ -280,8 +280,12 @@ def __init__(self: TileableType, *args, **kwargs):
 
         super().__init__(*args, **kwargs)
 
-        if hasattr(self, "_chunks") and self._chunks:
-            self._chunks = sorted(self._chunks, key=attrgetter("index"))
+        try:
+            chunks = self._chunks
+            if chunks:
+                self._chunks = sorted(chunks, key=attrgetter("index"))
+        except AttributeError:  # pragma: no cover
+            pass
 
         self._entities = WeakSet()
         self._executed_sessions = []
@@ -331,11 +335,9 @@ def is_coarse(self):
             return True
         return False
 
-    @enter_mode(build=True)
     def attach(self, entity):
         self._entities.add(entity)
 
-    @enter_mode(build=True)
     def detach(self, entity):
         self._entities.discard(entity)
 
@@ -345,10 +347,11 @@ class Tileable(Entity):
 
     def __init__(self, data: TileableType = None, **kw):
         super().__init__(data=data, **kw)
-        if self._data is not None:
-            self._data.attach(self)
-            if self._data.op.create_view:
-                entity_view_handler.add_observer(self._data.inputs[0], self)
+        data = self._data
+        if data is not None:
+            data.attach(self)
+            if data.op.create_view:
+                entity_view_handler.add_observer(data.inputs[0], self)
 
     def __copy__(self):
         return self._view()
@@ -412,11 +415,9 @@ def ndim(self):
 
     def __len__(self):
         try:
-            return self.shape[0]
-        except IndexError:
-            if is_build_mode():
-                return 0
-            raise TypeError("len() of unsized object")
+            return int(self.shape[0])
+        except (IndexError, ValueError):  # pragma: no cover
+            return 0
 
     @property
     def shape(self):

mars/core/mode.py

@@ -91,5 +91,6 @@ def enter_mode(kernel=None, build=None):
         "kernel": kernel,
         "build": build,
     }
+    mode_name_to_value = {k: v for k, v in mode_name_to_value.items() if v is not None}
 
     return _EnterModeFuncWrapper(mode_name_to_value)

mars/core/operand/base.py

@@ -127,6 +127,7 @@ class Operand(Base, metaclass=OperandMetaclass):
     attr_tag = "attr"
     _init_update_key_ = False
     _output_type_ = None
+    _no_copy_attrs_ = Base._no_copy_attrs_ | {"scheduling_hint"}
 
     sparse = BoolField("sparse", default=False)
     device = Int32Field("device", default=None)

mars/core/operand/core.py

@@ -23,12 +23,11 @@
     UFuncTypeError = None
 
 from ...typing import TileableType, ChunkType, OperandType
-from ...utils import calc_data_size
+from ...utils import calc_data_size, tokenize
 from ..context import Context
 from ..mode import is_eager_mode
 from ..entity import (
     OutputType,
-    TILEABLE_TYPE,
     ExecutableTuple,
     get_chunk_types,
     get_tileable_types,
@@ -46,13 +45,11 @@ class TileableOperandMixin:
     def check_inputs(self, inputs: List[TileableType]):
         if not inputs:
             return
+
+        from ...dataframe.core import DATAFRAME_TYPE
+
         for inp in inputs:
-            if isinstance(inp, TILEABLE_TYPE):
-                i = inp.extra_params["_i"]
-                if not inp.op.output_types:
-                    continue
-                if inp.op.output_types[i] != OutputType.dataframe:
-                    continue
+            if isinstance(inp, DATAFRAME_TYPE):
                 dtypes = getattr(inp, "dtypes", None)
                 if dtypes is None:
                     raise ValueError(
@@ -62,24 +59,25 @@ def check_inputs(self, inputs: List[TileableType]):
 
     @classmethod
     def _check_if_gpu(cls, inputs: List[TileableType]):
-        if (
-            inputs is not None
-            and len(
-                [
-                    inp
-                    for inp in inputs
-                    if inp is not None and getattr(inp, "op", None) is not None
-                ]
-            )
-            > 0
-        ):
-            if all(inp.op.gpu is True for inp in inputs):
-                return True
-            elif all(inp.op.gpu is False for inp in inputs):
-                return False
+        if not inputs:
+            return None
+        true_num = 0
+        for inp in inputs:
+            op = getattr(inp, "op", None)
+            if op is None or op.gpu is None:
+                return None
+            true_num += int(op.gpu)
+        if true_num == len(inputs):
+            return True
+        elif true_num == 0:
+            return False
+        return None
+
+    def _tokenize_output(self, output_idx: int, **kw):
+        return tokenize(self._key, output_idx)
 
     def _create_chunk(self, output_idx: int, index: Tuple[int], **kw) -> ChunkType:
-        output_type = kw.pop("output_type", self._get_output_type(output_idx))
+        output_type = kw.pop("output_type", None) or self._get_output_type(output_idx)
         if not output_type:
             raise ValueError("output_type should be specified")
 
@@ -92,6 +90,11 @@ def _create_chunk(self, output_idx: int, index: Tuple[int], **kw) -> ChunkType:
         if output_type == OutputType.scalar:
             # tensor
             kw["order"] = "C_ORDER"
+
+        # key of output chunks may only contain keys for its output ids
+        if "_key" not in kw:
+            kw["_key"] = self._tokenize_output(output_idx, **kw)
+
         data = chunk_data_type(**kw)
         return chunk_type(data)
 
@@ -189,6 +192,7 @@ def _create_tileable(self, output_idx: int, **kw) -> TileableType:
 
         if isinstance(output_type, (list, tuple)):
             output_type = output_type[output_idx]
+
         tileable_type, tileable_data_type = get_tileable_types(output_type)
         kw["_i"] = output_idx
         kw["op"] = self
@@ -197,6 +201,11 @@ def _create_tileable(self, output_idx: int, **kw) -> TileableType:
             kw["order"] = "C_ORDER"
 
         kw = self._fill_nan_shape(kw)
+
+        # key of output chunks may only contain keys for its output ids
+        if "_key" not in kw:
+            kw["_key"] = self._tokenize_output(output_idx, **kw)
+
         data = tileable_data_type(**kw)
         return tileable_type(data)
 
@@ -207,12 +216,11 @@ def _new_tileables(
         if output_limit is None:
             output_limit = getattr(self, "output_limit")
 
-        self.check_inputs(inputs)
-        getattr(self, "_set_inputs")(inputs)
-        if getattr(self, "gpu", None) is None:
+        self._set_inputs(inputs)
+        if self.gpu is None:
             self.gpu = self._check_if_gpu(self._inputs)
         if getattr(self, "_key", None) is None:
-            getattr(self, "_update_key")()  # update key when inputs are set
+            self._update_key()  # update key when inputs are set
 
         tileables = []
         for j in range(output_limit):
@@ -222,7 +230,7 @@ def _new_tileables(
             tileable = self._create_tileable(j, **create_tensor_kw)
             tileables.append(tileable)
 
-        setattr(self, "outputs", tileables)
+        self.outputs = tileables
         if len(tileables) > 1:
             # for each output tileable, hold the reference to the other outputs
             # so that either no one or everyone are gc collected

mars/dataframe/arithmetic/core.py

@@ -434,7 +434,10 @@ def _calc_properties(cls, x1, x2=None, axis="columns"):
         ):
             x2_dtype = x2.dtype if hasattr(x2, "dtype") else type(x2)
             x2_dtype = get_dtype(x2_dtype)
-            dtype = infer_dtype(x1.dtype, x2_dtype, cls._operator)
+            if hasattr(cls, "return_dtype"):
+                dtype = cls.return_dtype
+            else:
+                dtype = infer_dtype(x1.dtype, x2_dtype, cls._operator)
             ret = {"shape": x1.shape, "dtype": dtype, "index_value": x1.index_value}
             if pd.api.types.is_scalar(x2) or (
                 hasattr(x2, "ndim") and (x2.ndim == 0 or x2.ndim == 1)

mars/dataframe/arithmetic/equal.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
+
 from ... import opcodes as OperandDef
 from ...utils import classproperty
 from .core import DataFrameBinopUfunc
@@ -24,6 +26,8 @@ class DataFrameEqual(DataFrameBinopUfunc):
     _func_name = "eq"
     _rfunc_name = "eq"
 
+    return_dtype = np.dtype(bool)
+
     @classproperty
     def _operator(self):
         return lambda lhs, rhs: lhs.eq(rhs)

mars/dataframe/arithmetic/greater.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
+
 from ... import opcodes as OperandDef
 from ...utils import classproperty
 from .core import DataFrameBinopUfunc
@@ -24,6 +26,8 @@ class DataFrameGreater(DataFrameBinopUfunc):
     _func_name = "gt"
     _rfunc_name = "lt"
 
+    return_dtype = np.dtype(bool)
+
     @classproperty
     def _operator(self):
         return lambda lhs, rhs: lhs.gt(rhs)