Kaggle TPU环境下joblib结合NeuralProphet多进程报错求助

在Kaggle TPU环境中执行joblib多进程任务时，NeuralProphet无法正常运行，报错如下：

_RemoteTraceback: 
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/site-packages/joblib/externals/loky/process_executor.py", line 463, in _process_worker
    r = call_item()
  File "/usr/local/lib/python3.10/site-packages/joblib/externals/loky/process_executor.py", line 291, in __call__
    return self.fn(*self.args, **self.kwargs)
  File "/usr/local/lib/python3.10/site-packages/joblib/parallel.py", line 598, in __call__
    return [func(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/joblib/parallel.py", line 598, in <listcomp>
    return [func(*args, **kwargs)
  File "/tmp/ipykernel_13/4008567848.py", line 22, in run_one
  File "/usr/local/lib/python3.10/site-packages/neuralprophet/forecaster.py", line 1062, in fit
    metrics_df = self._train(
  File "/usr/local/lib/python3.10/site-packages/neuralprophet/forecaster.py", line 2757, in _train
    self.trainer, checkpoint_callback = utils.configure_trainer(
  File "/usr/local/lib/python3.10/site-packages/neuralprophet/utils.py", line 1007, in configure_trainer
    return pl.Trainer(**config), checkpoint_callback
  File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/utilities/argparse.py", line 348, in insert_env_defaults
    return fn(self, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 513, in __init__
    self._setup_on_init()
  File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 546, in _setup_on_init
    setup._log_device_info(self)
  File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/setup.py", line 166, in _log_device_info
    rank_zero_info(f"TPU available: {TPUAccelerator.is_available()}, using: {num_tpu_cores} TPU cores")
  File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/accelerators/tpu.py", line 77, in is_available
    return FabricTPUAccelerator.is_available()
  File "/usr/local/lib/python3.10/site-packages/lightning_fabric/accelerators/tpu.py", line 62, in is_available
    return bool(_XLA_AVAILABLE) and _is_device_tpu()
  File "/usr/local/lib/python3.10/site-packages/lightning_fabric/accelerators/tpu.py", line 90, in wrapper
    proc.start()
  File "/usr/local/lib/python3.10/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/usr/local/lib/python3.10/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/usr/local/lib/python3.10/site-packages/joblib/externals/loky/backend/process.py", line 45, in _Popen
    return Popen(process_obj)
  File "/usr/local/lib/python3.10/site-packages/joblib/externals/loky/backend/popen_loky_posix.py", line 48, in __init__
    self._launch(process_obj)
  File "/usr/local/lib/python3.10/site-packages/joblib/externals/loky/backend/popen_loky_posix.py", line 128, in _launch
    pid = fork_exec(cmd_python, self._fds, env=process_obj.env)
AttributeError: 'Process' object has no attribute 'env'

The above exception was the direct cause of the following exception:

AttributeError                            Traceback (most recent call last)
Cell In[34], line 1
----> 1 p = Parallel(n_jobs=96)(jobs)

File /usr/local/lib/python3.10/site-packages/joblib/parallel.py:2007, in Parallel.__call__(self, iterable)
   2001 # The first item from the output is blank, but it makes the interpreter
   2002 # progress until it enters the Try/Except block of the generator and
   2003 # reaches the first `yield` statement. This starts the asynchronous
   2004 # dispatch of the tasks to the workers.
   2005 next(output)
-> 2007 return output if self.return_generator else list(output)

File /usr/local/lib/python3.10/site-packages/joblib/parallel.py:1650, in Parallel._get_outputs(self, iterator, pre_dispatch)
   1647     yield
   1649     with self._backend.retrieval_context():
-> 1650         yield from self._retrieve()
   1652 except GeneratorExit:
   1653     # The generator has been garbage collected before being fully
   1654     # consumed. This aborts the remaining tasks if possible and warn
   1655     # the user if necessary.
   1656     self._exception = True

File /usr/local/lib/python3.10/site-packages/joblib/parallel.py:1754, in Parallel._retrieve(self)
   1747 while self._wait_retrieval():
   1748 
   1749     # If the callback thread of a worker has signaled that its task
   1750     # triggered an exception, or if the retrieval loop has raised an
   1751     # exception (e.g. `GeneratorExit`), exit the loop and surface the
   1752     # worker traceback.
   1753     if self._aborting:
-> 1754         self._raise_error_fast()
   1755         break
   1757     # If the next job is not ready for retrieval yet, we just wait for
   1758     # async callbacks to progress.

File /usr/local/lib/python3.10/site-packages/joblib/parallel.py:1789, in Parallel._raise_error_fast(self)
   1785 # If this error job exists, immediately raise the error by
   1786 # calling get_result. This job might not exists if abort has been
   1787 # called directly or if the generator is gc'ed.
   1788 if error_job is not None:
-> 1789     error_job.get_result(self.timeout)

File /usr/local/lib/python3.10/site-packages/joblib/parallel.py:745, in BatchCompletionCallBack.get_result(self, timeout)
    739 backend = self.parallel._backend
    741 if backend.supports_retrieve_callback:
    742     # We assume that the result has already been retrieved by the
    743     # callback thread, and is stored internally. It's just waiting to
    744     # be returned.
-> 745     return self._return_or_raise()
    747 # For other backends, the main thread needs to run the retrieval step.
    748 try:

File /usr/local/lib/python3.10/site-packages/joblib/parallel.py:763, in BatchCompletionCallBack._return_or_raise(self)
    761 try:
    762     if self.status == TASK_ERROR:
-> 763         raise self._result
    764     return self._result
    765 finally:

AttributeError: 'Process' object has no attribute 'env'

使用的代码如下：

def run_one(train_, region):
    print('region:', region)
    # 忽略所有警告
    warnings.filterwarnings("ignore")
    set_log_level("ERROR")

    m = NeuralProphet(
        trend_global_local="local",
        season_global_local="local",
        yearly_seasonality="auto",
        weekly_seasonality="auto",
        daily_seasonality=False,
        epochs=40,
        n_changepoints=12
    )
    m.set_plotting_backend("plotly-static")
    m = m.add_country_holidays("MG")
    m.fit(train_, freq="D", progress=None, num_workers=20)
    
    save(m, region+'.np')
    
    future = m.make_future_dataframe(train_, periods=h)
    forecast = m.predict(future, decompose=False)
    return forecast[['ds', 'yhat1', 'ID']]

jobs = list()
for reg in train['region'].unique():
    train_ = train.loc[train['region'] == reg].drop(columns=['region','site'])
    jobs.append(delayed(run_one)(train_, reg))
p = Parallel(n_jobs=96)(jobs)

问题根源 #

错误AttributeError: 'Process' object has no attribute 'env'源于joblib的loky后端与PyTorch Lightning的TPU检测逻辑冲突：PyTorch Lightning检测TPU可用性时会创建子进程，但loky的进程对象没有env属性，导致调用失败。此外，TPU架构不适合这种多Python进程的并行方式，更适合单进程内的分布式训练。

解决方案 #

方案1：切换joblib到线程后端 #

避免创建新Python进程，绕开进程属性问题：

p = Parallel(n_jobs=96, backend="threading")(jobs)

注：线程后端受GIL限制，但NeuralProphet的PyTorch计算部分会释放GIL，实际效率影响有限。

方案2：强制使用CPU，跳过TPU检测 #

在run_one函数开头添加环境变量设置，或初始化NeuralProphet时指定CPU加速器：

# 方法1：设置环境变量
import os
os.environ["PL_TPU_ENABLED"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = ""

# 方法2：初始化NeuralProphet时指定
m = NeuralProphet(
    # 其他参数...
    trainer_kwargs={"accelerator": "cpu"}
)

方案3：利用NeuralProphet原生多时间序列支持 #

放弃joblib多进程，直接用NeuralProphet的多时间序列模式处理所有region数据，更贴合TPU的并行特性：

m = NeuralProphet(
    trend_global_local="local",
    season_global_local="local",
    yearly_seasonality="auto",
    weekly_seasonality="auto",
    daily_seasonality=False,
    epochs=40,
    n_changepoints=12
)
m.set_plotting_backend("plotly-static")
m = m.add_country_holidays("MG")
# 保留region作为分组列，直接训练
m.fit(train, freq="D", progress=None, num_workers=20, group_column="region")

# 批量生成所有region的预测
future = m.make_future_dataframe(train, periods=h)
forecast = m.predict(future, decompose=False)

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