Python多线程中os.environ环境变量异常问题排查与解决

问题描述 #

在Python 3.9中使用ThreadPool多线程池时，出现环境变量跨线程互相干扰的异常行为。预期每个线程独立拥有CUDA_VISIBLE_DEVICES环境变量值，但实际输出显示线程0的变量值被线程1修改。尝试深拷贝os.environ、重新初始化os.environ等方法均未解决问题。

复现代码如下：

import multiprocessing.pool
import os
import copy
import time

multiprocessing.set_start_method('fork')
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
WORKERS = 2


def method(index: int):
    print(f'* index: {index}; position: 1; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    os.environ["CUDA_VISIBLE_DEVICES"] = str(index)
    print(f'* index: {index}; position: 2; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    time.sleep(3)  # some task
    print(f'* index: {index}; position: 3; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')


pool = multiprocessing.pool.ThreadPool(processes=WORKERS)
pool.map(lambda index: method(index), list(range(WORKERS)))

原因分析 #

核心原因是线程共享进程的内存空间：

• os.environ本质是进程级的环境变量字典，同一进程内的所有线程都直接操作这个全局共享对象。
• 线程没有独立的环境变量空间，任何线程对os.environ的修改都会立刻影响其他线程；深拷贝os.environ仅生成字典副本，无法改变线程共享全局环境变量的本质。

可行解决方案 #

方案1：使用进程池替代线程池 #

每个进程拥有独立的内存空间和环境变量，修改不会互相干扰。示例代码：

import multiprocessing
import os
import time

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
WORKERS = 2


def method(index: int):
    print(f'* index: {index}; position: 1; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    os.environ["CUDA_VISIBLE_DEVICES"] = str(index)
    print(f'* index: {index}; position: 2; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    time.sleep(3)  # some task
    print(f'* index: {index}; position: 3; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')


if __name__ == "__main__":
    multiprocessing.set_start_method('fork')
    pool = multiprocessing.Pool(processes=WORKERS)
    pool.map(method, list(range(WORKERS)))
    pool.close()
    pool.join()

方案2：线程本地存储+临时修改恢复 #

如果必须使用线程池，可借助threading.local()存储线程专属的环境变量值，执行任务时临时修改os.environ，任务结束后恢复原值，避免跨线程干扰：

import multiprocessing.pool
import os
import time
import threading

# 线程本地存储，保存每个线程的初始环境变量
thread_local = threading.local()
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
WORKERS = 2


def method(index: int):
    # 保存当前线程的初始值
    if not hasattr(thread_local, 'original_cuda_env'):
        thread_local.original_cuda_env = os.environ["CUDA_VISIBLE_DEVICES"]
    
    print(f'* index: {index}; position: 1; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    # 临时修改环境变量
    os.environ["CUDA_VISIBLE_DEVICES"] = str(index)
    print(f'* index: {index}; position: 2; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    
    try:
        time.sleep(3)  # some task
        print(f'* index: {index}; position: 3; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    finally:
        # 恢复初始值
        os.environ["CUDA_VISIBLE_DEVICES"] = thread_local.original_cuda_env


pool = multiprocessing.pool.ThreadPool(processes=WORKERS)
pool.map(lambda index: method(index), list(range(WORKERS)))

方案3：上下文管理器封装环境变量修改 #

用上下文管理器封装环境变量的临时修改与恢复，确保线程安全：

import multiprocessing.pool
import os
import time
from contextlib import contextmanager

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
WORKERS = 2


@contextmanager
def temp_env_var(key, value):
    original_value = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if original_value is not None:
            os.environ[key] = original_value
        else:
            del os.environ[key]


def method(index: int):
    print(f'* index: {index}; position: 1; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    with temp_env_var("CUDA_VISIBLE_DEVICES", str(index)):
        print(f'* index: {index}; position: 2; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
        time.sleep(3)  # some task
        print(f'* index: {index}; position: 3; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')
    # 退出上下文后恢复原值
    print(f'* index: {index}; position: 4; env var value: {os.environ["CUDA_VISIBLE_DEVICES"]} *')


pool = multiprocessing.pool.ThreadPool(processes=WORKERS)
pool.map(lambda index: method(index), list(range(WORKERS)))

内容的提问来源于stack exchange，提问作者Ferran Agulló