Files
worker-vllm/src/engine_args.py
T

587 lines
21 KiB
Python

import ast
import os
import json
import logging
from typing import get_origin, get_args
from torch.cuda import device_count
from vllm import AsyncEngineArgs
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
from src.utils import convert_limit_mm_per_prompt
# Backward-compat: env var names users already know → engine arg name
ENV_ALIASES = {
"MODEL_NAME": "model",
"MODEL_REVISION": "revision",
"TOKENIZER_NAME": "tokenizer",
}
# Literal defaults from original worker (used when env/local do not set a value)
DEFAULT_ARGS = {
"disable_log_stats": False,
"enable_log_requests": False,
"gpu_memory_utilization": 0.95,
"pipeline_parallel_size": 1,
"tensor_parallel_size": 1,
"skip_tokenizer_init": False,
"tokenizer_mode": "auto",
"trust_remote_code": False,
"load_format": "auto",
"dtype": "auto",
"kv_cache_dtype": "auto",
"seed": 0,
"worker_use_ray": False,
"block_size": 16,
"enable_prefix_caching": False,
"disable_sliding_window": False,
"swap_space": 4,
"cpu_offload_gb": 0,
"max_num_seqs": 256,
"max_logprobs": 20,
"enforce_eager": False,
"max_seq_len_to_capture": 8192,
"disable_custom_all_reduce": False,
"tokenizer_pool_size": 0,
"tokenizer_pool_type": "ray",
"enable_lora": False,
"max_loras": 1,
"max_lora_rank": 16,
"enable_prompt_adapter": False,
"max_prompt_adapters": 1,
"max_prompt_adapter_token": 0,
"fully_sharded_loras": False,
"lora_extra_vocab_size": 256,
"lora_dtype": "auto",
"device": "auto",
"ray_workers_use_nsight": False,
"num_lookahead_slots": 0,
"scheduler_delay_factor": 0.0,
"guided_decoding_backend": "outlines",
"spec_decoding_acceptance_method": "rejection_sampler",
"stream_interval": 1,
}
def _resolve_field_type(field_type: type) -> type:
"""Resolve Optional/Union to the concrete type for conversion."""
origin = get_origin(field_type)
args = get_args(field_type) if hasattr(field_type, "__args__") else ()
if origin is not None:
# Optional[X] is Union[X, None]; X | None is UnionType
non_none = [a for a in args if a is not type(None)]
if non_none:
return non_none[0]
return field_type
def _convert_env_value_to_field_type(value: str, field_name: str, field_type: type):
"""Convert env var string to the type expected by AsyncEngineArgs for this field."""
val = value.strip() if isinstance(value, str) else value
if val in ("", "None", "none"):
args = get_args(field_type) if hasattr(field_type, "__args__") else ()
if type(None) in (args or ()):
return None
raise ValueError("empty value not allowed for non-optional field")
# Union[bool, str, ...]: only coerce to bool for unambiguous literals;
# otherwise preserve the string (e.g. hf_token="hf_abc..." must stay a str).
if get_origin(field_type) is not None:
union_types = [a for a in (get_args(field_type) or ()) if a is not type(None)]
if bool in union_types and str in union_types:
if str(val).lower() in ("true", "false", "1", "0", "yes", "no", "on", "off"):
return str(val).lower() in ("true", "1", "yes", "on")
return str(val)
effective_type = _resolve_field_type(field_type)
# bool
if effective_type is bool:
return str(val).lower() in ("true", "1", "yes", "on")
# int
if effective_type is int:
return int(val)
# float
if effective_type is float:
return float(val)
# str
if effective_type is str:
return str(val)
# dict, list, or complex (try JSON)
origin = get_origin(effective_type)
if effective_type in (dict, list) or origin in (dict, list):
try:
return json.loads(val)
except json.JSONDecodeError:
return val
# tuple (e.g. long_lora_scaling_factors) — comma-separated or JSON array
if effective_type is tuple or origin is tuple:
args = get_args(field_type) if hasattr(field_type, "__args__") else ()
elem_types = [a for a in args if a is not Ellipsis]
elem_type = elem_types[0] if elem_types else str
try:
parsed = json.loads(val)
if isinstance(parsed, list):
return tuple(elem_type(x) for x in parsed)
except (json.JSONDecodeError, TypeError):
pass
return tuple(elem_type(x.strip()) for x in str(val).split(",") if x.strip())
# For dataclass/complex types, try JSON then Python literal parsing to dict
try:
return json.loads(val)
except (json.JSONDecodeError, TypeError):
pass
try:
parsed = ast.literal_eval(val)
if isinstance(parsed, (dict, list)):
return parsed
except (ValueError, SyntaxError):
pass
# Fallback: try int, float, then str
try:
return int(val)
except ValueError:
pass
try:
return float(val)
except ValueError:
pass
return str(val)
def _get_args_from_env_auto_discover() -> dict:
"""Auto-discover engine args from env vars using UPPERCASED field names.
For every field in AsyncEngineArgs, check os.getenv(FIELD_NAME).
E.g. MAX_MODEL_LEN=4096 -> max_model_len=4096.
Uses same type conversion as before; supports all vLLM engine args without manual listing.
"""
args = {}
valid_fields = AsyncEngineArgs.__dataclass_fields__
for field_name, field in valid_fields.items():
env_key = field_name.upper()
value = os.environ.get(env_key)
if value is None:
continue
try:
args[field_name] = _convert_env_value_to_field_type(
value, field_name, field.type
)
except (ValueError, TypeError, json.JSONDecodeError) as e:
logging.warning(
"Skip env %s=%r: %s", env_key, value, e
)
return args
def _apply_env_aliases(args: dict) -> None:
"""Apply ENV_ALIASES: if MODEL_NAME etc. are set, set the target engine arg."""
valid_fields = AsyncEngineArgs.__dataclass_fields__
for alias, target in ENV_ALIASES.items():
value = os.environ.get(alias)
if value is None or target not in valid_fields:
continue
try:
args[target] = _convert_env_value_to_field_type(
value, target, valid_fields[target].type
)
except (ValueError, TypeError, json.JSONDecodeError) as e:
logging.warning("Skip env alias %s=%r: %s", alias, value, e)
def get_speculative_config():
"""Build speculative decoding configuration from environment variables.
Supports two modes:
1. Full JSON config via SPECULATIVE_CONFIG env var
2. Individual env vars for common settings
"""
# Option 1: Full JSON configuration
spec_config_json = os.getenv('SPECULATIVE_CONFIG')
if spec_config_json:
try:
config = json.loads(spec_config_json)
logging.info(f"Using speculative config from SPECULATIVE_CONFIG: {config}")
return config
except json.JSONDecodeError as e:
logging.error(f"Failed to parse SPECULATIVE_CONFIG JSON: {e}")
return None
# Option 2: Build config from individual environment variables
spec_method = os.getenv('SPECULATIVE_METHOD')
spec_model = os.getenv('SPECULATIVE_MODEL')
_num_spec_tokens = os.getenv('NUM_SPECULATIVE_TOKENS')
_ngram_max = os.getenv('NGRAM_PROMPT_LOOKUP_MAX')
_ngram_min = os.getenv('NGRAM_PROMPT_LOOKUP_MIN')
# Convert numeric vars to int so '0' (hub.json default) is treated as unset
num_spec_tokens = (int(_num_spec_tokens) or None) if _num_spec_tokens else None
ngram_max = (int(_ngram_max) or None) if _ngram_max else None
ngram_min = (int(_ngram_min) or None) if _ngram_min else None
if not any([spec_method, spec_model, ngram_max]):
return None
config = {}
# Determine method
if spec_method:
config['method'] = spec_method
elif ngram_max and not spec_model:
config['method'] = 'ngram'
elif spec_model:
model_lower = spec_model.lower()
if 'eagle3' in model_lower:
config['method'] = 'eagle3'
elif 'eagle' in model_lower:
config['method'] = 'eagle'
elif 'medusa' in model_lower:
config['method'] = 'medusa'
else:
config['method'] = 'draft_model'
if spec_model:
config['model'] = spec_model
if num_spec_tokens:
config['num_speculative_tokens'] = num_spec_tokens
if ngram_max:
config['prompt_lookup_max'] = ngram_max
if ngram_min:
config['prompt_lookup_min'] = ngram_min
draft_tp = os.getenv('SPECULATIVE_DRAFT_TENSOR_PARALLEL_SIZE')
if draft_tp:
config['draft_tensor_parallel_size'] = int(draft_tp)
spec_max_len = os.getenv('SPECULATIVE_MAX_MODEL_LEN')
if spec_max_len:
config['max_model_len'] = int(spec_max_len)
disable_batch = os.getenv('SPECULATIVE_DISABLE_BY_BATCH_SIZE')
if disable_batch:
config['disable_by_batch_size'] = int(disable_batch)
spec_quant = os.getenv('SPECULATIVE_QUANTIZATION')
if spec_quant:
config['quantization'] = spec_quant
spec_revision = os.getenv('SPECULATIVE_MODEL_REVISION')
if spec_revision:
config['revision'] = spec_revision
spec_eager = os.getenv('SPECULATIVE_ENFORCE_EAGER')
if spec_eager:
config['enforce_eager'] = spec_eager.lower() == 'true'
if config:
logging.info(f"Built speculative config from env vars: {config}")
return config
return None
def _resolve_max_model_len(model, trust_remote_code=False, revision=None):
"""Resolve max_model_len from the model's HuggingFace config."""
try:
from transformers import AutoConfig
config = AutoConfig.from_pretrained(
model,
trust_remote_code=trust_remote_code,
revision=revision,
)
for attr in ('max_position_embeddings', 'n_positions', 'max_seq_len', 'seq_length'):
val = getattr(config, attr, None)
if val is not None:
logging.info(f"Resolved max_model_len={val} from model config ({attr})")
return val
except Exception as e:
logging.warning(f"Could not resolve max_model_len from model config: {e}")
return None
def _local_args_to_engine_args(local: dict) -> dict:
"""Map local args (e.g. from /local_model_args.json) to engine arg names and filter."""
valid = AsyncEngineArgs.__dataclass_fields__
out = {}
for k, v in local.items():
target = ENV_ALIASES.get(k, k.lower().replace("-", "_"))
if target not in valid or v in (None, "", "None"):
continue
out[target] = v
return out
def _sanitize_hf_overrides(hf_overrides: dict) -> dict | None:
"""Strip rope_scaling from hf_overrides sub-configs if vLLM rejects them.
Older vLLM (<0.7) required explicit mrope rope_scaling in hf_overrides for
models like Qwen2-VL. Newer vLLM auto-detects mrope and raises a ValueError
in patch_rope_scaling_dict when it finds conflicting rope_type values. Strip
the offending rope_scaling so the model loads with its native config.
"""
if not isinstance(hf_overrides, dict):
return hf_overrides
try:
from vllm.transformers_utils.config import patch_rope_scaling_dict
except ImportError:
return hf_overrides
import copy
cleaned = {}
changed = False
for key, value in hf_overrides.items():
if isinstance(value, dict) and "rope_scaling" in value:
rope_scaling = value.get("rope_scaling")
if isinstance(rope_scaling, dict):
try:
patch_rope_scaling_dict(copy.deepcopy(rope_scaling))
except (ValueError, Exception) as e:
logging.warning(
"Stripping hf_overrides['%s']['rope_scaling'] because vLLM "
"rejected it (%s). Newer vLLM auto-detects rope scaling from "
"the model config.", key, e
)
stripped = {k: v for k, v in value.items() if k != "rope_scaling"}
cleaned[key] = stripped if stripped else None
changed = True
continue
cleaned[key] = value
if not changed:
return hf_overrides
result = {k: v for k, v in cleaned.items() if v is not None}
return result or None
def _resolve_cached_model_path(model_name: str) -> str:
"""Return a local snapshot path when the HF cache was stored with lowercase names.
Some model stores (e.g. RunPod pre-cached volumes) normalize repo IDs to
lowercase. HuggingFace Hub stores caches as
``models--{org}--{model}/snapshots/{hash}/`` preserving the original casing,
so MODEL_NAME=Qwen/Qwen2.5-Coder-32B-Instruct-AWQ will miss a cache stored
as ``models--qwen--qwen2.5-coder-32b-instruct-awq/``.
If the exact-case cache directory is absent but a lowercase variant exists,
the latest snapshot path is returned so vLLM loads from disk rather than
attempting a redundant download.
"""
if os.path.isabs(model_name):
return model_name
cache_dir = (
os.getenv("HUGGINGFACE_HUB_CACHE")
or os.getenv("HF_HOME")
or os.path.expanduser("~/.cache/huggingface/hub")
)
folder_name = f"models--{model_name.replace('/', '--')}"
if os.path.isdir(os.path.join(cache_dir, folder_name)):
return model_name
lower_dir = os.path.join(cache_dir, folder_name.lower())
if not os.path.isdir(lower_dir):
return model_name
snapshots_dir = os.path.join(lower_dir, "snapshots")
if not os.path.isdir(snapshots_dir):
return model_name
try:
snapshots = sorted(os.listdir(snapshots_dir))
except OSError:
return model_name
if not snapshots:
return model_name
resolved = os.path.join(snapshots_dir, snapshots[-1])
logging.info(
"MODEL_NAME %r not found at original casing in HF cache; "
"resolved to lowercase cached snapshot at %r",
model_name, resolved,
)
return resolved
def get_local_args():
"""
Retrieve local arguments from a JSON file.
Returns:
dict: Local arguments.
"""
if not os.path.exists("/local_model_args.json"):
return {}
with open("/local_model_args.json", "r") as f:
local_args = json.load(f)
if local_args.get("MODEL_NAME") is None:
logging.warning("Model name not found in /local_model_args.json. There maybe was a problem when baking the model in.")
logging.info(f"Using baked in model with args: {local_args}")
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_HUB_OFFLINE"] = "1"
return local_args
def get_engine_args():
# Start with worker custom defaults (only where we differ from vLLM)
args = dict(DEFAULT_ARGS)
# Auto-discover: every AsyncEngineArgs field from env UPPERCASED (e.g. MAX_MODEL_LEN)
args.update(_get_args_from_env_auto_discover())
# Backward-compat aliases (MODEL_NAME → model, etc.)
_apply_env_aliases(args)
# Local baked-in model overrides
local = get_local_args()
if local:
args.update(_local_args_to_engine_args(local))
# Filter to valid engine args and drop sentinel empty values
valid_fields = AsyncEngineArgs.__dataclass_fields__
args = {
k: v for k, v in args.items()
if k in valid_fields and v not in (None, "", "None")
}
# Special conversion for limit_mm_per_prompt (e.g. "image=1,video=0")
limit_mm_env = os.getenv("LIMIT_MM_PER_PROMPT")
if limit_mm_env is not None:
args["limit_mm_per_prompt"] = convert_limit_mm_per_prompt(limit_mm_env)
# if args.get("TENSORIZER_URI"): TODO: add back once tensorizer is ready
# args["load_format"] = "tensorizer"
# args["model_loader_extra_config"] = TensorizerConfig(tensorizer_uri=args["TENSORIZER_URI"], num_readers=None)
# logging.info(f"Using tensorized model from {args['TENSORIZER_URI']}")
if "hf_overrides" in args:
sanitized = _sanitize_hf_overrides(args["hf_overrides"])
if sanitized:
args["hf_overrides"] = sanitized
else:
del args["hf_overrides"]
if args.get("load_format") == "bitsandbytes":
args["quantization"] = args["load_format"]
# Set tensor parallel size and max parallel loading workers if more than 1 GPU is available
num_gpus = device_count()
if num_gpus > 1:
args["tensor_parallel_size"] = num_gpus
args["max_parallel_loading_workers"] = None
if os.getenv("MAX_PARALLEL_LOADING_WORKERS"):
logging.warning("Overriding MAX_PARALLEL_LOADING_WORKERS with None because more than 1 GPU is available.")
# LMCache requires HMA to be disabled
try:
_kv_transfer = args.get("kv_transfer_config")
if isinstance(_kv_transfer, str):
parsed = None
try:
parsed = json.loads(_kv_transfer)
except (json.JSONDecodeError, TypeError):
pass
if parsed is None:
try:
result = ast.literal_eval(_kv_transfer)
if isinstance(result, dict):
parsed = result
except (ValueError, SyntaxError):
pass
if parsed is not None:
_kv_transfer = parsed
args["kv_transfer_config"] = _kv_transfer
_kv_offload = args.get("kv_offloading_backend")
lmcache_via_offload = _kv_offload == "lmcache"
lmcache_via_transfer = (
isinstance(_kv_transfer, dict)
and isinstance(_kv_transfer.get("kv_connector"), str)
and "lmcache" in _kv_transfer.get("kv_connector", "").lower()
)
lmcache_detected = lmcache_via_offload or lmcache_via_transfer
if lmcache_detected:
current = args.get("disable_hybrid_kv_cache_manager")
if current is False:
logging.warning(
"disable_hybrid_kv_cache_manager=False conflicts with LMCache; "
"overriding to True (HMA must be disabled when using LMCache)"
)
args["disable_hybrid_kv_cache_manager"] = True
elif current is None:
args["disable_hybrid_kv_cache_manager"] = True
logging.info("LMCache detected: automatically setting disable_hybrid_kv_cache_manager=True")
except Exception as e:
logging.error(
"Failed to check LMCache configuration: %s",
e,
exc_info=True
)
# Deprecated env args backwards compatibility
if args.get("kv_cache_dtype") == "fp8_e5m2":
args["kv_cache_dtype"] = "fp8"
logging.warning("Using fp8_e5m2 is deprecated. Please use fp8 instead.")
if os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE"):
args["max_seq_len_to_capture"] = int(os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE"))
logging.warning("Using MAX_CONTEXT_LEN_TO_CAPTURE is deprecated. Please use MAX_SEQ_LEN_TO_CAPTURE instead.")
# if "gemma-2" in args.get("model", "").lower():
# os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
# logging.info("Using FLASHINFER for gemma-2 model.")
# Set max_num_batched_tokens to max_model_len for unlimited batching.
# vLLM defaults max_num_batched_tokens to 2048 when None, which is too low.
if args.get("max_model_len") == 0:
args["max_model_len"] = None
if args.get("max_num_batched_tokens") == 0:
args["max_num_batched_tokens"] = None
if args.get("max_num_batched_tokens") is None:
max_model_len = args.get("max_model_len")
if max_model_len is None:
max_model_len = _resolve_max_model_len(
args.get("model"),
trust_remote_code=args.get("trust_remote_code", False),
revision=args.get("revision"),
)
if max_model_len is not None:
args["max_num_batched_tokens"] = max_model_len
logging.info(f"Setting max_num_batched_tokens to {max_model_len}")
# VLLM_ATTENTION_BACKEND is deprecated, migrate to attention_backend
if os.getenv('VLLM_ATTENTION_BACKEND'):
logging.warning(
"VLLM_ATTENTION_BACKEND env var is deprecated. "
"Use ATTENTION_BACKEND instead (maps to --attention-backend CLI arg)."
)
if not args.get('attention_backend'):
args['attention_backend'] = os.getenv('VLLM_ATTENTION_BACKEND')
# DISABLE_LOG_REQUESTS is deprecated, use ENABLE_LOG_REQUESTS instead
if os.getenv('DISABLE_LOG_REQUESTS'):
logging.warning(
"DISABLE_LOG_REQUESTS env var is deprecated. "
"Use ENABLE_LOG_REQUESTS instead (default: False)."
)
# Honor old behavior: if DISABLE_LOG_REQUESTS=true, don't enable logging
if os.getenv('DISABLE_LOG_REQUESTS', 'False').lower() == 'true':
args['enable_log_requests'] = False
# Add speculative decoding configuration if present
speculative_config = get_speculative_config()
if speculative_config:
args["speculative_config"] = speculative_config
# Resolve lowercase HF cache paths (FDE-174)
if args.get("model"):
args["model"] = _resolve_cached_model_path(args["model"])
return AsyncEngineArgs(**args)