diff --git a/.runpod/README.md b/.runpod/README.md index f0f4b04..f3ae53b 100644 --- a/.runpod/README.md +++ b/.runpod/README.md @@ -28,6 +28,8 @@ All behaviour is controlled through environment variables: | `OPENAI_SERVED_MODEL_NAME_OVERRIDE` | Override served model name in API | | String | | `MAX_CONCURRENCY` | Maximum concurrent requests | 300 | Integer | +**Pass any vLLM engine arg** not listed above by setting an env var with the **UPPERCASED** field name (e.g. `MAX_MODEL_LEN=4096`, `ENABLE_CHUNKED_PREFILL=true`). The worker auto-discovers all `AsyncEngineArgs` fields from env. See the [vLLM engine args docs](https://docs.vllm.ai/en/latest/configuration/engine_args) for all available options. + For complete configuration options, see the [full configuration documentation](https://github.com/runpod-workers/worker-vllm/blob/main/docs/configuration.md). ## API Usage diff --git a/README.md b/README.md index c64180e..78b6181 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,16 @@ Configure worker-vllm using environment variables: | `OPENAI_SERVED_MODEL_NAME_OVERRIDE` | Override served model name in API | | String | | `MAX_CONCURRENCY` | Maximum concurrent requests | 30 | Integer | +**Pass any vLLM engine arg** not listed above by setting an environment variable with the **UPPERCASED** field name (same names vLLM uses). The worker auto-discovers all `AsyncEngineArgs` fields from env. For example: + +| Environment Variable | vLLM Engine Arg | Example Value | +| ------------------------- | ------------------------ | ------------- | +| `MAX_MODEL_LEN` | `max_model_len` | `4096` | +| `ENFORCE_EAGER` | `enforce_eager` | `true` | +| `ENABLE_CHUNKED_PREFILL` | `enable_chunked_prefill` | `true` | + +Any env var whose name matches a valid `AsyncEngineArgs` field (uppercased) is applied automatically. Backward-compat aliases: `MODEL_NAME`, `TOKENIZER_NAME`, `MAX_CONTEXT_LEN_TO_CAPTURE`. This lets you configure any vLLM option without waiting for explicit worker support. + For the complete list of all available environment variables, examples, and detailed descriptions: **[Configuration](docs/configuration.md)** ## Option 2: Build Docker Image with Model Inside diff --git a/docs/configuration.md b/docs/configuration.md index 323fce3..fb21f9d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -156,6 +156,29 @@ The way this works is that the first request will have a batch size of `DEFAULT_ | `DISABLE_LOGGING_REQUEST` | False | `bool` | Disable logging requests. | | `MAX_LOG_LEN` | None | `int` | Max number of prompt characters or prompt ID numbers being printed in log. | +## UPPERCASED env vars: Pass any engine arg + +Any vLLM `AsyncEngineArgs` field can be set via an environment variable using the **UPPERCASED** field name (the same names vLLM uses). The worker auto-discovers all fields from env — no prefix. + +**Format:** `=` (e.g. `MAX_MODEL_LEN=4096`) + +**Examples:** + +| Environment Variable | vLLM Engine Arg | Value Example | +| ------------------------ | ------------------------ | ------------- | +| `MAX_MODEL_LEN` | `max_model_len` | `4096` | +| `ENFORCE_EAGER` | `enforce_eager` | `true` | +| `ENABLE_CHUNKED_PREFILL` | `enable_chunked_prefill` | `true` | +| `NUM_SCHEDULER_STEPS` | `num_scheduler_steps` | `8` | +| `TOKENIZER_POOL_SIZE` | `tokenizer_pool_size` | `4` | + +**Backward-compat aliases:** `MODEL_NAME` → `model`, `TOKENIZER_NAME` → `tokenizer`, `MAX_CONTEXT_LEN_TO_CAPTURE` → `max_seq_len_to_capture`, `MODEL_REVISION` → `revision`. + +**Notes:** +- Only valid `AsyncEngineArgs` fields are applied. Unknown keys are silently ignored. +- Values are automatically cast to the correct type (`int`, `float`, `bool`, `str`, or JSON for `dict`/`list`/`tuple`). +- For a full list of available engine args, see the [vLLM AsyncEngineArgs documentation](https://docs.vllm.ai/en/latest/configuration/engine_args/). + ## Docker Build Arguments These variables are used when building custom Docker images with models baked in: diff --git a/src/engine_args.py b/src/engine_args.py index ef58b3d..7d656e3 100644 --- a/src/engine_args.py +++ b/src/engine_args.py @@ -1,106 +1,169 @@ 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 -RENAME_ARGS_MAP = { +# Backward-compat: env var names users already know → engine arg name +ENV_ALIASES = { "MODEL_NAME": "model", "MODEL_REVISION": "revision", "TOKENIZER_NAME": "tokenizer", - "MAX_CONTEXT_LEN_TO_CAPTURE": "max_seq_len_to_capture" } +# Literal defaults from original worker (used when env/local do not set a value) DEFAULT_ARGS = { - "disable_log_stats": os.getenv('DISABLE_LOG_STATS', 'False').lower() == 'true', - # disable_log_requests is deprecated, use enable_log_requests instead - "enable_log_requests": os.getenv('ENABLE_LOG_REQUESTS', 'False').lower() == 'true', - "gpu_memory_utilization": float(os.getenv('GPU_MEMORY_UTILIZATION', 0.95)), - "pipeline_parallel_size": int(os.getenv('PIPELINE_PARALLEL_SIZE', 1)), - "tensor_parallel_size": int(os.getenv('TENSOR_PARALLEL_SIZE', 1)), - "served_model_name": os.getenv('SERVED_MODEL_NAME', None), - "tokenizer": os.getenv('TOKENIZER', None), - "skip_tokenizer_init": os.getenv('SKIP_TOKENIZER_INIT', 'False').lower() == 'true', - "tokenizer_mode": os.getenv('TOKENIZER_MODE', 'auto'), - "trust_remote_code": os.getenv('TRUST_REMOTE_CODE', 'False').lower() == 'true', - "download_dir": os.getenv('DOWNLOAD_DIR', None), - "load_format": os.getenv('LOAD_FORMAT', 'auto'), - "config_format": os.getenv('CONFIG_FORMAT', 'auto'), - "dtype": os.getenv('DTYPE', 'auto'), - "kv_cache_dtype": os.getenv('KV_CACHE_DTYPE', 'auto'), - "quantization_param_path": os.getenv('QUANTIZATION_PARAM_PATH', None), - "seed": int(os.getenv('SEED', 0)), - "max_model_len": int(os.getenv('MAX_MODEL_LEN', 0)) or None, - "worker_use_ray": os.getenv('WORKER_USE_RAY', 'False').lower() == 'true', - "distributed_executor_backend": os.getenv('DISTRIBUTED_EXECUTOR_BACKEND', None), - "max_parallel_loading_workers": int(os.getenv('MAX_PARALLEL_LOADING_WORKERS', 0)) or None, - "block_size": int(os.getenv('BLOCK_SIZE', 16)), - "enable_prefix_caching": os.getenv('ENABLE_PREFIX_CACHING', 'False').lower() == 'true', - "disable_sliding_window": os.getenv('DISABLE_SLIDING_WINDOW', 'False').lower() == 'true', - # attention_backend replaces deprecated VLLM_ATTENTION_BACKEND env var - "attention_backend": os.getenv('ATTENTION_BACKEND', None), - # Enabled by default for improved throughput. Set to False to disable if experiencing issues - "async_scheduling": None if os.getenv('ASYNC_SCHEDULING') is None else os.getenv('ASYNC_SCHEDULING', 'True').lower() == 'true', - # Controls how often to yield streaming results - "stream_interval": int(os.getenv('STREAM_INTERVAL', 1)), - "swap_space": int(os.getenv('SWAP_SPACE', 4)), # GiB - "cpu_offload_gb": int(os.getenv('CPU_OFFLOAD_GB', 0)), # GiB - # vLLM defaults None to 2048; keep 0 as None to let vLLM auto-calculate - "max_num_batched_tokens": int(os.getenv('MAX_NUM_BATCHED_TOKENS', 0)) or None, - "max_num_seqs": int(os.getenv('MAX_NUM_SEQS', 256)), - "max_logprobs": int(os.getenv('MAX_LOGPROBS', 20)), # Default value for OpenAI Chat Completions API - "revision": os.getenv('REVISION', None), - "code_revision": os.getenv('CODE_REVISION', None), - "rope_scaling": os.getenv('ROPE_SCALING', None), - "rope_theta": float(os.getenv('ROPE_THETA', 0)) or None, - "tokenizer_revision": os.getenv('TOKENIZER_REVISION', None), - "quantization": os.getenv('QUANTIZATION', None), - "enforce_eager": os.getenv('ENFORCE_EAGER', 'False').lower() == 'true', - "max_context_len_to_capture": int(os.getenv('MAX_CONTEXT_LEN_TO_CAPTURE', 0)) or None, - "max_seq_len_to_capture": int(os.getenv('MAX_SEQ_LEN_TO_CAPTURE', 8192)), - "disable_custom_all_reduce": os.getenv('DISABLE_CUSTOM_ALL_REDUCE', 'False').lower() == 'true', - "tokenizer_pool_size": int(os.getenv('TOKENIZER_POOL_SIZE', 0)), - "tokenizer_pool_type": os.getenv('TOKENIZER_POOL_TYPE', 'ray'), - "tokenizer_pool_extra_config": os.getenv('TOKENIZER_POOL_EXTRA_CONFIG', None), - "enable_lora": os.getenv('ENABLE_LORA', 'False').lower() == 'true', - "max_loras": int(os.getenv('MAX_LORAS', 1)), - "max_lora_rank": int(os.getenv('MAX_LORA_RANK', 16)), - "enable_prompt_adapter": os.getenv('ENABLE_PROMPT_ADAPTER', 'False').lower() == 'true', - "max_prompt_adapters": int(os.getenv('MAX_PROMPT_ADAPTERS', 1)), - "max_prompt_adapter_token": int(os.getenv('MAX_PROMPT_ADAPTER_TOKEN', 0)), - "fully_sharded_loras": os.getenv('FULLY_SHARDED_LORAS', 'False').lower() == 'true', - "lora_extra_vocab_size": int(os.getenv('LORA_EXTRA_VOCAB_SIZE', 256)), - "long_lora_scaling_factors": tuple(map(float, os.getenv('LONG_LORA_SCALING_FACTORS', '').split(','))) if os.getenv('LONG_LORA_SCALING_FACTORS') else None, - "lora_dtype": os.getenv('LORA_DTYPE', 'auto'), - "max_cpu_loras": int(os.getenv('MAX_CPU_LORAS', 0)) or None, - "device": os.getenv('DEVICE', 'auto'), - "ray_workers_use_nsight": os.getenv('RAY_WORKERS_USE_NSIGHT', 'False').lower() == 'true', - "num_gpu_blocks_override": int(os.getenv('NUM_GPU_BLOCKS_OVERRIDE', 0)) or None, - "num_lookahead_slots": int(os.getenv('NUM_LOOKAHEAD_SLOTS', 0)), - "model_loader_extra_config": os.getenv('MODEL_LOADER_EXTRA_CONFIG', None), - "ignore_patterns": os.getenv('IGNORE_PATTERNS', None), - "preemption_mode": os.getenv('PREEMPTION_MODE', None), - "scheduler_delay_factor": float(os.getenv('SCHEDULER_DELAY_FACTOR', 0.0)), - "enable_chunked_prefill": os.getenv('ENABLE_CHUNKED_PREFILL', None), - "guided_decoding_backend": os.getenv('GUIDED_DECODING_BACKEND', 'outlines'), - "speculative_model": os.getenv('SPECULATIVE_MODEL', None), - "speculative_draft_tensor_parallel_size": int(os.getenv('SPECULATIVE_DRAFT_TENSOR_PARALLEL_SIZE', 0)) or None, - "enable_expert_parallel": bool(os.getenv('ENABLE_EXPERT_PARALLEL', 'False').lower() == 'true'), - "num_speculative_tokens": int(os.getenv('NUM_SPECULATIVE_TOKENS', 0)) or None, - "speculative_max_model_len": int(os.getenv('SPECULATIVE_MAX_MODEL_LEN', 0)) or None, - "speculative_disable_by_batch_size": int(os.getenv('SPECULATIVE_DISABLE_BY_BATCH_SIZE', 0)) or None, - "ngram_prompt_lookup_max": int(os.getenv('NGRAM_PROMPT_LOOKUP_MAX', 0)) or None, - "ngram_prompt_lookup_min": int(os.getenv('NGRAM_PROMPT_LOOKUP_MIN', 0)) or None, - "spec_decoding_acceptance_method": os.getenv('SPEC_DECODING_ACCEPTANCE_METHOD', 'rejection_sampler'), - "typical_acceptance_sampler_posterior_threshold": float(os.getenv('TYPICAL_ACCEPTANCE_SAMPLER_POSTERIOR_THRESHOLD', 0)) or None, - "typical_acceptance_sampler_posterior_alpha": float(os.getenv('TYPICAL_ACCEPTANCE_SAMPLER_POSTERIOR_ALPHA', 0)) or None, - "qlora_adapter_name_or_path": os.getenv('QLORA_ADAPTER_NAME_OR_PATH', None), - "disable_logprobs_during_spec_decoding": os.getenv('DISABLE_LOGPROBS_DURING_SPEC_DECODING', None), - "otlp_traces_endpoint": os.getenv('OTLP_TRACES_ENDPOINT', None), + "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") + 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()) + # 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. @@ -191,6 +254,7 @@ def get_speculative_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: @@ -209,25 +273,19 @@ def _resolve_max_model_len(model, trust_remote_code=False, revision=None): logging.warning(f"Could not resolve max_model_len from model config: {e}") return None -limit_mm_env = os.getenv('LIMIT_MM_PER_PROMPT') -if limit_mm_env is not None: - DEFAULT_ARGS["limit_mm_per_prompt"] = convert_limit_mm_per_prompt(limit_mm_env) -def match_vllm_args(args): - """Rename args to match vllm by: - 1. Renaming keys to lower case - 2. Renaming keys to match vllm - 3. Filtering args to match vllm's AsyncEngineArgs +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 - Args: - args (dict): Dictionary of args - Returns: - dict: Dictionary of args with renamed keys - """ - renamed_args = {RENAME_ARGS_MAP.get(k, k): v for k, v in args.items()} - matched_args = {k: v for k, v in renamed_args.items() if k in AsyncEngineArgs.__dataclass_fields__} - return {k: v for k, v in matched_args.items() if v not in [None, "", "None"]} def get_local_args(): """ Retrieve local arguments from a JSON file. @@ -250,23 +308,36 @@ def get_local_args(): return local_args def get_engine_args(): - # Start with default args - args = DEFAULT_ARGS - - # Get env args that match keys in AsyncEngineArgs - args.update(os.environ) - - # Get local args if model is baked in and overwrite env args - args.update(get_local_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']}") - - - # Rename and match to vllm args - args = match_vllm_args(args) if args.get("load_format") == "bitsandbytes": args["quantization"] = args["load_format"]