Repository: pytorch
Issue: os.environ modification is not thread safe
Labels: triaged,module: multithreading

### 🐛 Describe the bug

Modifications to os.environ affect a process global data structure (the environment struct), and are therefore not thread safe.

Here is a quick and dirty grep for environ modification, excluding results from torch/_testing.py. Some of these are OK to ignore as they are top-level, I haven't carefully audited them yet.

```
$ git grep "os.environ.*\s=\s" torch
torch/__init__.py:                    os.environ["PATH"] = ";".join(dll_paths + [os.environ["PATH"]])
torch/__init__.py:            os.environ["DISABLE_CUPTI_LAZY_REINIT"] = "1"
torch/__init__.py:            os.environ["TEARDOWN_CUPTI"] = "0"
torch/_inductor/autotune_process.py:    os.environ[CUDA_VISIBLE_DEVICES] = str(device)
torch/_inductor/autotune_process.py:            os.environ[CUDA_VISIBLE_DEVICES] = current
torch/_inductor/codecache.py:        os.environ["_TORCHINDUCTOR_PYOBJECT_TENSOR_DATA_PTR"] = str(
torch/_inductor/cpp_builder.py:            os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
torch/_inductor/cpp_builder.py:                os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
torch/_inductor/cpp_builder.py:        os.environ["CUDA_HOME"] = build_paths.sdk_home
torch/_inductor/cpp_builder.py:        os.environ["CUDA_HOME"] = build_paths.sdk_home
torch/_inductor/fx_passes/numeric_utils.py:os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
torch/_inductor/runtime/cache_dir_utils.py:        os.environ["TORCH_COMPILE_CACHE_DIR"] = cache_dir
torch/_inductor/runtime/compile_tasks.py:        os.environ["TRITON_PTXAS_PATH"] = ptxas_path
torch/_inductor/runtime/triton_heuristics.py:            os.environ["TRITON_CACHE_DIR"] = os.path.join(
torch/backends/xeon/run_cpu.py:                    os.environ["LD_PRELOAD"] = os.pathsep.join(
torch/backends/xeon/run_cpu.py:            os.environ[env_name] = env_value
torch/backends/xeon/run_cpu.py:            os.environ["LD_PRELOAD"] = ":".join(lst_valid)
torch/backends/xeon/run_cpu.py:            os.environ["LD_PRELOAD"] = ""
torch/cuda/__init__.py:            os.environ["CUDA_MODULE_LOADING"] = "LAZY"
torch/cuda/__init__.py:        >>> os.environ['CUDA_PROFILE'] = '1'
torch/distributed/elastic/multiprocessing/api.py:        os.environ[k] = v
torch/distributed/run.py:        os.environ["OMP_NUM_THREADS"] = str(omp_num_threads)
torch/profiler/profiler.py:                    os.environ["DISABLE_CUPTI_LAZY_REINIT"] = "1"
torch/profiler/profiler.py:                    os.environ["TEARDOWN_CUPTI"] = "0"
```

Also, using mock.patch to modify the environ in a scoped way:

```
$ git grep "mock.patch.*os.environ" torch
torch/_inductor/utils.py:            with mock.patch.dict(os.environ, {"TRITON_CACHE_DIR": triton_cache_dir}):
torch/testing/_internal/logging_utils.py:    settings_patch = unittest.mock.patch.dict(os.environ, {"TORCH_LOGS": settings})
torch/testing/_internal/logging_utils.py:            unittest.mock.patch.dict(os.environ, {"___LOG_TESTING": ""})
```

also these two harder to find sites:

```
torch/_inductor/aoti_eager.py:        os.environ,
torch/_inductor/utils.py:            os.environ, {"TORCHINDUCTOR_CACHE_DIR": inductor_cache_dir}
```

Environment variables may be modified for several reasons:

* To modify the environment variable that will be passed on subprocess invocation. In this case, the environment setting can be eliminated in favor of explicitly plumbing the environment variable changes to the subprocess invocations (possibly using TLS, if manual plumbing is not feasible
* To affect the behavior of preexisting code (potentially in a library PyTorch links against) which only consults the environment variable. If the envvar read happens in third party code (e.g., as in KMP) is no easy way to solve this, we can only patch environment on a global basis.

## Case study: `TORCHINDUCTOR_CACHE_DIR` patching

torch/_inductor/aoti_eager.py patches environment to temporarily modify the cache directory location for AOTInductor lowering. This is because the inductor cache dir is currently only via an environment variable.  If we introduce TLS that can be patched to modify this value temporarily, we must also take care to ensure that AOTI isn't relying on this value to be propagated to subprocess calls to the compiler. From this example, it is also clear that the local patching of the variable must *override* the environment variable.

### Versions

main

cc @albanD