owen/worker-vllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update changes

Browse Source
This commit is contained in:
velaraptor-runpod
2026-02-13 03:16:43 -06:00
parent ee961ad28d
commit fefdbe21a9
3 changed files with 26 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+1 -2
View File
@@ -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 \
src/engine.py
+20 -19
View File
@@ -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
src/engine_args.py
+5 -34
View File
@@ -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