Jupyter Notebook中通过Pybind11运行CUDA加速oxDNA代码遗留僵尸进程求助

问题场景 #

• 主线程模拟后残留GPU进程：用oxpy.Context()运行模拟完成后，GPU上遗留僵尸进程，重启Notebook内核无法清除，必须手动执行kill -9终止。对应代码：

  import oxpy
  with oxpy.Context():
      manager = oxpy.OxpyManager('sim0/input')
      manager.run_complete()
  

• 多进程终止时残留进程：使用multiprocessing.Pool运行模拟，中途调用Pool.terminate()会遗留GPU进程；且需两次触发KeyboardInterrupt才能退出内核循环，模拟正常结束则无此问题。对应代码：

  import oxpy
  import multiprocessing as mp
  
  def launch(inp):
      print("running:", inp)
      with oxpy.Context():
          manager = oxpy.OxpyManager(inp)
          manager.run_complete()
  
  pool = mp.Pool(1)
  
  try:
      response = [pool.apply_async(launch, (i,)) for i in ['sim0/input']]
      for i, r in enumerate(response):
          try:
              r.get(100000000)
          except KeyboardInterrupt:
              print("Caught KeyboardInterrupt1, terminating workers")
              pool.terminate()
              break
          except Exception as e:
              print("Error in simulation", i)
              print("Error was:", e)
              continue
          
  except KeyboardInterrupt:
      print("Caught KeyboardInterrupt2, terminating workers")
      pool.terminate()
      pool.close()
      
  pool.close()
  

解决方法 #

1. 显式重置CUDA设备 #

oxpy基于CUDA实现，模拟结束后可能未彻底释放GPU资源。在模拟完成或异常终止后，强制重置CUDA设备：

import oxpy
import torch

with oxpy.Context():
    manager = oxpy.OxpyManager('sim0/input')
    manager.run_complete()
# 清理GPU缓存并重置设备
torch.cuda.empty_cache()
torch.cuda.device_reset()

如果未安装PyTorch，可使用pycuda直接操作：

import oxpy
import pycuda.driver as cuda

with oxpy.Context():
    manager = oxpy.OxpyManager('sim0/input')
    manager.run_complete()

# 重置CUDA上下文
cuda.init()
device = cuda.Device(0)
current_ctx = cuda.Context.get_current()
if current_ctx:
    current_ctx.pop()
cuda.Context.destroy_all()

2. 优雅终止多进程 #

Pool.terminate()会强制杀死子进程，导致oxpy无法正常清理GPU资源。改用close() + join()的优雅终止方式，同时在子进程中添加资源清理逻辑：

import oxpy
import multiprocessing as mp
import signal

def launch(inp):
    # 为子进程注册信号处理，确保收到终止信号时清理资源
    def sigterm_handler(signum, frame):
        # 显式清理GPU缓存
        try:
            import torch
            torch.cuda.empty_cache()
        except:
            pass
        exit(0)
    signal.signal(signal.SIGTERM, sigterm_handler)
    
    try:
        print("running:", inp)
        with oxpy.Context():
            manager = oxpy.OxpyManager(inp)
            manager.run_complete()
    finally:
        # 无论模拟成功或失败，都清理GPU缓存
        try:
            import torch
            torch.cuda.empty_cache()
        except:
            pass

pool = mp.Pool(1)

try:
    response = [pool.apply_async(launch, (i,)) for i in ['sim0/input']]
    for i, r in enumerate(response):
        try:
            r.get(100000000)
        except KeyboardInterrupt:
            print("Caught KeyboardInterrupt, terminating workers gracefully")
            pool.close()
            pool.join()
            break
        except Exception as e:
            print("Error in simulation", i)
            print("Error was:", e)
            continue
except KeyboardInterrupt:
    print("Caught KeyboardInterrupt, terminating workers gracefully")
    pool.close()
    pool.join()

3. 手动管理oxpy上下文 #

避免依赖with语句的自动清理，手动创建并销毁oxpy上下文：

import oxpy

ctx = oxpy.Context()
try:
    manager = oxpy.OxpyManager('sim0/input')
    manager.run_complete()
finally:
    # 若oxpy.Context提供destroy方法，显式调用
    if hasattr(ctx, 'destroy'):
        ctx.destroy()
    del ctx
    # 额外清理GPU资源
    import torch
    torch.cuda.empty_cache()

4. 用独立进程运行模拟 #

在Notebook中使用魔法命令将模拟代码作为独立进程执行，彻底隔离内核进程与GPU资源：

%%bash
python -c "
import oxpy
with oxpy.Context():
    manager = oxpy.OxpyManager('sim0/input')
    manager.run_complete()
"

核心原因 #

GPU僵尸进程源于CUDA上下文未被正确销毁，或设备内存未彻底释放。oxpy的Pybind11封装在Jupyter交互式环境中，受内核进程生命周期管理影响，无法自动回收所有GPU资源。显式重置设备、优雅终止进程，可强制清理残留资源。

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