[hardware][rocm] allow rocm to override default env var (#7926)

This commit is contained in:
youkaichao
2024-08-27 19:50:06 -07:00
committed by GitHub
parent fab5f53e2d
commit bc6e42a9b1
2 changed files with 14 additions and 2 deletions

View File

@@ -1,10 +1,21 @@
import os
from functools import lru_cache
from typing import Tuple
import torch
from vllm.logger import init_logger
from .interface import Platform, PlatformEnum
logger = init_logger(__name__)
if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD", None) in ["fork", None]:
logger.warning("`fork` method is not supported by ROCm. "
"VLLM_WORKER_MULTIPROC_METHOD is overridden to"
" `spawn` instead.")
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
class RocmPlatform(Platform):
_enum = PlatformEnum.ROCM