From 3a8d2decfd7474aa11518a0cbcf3f6394355d780 Mon Sep 17 00:00:00 2001 From: mags0ft Date: Sat, 15 Nov 2025 14:03:10 +0100 Subject: [PATCH] change project to work with llama.cpp instead of Ollama --- .github/workflows/build-and-push.yml | 57 ---- .github/workflows/check-n8n-new-release.yml | 93 ------ .../workflows/create-new-release-on-push.yml | 87 ------ .gitignore | 218 +++++++++++++- .runpod/hub.json | 78 +++-- .runpod/tests.json | 44 +-- Dockerfile | 14 +- LICENSE | 11 + README.md | 44 ++- embed_model/Dockerfile | 11 - embed_model/README.md | 3 - embed_model/preload_model.sh | 36 --- src/engine.py | 276 +++++++++++++++--- src/handler.py | 41 ++- src/requirements.txt | 2 +- src/start.sh | 54 ++-- src/test_input.json | 6 +- src/utils.py | 23 +- test_inputs/default_runpodio.json | 5 - test_inputs/openai_completion.json | 10 - test_inputs/openai_get_models.json | 6 - 21 files changed, 622 insertions(+), 497 deletions(-) delete mode 100644 .github/workflows/build-and-push.yml delete mode 100644 .github/workflows/check-n8n-new-release.yml delete mode 100644 .github/workflows/create-new-release-on-push.yml create mode 100644 LICENSE delete mode 100644 embed_model/Dockerfile delete mode 100644 embed_model/README.md delete mode 100644 embed_model/preload_model.sh delete mode 100644 test_inputs/default_runpodio.json delete mode 100644 test_inputs/openai_completion.json delete mode 100644 test_inputs/openai_get_models.json diff --git a/.github/workflows/build-and-push.yml b/.github/workflows/build-and-push.yml deleted file mode 100644 index 05fc485..0000000 --- a/.github/workflows/build-and-push.yml +++ /dev/null @@ -1,57 +0,0 @@ -name: 📦 Build custom Docker image and push to Docker Hub - -# Workflow triggered manually (GitHub UI) or programmatically by the check-ollama-new-release workflow -on: - workflow_call: - inputs: - ollama_version: - description: 'ollama version to build. Example: 1.20.0' - required: true - type: string - next_version_suffix: - description: 'Next version suffix to compare with the new release. Example: 0' - required: false - type: string - workflow_dispatch: - inputs: - ollama_version: - description: 'ollama version to build. Example: 1.20.0' - required: true - type: string - next_version_suffix: - description: 'Next version suffix to compare with the new release. Example: 0' - required: false - type: string - -jobs: - build-push-docker-image: - name: 🐳 Build and push Docker image - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - uses: actions/checkout@v4 - - uses: docker/setup-qemu-action@v3 - - uses: docker/setup-buildx-action@v3 - - uses: docker/login-action@v3 - with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} - - uses: docker/build-push-action@v6 - with: - context: ./ - build-args: OLLAMA_VERSION=${{ inputs.ollama_version }} - push: true - tags: svenbrnn/runpod-ollama:${{ inputs.ollama_version }}${{ inputs.next_version_suffix }},svenbrnn/runpod-ollama:latest - - create-release: - name: 🏷️ Create GitHub Release - needs: build-push-docker-image - runs-on: ubuntu-latest - timeout-minutes: 1 - permissions: - contents: write - steps: - - uses: ncipollo/release-action@v1 - with: - commit: master - tag: ollama@${{ inputs.ollama_version }}${{ inputs.next_version_suffix }} \ No newline at end of file diff --git a/.github/workflows/check-n8n-new-release.yml b/.github/workflows/check-n8n-new-release.yml deleted file mode 100644 index 2251841..0000000 --- a/.github/workflows/check-n8n-new-release.yml +++ /dev/null @@ -1,93 +0,0 @@ -name: 👀 Check for new ollama stable releases and builds a new custom image release if needed - -on: - schedule: - - cron: '0 8 * * *' # Every day at 8am - workflow_dispatch: - -jobs: - check-new-release: - runs-on: ubuntu-latest - - timeout-minutes: 5 - - # Output mapping needed to share data between jobs. Separated jobs to reuse the build-and-push.yml workflow. - outputs: - has_new_release: ${{ steps.new_release_check.outputs.has_new_release }} - version_to_release: ${{ steps.ollama_version.outputs.version_to_release }} - - steps: - - name: 🆕 Get the latest ollama release - id: custom_image_release - uses: dvershinin/lastversion-action@main - with: - repository: 'SvenBrnn/runpod-worker-ollama' - - - name: 🆕 Get the latest ollama release - id: ollama_release - uses: dvershinin/lastversion-action@main - with: - repository: 'ollama/ollama' - - - name: 🤲 Check if it is a newer ollama release - id: new_release_check - run: | - echo "👀 Latest Codely custom image release: ${{ steps.custom_image_release.outputs.last_version }}" - echo "🆕 Latest ollama release: ${{ steps.ollama_release.outputs.last_version }}" - - custom_version_orig=$(echo "${{ steps.custom_image_release.outputs.last_version }}" | sed 's/ollama@//g') - custom_version=$(echo "$custom_version_orig" | sed 's/-[0-9]*//g') # Remove the suffix - ollama_version=$(echo "${{ steps.ollama_release.outputs.last_version }}" | sed 's/ollama@//g') - - latest_version=$(echo -e "$custom_version\n$ollama_version" | sort -V | tail -n1) - - if [ "$custom_version" != "$latest_version" ]; then - echo "has_new_release=true" >> $GITHUB_OUTPUT - echo "🤩 New release detected! Releasing new custom image version in 1, 2…" - else - echo "has_new_release=false" >> $GITHUB_OUTPUT - echo "👍 Release has not changed or it is previous to the current one. Doing nothing." - fi - - - name: 🔡 Remove 'ollama@' from release name to get the version - id: ollama_version - if: steps.new_release_check.outputs.has_new_release == 'true' - run: | - version_to_release=$(echo "${{ steps.ollama_release.outputs.last_version }}" | sed 's/ollama@//g') - echo "version_to_release=$version_to_release" >> $GITHUB_OUTPUT - - patch-dockerfile-and-push: - needs: check-new-release - if: ${{ needs.check-new-release.outputs.has_new_release == 'true' }} - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - name: Patch Dockerfile with new OLLAMA_VERSION - run: | - sed -i '1s/^ARG OLLAMA_VERSION=.*/ARG OLLAMA_VERSION=${{ needs.check-new-release.outputs.version_to_release }}/' Dockerfile - - - name: Commit and push changes - run: | - git config user.name "github-actions[bot]" - git config user.email "github-actions[bot]@users.noreply.github.com" - git add Dockerfile - if ! git diff --cached --quiet; then - git commit -m "chore: bump OLLAMA_VERSION to ${{ needs.check-new-release.outputs.version_to_release }}" - git push - else - echo "No changes to commit." - fi - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - - build-and-push-new-docker-image: - needs: - - check-new-release - - patch-dockerfile-and-push - if: ${{ needs.check-new-release.outputs.has_new_release == 'true' }} - uses: ./.github/workflows/build-and-push.yml - secrets: inherit - with: - ollama_version: ${{ needs.check-new-release.outputs.version_to_release }} diff --git a/.github/workflows/create-new-release-on-push.yml b/.github/workflows/create-new-release-on-push.yml deleted file mode 100644 index 3f59e5a..0000000 --- a/.github/workflows/create-new-release-on-push.yml +++ /dev/null @@ -1,87 +0,0 @@ -name: Create new release on push - -on: - push: - branches: - - master - -jobs: - check-new-release: - runs-on: ubuntu-latest - - timeout-minutes: 5 - - # Output mapping needed to share data between jobs. Separated jobs to reuse the build-and-push.yml workflow. - outputs: - has_new_release: ${{ steps.new_release_check.outputs.has_new_release }} - version_to_release: ${{ steps.ollama_version.outputs.version_to_release }} - next_version_suffix: ${{ steps.new_release_check.outputs.next_version_suffix }} - - steps: - - name: 👀 Get the latest ollama-serverless stable release - id: custom_image_release - uses: pozetroninc/github-action-get-latest-release@v0.8.0 - with: - repository: SvenBrnn/runpod-worker-ollama - excludes: prerelease,draft - - - name: 🆕 Get the latest ollama release - id: ollama_release - uses: dvershinin/lastversion-action@main - with: - repository: 'ollama/ollama' - - - name: 🤲 Check if it is a newer ollama release - id: new_release_check - run: | - echo "👀 Latest Codely custom image release: ${{ steps.custom_image_release.outputs.release }}" - echo "🆕 Latest ollama release: ${{ steps.ollama_release.outputs.last_version }}" - - custom_version_orig=$(echo "${{ steps.custom_image_release.outputs.release }}" | sed 's/ollama@//g') - custom_version=$(echo "$custom_version_orig" | sed 's/-[0-9]*//g') # Remove the suffix - ollama_version=$(echo "${{ steps.ollama_release.outputs.last_version }}" | sed 's/ollama@//g') - - latest_version=$(echo -e "$custom_version\n$ollama_version" | sort -V | tail -n1) - echo "has_new_release=true" >> $GITHUB_OUTPUT - - # Debug print - echo "custom_version_orig=$custom_version_orig" - echo "custom_version=$custom_version" - echo "ollama_version=$ollama_version" - - # Calculate the next_version_suffix to compare with the new release - # If the custom_version does not equal the latest_version, then it is a new release and we don't need a suffix - if [ "$custom_version" != "$ollama_version" ]; then - echo "next_version_suffix=" >> $GITHUB_OUTPUT - else - # If the custom_version equals the latest_version, then it is a new release and we need to increment the suffix - # If there is no suffix present we start with 1 - # Check if we have a - first - has_suffix=$(echo "$custom_version_orig" | sed 's/[^-]//g') - if [ -z "$has_suffix" ]; then - suffix="-1" - else - suffix=$(echo "$custom_version_orig" | grep -o '[0-9]*$') - suffix="-$((suffix + 1))" - fi - - echo "next_version_suffix=$suffix" - echo "next_version_suffix=$suffix" >> $GITHUB_OUTPUT - fi - - - name: 🔡 Remove 'ollama@' from release name to get the version - id: ollama_version - if: steps.new_release_check.outputs.has_new_release == 'true' - run: | - version_to_release=$(echo "${{ steps.ollama_release.outputs.last_version }}" | sed 's/ollama@//g') - echo "version_to_release=$version_to_release" >> $GITHUB_OUTPUT - echo "next_version_suffix=${{ steps.new_release_check.outputs.next_version_suffix }}" >> $GITHUB_OUTPUT - - build-and-push-new-docker-image: - needs: check-new-release - if: ${{ needs.check-new-release.outputs.has_new_release == 'true' }} - uses: ./.github/workflows/build-and-push.yml - secrets: inherit - with: - ollama_version: ${{ needs.check-new-release.outputs.version_to_release }} - next_version_suffix: ${{ needs.check-new-release.outputs.next_version_suffix }} diff --git a/.gitignore b/.gitignore index 090a1f0..e15106e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,216 @@ -.idea -.DS_Store +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +# Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +# poetry.lock +# poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +# pdm.lock +# pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +# pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# Redis +*.rdb +*.aof +*.pid + +# RabbitMQ +mnesia/ +rabbitmq/ +rabbitmq-data/ + +# ActiveMQ +activemq-data/ + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +# .idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ + +# Streamlit +.streamlit/secrets.toml diff --git a/.runpod/hub.json b/.runpod/hub.json index 4aca0f3..bf87c42 100644 --- a/.runpod/hub.json +++ b/.runpod/hub.json @@ -1,46 +1,36 @@ { - "title": "Runpod Worker Ollama", - "description": "A serverless Ollama Worker for Runpod", - "type": "serverless", - "category": "language", - "iconUrl": "https://ollama.com/public/ollama.png", - "config": { - "runsOn": "GPU", - "gpuCount": 1, - "gpuIds": "AMPERE_16,AMPERE_24,ADA_24", - "containerDiskInGb": 20, - "presets": [], - "env": [ - { - "key": "OLLAMA_MODEL_NAME", - "input": { - "name": "Model Name", - "type": "string", - "description": "Name of a model to preload", - "default": "phi3", - "advanced": false - } - }, - { - "key": "MAX_CONCURRENCY", - "input": { - "name": "Max Concurrency", - "type": "number", - "description": "Maximum number of concurrent requests to handle (default: 8)", - "default": 8, - "advanced": true - } - }, - { - "key": "OLLAMA_NUM_PARALLEL", - "input": { - "name": "Parallel Requests", - "type": "string", - "description": "Maximum number of concurrent requests to handle (default: 4 or 1)", - "default": "", - "advanced": true - } - } - ] - } + "title": "llama.cpp inference", + "description": "Run llama.cpp inference using serverless RunPod workers!", + "type": "serverless", + "category": "language", + "iconUrl": "https://raw.githubusercontent.com/ggml-org/llama.cpp/master/media/llama1-icon-transparent.png", + "config": { + "runsOn": "GPU", + "gpuCount": 1, + "gpuIds": "AMPERE_16,AMPERE_24,ADA_24", + "containerDiskInGb": 32, + "presets": [], + "env": [ + { + "key": "LLAMA_SERVER_CMD_ARGS", + "input": { + "name": "Model Name", + "type": "string", + "description": "Launch command line arguments (argv) for the llama-server binary. Do not define the port.", + "default": "-hf unsloth/Mistral-Small-3.2-24B-Instruct-2506-GGUF:Q4_K_M -ctx_size 4096", + "advanced": false + } + }, + { + "key": "MAX_CONCURRENCY", + "input": { + "name": "Maximum Concurrency", + "type": "number", + "description": "Maximum number of concurrent requests to handle (default: 8).", + "default": 8, + "advanced": true + } + } + ] + } } \ No newline at end of file diff --git a/.runpod/tests.json b/.runpod/tests.json index 860a40b..5fcc642 100644 --- a/.runpod/tests.json +++ b/.runpod/tests.json @@ -1,24 +1,24 @@ { - "tests": [ - { - "name": "execute_a_prompt", - "input": { - "prompt": "Say: Hallo World!" - }, - "timeout": 120000 - } - ], - "config": { - "gpuTypeId": "NVIDIA GeForce RTX 4090", - "gpuCount": 1, - "env": [ - { - "key": "OLLAMA_MODEL_NAME", - "value": "phi3" - } + "tests": [ + { + "name": "execute_a_prompt", + "input": { + "prompt": "Hi! Who are you?" + }, + "timeout": 120000 + } ], - "allowedCudaVersions": [ - "12.8" - ] - } -} + "config": { + "gpuTypeId": "NVIDIA GeForce RTX 4090", + "gpuCount": 1, + "env": [ + { + "key": "LLAMA_SERVER_CMD_ARGS", + "value": "-hf unsloth/Mistral-Small-3.2-24B-Instruct-2506-GGUF:Q4_K_M -ctx_size 4096" + } + ], + "allowedCudaVersions": [ + "12.8" + ] + } +} \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index b9cda05..afeb55f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,5 @@ -ARG OLLAMA_VERSION=0.12.11 - -# Use an official base${OLLAMA_VERSION} image with your desired version -FROM ollama/ollama:${OLLAMA_VERSION} +# Use an official ggml-org/llama.cpp image as the base image +FROM ghcr.io/ggml-org/llama.cpp:server-cuda ENV PYTHONUNBUFFERED=1 @@ -31,15 +29,11 @@ RUN apt-get update --yes --quiet && DEBIAN_FRONTEND=noninteractive apt-get insta # Set the working directory WORKDIR /work -# Add my src as /work +# Add ./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 /work/start.sh - # Set the entrypoint -ENTRYPOINT ["/bin/sh", "-c", "/work/start.sh"] \ No newline at end of file +ENTRYPOINT ["/bin/sh", "-c", "/work/start.sh"] diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..8553d3b --- /dev/null +++ b/LICENSE @@ -0,0 +1,11 @@ +LICENSE + +NOTE: This project is a fork of the original repository at https://github.com/svenbrnn/runpod-worker-ollama. + +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: + +- 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. +- You must reference the original repository at https://github.com/svenbrnn/runpod-worker-ollama. + +For more details, see the license (https://creativecommons.org/licenses/by/4.0/). diff --git a/README.md b/README.md index 62db75e..e13891d 100644 --- a/README.md +++ b/README.md @@ -1,36 +1,30 @@ -# Runpod serverless runner for ollama +

+ llama.cpp logo +

-## How to use +# Serverless llama.cpp inference worker for RunPod -Start a runpod serverless with the docker container ``svenbrnn/runpod-ollama:latest``. Set ``OLLAMA_MODEL_NAME`` environment to a model from ollama.com to automatically download a model. -A mounted volume will be automatically used. +This repository contains a serverless inference worker for running llama.cpp models on RunPod. It uses the `llama-server` image to provide an API for interacting with the models. +The following OpenAI API endpoints are supported: -[![RunPod](https://api.runpod.io/badge/SvenBrnn/runpod-worker-ollama)](https://www.runpod.io/console/hub/SvenBrnn/runpod-worker-ollama) +- `v1/models` +- `v1/chat/completions` +- `v1/completions` -## Environment variables +Streaming responses is also supported. -| Variable Name | Description | Default Value | -|---------------------|------------------------------------------|---------------------| -| `OLLAMA_MODEL_NAME` | The name of the model to download | NULL | +## Setup -## Test requests for runpod.io console +For the setup to work best, it is recommended to use a network volume attached to all workers which stores the model GGUFs and then reference those files in the launch arguments. +Make sure your RunPod worker has access to the network volume, i.e. is located in the correct data center. -See the [test_inputs](./test_inputs) directory for example test requests. +## Configuration +The worker can be configured via environment variables set in the RunPod hub configuration: -## Streaming +- `LLAMA_SERVER_CMD_ARGS`: Command line arguments (argv) for the `llama-server` binary. Example: `-hf /path/to/model.gguf:Q4_K_M -ctx_size 4096`. **IMPORTANT**: Please do not define the port argument here, as the worker will always use port `3098` automatically. +- `MAX_CONCURRENCY`: Maximum number of concurrent requests the worker can handle. Default is `8`. -Streaming for openai requests are fully working. +## License -## Preload model into the docker image - -See the [embed_model](./embed_model/) directory for instructions. - -## 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-worker-ollama](https://github.com/svenbrnn/runpod-worker-ollama). - -For more details, see the [license](https://creativecommons.org/licenses/by/4.0/). \ No newline at end of file +Please see the [LICENSE](./LICENSE) file for more information. diff --git a/embed_model/Dockerfile b/embed_model/Dockerfile deleted file mode 100644 index 4f93373..0000000 --- a/embed_model/Dockerfile +++ /dev/null @@ -1,11 +0,0 @@ -FROM svenbrnn/runpod-ollama:0.5.7 - -ARG MODEL_NAMES -ENV MODEL_NAMES=$MODEL_NAMES -ADD preload_model.sh /preload_model.sh -RUN apt-get update && apt-get install bash -RUN chmod +x /preload_model.sh && /preload_model.sh - -# Copy the model to the volume -FROM svenbrnn/runpod-ollama:0.5.7 -COPY --from=0 /runpod-volume /runpod-volume diff --git a/embed_model/README.md b/embed_model/README.md deleted file mode 100644 index a2536c8..0000000 --- a/embed_model/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Embed model into runpod-worker-ollama - -Run ``docker build --build-arg MODEL_NAMES=",..." -t /: .`` diff --git a/embed_model/preload_model.sh b/embed_model/preload_model.sh deleted file mode 100644 index 5e503cb..0000000 --- a/embed_model/preload_model.sh +++ /dev/null @@ -1,36 +0,0 @@ -#!/bin/bash - -# Start the Ollama server in the background -echo "Starting Ollama server to preload models: $MODEL_NAMES" -ollama serve & - -# Capture the PID of the Ollama server -OLLAMA_PID=$! - -# Wait for the server to be ready (adjust if necessary) -echo "Waiting for Ollama server to start..." -sleep 5 - -# Split the comma-separated model names into an array -IFS=',' read -r -a MODELS <<< "$MODEL_NAMES" - -# Loop through each model and pull it -for MODEL_NAME in "${MODELS[@]}"; do - echo "Pulling model: $MODEL_NAME" - if ollama pull "$MODEL_NAME"; then - echo "Successfully pulled model: $MODEL_NAME" - else - echo "Failed to pull model: $MODEL_NAME" - kill $OLLAMA_PID - exit 1 - fi -done - -# Stop the Ollama server -echo "Stopping Ollama server..." -kill $OLLAMA_PID - -# Wait for the server to terminate -wait $OLLAMA_PID - -echo "Model preloading complete." diff --git a/src/engine.py b/src/engine.py index e1a1393..2830446 100644 --- a/src/engine.py +++ b/src/engine.py @@ -1,3 +1,25 @@ +""" +Module engine + +This module provides lightweight adapters that present an OpenAI-compatible +interface for local llama.cpp. It contains two primary classes: + +- LlamaCPPEngine: A front-facing adapter that accepts a JobInput and transforms + it into OpenAI-compatible routes and payloads. It normalizes prompt vs chat + input and delegates actual interaction to an OpenAI-style engine. +- LlamaCPPOpenAIEngine: A concrete engine that uses the OpenAI client (pointing + at a locally hosted inference service) to list models and create completions + or chat completions. It supports both non-streaming and streaming responses. + +Environment: +- The OpenAI client is configured to point at a local base_url in this script. + +Typical usage: +- Construct LlamaCPPEngine() and call its async generate() with a JobInput. +- The generator yields either dicts (non-stream responses) or streaming chunks + formatted as strings when stream=True. +""" + import json import os @@ -6,61 +28,165 @@ from openai import OpenAI from utils import JobInput client = OpenAI( - base_url='http://localhost:11434/v1/', - - # required but ignored - api_key='ollama', + base_url="http://localhost:3098/v1/", + api_key="", ) -class OllamaEngine: + +class LlamaCPPEngine: + """ + Adapter that prepares JobInput for an OpenAI-style engine and yields + responses. + + This class normalizes incoming JobInput.llm_input into an OpenAI-style + payload. If llm_input is a string, it prepares a completions request; if + it's a list (messages), it prepares a chat.completions request. It then + delegates the actual generation to a lower-level OpenAI-compatible engine + (LlamaCPPOpenAIEngine). + + Attributes: + None beyond temporary per-call variables — the class is stateless with + respect to client configuration. + """ + def __init__(self): + """ + Initialize the LlamaCPPEngine. + + Loads environment variables (via python-dotenv). This constructor does + not perform network calls and is intentionally lightweight to allow + quick instantiation in worker processes. + + Example: + engine = LlamaCPPEngine() + """ + load_dotenv() - print ("OllamaEngine initialized") + print("Llama.cpp engine initialized") async def generate(self, job_input): - # Get model from OLLAMA_MODEL_NAME defauting to llama3.2:1b - model = os.getenv("OLLAMA_MODEL_NAME", "llama3.2:1b") + """ + Asynchronously generate responses for a given JobInput. - # Depending if prompt is a string or a list, we need to handle it differently and send it to the OpenAI API + The method inspects job_input.llm_input and builds an OpenAI-compatible + JobInput object that routes to either "/v1/completions" or + "/v1/chat/completions". It then delegates to LlamaCPPOpenAIEngine to + perform the actual request and yields each response chunk produced by + that engine. + + Args: + job_input (utils.JobInput): Input job object that contains at \ + least: + - llm_input: either a prompt string (for completions) or a list + of chat messages (for chat completion). + - stream: boolean indicating whether a streaming response is + desired. + + Yields: + dict or str: Each yielded item is either: + - a dict representing a non-streaming response, + - a string that contains streaming "data: ..." chunks for + stream=True, + - or an error dict with an "error" key. + + Notes: + - This method does not validate the full schema of job_input; it + relies on upstream code to supply a correctly shaped JobInput. + - The default model selection uses the first model listed by the + server upon requesting /v1/models. + """ + + openAIEngine = LlamaCPPOpenAIEngine() + + # Get model to use (defaults to first model in list of models) + model = client.models.list().data[0].id + + # 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 + # Build 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 + # Build 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) - 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) + generate = openAIEngine.generate(openAIjob) # Yield the response from the OpenAI API async for batch in generate: yield batch -class OllamaOpenAiEngine(OllamaEngine): + +class LlamaCPPOpenAIEngine(LlamaCPPEngine): + """ + Concrete OpenAI-compatible engine that uses the OpenAI client to perform + model listing and completions against a locally hosted inference service. + + This class expects the module-level `client` object to be configured to + point to a compatible service (llama.cpp). + that implements the same interface as the official OpenAI Python client. + + Behavior: + - For route "/v1/models": lists available models and yields a dict + describing them. + - For "/v1/chat/completions" or "/v1/completions": performs a request + and yields either a single response (non-stream) or streaming chunks. + """ + def __init__(self): + """ + Initialize the LlamaCPPOpenAIEngine. + + Loads environment variables and prints an initialization message. No + network calls are made here. + + Example: + engine = LlamaCPPOpenAIEngine() + """ + load_dotenv() - print ("OllamaOpenAiEngine initialized") + print("LlamaCPPOpenAIEngine initialized") async def generate(self, job_input): - print("Generating response for job_input:", job_input) + """ + Dispatch the provided OpenAI-style JobInput to the appropriate handler. + + Args: + job_input (JobInput): JobInput with at least: + - openai_route: str, one of "/v1/models", + "/v1/chat/completions", "/v1/completions" + - openai_input: dict containing the parameters for the chosen + route. + + Yields: + dict or str: Each yielded value can be: + - A dictionary representing the full response (non-streaming). + - A sequence of strings for streaming responses (prefixed with + "data: "). + - An error dict if the route is invalid or an exception occurs. + """ # Dump job_input to console openai_input = job_input.openai_input @@ -70,24 +196,83 @@ class OllamaOpenAiEngine(OllamaEngine): # 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"): + 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"} + yield {"error": "invalid route"} async def _handle_model_request(self): + """ + Handle a model-listing request using the OpenAI client. + + This method calls client.models.list() and yields a single dict that + mirrors the structure of OpenAI's model list responses: + {"object": "list", "data": []} + + Yields: + dict: The model list as a dictionary, or an error dict on failure. + + Raises: + No exceptions are propagated; errors are converted to + `{"error": str(e)}`. + """ + 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]} + + 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): + async def _handle_chat_or_completion_request( + self, openai_input, chat=False + ): + """ + Handle a chat or completion request and yield responses or streaming + chunks. + + This method chooses between client.chat.completions.create and + client.completions.create based on the 'chat' flag. If the input + requests non-streaming behavior (openai_input.get("stream") is falsy), + it yields a single dict representation of the response. If streaming is + requested, it iterates over the streaming response and yields + JSON-formatted chunks prefixed with "data: ", matching common SSE-like + streaming conventions. + + Args: + openai_input (dict): Parameters to pass to the client's create + method. + chat (bool): If True, call the chat completion endpoint; otherwise, + call the standard completion endpoint. + + Yields: + dict or str: + - For non-stream: yields the full response as a dict. + - For stream: yields strings of the form "data: \\n\\n" + for each chunk, and finally "data: [DONE]". + - On error: yields {"error": ""}. + + Notes: + - Exceptions are caught and yielded as error dictionaries so the + async generator consumer can handle them without dealing with + exceptions. + - The exact behavior depends on the local client implementation + being OpenAI-compatible. + """ + try: - # Call openai.chat.completions.create or openai.completions.create based on the route + # Call openai.chat.completions.create or openai.completions.create + # based on the route if chat: response = client.chat.completions.create(**openai_input) else: @@ -99,11 +284,12 @@ class OllamaOpenAiEngine(OllamaEngine): 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: " + 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 + yield {"error": str(e)} diff --git a/src/handler.py b/src/handler.py index 26354fe..c20b787 100644 --- a/src/handler.py +++ b/src/handler.py @@ -1,36 +1,45 @@ +""" +Runpod handler for processing jobs using LlamaCPP or OpenAI engines. This +module defines an asynchronous handler function that receives job inputs, +instantiates the appropriate engine based on the job input, and yields +generated output in a streaming fashion. +""" + +from typing import Any import runpod import os from utils import JobInput -from engine import OllamaEngine, OllamaOpenAiEngine +from engine import LlamaCPPEngine, LlamaCPPOpenAIEngine +# set max concurrency from environment variable or default DEFAULT_MAX_CONCURRENCY = 8 + max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY)) -async def handler(job: any): - # Just dump the whole input to the console and then return an {"ok": True} response - print('Job:', job) + +async def handler(job: Any): + """ + Asynchronous handler function for processing jobs. It receives a job + dictionary, extracts the input, determines the appropriate engine to use + (LlamaCPP or OpenAI), and yields generated output in a streaming manner. + """ job_input = JobInput(job["input"]) - engine_class = OllamaOpenAiEngine if job_input.openai_route else OllamaEngine - engine = engine_class() # Instantiate the engine + engine_class = ( + LlamaCPPOpenAIEngine if job_input.openai_route else LlamaCPPEngine + ) + engine = engine_class() - job = engine.generate(job_input) # Call generate with job_input + job = engine.generate(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, - "concurrency_modifier": lambda x: max_concurrency, + "concurrency_modifier": lambda _x: max_concurrency, "return_aggregate_stream": True, } -) \ No newline at end of file +) diff --git a/src/requirements.txt b/src/requirements.txt index bc90836..47eb4ac 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -1,4 +1,4 @@ runpod python-dotenv openai -orjson==3.10.14 \ No newline at end of file +orjson==3.10.14 diff --git a/src/start.sh b/src/start.sh index afd6d09..a889fb1 100644 --- a/src/start.sh +++ b/src/start.sh @@ -1,40 +1,54 @@ #!/bin/bash +# This script starts the llama-server with the command line arguments +# specified in the environment variable LLAMA_SERVER_CMD_ARGS, ensuring +# that the server listens on port 3098. It also starts the handler.py +# script after the server is up and running. + cleanup() { echo "Cleaning up..." - pkill -P $$ # Kill all child processes of the current script + pkill -P $$ # kill all child processes of the current script exit 0 } -# Trap exit signals and call the cleanup function +# check if the substring /workspace is in LLAMA_SERVER_CMD_ARGS +if [[ "$LLAMA_SERVER_CMD_ARGS" != *"/workspace"* ]]; then + echo "Tip: For reduced downloads and faster startup times, consider using a model stored in a network volume mounted to /workspace." +fi + +# check if the substring -port is in LLAMA_SERVER_CMD_ARGS +if [[ "$LLAMA_SERVER_CMD_ARGS" != *"-port"* ]]; then + echo "Error: You must not define -port in LLAMA_SERVER_CMD_ARGS, as port 3098 is required." + exit 1 +fi + +# trap exit signals and call the cleanup function trap cleanup SIGINT SIGTERM -# Kill any existing ollama processes -pgrep ollama | xargs kill +# kill any existing llama-server processes +pgrep llama-server | 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 +# we have a string with all the command line arguments in the env var LLAMA_SERVER_CMD_ARGS; +# it contains a.e. "-hf modelname -ctx_size 4096". + +# We need to pass these arguments to llama-server verbatim. +llama-server $LLAMA_SERVER_CMD_ARGS -port 3098 2>&1 | tee llama.server.log & + +LLAMA_SERVER_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 + echo "Checking if llama-server is done initializing..." + + if cat llama.server.log | grep -q "listening"; then + return 0 # success else - return 1 # Failure + return 1 # failure fi } -# Wait for the server to start +# 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 "$OLLAMA_MODEL_NAME" ]; then - echo "No model name provided. Skipping model pull..." -else - echo "Pulling model $OLLAMA_MODEL_NAME..." - ollama pull $OLLAMA_MODEL_NAME -fi -python -u handler.py $1 \ No newline at end of file +python -u handler.py $1 diff --git a/src/test_input.json b/src/test_input.json index 484acbd..c6294a9 100644 --- a/src/test_input.json +++ b/src/test_input.json @@ -1,5 +1,5 @@ { - "input": { - "prompt": "How are you?" - } + "input": { + "prompt": "How are you?" + } } \ No newline at end of file diff --git a/src/utils.py b/src/utils.py index dcdb385..ebfec0e 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,6 +1,27 @@ +""" +This module contains utility classes and functions for handling job inputs. +""" + + class JobInput: + """ + Class to parse and store job input data. It extracts fields such as + llm_input, stream, openai_route, and openai_input from the provided job + dictionary. + """ + def __init__(self, job): + """ + Initialize the JobInput instance by parsing the job dictionary. + + Default values: + - llm_input: job["messages"] if present, else job["prompt"] + - stream: False + - openai_route: None + - openai_input: None + """ + 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 + self.openai_input = job.get("openai_input") diff --git a/test_inputs/default_runpodio.json b/test_inputs/default_runpodio.json deleted file mode 100644 index c6294a9..0000000 --- a/test_inputs/default_runpodio.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "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 deleted file mode 100644 index eeda987..0000000 --- a/test_inputs/openai_completion.json +++ /dev/null @@ -1,10 +0,0 @@ -{ - "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 deleted file mode 100644 index 4cdb9cd..0000000 --- a/test_inputs/openai_get_models.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "input": { - "openai_route": "/v1/models", - "openai_input": {} - } -} \ No newline at end of file