Optimize import speed for Mars package (#3022)

Author: wjsi

benchmarks/asv_bench/benchmarks/import.py

@@ -0,0 +1,40 @@
+# Copyright 1999-2022 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.
+
+import subprocess
+import sys
+
+# make sure necessary pyc files generated
+import mars.dataframe as md
+import mars.tensor as mt
+
+del md, mt
+
+
+class ImportPackageSuite:
+    """
+    Benchmark that times performance of chunk graph builder
+    """
+
+    def time_import_mars(self):
+        proc = subprocess.Popen([sys.executable, "-c", "import mars"])
+        proc.wait(120)
+
+    def time_import_mars_tensor(self):
+        proc = subprocess.Popen([sys.executable, "-c", "import mars.tensor"])
+        proc.wait(120)
+
+    def time_import_mars_dataframe(self):
+        proc = subprocess.Popen([sys.executable, "-c", "import mars.dataframe"])
+        proc.wait(120)

benchmarks/asv_bench/benchmarks/serialize.py

@@ -46,6 +46,10 @@
 from mars.utils import tokenize
 
 
+# do warmup
+serialize(None)
+
+
 class SerializableChild(Serializable):
     str_field = StringField("str_field")
     int_field = Int64Field("int_field")

benchmarks/asv_bench/benchmarks/tokenize.py

@@ -47,6 +47,9 @@ class MockOperand(Operand, TensorOperandMixin):
 
 class TokenizeOperandSuite:
     def setup(self):
+        # do some warm up
+        tokenize(None)
+
         chunks = []
         for idx in range(1000):
             op = MockOperand(

mars/_utils.pxd

@@ -17,6 +17,7 @@ cdef class TypeDispatcher:
     cdef dict _handlers
     cdef dict _lazy_handlers
     cdef dict _inherit_handlers
+    cpdef object __weakref__
 
     cpdef void register(self, object type_, object handler)
     cpdef void unregister(self, object type_)

mars/_utils.pyx

@@ -22,6 +22,7 @@ import uuid
 from datetime import date, datetime, timedelta, tzinfo
 from enum import Enum
 from functools import lru_cache, partial
+from weakref import WeakSet
 
 import numpy as np
 import pandas as pd
@@ -31,17 +32,13 @@ try:
     from pandas.tseries.offsets import Tick as PDTick
 except ImportError:
     PDTick = None
-try:
-    from sqlalchemy.sql import Selectable as SASelectable
-    from sqlalchemy.sql.sqltypes import TypeEngine as SATypeEngine
-except ImportError:
-    SASelectable, SATypeEngine = None, None
 
 from .lib.mmh3 import hash as mmh_hash, hash_bytes as mmh_hash_bytes, \
     hash_from_buffer as mmh3_hash_from_buffer
 
 cdef bint _has_cupy = bool(pkgutil.find_loader('cupy'))
 cdef bint _has_cudf = bool(pkgutil.find_loader('cudf'))
+cdef bint _has_sqlalchemy = bool(pkgutil.find_loader('sqlalchemy'))
 
 
 cpdef str to_str(s, encoding='utf-8'):
@@ -83,29 +80,41 @@ cpdef unicode to_text(s, encoding='utf-8'):
         raise TypeError(f"Could not convert from {s} to unicode.")
 
 
+_type_dispatchers = WeakSet()
+
+
 cdef class TypeDispatcher:
     def __init__(self):
         self._handlers = dict()
         self._lazy_handlers = dict()
         # store inherited handlers to facilitate unregistering
         self._inherit_handlers = dict()
 
+        _type_dispatchers.add(self)
+
     cpdef void register(self, object type_, object handler):
         if isinstance(type_, str):
             self._lazy_handlers[type_] = handler
+        elif isinstance(type_, tuple):
+            for t in type_:
+                self.register(t, handler)
         else:
             self._handlers[type_] = handler
 
     cpdef void unregister(self, object type_):
-        self._lazy_handlers.pop(type_, None)
-        self._handlers.pop(type_, None)
-        self._inherit_handlers.clear()
+        if isinstance(type_, tuple):
+            for t in type_:
+                self.unregister(t)
+        else:
+            self._lazy_handlers.pop(type_, None)
+            self._handlers.pop(type_, None)
+            self._inherit_handlers.clear()
 
     cdef _reload_lazy_handlers(self):
         for k, v in self._lazy_handlers.items():
             mod_name, obj_name = k.rsplit('.', 1)
             mod = importlib.import_module(mod_name, __name__)
-            self._handlers[getattr(mod, obj_name)] = v
+            self.register(getattr(mod, obj_name), v)
         self._lazy_handlers = dict()
 
     cpdef get_handler(self, object type_):
@@ -134,6 +143,11 @@ cdef class TypeDispatcher:
     def __call__(self, object obj, *args, **kwargs):
         return self.get_handler(type(obj))(obj, *args, **kwargs)
 
+    @staticmethod
+    def reload_all_lazy_handlers():
+        for dispatcher in _type_dispatchers:
+            (<TypeDispatcher>dispatcher)._reload_lazy_handlers()
+
 
 cdef inline build_canonical_bytes(tuple args, kwargs):
     if kwargs:
@@ -376,10 +390,13 @@ if _has_cudf:
 
 if PDTick is not None:
     tokenize_handler.register(PDTick, tokenize_pandas_tick)
-if SATypeEngine is not None:
-    tokenize_handler.register(SATypeEngine, tokenize_sqlalchemy_data_type)
-if SASelectable is not None:
-    tokenize_handler.register(SASelectable, tokenize_sqlalchemy_selectable)
+if _has_sqlalchemy:
+    tokenize_handler.register(
+        "sqlalchemy.sql.sqltypes.TypeEngine", tokenize_sqlalchemy_data_type
+    )
+    tokenize_handler.register(
+        "sqlalchemy.sql.Selectable", tokenize_sqlalchemy_selectable
+    )
 
 cpdef register_tokenizer(cls, handler):
     tokenize_handler.register(cls, handler)

mars/contrib/dask/tests/test_dask.py

@@ -17,8 +17,8 @@
 from ....utils import lazy_import
 from .. import convert_dask_collection, mars_scheduler
 
-dask_installed = lazy_import("dask", globals=globals()) is not None
-mimesis_installed = lazy_import("mimesis", globals=globals()) is not None
+dask_installed = lazy_import("dask") is not None
+mimesis_installed = lazy_import("mimesis") is not None
 
 
 @pytest.mark.skipif(not dask_installed, reason="dask not installed")

mars/core/entrypoints.py

@@ -16,8 +16,6 @@
 import warnings
 import functools
 
-from pkg_resources import iter_entry_points
-
 logger = logging.getLogger(__name__)
 
 
@@ -28,6 +26,8 @@ def init_extension_entrypoints():
     """Execute all `mars_extensions` entry points with the name `init`
     If extensions have already been initialized, this function does nothing.
     """
+    from pkg_resources import iter_entry_points
+
     for entry_point in iter_entry_points("mars_extensions", "init"):
         logger.info("Loading extension: %s", entry_point)
         try:

mars/dataframe/base/_duplicate.py

@@ -26,7 +26,7 @@
 from ..initializer import DataFrame as asdataframe
 from ..operands import DataFrameOperandMixin, DataFrameShuffleProxy
 
-cudf = lazy_import("cudf", globals=globals())
+cudf = lazy_import("cudf")
 
 
 class DuplicateOperand(MapReduceOperand, DataFrameOperandMixin):

mars/dataframe/base/drop_duplicates.py

@@ -28,7 +28,7 @@
 )
 from ._duplicate import DuplicateOperand, validate_subset
 
-cudf = lazy_import("cudf", globals=globals())
+cudf = lazy_import("cudf")
 
 
 class DataFrameDropDuplicates(DuplicateOperand):

mars/dataframe/base/memory_usage.py

@@ -27,7 +27,7 @@
 from ..core import IndexValue
 from ..utils import parse_index
 
-cudf = lazy_import("cudf", globals=globals())
+cudf = lazy_import("cudf")
 
 
 class DataFrameMemoryUsage(DataFrameOperand, DataFrameOperandMixin):