DEBUG Subtile execution hanging

In some cases, the execution of the program hangs.
Is it due to the thread teardown done via Scheduler.finish?
Co-authored-by: Xavier Thompson <xavier.thompson@nexedi.com>

kdtree.pyx

@@ -75,6 +75,10 @@ cdef extern from *:
                 I n_features
     ) nogil except +
 
+cdef extern from "<sys/syscall.h>" nogil:
+    int SYS_gettid
+    long syscall(long)
+
 cdef I_t find_node_split_dim(D_t* data,
                              I_t* node_indices,
                              I_t n_features,
@@ -515,6 +519,30 @@ cdef cypclass Node activable:
             counter=counter,
         )
 
+# TODO (jjerphan): integrate this where needed once the
+#       compilation problem with Cython has been resolved
+cdef void prange_workaround(
+    KDTree tree,
+    NeighborsHeaps heaps,
+    I_t n_query,
+    I_t num_threads,
+    I_t n_features,
+    D_t * _query_points_ptr,
+) nogil:
+    cdef:
+        D_t rdist_lower_bound = 0
+        I_t i
+
+    printf("pre-prange tid: %d\n", syscall(SYS_gettid))
+
+    for i in prange(n_query, schedule='static', num_threads=num_threads):
+        printf("i, tid: %d\n", syscall(SYS_gettid))
+        printf("%d / %d: start\n", i, n_query)
+        rdist_lower_bound = tree.min_rdist(0, _query_points_ptr + i * n_features)
+        tree._query_single_depthfirst(0, _query_points_ptr, i, heaps, rdist_lower_bound)
+        printf("%d / %d: end\n", i, n_query)
+    # printf("end prange\n")
+
 
 cdef cypclass KDTree:
     """A KDTree based on asynchronous and parallel computations.
@@ -613,13 +641,9 @@ cdef cypclass KDTree:
 
         # Waiting for the tree construction to end
         # Somewhat similar to a thread barrier
-        while initialised < self._n_samples:
-            initialised = counter.value(NULL).getIntResult()
-
-        counter.reset(NULL)
+        scheduler.finish()
 
     void __dealloc__(self):
-        scheduler.finish()
         free(self._indices_ptr)
         free(self._node_data_ptr)
         free(self._node_bounds_ptr)
@@ -699,9 +723,7 @@ cdef cypclass KDTree:
                 n_neighbors
             )
 
-        for i in prange(n_query, schedule='static', num_threads=num_threads):
-            rdist_lower_bound = self.min_rdist(0, _query_points_ptr + i * n_features)
-            self._query_single_depthfirst(0, _query_points_ptr, i, heaps, rdist_lower_bound)
+        prange_workaround(self, heaps, n_query, num_threads, n_features, _query_points_ptr)
 
         heaps.sort()
 

script.py

+import kdtree
+import numpy as np
+
+if __name__ == "__main__":
+    X = np.load("X.npy")
+    tree = kdtree.KDTree(X, leaf_size=256)
+    print("done")
+    n_query = 2 # X.shape[0] // 100000
+    k = 1
+    Y = np.random.rand(n_query, X.shape[1])
+
+    knn_indices = np.zeros((n_query, k), dtype=np.int32)
+    knn_distances = np.zeros((n_query, k), dtype=np.float64)
+    tree.query(Y, knn_indices, knn_distances)
+