32 lines
938 B
Python
32 lines
938 B
Python
import runpod
|
|
from utils import JobInput
|
|
from engine import OllamaEngine, OllamaOpenAiEngine
|
|
|
|
|
|
async def handler(job: any):
|
|
# Just dump the whole input to the console and then return an {"ok": True} response
|
|
print('Job:', job)
|
|
|
|
job_input = JobInput(job["input"])
|
|
engine_class = OllamaOpenAiEngine if job_input.openai_route else OllamaEngine
|
|
engine = engine_class() # Instantiate the engine
|
|
|
|
job = engine.generate(job_input) # Call generate with job_input
|
|
|
|
async for batch in job:
|
|
yield batch
|
|
|
|
# Original code from vllm runpod_wrapper.py
|
|
#async def handler(job):
|
|
# job_input = JobInput(job["input"])
|
|
# engine = OpenAIvLLMEngine if job_input.openai_route else vllm_engine
|
|
# results_generator = engine.generate(job_input)
|
|
# async for batch in results_generator:
|
|
# yield batch
|
|
|
|
runpod.serverless.start(
|
|
{
|
|
"handler": handler,
|
|
"return_aggregate_stream": True,
|
|
}
|
|
) |