Repository: mars-project/mars

Issue:
[BUG] KMeans hangs for local cluster
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

KMeans may hang for local cluster.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and protobuf
4. Full stack of the error.
5. Minimized code to reproduce the error.

The code to reproduce.

```python
import mars.tensor as mt
from mars.learn.cluster import KMeans
from mars.deploy.local import new_cluster

cluster = new_cluster()

X = mt.array([[1, 2], [1, 4], [1, 0],
                        [10, 2], [10, 4], [10, 0]])
kmeans = KMeans(n_clusters=2, random_state=0, init='k-means++').fit(X)
```





Discussion:

Resolving patch:
diff --git a/mars/learn/cluster/_kmeans.py b/mars/learn/cluster/_kmeans.py
--- a/mars/learn/cluster/_kmeans.py
+++ b/mars/learn/cluster/_kmeans.py
@@ -183,7 +183,7 @@ def k_means(X, n_clusters, sample_weight=None, init='k-means||',
 def _kmeans_single_elkan(X, sample_weight, n_clusters, max_iter=300,
                          init='k-means++', verbose=False, x_squared_norms=None,
                          random_state=None, tol=1e-4, oversampling_factor=2,
-                         init_iter=5, session=None, run_kwargs=None):
+                         init_iter=5, X_mean=None, session=None, run_kwargs=None):
     random_state = check_random_state(random_state)
     sample_weight = _check_normalize_sample_weight(sample_weight, X)
 
@@ -195,9 +195,11 @@ def _kmeans_single_elkan(X, sample_weight, n_clusters, max_iter=300,
 
     # execute X, centers and tol first
     tol = mt.asarray(tol)
-    mt.ExecutableTuple([X, sample_weight, centers, x_squared_norms, tol]).execute(
-        session=session, **(run_kwargs or dict()))
-    tol = tol.fetch()
+    to_run = [X, sample_weight, centers, x_squared_norms, tol]
+    if X_mean is not None:
+        to_run.append(X_mean)
+    mt.ExecutableTuple(to_run).execute(session=session, **(run_kwargs or dict()))
+    tol = tol.fetch(session=session)
 
     if verbose:
         print('Initialization complete')
@@ -237,9 +239,9 @@ def _kmeans_single_elkan(X, sample_weight, n_clusters, max_iter=300,
         mt.ExecutableTuple(to_runs).execute(session=session, **(run_kwargs or dict()))
 
         if verbose:
-            print("Iteration {0}, inertia {1}" .format(i, inertia.fetch()))
+            print("Iteration {0}, inertia {1}" .format(i, inertia.fetch(session=session)))
 
-        center_shift_tot = center_shift_tot.fetch()
+        center_shift_tot = center_shift_tot.fetch(session=session)
         if center_shift_tot <= tol:
             if verbose:  # pragma: no cover
                 print("Converged at iteration {0}: "
@@ -267,7 +269,7 @@ def _kmeans_single_elkan(X, sample_weight, n_clusters, max_iter=300,
 def _kmeans_single_lloyd(X, sample_weight, n_clusters, max_iter=300,
                          init='k-means++', verbose=False, x_squared_norms=None,
                          random_state=None, tol=1e-4, oversampling_factor=2,
-                         init_iter=5, session=None, run_kwargs=None):
+                         init_iter=5, X_mean=None, session=None, run_kwargs=None):
     random_state = check_random_state(random_state)
     sample_weight = _check_normalize_sample_weight(sample_weight, X)
 
@@ -279,9 +281,11 @@ def _kmeans_single_lloyd(X, sample_weight, n_clusters, max_iter=300,
 
     # execute X, centers and tol first
     tol = mt.asarray(tol)
-    mt.ExecutableTuple([X, centers, x_squared_norms, tol]).execute(
-        session=session, **(run_kwargs or dict()))
-    tol = tol.fetch()
+    to_run = [X, centers, x_squared_norms, tol]
+    if X_mean is not None:
+        to_run.append(X_mean)
+    mt.ExecutableTuple(to_run).execute(session=session, **(run_kwargs or dict()))
+    tol = tol.fetch(session=session)
 
     if verbose:  # pragma: no cover
         print("Initialization complete")
@@ -308,9 +312,9 @@ def _kmeans_single_lloyd(X, sample_weight, n_clusters, max_iter=300,
         mt.ExecutableTuple(to_runs).execute(session=session, **(run_kwargs or dict()))
 
         if verbose:  # pragma: no cover
-            print("Iteration {0}, inertia {1}" .format(i, inertia.fetch()))
+            print("Iteration {0}, inertia {1}" .format(i, inertia.fetch(session=session)))
 
-        center_shift_tot = center_shift_tot.fetch()
+        center_shift_tot = center_shift_tot.fetch(session=session)
         if center_shift_tot <= tol:
             if verbose:  # pragma: no cover
                 print("Converged at iteration {0}: "
@@ -760,6 +764,7 @@ def fit(self, X, y=None, sample_weight=None, session=None, run_kwargs=None):
                 n_init = 1
 
         # subtract of mean of x for more accurate distance computations
+        X_mean = None
         if not X.issparse():
             X_mean = X.mean(axis=0)
             # The copy was already done above
@@ -800,8 +805,9 @@ def fit(self, X, y=None, sample_weight=None, session=None, run_kwargs=None):
                 init=init, verbose=self.verbose, tol=tol,
                 x_squared_norms=x_squared_norms, random_state=seed,
                 oversampling_factor=self.oversampling_factor,
-                init_iter=self.init_iter)
-            inertia = inertia.fetch()
+                init_iter=self.init_iter, X_mean=X_mean,
+                session=session, run_kwargs=run_kwargs)
+            inertia = inertia.fetch(session=session)
             # determine if these results are the best so far
             if best_inertia is None or inertia < best_inertia:
                 best_labels = labels
@@ -815,7 +821,7 @@ def fit(self, X, y=None, sample_weight=None, session=None, run_kwargs=None):
             best_centers += X_mean
             best_centers.execute(session=session, **(run_kwargs or dict()))
 
-        distinct_clusters = len(set(best_labels.fetch()))
+        distinct_clusters = len(set(best_labels.fetch(session=session)))
         if distinct_clusters < self.n_clusters:  # pragma: no cover
             warnings.warn(
                 "Number of distinct clusters ({}) found smaller than "
diff --git a/mars/promise.py b/mars/promise.py
--- a/mars/promise.py
+++ b/mars/promise.py
@@ -202,7 +202,8 @@ def _internal_step_next(self, *args, **kwargs):
                 if acceptor and acceptor._accept_handler:
                     # remove the handler in promise in case that
                     # function closure is not freed
-                    handler, acceptor._accept_handler = acceptor._accept_handler, None
+                    handler, acceptor._accept_handler, acceptor._reject_handler = \
+                        acceptor._accept_handler, None, None
                     next_call = handler(*args, **kwargs)
                 else:
                     acceptor._accepted = accept
@@ -214,7 +215,8 @@ def _internal_step_next(self, *args, **kwargs):
                 if rejecter and rejecter._reject_handler:
                     # remove the handler in promise in case that
                     # function closure is not freed
-                    handler, rejecter._reject_handler = rejecter._reject_handler, None
+                    handler, rejecter._reject_handler, rejecter._accept_handler = \
+                        rejecter._reject_handler, None, None
                     next_call = handler(*args, **kwargs)
                 else:
                     rejecter._accepted = accept
@@ -470,6 +472,7 @@ def get_promise(self, promise_id):
         return obj()
 
     def delete_promise(self, promise_id):
+        _promise_pool.pop(promise_id, None)
         if promise_id not in self._promises:
             return
         ref_key = self._promise_ref_keys[promise_id]
@@ -523,25 +526,29 @@ def handle_promise(self, promise_id, *args, **kwargs):
         Callback entry for promise results
         :param promise_id: promise key
         """
-        p = self.get_promise(promise_id)
-        if p is None:
-            logger.warning('Promise %r reentered or not registered in %s', promise_id, self.uid)
-            return
-        self.get_promise(promise_id).step_next([args, kwargs])
-        self.delete_promise(promise_id)
+        try:
+            p = self.get_promise(promise_id)
+            if p is None:
+                logger.warning('Promise %r reentered or not registered in %s', promise_id, self.uid)
+                return
+            self.get_promise(promise_id).step_next([args, kwargs])
+        finally:
+            self.delete_promise(promise_id)
 
     def handle_promise_timeout(self, promise_id):
         """
         Callback entry for promise timeout
         :param promise_id: promise key
         """
-        p = self.get_promise(promise_id)
-        if p is None or p._accepted is not None:
-            # skip promises that are already finished
-            return
+        try:
+            p = self.get_promise(promise_id)
+            if p is None or p._accepted is not None:
+                # skip promises that are already finished
+                return
 
-        self.delete_promise(promise_id)
-        p.step_next([build_exc_info(PromiseTimeout), dict(_accept=False)])
+            p.step_next([build_exc_info(PromiseTimeout), dict(_accept=False)])
+        finally:
+            self.delete_promise(promise_id)
 
 
 def all_(promises):


Test patch:
diff --git a/mars/deploy/local/tests/test_cluster.py b/mars/deploy/local/tests/test_cluster.py
--- a/mars/deploy/local/tests/test_cluster.py
+++ b/mars/deploy/local/tests/test_cluster.py
@@ -32,6 +32,10 @@
     import h5py
 except ImportError:
     h5py = None
+try:
+    import sklearn
+except ImportError:
+    sklearn = None
 
 from mars import tensor as mt
 from mars import dataframe as md
@@ -1111,8 +1115,11 @@ def f5(t, x):
             expected = (raw[raw < 0.5] * 3).sum()
             self.assertAlmostEqual(result, expected)
 
+    @unittest.skipIf(sklearn is None, 'sklearn not installed')
     def testLearnInLocalCluster(self, *_):
+        from mars.learn.cluster import KMeans
         from mars.learn.neighbors import NearestNeighbors
+        from sklearn.cluster import KMeans as SK_KMEANS
         from sklearn.neighbors import NearestNeighbors as SkNearestNeighbors
 
         with new_cluster(scheduler_n_process=2, worker_n_process=3, shared_memory='20M') as cluster:
@@ -1134,3 +1141,10 @@ def testLearnInLocalCluster(self, *_):
             result = [r.fetch() for r in ret]
             np.testing.assert_almost_equal(result[0], expected[0])
             np.testing.assert_almost_equal(result[1], expected[1])
+
+            raw = np.array([[1, 2], [1, 4], [1, 0],
+                            [10, 2], [10, 4], [10, 0]])
+            X = mt.array(raw)
+            kmeans = KMeans(n_clusters=2, random_state=0, init='k-means++').fit(X)
+            sk_km_elkan = SK_KMEANS(n_clusters=2, random_state=0, init='k-means++').fit(raw)
+            np.testing.assert_allclose(kmeans.cluster_centers_, sk_km_elkan.cluster_centers_)