Python函数执行完成与返回间存在显著延迟问题排查

我遇到了一个难以诊断的问题：Python函数在实际执行完成后，需要很长时间才会返回。以下是简化后的最小复现代码：

from sklearn.neighbors import KDTree
import time

def mwe(n):
    start = time.time()
    a = KDTree([[1]]).query_radius([[1]]*int(n), 1)
    print(f'Function completed in: {time.time()-start:.2f} seconds')

start = time.time()
mwe(1e6)
print(f'Function returned after {time.time()-start:.1f} seconds')

start = time.time()
mwe(1e7)
print(f'Function returned after {time.time()-start:.1f} seconds')

输出结果：

Function completed in: 0.36 seconds
Function returned after 4.4 seconds

Function completed in: 3.75 seconds
Function returned after 44.0 seconds

起初我认为这是垃圾回收问题，但添加gc.disable()后没有任何效果。而且如果是垃圾回收问题，应该会在性能分析结果中明显体现，但cProfile甚至没有识别到该问题的存在：

645 function calls (636 primitive calls) in 2.701 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        2    1.069    0.534    1.076    0.538 {method 'query_radius' of 'sklearn.neighbors._ball_tree.BinaryTree' objects}
    32/23    0.919    0.029    1.059    0.046 {built-in method numpy.core._multiarray_umath.implement_array_function}
        1    0.570    0.570    2.701    2.701 test_return_delay.py:6(test_loop)
        4    0.135    0.034    0.135    0.034 {method 'repeat' of 'numpy.ndarray' objects}
        7    0.004    0.001    0.004    0.001 {method 'reduce' of 'numpy.ufunc' objects}
        5    0.002    0.000    0.002    0.000 {built-in method numpy.arange}
        2    0.000    0.000    0.000    0.000 {built-in method builtins.print}
        2    0.000    0.000    0.246    0.123 index_tricks.py:322(__getitem__)
        4    0.000    0.000    0.005    0.001 validation.py:96(_assert_all_finite)
        3    0.000    0.000    0.000    0.000 function_base.py:23(linspace)
        4    0.000    0.000    0.005    0.001 validation.py:629(check_array)
        4    0.000    0.000    0.000    0.000 _array_api.py:168(_asarray_with_order)
        4    0.000    0.000    0.000    0.000 numerictypes.py:573(_can_coerce_all)
        4    0.000    0.000    0.000    0.000 validation.py:320(_num_samples)
       47    0.000    0.000    0.000    0.000 {built-in method builtins.isinstance}
        4    0.000    0.000    0.135    0.034 <__array_function__ internals>:177(repeat)
...

如果以模块形式运行cProfile分析整个脚本（即python -m cProfile -s time test_return_delay.py），虽然能显示存在问题，但我无法从中定位到问题根源：

788549 function calls (774939 primitive calls) in 12.372 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1   11.257   11.257   11.257   11.257 {method 'enable' of '_lsprof.Profiler' objects}
      808    0.173    0.000    0.173    0.000 {built-in method io.open_code}
     3514    0.124    0.000    0.124    0.000 {built-in method nt.stat}
      157    0.094    0.001    0.097    0.001 {built-in method _imp.create_dynamic}
      808    0.076    0.000    0.076    0.000 {built-in method marshal.loads}
      808    0.045    0.000    0.045    0.000 {method 'read' of '_io.BufferedReader' objects}
     6273    0.033    0.000    0.051    0.000 <frozen importlib._bootstrap_external>:96(_path_join)
   1066/1    0.025    0.000   12.373   12.373 {built-in method builtins.exec}
      391    0.025    0.000    0.067    0.000 doccer.py:13(docformat)
1701/1692    0.020    0.000    0.086    0.000 {built-in method builtins.__build_class__}
  157/126    0.016    0.000    0.084    0.001 {built-in method _imp.exec_dynamic}
      399    0.011    0.000    0.019    0.000 doccer.py:180(indentcount_lines)
     1359    0.011    0.000    0.161    0.000 <frozen importlib._bootstrap_external>:1536(find_spec)
     9043    0.009    0.000    0.009    0.000 {method 'splitlines' of 'str' objects}
  499/100    0.009    0.000    0.025    0.000 sre_parse.py:494(_parse)
      809    0.009    0.000    0.009    0.000 {method '__exit__' of '_io._IOBase' objects}
    17871    0.009    0.000    0.009    0.000 {built-in method builtins.getattr}
      808    0.008    0.000    0.367    0.000 <frozen importlib._bootstrap_external>:950(get_code)
   100236    0.008    0.000    0.008    0.000 {method 'append' of 'list' objects}
 1013/842    0.007    0.000    0.337    0.000 <frozen importlib._bootstrap>:921(_find_spec)
   1040/2    0.007    0.000    1.116    0.558 <frozen importlib._bootstrap>:1022(_find_and_load)
     9075    0.007    0.000    0.007    0.000 {method 'expandtabs' of 'str' objects}
      671    0.007    0.000    0.020    0.000 inspect.py:2280(_signature_from_function)
        1    0.006    0.006    0.006    0.006 memory.py:371(MemorizedFunc)
...

请问问题可能出在哪里，或者我该如何进一步调试？

问题根源 #

延迟返回的核心原因是sklearn.neighbors.KDTree的query_radius方法返回的是数组组成的嵌套结构，而非像scipy.spatial.KDTree那样返回嵌套列表。当返回结果包含百万甚至千万级元素时，Python需要将C++层面生成的数组结构逐个转换为Python对象，这个过程会产生大量小对象，耗时极长且无法被cProfile准确捕获——因为大部分开销发生在C到Python的对象转换阶段，不属于Python函数调用范畴。

调试建议 #

• 对比返回对象结构：打印type(a)及单个元素的类型，和scipy版本的结果做对比，能快速定位差异。
• 跟踪内存分配：用tracemalloc观察函数返回后的内存变化，会发现大量小对象被分配，佐证对象转换的开销。
• 单独测试转换耗时：将query_radius的结果赋值后，单独测试把它转为列表的时间，可明确看到这部分的耗时占比。

修复方向 #

调整KDTree及其他近邻类的查询方法，使其返回嵌套列表而非数组嵌套结构，对齐scipy.spatial.KDTree的行为，避免大量小数组对象的创建与转换开销。

内容的提问来源于stack exchange，提问作者brandonsmithj