Repository: mars-project/mars

Issue:
[BUG] `build_faiss_index` failed for some index type
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

`build_faiss_index` failed for some index type.

**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.

```
In [1]: from sklearn.datasets import make_classification                        

In [2]: x, y = make_classification()                                            

In [3]: import mars.tensor as mt                                                
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
  return f(*args, **kwds)
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
  return f(*args, **kwds)

In [4]: x = mt.tensor(x, chunk_size=20)                                         

In [5]: x.shape                                                                 
Out[5]: (100, 20)

In [6]: y = mt.tensor(y, chunk_size=20)                                         

In [7]: y.shape                                                                 
Out[7]: (100,)

/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
  return f(*args, **kwds)
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
  return f(*args, **kwds)
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
  return f(*args, **kwds)
In [8]: from mars.learn.neighbors._faiss import build_faiss_index               

In [39]: index = build_faiss_index(x, index_name='PCAR6,IVF8_HNSW32,SQ8', n_samp
    ...: le=10)                                                                 

In [40]: index.execute()                                                        
WARNING clustering 10 points to 8 centroids: please provide at least 312 training points
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-40-bd4069985a62> in <module>
----> 1 index.execute()

~/Workspace/mars/mars/core.py in execute(self, session, **kw)
    373 
    374         if wait:
--> 375             return run()
    376         else:
    377             # leverage ThreadPoolExecutor to submit task,

~/Workspace/mars/mars/core.py in run()
    368         def run():
    369             # no more fetch, thus just fire run
--> 370             session.run(self, **kw)
    371             # return Tileable or ExecutableTuple itself
    372             return self

~/Workspace/mars/mars/session.py in run(self, *tileables, **kw)
    498         tileables = tuple(mt.tensor(t) if not isinstance(t, (Entity, Base)) else t
    499                           for t in tileables)
--> 500         result = self._sess.run(*tileables, **kw)
    501 
    502         for t in tileables:

~/Workspace/mars/mars/session.py in run(self, *tileables, **kw)
    106             # set number of running cores
    107             self.context.set_ncores(kw['n_parallel'])
--> 108             res = self._executor.execute_tileables(tileables, **kw)
    109             return res
    110 

~/Workspace/mars/mars/utils.py in _inner(*args, **kwargs)
    437         def _inner(*args, **kwargs):
    438             with self:
--> 439                 return func(*args, **kwargs)
    440 
    441         return _inner

~/Workspace/mars/mars/executor.py in execute_tileables(self, tileables, fetch, n_parallel, n_thread, print_progress, mock, compose, name)
    879                                    n_parallel=n_parallel or n_thread,
    880                                    print_progress=print_progress, mock=mock,
--> 881                                    chunk_result=chunk_result)
    882 
    883                 # update shape of tileable and its chunks whatever it's successful or not

~/Workspace/mars/mars/executor.py in execute_graph(self, graph, keys, n_parallel, print_progress, mock, no_intermediate, compose, retval, chunk_result)
    691             print_progress=print_progress, mock=mock, mock_max_memory=self._mock_max_memory,
    692             fetch_keys=fetch_keys, no_intermediate=no_intermediate)
--> 693         res = graph_execution.execute(retval)
    694         self._mock_max_memory = max(self._mock_max_memory, graph_execution._mock_max_memory)
    695         if mock:

~/Workspace/mars/mars/executor.py in execute(self, retval)
    572         # wait until all the futures completed
    573         for future in executed_futures:
--> 574             future.result()
    575 
    576         if retval:

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in result(self, timeout)
    433                 raise CancelledError()
    434             elif self._state == FINISHED:
--> 435                 return self.__get_result()
    436             else:
    437                 raise TimeoutError()

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in __get_result(self)
    382     def __get_result(self):
    383         if self._exception:
--> 384             raise self._exception
    385         else:
    386             return self._result

~/miniconda3/lib/python3.7/concurrent/futures/thread.py in run(self)
     55 
     56         try:
---> 57             result = self.fn(*self.args, **self.kwargs)
     58         except BaseException as exc:
     59             self.future.set_exception(exc)

~/Workspace/mars/mars/utils.py in _inner(*args, **kwargs)
    437         def _inner(*args, **kwargs):
    438             with self:
--> 439                 return func(*args, **kwargs)
    440 
    441         return _inner

~/Workspace/mars/mars/executor.py in _execute_operand(self, op)
    444             # so we pass the first operand's first output to Executor.handle
    445             first_op = ops[0]
--> 446             self.handle_op(first_op, results, self._mock)
    447 
    448             # update maximal memory usage during execution

~/Workspace/mars/mars/executor.py in handle_op(self, *args, **kw)
    376 
    377     def handle_op(self, *args, **kw):
--> 378         return Executor.handle(*args, **kw)
    379 
    380     def _order_starts(self):

~/Workspace/mars/mars/executor.py in handle(cls, op, results, mock)
    642                 # The `UFuncTypeError` was introduced by numpy#12593 since v1.17.0.
    643                 try:
--> 644                     return runner(results, op)
    645                 except UFuncTypeError as e:
    646                     raise TypeError(str(e)).with_traceback(sys.exc_info()[2]) from None

~/Workspace/mars/mars/learn/neighbors/_faiss.py in execute(cls, ctx, op)
    340             cls._execute_map(ctx, op)
    341         elif op.stage == OperandStage.agg:
--> 342             cls._execute_agg(ctx, op)
    343         else:
    344             assert op.stage is None

~/Workspace/mars/mars/learn/neighbors/_faiss.py in _execute_agg(cls, ctx, op)
    327                 index = _load_index(ctx, op, index, device_id)
    328                 indexes.append(index)
--> 329                 assert hasattr(index, 'merge_from')
    330                 if merged_index is None:
    331                     merged_index = index

AssertionError:

```




Discussion:

Resolving patch:
diff --git a/mars/learn/neighbors/_faiss.py b/mars/learn/neighbors/_faiss.py
--- a/mars/learn/neighbors/_faiss.py
+++ b/mars/learn/neighbors/_faiss.py
@@ -186,6 +186,7 @@ def _tile_chunks(cls, op, in_tensor, faiss_index, n_sample):
         # 2. distributions of chunks are the same, in not,
         #    train separately on each chunk data
         need_sample_train = not faiss_index_.is_trained and op.same_distribution
+        need_merge_index = hasattr(faiss_index_, 'merge_from') if need_sample_train else False
 
         train_chunk = None
         if need_sample_train:
@@ -217,7 +218,7 @@ def _tile_chunks(cls, op, in_tensor, faiss_index, n_sample):
             build_index_chunks.append(build_index_chunk)
 
         out_chunks = []
-        if need_sample_train:
+        if need_merge_index:
             assert op.n_sample is not None
             # merge all indices into one, do only when trained on sample data
             out_chunk_op = op.copy().reset_key()
@@ -238,6 +239,7 @@ def _execute_one_chunk(cls, ctx, op):
             [ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True)
 
         with device(device_id):
+            inp = inp.astype(np.float32, copy=False)
             # create index
             index = faiss.index_factory(inp.shape[1], op.faiss_index,
                                         op.faiss_metric_type)
@@ -262,7 +264,6 @@ def _execute_one_chunk(cls, ctx, op):
             # add vectors to index
             if device_id >= 0:  # pragma: no cover
                 # gpu
-                inp = inp.astype(np.float32, copy=False)
                 index.add_c(inp.shape[0], _swig_ptr_from_cupy_float32_array(inp))
             else:
                 index.add(inp)


Test patch:
diff --git a/mars/learn/neighbors/tests/test_faiss.py b/mars/learn/neighbors/tests/test_faiss.py
--- a/mars/learn/neighbors/tests/test_faiss.py
+++ b/mars/learn/neighbors/tests/test_faiss.py
@@ -103,6 +103,16 @@ def testManualBuildFaissIndex(self):
             self.assertIsInstance(ind, faiss.IndexIVFFlat)
             self.assertEqual(ind.ntotal, 10)
 
+        # test more index type
+        index = build_faiss_index(X, 'PCAR6,IVF8_HNSW32,SQ8', 10, random_state=0,
+                                  return_index_type='object')
+        faiss_index = self.executor.execute_tileable(index)
+
+        self.assertEqual(len(faiss_index), 5)
+        for ind in faiss_index:
+            self.assertIsInstance(ind, faiss.IndexPreTransform)
+            self.assertEqual(ind.ntotal, 10)
+
         # test one chunk, train
         X = mt.tensor(x, chunk_size=50)
         index = build_faiss_index(X, 'IVF30,Flat', 30, random_state=0,
@@ -148,6 +158,15 @@ def testFaissQuery(self):
                 np.testing.assert_array_equal(indices, expected_indices.fetch())
                 np.testing.assert_almost_equal(distance, expected_distance.fetch())
 
+                # test other index
+                X2 = X.astype(np.float64)
+                Y2 = y.astype(np.float64)
+                faiss_index = build_faiss_index(X2, 'PCAR6,IVF8_HNSW32,SQ8', 10,
+                                                random_state=0, return_index_type='object')
+                d, i = faiss_query(faiss_index, Y2, 5, nprobe=10)
+                # test execute only
+                self.executor.execute_tensors([d, i])
+
     def testGenIndexStringAndSampleCount(self):
         d = 32