env_override.py 798 Bytes
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
import os

import torch

7
8
9
10
from vllm.logger import init_logger

logger = init_logger(__name__)

11
12
13
14
15
16
17
# set some common config/environment variables that should be set
# for all processes created by vllm and all processes
# that interact with vllm workers.
# they are executed whenever `import vllm` is called.

# see https://github.com/vllm-project/vllm/pull/15951
# it avoids unintentional cuda initialization from torch.cuda.is_available()
18
os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
19
20

# see https://github.com/vllm-project/vllm/issues/10480
21
os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
22
23
# see https://github.com/vllm-project/vllm/issues/10619
torch._inductor.config.compile_threads = 1