[BACKPORT] Add mars.learn.cluster.KMeans support (#1426) (#1428)

Author: Xuye (Chris) Qin

.codecov.yml

@@ -28,6 +28,10 @@ ignore:
   - "mars/lib/uhashring"
   - "mars/serialize/protos"
   - "mars/learn/contrib/xgboost/tracker.py"
+  - "mars/learn/cluster/_k_means_fast.*"
+  - "mars/learn/cluster/_k_means_elkan.pyx"
+  - "mars/learn/cluster/_k_means_lloyd.pyx"
+  - "mars/learn/utils/_cython_blas.*"
   - "mars/tensor/einsum/einsumfunc.py"
   - "**/*.html"
   - "**/*.pxd"

.coveragerc

@@ -19,6 +19,10 @@ omit =
     mars/lib/uhashring/*
     mars/serialize/protos/*
     mars/learn/contrib/xgboost/tracker.py
+    mars/learn/cluster/_k_means_fast.*
+    mars/learn/cluster/_k_means_elkan.pyx
+    mars/learn/cluster/_k_means_lloyd.pyx
+    mars/learn/utils/_cython_blas.*
     mars/tensor/einsum/einsumfunc.py
     *.html
     *.pxd

.coveragerc-threaded

@@ -22,6 +22,10 @@ omit =
     mars/lib/uhashring/*
     mars/serialize/protos/*
     mars/learn/contrib/xgboost/tracker.py
+    mars/learn/cluster/_k_means_fast.*
+    mars/learn/cluster/_k_means_elkan.pyx
+    mars/learn/cluster/_k_means_lloyd.pyx
+    mars/learn/utils/_cython_blas.*
     mars/tensor/einsum/einsumfunc.py
     *.html
     *.pxd

.github/codecov-upstream.yml

@@ -30,6 +30,10 @@ ignore:
   - "mars/lib/uhashring"
   - "mars/serialize/protos"
   - "mars/learn/contrib/xgboost/tracker.py"
+  - "mars/learn/cluster/_k_means_fast.*"
+  - "mars/learn/cluster/_k_means_elkan.pyx"
+  - "mars/learn/cluster/_k_means_lloyd.pyx"
+  - "mars/learn/utils/_cython_blas.*"
   - "mars/tensor/einsum/einsumfunc.py"
   - "**/*.html"
   - "**/*.pxd"

.github/workflows/ci.yml

@@ -195,7 +195,7 @@ jobs:
 
       - name: Collect coverage data
         run: |
-          pip install numpy cython gevent coverage
+          pip install numpy scipy cython gevent coverage
           pip install codecov
           CYTHON_TRACE=1 python setup.py sdist
           export DEFAULT_VENV=$VIRTUAL_ENV

.gitignore

@@ -82,6 +82,8 @@ mars/optimizes/**/*.c*
 mars/scheduler/*.c*
 mars/serialize/*.c*
 mars/worker/*.c*
+mars/learn/cluster/*.c*
+mars/learn/utils/*.c*
 
 # files built from protobuf files
 mars/opcodes.py

docs/source/learn/generated/mars.learn.cluster.KMeans.rst

@@ -0,0 +1,30 @@
+mars.learn.cluster.KMeans
+=========================
+
+.. currentmodule:: mars.learn.cluster
+
+.. autoclass:: KMeans
+
+   
+   .. automethod:: __init__
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+      ~KMeans.__init__
+      ~KMeans.fit
+      ~KMeans.fit_predict
+      ~KMeans.fit_transform
+      ~KMeans.get_params
+      ~KMeans.predict
+      ~KMeans.score
+      ~KMeans.set_params
+      ~KMeans.transform
+   
+   
+
+   
+   
+   

docs/source/learn/generated/mars.learn.cluster.k_means.rst

@@ -0,0 +1,6 @@
+mars.learn.cluster.k\_means
+===========================
+
+.. currentmodule:: mars.learn.cluster
+
+.. autofunction:: k_means

mars/_version.py

@@ -15,7 +15,7 @@
 import subprocess
 import os
 
-version_info = (0, 4, 3)
+version_info = (0, 4, 4)
 _num_index = max(idx if isinstance(v, int) else 0
                  for idx, v in enumerate(version_info))
 __version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \

mars/executor.py

@@ -28,7 +28,7 @@
 import numpy as np
 
 from .operands import Fetch, ShuffleProxy
-from .graph import DAG, DirectedGraph
+from .graph import DAG
 from .config import options
 from .tiles import IterativeChunkGraphBuilder, ChunkGraphBuilder, get_tiled
 from .optimizes.runtime.core import RuntimeOptimizer
@@ -712,7 +712,7 @@ def _on_tile_success(before_tile_data, after_tile_data):
         chunk_result = self._chunk_result.copy()
         tileable_graph_builder = TileableGraphBuilder()
         tileable_graph = tileable_graph_builder.build([tileable])
-        chunk_graph_builder = ChunkGraphBuilder(graph_cls=DirectedGraph, compose=compose,
+        chunk_graph_builder = ChunkGraphBuilder(compose=compose,
                                                 on_tile_success=_on_tile_success)
         chunk_graph = chunk_graph_builder.build([tileable], tileable_graph=tileable_graph)
         ret = self.execute_graph(chunk_graph, result_keys, n_parallel=n_parallel or n_thread,
@@ -774,6 +774,7 @@ def execute_tileables(self, tileables, fetch=True, n_parallel=None, n_thread=Non
         tileable_data_to_chunks = weakref.WeakKeyDictionary()
 
         node_to_fetch = weakref.WeakKeyDictionary()
+        skipped_tileables = set()
 
         def _generate_fetch_tileable(node):
             # Attach chunks to fetch tileables to skip tile.
@@ -786,6 +787,22 @@ def _generate_fetch_tileable(node):
 
             return node
 
+        def _skip_executed_tileables(inps):
+            # skip the input that executed, and not gc collected
+            new_inps = []
+            for inp in inps:
+                if inp.key in self.stored_tileables:
+                    try:
+                        get_tiled(inp)
+                    except KeyError:
+                        new_inps.append(inp)
+                    else:
+                        skipped_tileables.add(inp)
+                        continue
+                else:
+                    new_inps.append(inp)
+            return new_inps
+
         def _generate_fetch_if_executed(nd):
             # node processor that if the node is executed
             # replace it with a fetch node
@@ -830,7 +847,8 @@ def _get_tileable_graph_builder(**kwargs):
         with self._gen_local_context(chunk_result):
             # build tileable graph
             tileable_graph_builder = _get_tileable_graph_builder(
-                node_processor=_generate_fetch_tileable)
+                node_processor=_generate_fetch_tileable,
+                inputs_selector=_skip_executed_tileables)
             tileable_graph = tileable_graph_builder.build(tileables)
             chunk_graph_builder = IterativeChunkGraphBuilder(
                 graph_cls=DAG, node_processor=_generate_fetch_if_executed,

mars/learn/__init__.py

@@ -17,6 +17,7 @@
 # see https://github.com/pytorch/pytorch/issues/2575
 from .contrib import pytorch, tensorflow, xgboost, lightgbm
 from .metrics import pairwise
+from . import cluster
 from . import preprocessing
 from . import neighbors
 from . import utils

mars/learn/cluster/__init__.py

@@ -0,0 +1,35 @@
+# Copyright 1999-2020 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.
+
+try:
+    from ._kmeans import KMeans, k_means
+
+    def _install():
+        from ._k_means_common import KMeansInertia, KMeansRelocateEmptyClusters
+        from ._k_means_elkan_iter import KMeansElkanInitBounds, \
+            KMeansElkanUpdate, KMeansElkanPostprocess
+        from ._k_means_init import KMeansPlusPlusInit
+        from ._k_means_lloyd_iter import KMeansLloydUpdate, KMeansLloydPostprocess
+
+        del KMeansInertia, KMeansRelocateEmptyClusters, KMeansElkanInitBounds, \
+            KMeansElkanUpdate, KMeansElkanPostprocess, KMeansPlusPlusInit, \
+            KMeansLloydUpdate, KMeansLloydPostprocess
+
+
+    _install()
+    del _install
+except ImportError:
+    KMeans = None
+    k_means = None
+