c45ac42acd
Release / release (push) Waiting to run
* VLLM upgrade to 0.12.0 and compatibility fixes * MAX_NUM_BATCHED_TOKENS fix and CUDA tester * Sys kill worker instead of marking as failed * upgrade to vllm 0.12.0 * Update to vllm 0.15.0 and lora fix * Update for HUB and removal of deprected env variables * reverted docker-bake changes * removed leftovers * Update src/handler.py Co-authored-by: Dj Isaac <contact@dejaydev.com> * Update src/utils.py Co-authored-by: Dj Isaac <contact@dejaydev.com> * Update src/handler.py Co-authored-by: Dj Isaac <contact@dejaydev.com> * Clean up of docs and comments in code * nit: lowercase p * nit: lowercase p --------- Co-authored-by: Dj Isaac <contact@dejaydev.com> Co-authored-by: chrisvela <chris.vela@runpod.io>
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
import sys
|
|
import multiprocessing
|
|
import traceback
|
|
import runpod
|
|
from runpod import RunPodLogger
|
|
|
|
log = RunPodLogger()
|
|
|
|
vllm_engine = None
|
|
openai_engine = None
|
|
|
|
|
|
async def handler(job):
|
|
try:
|
|
from utils import JobInput
|
|
job_input = JobInput(job["input"])
|
|
engine = openai_engine if job_input.openai_route else vllm_engine
|
|
results_generator = engine.generate(job_input)
|
|
async for batch in results_generator:
|
|
yield batch
|
|
except Exception as e:
|
|
error_str = str(e)
|
|
full_traceback = traceback.format_exc()
|
|
|
|
log.error(f"Error during inference: {error_str}")
|
|
log.error(f"Full traceback:\n{full_traceback}")
|
|
|
|
# CUDA errors = worker is broken, exit to let RunPod spin up a healthy one
|
|
if "CUDA" in error_str or "cuda" in error_str:
|
|
log.error("Terminating worker due to CUDA/GPU error")
|
|
sys.exit(1)
|
|
|
|
yield {"error": error_str}
|
|
|
|
|
|
# Only run in main process to prevent re-initialization when vLLM spawns worker subprocesses
|
|
if __name__ == "__main__" or multiprocessing.current_process().name == "MainProcess":
|
|
|
|
try:
|
|
from engine import vLLMEngine, OpenAIvLLMEngine
|
|
|
|
vllm_engine = vLLMEngine()
|
|
openai_engine = OpenAIvLLMEngine(vllm_engine)
|
|
log.info("vLLM engines initialized successfully")
|
|
except Exception as e:
|
|
log.error(f"Worker startup failed: {e}\n{traceback.format_exc()}")
|
|
sys.exit(1)
|
|
|
|
runpod.serverless.start(
|
|
{
|
|
"handler": handler,
|
|
"concurrency_modifier": lambda x: vllm_engine.max_concurrency if vllm_engine else 1,
|
|
"return_aggregate_stream": True,
|
|
}
|
|
)
|