commit 1462e142c961adf78d862189994b1a1ef5c6a437 Author: SvenBrnn Date: Sat Jan 18 18:01:22 2025 +0100 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..090a1f0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.idea +.DS_Store diff --git a/README.md b/README.md new file mode 100644 index 0000000..fbf4789 --- /dev/null +++ b/README.md @@ -0,0 +1,30 @@ +# Runpod serverless runner for ollama + +## How to use + +Start a runpod serverless with the docker container ``svenbrnn/runpod-ollama:0.5.7``. Set ``MODEL_NAME`` environment to a model from ollama.com to automatically download a model. +A mounted volume will be automatically used. + +## Environment variables + +| Variable Name | Description | Default Value | +|---------------|------------------------------------------|---------------------| +| `MODEL_NAME` | The name of the model to download | NULL | + +## Test requests for runpod.io console + +See the [test_inputs](./test_inputs) directory for example test requests. + + +## Streaming + +Streaming for openai requests are fully working. + +## Licence + +This project is licensed under the Creative Commons Attribution 4.0 International License. You are free to use, share, and adapt the material for any purpose, even commercially, under the following terms: + +- **Attribution**: You must give appropriate credit, provide a link to the license, and indicate if changes were made. You may do so in any reasonable manner, but not in any way that suggests the licensor endorses you or your use. +- **Reference**: You must reference the original repository at [https://github.com/svenbrnn/runpod-ollama](https://github.com/svenbrnn/runpod-ollama). + +For more details, see the [license](https://creativecommons.org/licenses/by/4.0/). \ No newline at end of file diff --git a/test_inputs/default_runpodio.json b/test_inputs/default_runpodio.json new file mode 100644 index 0000000..c6294a9 --- /dev/null +++ b/test_inputs/default_runpodio.json @@ -0,0 +1,5 @@ +{ + "input": { + "prompt": "How are you?" + } +} \ No newline at end of file diff --git a/test_inputs/openai_completion.json b/test_inputs/openai_completion.json new file mode 100644 index 0000000..eeda987 --- /dev/null +++ b/test_inputs/openai_completion.json @@ -0,0 +1,10 @@ +{ + "input": { + "openai_route": "/v1/completions", + "openai_input": { + "model": "llama3.2:1b", + "prompt": "How are you?", + "max_tokens": 60 + } + } +} \ No newline at end of file diff --git a/test_inputs/openai_get_models.json b/test_inputs/openai_get_models.json new file mode 100644 index 0000000..4cdb9cd --- /dev/null +++ b/test_inputs/openai_get_models.json @@ -0,0 +1,6 @@ +{ + "input": { + "openai_route": "/v1/models", + "openai_input": {} + } +} \ No newline at end of file diff --git a/wrapper/Dockerfile b/wrapper/Dockerfile new file mode 100644 index 0000000..562efdd --- /dev/null +++ b/wrapper/Dockerfile @@ -0,0 +1,46 @@ +ARG OLLAMA_VERSION=0.5.7 + +# Use an official base${OLLAMA_VERSION} image with your desired version +FROM ollama/ollama:${OLLAMA_VERSION} + +ENV PYTHONUNBUFFERED=1 + +# Set up the working directory +WORKDIR / + +RUN apt-get update --yes --quiet && DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \ + software-properties-common \ + gpg-agent \ + build-essential apt-utils \ + && apt-get install --reinstall ca-certificates \ + && add-apt-repository --yes ppa:deadsnakes/ppa && apt update --yes --quiet \ + && DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \ + python3.11 \ + python3.11-dev \ + python3.11-distutils \ + python3.11-lib2to3 \ + python3.11-gdbm \ + python3.11-tk \ + curl \ + pip && \ + update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 999 && \ + update-alternatives --config python3 && ln -s /usr/bin/python3 /usr/bin/python && \ + pip install --upgrade pip && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +# Set the working directory +WORKDIR /work + +# Add my src as /work +ADD src /work + +# Set defaut ollama models directory to /runpod-volume where runpod will mount the volume by default +ENV OLLAMA_MODELS="/runpod-volume" + +# Install runpod and its dependencies +RUN pip install -r requirements.txt && \ + chmod +x start.sh + +# Set the entrypoint +ENTRYPOINT ["/bin/sh", "-c", "/work/start.sh"] \ No newline at end of file diff --git a/wrapper/src/engine.py b/wrapper/src/engine.py new file mode 100644 index 0000000..3cdf5c0 --- /dev/null +++ b/wrapper/src/engine.py @@ -0,0 +1,109 @@ +import json +import os + +from dotenv import load_dotenv +from openai import OpenAI +from utils import JobInput + +client = OpenAI( + base_url='http://localhost:11434/v1/', + + # required but ignored + api_key='ollama', +) + +class OllamaEngine: + def __init__(self): + load_dotenv() + print ("OllamaEngine initialized") + + async def generate(self, job_input): + # Get model from MODEL_NAME defauting to llama3.2:1b + model = os.getenv("MODEL_NAME", "llama3.2:1b") + + # Depending if prompt is a string or a list, we need to handle it differently and send it to the OpenAI API + if isinstance(job_input.llm_input, str): + # Buid new JobInput object with the OpenAI route and input + openAiJob = JobInput({ + "openai_route": "/v1/completions", + "openai_input": { + "model": model, + "prompt": job_input.llm_input, + "stream": job_input.stream + } + }) + else: + # Buid new JobInput object with the OpenAI route and input + openAiJob = JobInput({ + "openai_route": "/v1/chat/completions", + "openai_input": { + "model": model, + "messages": job_input.llm_input, + "stream": job_input.stream + } + }) + + print ("Generating response for job_input:", job_input) + print ("OpenAI job:", openAiJob) + + # Create a generator that will yield the response from the OpenAI API + openAIEngine = OllamaOpenAiEngine() + generate = openAIEngine.generate(openAiJob) + + # Yield the response from the OpenAI API + async for batch in generate: + yield batch + +class OllamaOpenAiEngine(OllamaEngine): + def __init__(self): + load_dotenv() + print ("OllamaOpenAiEngine initialized") + + async def generate(self, job_input): + print("Generating response for job_input:", job_input) + + # Dump job_input to console + openai_input = job_input.openai_input + + # for now e just mock the response + if job_input.openai_route == "/v1/models": + # Async response + async for response in self._handle_model_request(): + yield response + elif job_input.openai_route in ["/v1/chat/completions", "/v1/completions"]: + async for response in self._handle_chat_or_completion_request(openai_input, chat=job_input.openai_route == "/v1/chat/completions"): + yield response + else: + yield {"error": "Invalid route"} + + async def _handle_model_request(self): + try: + response = client.models.list() + # build a json response from the response object + # SyncPage[Model](data=[Model(id='llama3.2:1b', created=1737206544, object='model', owned_by='library')], object='list')\n + yield {"object": "list", "data": [model.to_dict() for model in response.data]} + except Exception as e: + yield {"error": str(e)} + + async def _handle_chat_or_completion_request(self, openai_input, chat=False): + try: + # Call openai.chat.completions.create or openai.completions.create based on the route + if chat: + response = client.chat.completions.create(**openai_input) + else: + response = client.completions.create(**openai_input) + + # If streaming is False, we can just return the response + if not openai_input.get("stream", False): + yield response.to_dict() + return + + for chunk in response: + # Log message to console + print("Message:", chunk) + # Return json of the chunk without any line breaks + yield "data: " + json.dumps(chunk.to_dict(), separators=(',', ':')) + "\n\n" + + yield "data: [DONE]" + except Exception as e: + yield {"error": str(e)} \ No newline at end of file diff --git a/wrapper/src/requirements.txt b/wrapper/src/requirements.txt new file mode 100644 index 0000000..bc90836 --- /dev/null +++ b/wrapper/src/requirements.txt @@ -0,0 +1,4 @@ +runpod +python-dotenv +openai +orjson==3.10.14 \ No newline at end of file diff --git a/wrapper/src/runpod_wrapper.py b/wrapper/src/runpod_wrapper.py new file mode 100644 index 0000000..d0bcfd7 --- /dev/null +++ b/wrapper/src/runpod_wrapper.py @@ -0,0 +1,32 @@ +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, + } +) \ No newline at end of file diff --git a/wrapper/src/start.sh b/wrapper/src/start.sh new file mode 100644 index 0000000..1ce9fd7 --- /dev/null +++ b/wrapper/src/start.sh @@ -0,0 +1,40 @@ +#!/bin/bash + +cleanup() { + echo "Cleaning up..." + pkill -P $$ # Kill all child processes of the current script + exit 0 +} + +# Trap exit signals and call the cleanup function +trap cleanup SIGINT SIGTERM + +# Kill any existing ollama processes +pgrep ollama | xargs kill + +# Start the ollama server and log its output +ollama serve 2>&1 | tee ollama.server.log & +OLLAMA_PID=$! # Store the process ID (PID) of the background command + +check_server_is_running() { + echo "Checking if server is running..." + if cat ollama.server.log | grep -q "Listening"; then + return 0 # Success + else + return 1 # Failure + fi +} + +# Wait for the server to start +while ! check_server_is_running; do + sleep 5 +done +# IF $MODEL_NAME is set, make sure to pull the model, else just skip +if [ -z "$MODEL_NAME" ]; then + echo "No model name provided. Skipping model pull..." +else + echo "Pulling model $MODEL_NAME..." + ollama pull $MODEL_NAME +fi + +python -u runpod_wrapper.py $1 \ No newline at end of file diff --git a/wrapper/src/test_input.json b/wrapper/src/test_input.json new file mode 100644 index 0000000..484acbd --- /dev/null +++ b/wrapper/src/test_input.json @@ -0,0 +1,5 @@ +{ + "input": { + "prompt": "How are you?" + } +} \ No newline at end of file diff --git a/wrapper/src/utils.py b/wrapper/src/utils.py new file mode 100644 index 0000000..dcdb385 --- /dev/null +++ b/wrapper/src/utils.py @@ -0,0 +1,6 @@ +class JobInput: + def __init__(self, job): + self.llm_input = job.get("messages", job.get("prompt")) + self.stream = job.get("stream", False) + self.openai_route = job.get("openai_route") + self.openai_input = job.get("openai_input") \ No newline at end of file