From fefdbe21a9e8f239b5cfe42b6f9c96ab455029e8 Mon Sep 17 00:00:00 2001 From: velaraptor-runpod Date: Fri, 13 Feb 2026 03:16:43 -0600 Subject: [PATCH] update changes --- Dockerfile | 3 +-- src/engine.py | 39 ++++++++++++++++++++------------------- src/engine_args.py | 39 +++++---------------------------------- 3 files changed, 26 insertions(+), 55 deletions(-) diff --git a/Dockerfile b/Dockerfile index c2128b3..9883925 100644 --- a/Dockerfile +++ b/Dockerfile @@ -45,13 +45,12 @@ ENV MODEL_NAME=$MODEL_NAME \ ENV PYTHONPATH="/:/vllm-workspace" -RUN if [ -n "${VLLM_NIGHTLY}" ]; then \ +RUN if [ "${VLLM_NIGHTLY}" = "true" ]; then \ pip install -U vllm --pre --index-url https://pypi.org/simple --extra-index-url https://wheels.vllm.ai/nightly && \ apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/* && \ pip install git+https://github.com/huggingface/transformers.git; \ fi - COPY src /src RUN --mount=type=secret,id=HF_TOKEN,required=false \ if [ -f /run/secrets/HF_TOKEN ]; then \ diff --git a/src/engine.py b/src/engine.py index 67c903b..84b061b 100644 --- a/src/engine.py +++ b/src/engine.py @@ -26,20 +26,20 @@ class vLLMEngine: load_dotenv() # For local development self.engine_args = get_engine_args() logging.info(f"Engine args: {self.engine_args}") - + # Initialize vLLM engine first self.llm = self._initialize_llm() if engine is None else engine.llm - + # Only create custom tokenizer wrapper if not using mistral tokenizer mode # For mistral models, let vLLM handle tokenizer initialization if self.engine_args.tokenizer_mode != 'mistral': - self.tokenizer = TokenizerWrapper(self.engine_args.tokenizer or self.engine_args.model, - self.engine_args.tokenizer_revision, + self.tokenizer = TokenizerWrapper(self.engine_args.tokenizer or self.engine_args.model, + self.engine_args.tokenizer_revision, self.engine_args.trust_remote_code) else: # For mistral models, we'll get the tokenizer from vLLM later self.tokenizer = None - + self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY)) self.default_batch_size = int(os.getenv("DEFAULT_BATCH_SIZE", DEFAULT_BATCH_SIZE)) self.batch_size_growth_factor = int(os.getenv("BATCH_SIZE_GROWTH_FACTOR", DEFAULT_BATCH_SIZE_GROWTH_FACTOR)) @@ -67,7 +67,7 @@ class vLLMEngine: self.has_chat_template = bool(self.tokenizer.chat_template) or bool(self.custom_chat_template) if self.custom_chat_template and isinstance(self.custom_chat_template, str): self.tokenizer.chat_template = self.custom_chat_template - + def apply_chat_template(self, input): if isinstance(input, list): if not self.has_chat_template: @@ -78,11 +78,11 @@ class vLLMEngine: input = [{"role": "user", "content": input}] else: raise ValueError("Input must be a string or a list of messages") - + return self.tokenizer.apply_chat_template( input, tokenize=False, add_generation_prompt=True ) - + return MinimalTokenizerWrapper(tokenizer) except Exception as e: logging.error(f"Failed to create fallback tokenizer: {e}") @@ -90,7 +90,7 @@ class vLLMEngine: def dynamic_batch_size(self, current_batch_size, batch_size_growth_factor): return min(current_batch_size*batch_size_growth_factor, self.default_batch_size) - + async def generate(self, job_input: JobInput): try: async for batch in self._generate_vllm( @@ -118,11 +118,11 @@ class vLLMEngine: batch = { "choices": [{"tokens": []} for _ in range(n_responses)], } - + max_batch_size = batch_size or self.default_batch_size batch_size_growth_factor, min_batch_size = batch_size_growth_factor or self.batch_size_growth_factor, min_batch_size or self.min_batch_size batch_size = BatchSize(max_batch_size, min_batch_size, batch_size_growth_factor) - + async for request_output in results_generator: if is_first_output: # Count input tokens only once @@ -242,14 +242,14 @@ class OpenAIvLLMEngine(vLLMEngine): lora_modules=self.lora_adapters, ) await self.serving_models.init_static_loras() - + # Get chat template from vLLM tokenizer if available chat_template = None if self.tokenizer and hasattr(self.tokenizer, 'tokenizer'): chat_template = self.tokenizer.tokenizer.chat_template - + self.chat_engine = OpenAIServingChat( - engine_client=self.llm, + engine_client=self.llm, models=self.serving_models, response_role=self.response_role, request_logger=None, @@ -290,11 +290,11 @@ class OpenAIvLLMEngine(vLLMEngine): yield response else: yield create_error_response("Invalid route").model_dump() - + async def _handle_model_request(self): models = await self.serving_models.show_available_models() return models.model_dump() - + async def _handle_chat_or_completion_request(self, openai_request: JobInput): if openai_request.openai_route == "/v1/chat/completions": request_class = ChatCompletionRequest @@ -302,7 +302,7 @@ class OpenAIvLLMEngine(vLLMEngine): elif openai_request.openai_route == "/v1/completions": request_class = CompletionRequest generator_function = self.completion_engine.create_completion - + try: request = request_class( **openai_request.openai_input @@ -310,7 +310,7 @@ class OpenAIvLLMEngine(vLLMEngine): except Exception as e: yield create_error_response(str(e)).model_dump() return - + dummy_request = DummyRequest() response_generator = await generator_function(request, raw_request=dummy_request) @@ -320,7 +320,7 @@ class OpenAIvLLMEngine(vLLMEngine): batch = [] batch_token_counter = 0 batch_size = BatchSize(self.default_batch_size, self.min_batch_size, self.batch_size_growth_factor) - + async for chunk_str in response_generator: if "data" in chunk_str: if self.raw_openai_output: @@ -342,3 +342,4 @@ class OpenAIvLLMEngine(vLLMEngine): if self.raw_openai_output: batch = "".join(batch) yield batch + \ No newline at end of file diff --git a/src/engine_args.py b/src/engine_args.py index 2a4da0e..e2a5c5b 100644 --- a/src/engine_args.py +++ b/src/engine_args.py @@ -38,6 +38,9 @@ DEFAULT_ARGS = { "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": os.getenv('ATTENTION_BACKEND', None), + "async_scheduling": os.getenv('ASYNC_SCHEDULING', 'False').lower() == 'true', + "stream_interval": float(os.getenv('STREAM_INTERVAL', 0)), "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 @@ -73,29 +76,14 @@ DEFAULT_ARGS = { "enable_expert_parallel": bool(os.getenv('ENABLE_EXPERT_PARALLEL', 'False').lower() == 'true'), "qlora_adapter_name_or_path": os.getenv('QLORA_ADAPTER_NAME_OR_PATH', None), "otlp_traces_endpoint": os.getenv('OTLP_TRACES_ENDPOINT', None), - "attention_backend": os.getenv('ATTENTION_BACKEND', None), - "async_scheduling": os.getenv('ASYNC_SCHEDULING', 'False').lower() == 'true', - "stream_interval": float(os.getenv('STREAM_INTERVAL', 0)), } - def get_speculative_config(): - """ - Build speculative decoding configuration from environment variables. + """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 - - Speculative Methods: - - "draft_model": Use a smaller draft model for speculation - - "ngram": Use n-gram based prompt lookup (no additional model needed) - - "eagle" / "eagle3": Use EAGLE-based speculation - - "medusa": Use Medusa heads for speculation - - "mlp_speculator": Use MLP-based speculator - - Returns: - dict | None: Speculative config dictionary or None if not configured """ # Option 1: Full JSON configuration spec_config_json = os.getenv('SPECULATIVE_CONFIG') @@ -109,15 +97,12 @@ def get_speculative_config(): return None # Option 2: Build config from individual environment variables - spec_method = os.getenv('SPECULATIVE_METHOD') # ngram, draft_model, eagle, eagle3, medusa, mlp_speculator + spec_method = os.getenv('SPECULATIVE_METHOD') spec_model = os.getenv('SPECULATIVE_MODEL') num_spec_tokens = os.getenv('NUM_SPECULATIVE_TOKENS') - - # N-gram specific settings ngram_max = os.getenv('NGRAM_PROMPT_LOOKUP_MAX') ngram_min = os.getenv('NGRAM_PROMPT_LOOKUP_MIN') - # Check if any speculative decoding is configured if not any([spec_method, spec_model, ngram_max]): return None @@ -129,7 +114,6 @@ def get_speculative_config(): elif ngram_max and not spec_model: config['method'] = 'ngram' elif spec_model: - # Auto-detect method based on model name if not specified model_lower = spec_model.lower() if 'eagle3' in model_lower: config['method'] = 'eagle3' @@ -140,46 +124,35 @@ def get_speculative_config(): else: config['method'] = 'draft_model' - # Model configuration if spec_model: config['model'] = spec_model - - # Number of speculative tokens if num_spec_tokens: config['num_speculative_tokens'] = int(num_spec_tokens) - - # N-gram settings if ngram_max: config['prompt_lookup_max'] = int(ngram_max) if ngram_min: config['prompt_lookup_min'] = int(ngram_min) - # Draft model tensor parallel size draft_tp = os.getenv('SPECULATIVE_DRAFT_TENSOR_PARALLEL_SIZE') if draft_tp: config['draft_tensor_parallel_size'] = int(draft_tp) - # Max model length for draft spec_max_len = os.getenv('SPECULATIVE_MAX_MODEL_LEN') if spec_max_len: config['max_model_len'] = int(spec_max_len) - # Disable by batch size disable_batch = os.getenv('SPECULATIVE_DISABLE_BY_BATCH_SIZE') if disable_batch: config['disable_by_batch_size'] = int(disable_batch) - # Draft model quantization spec_quant = os.getenv('SPECULATIVE_QUANTIZATION') if spec_quant: config['quantization'] = spec_quant - # Draft model revision spec_revision = os.getenv('SPECULATIVE_MODEL_REVISION') if spec_revision: config['revision'] = spec_revision - # Enforce eager mode for draft model spec_eager = os.getenv('SPECULATIVE_ENFORCE_EAGER') if spec_eager: config['enforce_eager'] = spec_eager.lower() == 'true' @@ -209,7 +182,6 @@ def match_vllm_args(args): 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. @@ -231,7 +203,6 @@ def get_local_args(): os.environ["HF_HUB_OFFLINE"] = "1" return local_args - def get_engine_args(): # Start with default args args = DEFAULT_ARGS