TimeoutCallback Python

Summary: facebookresearch#3351

Differential Revision: D56856744

faiss/impl/AuxIndexStructures.cpp

@@ -255,4 +255,10 @@ bool TimeoutCallback::want_interrupt() {
     return false;
 }
 
+void TimeoutCallback::reset(double timeout_in_seconds) {
+    auto tc(new faiss::TimeoutCallback());
+    faiss::InterruptCallback::instance.reset(tc);
+    tc->set_timeout(timeout_in_seconds);
+}
+
 } // namespace faiss

faiss/impl/AuxIndexStructures.h

@@ -166,6 +166,7 @@ struct TimeoutCallback : InterruptCallback {
     double timeout;
     bool want_interrupt() override;
     void set_timeout(double timeout_in_seconds);
+    static void reset(double timeout_in_seconds);
 };
 
 /// set implementation optimized for fast access.

faiss/python/__init__.py

@@ -316,3 +316,14 @@ def deserialize_index_binary(data):
     reader = VectorIOReader()
     copy_array_to_vector(data, reader.data)
     return read_index_binary(reader)
+
+
+class TimeoutGuard:
+    def __init__(self, timeout_in_seconds: float):
+        self.timeout = timeout_in_seconds
+
+    def __enter__(self):
+        TimeoutCallback.reset(self.timeout)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        PythonInterruptCallback.reset()

faiss/python/swigfaiss.swig

@@ -1041,7 +1041,9 @@ PyObject *swig_ptr (PyObject *a)
     PyErr_SetString(PyExc_ValueError, "did not recognize array type");
     return NULL;
 }
+%}
 
+%inline %{
 
 struct PythonInterruptCallback: faiss::InterruptCallback {
 
@@ -1056,15 +1058,18 @@ struct PythonInterruptCallback: faiss::InterruptCallback {
         return err == -1;
     }
 
+    static void reset() {
+        faiss::InterruptCallback::instance.reset(new PythonInterruptCallback());
+    }
 };
+
 %}
 
 %init %{
     /* needed, else crash at runtime */
     import_array();
 
-    faiss::InterruptCallback::instance.reset(new PythonInterruptCallback());
-
+    PythonInterruptCallback::reset();
 %}
 
 // return a pointer usable as input for functions that expect pointers

tests/test_callback.cpp

@@ -31,9 +31,7 @@ TEST(TestCallback, timeout) {
 
     faiss::Clustering kmeans(d, k, cp);
 
-    auto tc(new faiss::TimeoutCallback());
-    faiss::InterruptCallback::instance.reset(tc);
-    tc->set_timeout(0.010);
+    faiss::TimeoutCallback::reset(0.010);
     EXPECT_THROW(kmeans.train(n, vecs.data(), *index), faiss::FaissException);
     delete index;
 }

tests/test_callback_py.py

@@ -0,0 +1,32 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+import numpy as np
+import faiss
+
+
+class TestCallbackPy(unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+
+    def test_timeout(self) -> None:
+        n = 1_000_000
+        k = 100
+        d = 128
+        niter = 10
+
+        x = np.random.rand(n, d).astype('float32')
+        index = faiss.IndexFlat(d)
+
+        cp = faiss.ClusteringParameters()
+        cp.niter = niter
+        cp.verbose = False
+
+        kmeans = faiss.Clustering(d, k, cp)
+
+        with self.assertRaises(RuntimeError):
+            with faiss.TimeoutGuard(0.010):
+                kmeans.train(x, index)