为何在子进程中训练XGBoost模型无法终止？

核心问题 #

你的代码里有两个关键问题导致子进程超时：

1. 样本维度不匹配：my_function中pol.fit的特征矩阵是(5,1500)（仅5个样本），但标签是(50,1)（50个样本），XGBoost要求两者样本数必须一致，这会直接抛出ValueError，导致函数提前终止，子进程根本没机会把结果放进队列，主进程就一直等待直到超时。你看到Started打印，只是因为函数开头的代码执行了，后面的fit报错后函数直接崩了。
2. 异常处理缺失：子进程里没捕获函数可能抛出的异常，主进程也没处理这种情况，只会傻等队列结果。另外Empty和multiprocessing.TimeoutError没正确导入，会导致捕获异常时出错。

修复后的代码 #

from multiprocessing import Process, Queue, TimeoutError
import time
import numpy as np
import xgboost
from queue import Empty


def run_process_timeout_wrapper(function, args, timeout):
    def foo(n, out_q, err_q):
        try:
            res = function(*n)
            out_q.put(res)
        except Exception as e:
            err_q.put(e)

    result_q = Queue()
    error_q = Queue()
    p = Process(target=foo, args=(args, result_q, error_q))
    p.start()

    try:
        # 先检查子进程是否抛出异常
        if not error_q.empty():
            raise error_q.get()
        # 获取函数执行结果
        x = result_q.get(timeout=timeout)
    except Empty:
        p.terminate()
        raise TimeoutError(f"Timed out after waiting for {timeout}s")
    except Exception as e:
        p.terminate()
        raise e
    finally:
        # 确保子进程正常结束，避免僵尸进程
        p.join()

    return x


def my_function(fun):
    print("Started")
    t1 = time.time()
    pol = xgboost.XGBRegressor()
    # 修复：特征和标签样本数统一为50个
    pol.fit(np.random.rand(50, 1500), np.random.rand(50, 1))
    print("Took ", time.time() - t1)
    pol.predict(np.random.rand(2, 1500))
    return 5


if __name__ == '__main__':
    t1 = time.time()
    pol = xgboost.XGBRegressor()
    pol.fit(np.random.rand(50, 150000), np.random.rand(50, 1))
    print("Took ", time.time() - t1)

    my_function(None)

    t1 = time.time()
    try:
        res = run_process_timeout_wrapper(my_function, (None,), 160)
        print("Res ", res, " Time ", time.time() - t1)
    except Exception as e:
        print(f"Error: {e}, Time taken: {time.time() - t1}")

主要修改说明 #

• 修正样本维度：把my_function里的特征矩阵改成(50,1500)，和标签样本数匹配，确保fit能正常执行完成。
• 新增异常队列：子进程中捕获函数的所有异常，通过error_q传递给主进程，这样能直接看到子进程里的报错，不用瞎猜。
• 完善导入与异常处理：正确导入queue.Empty和multiprocessing.TimeoutError，主进程先检查异常再获取结果，避免无意义的等待。
• 优化子进程管理：用finally块确保子进程执行join，避免产生僵尸进程；异常时及时终止子进程并抛出错误。

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