feat: Add Kokoros backend #9212
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,3 +1,6 @@ | ||
| [submodule "docs/themes/hugo-theme-relearn"] | ||
| path = docs/themes/hugo-theme-relearn | ||
| url = https://github.com/McShelby/hugo-theme-relearn.git | ||
| [submodule "backend/rust/kokoros/sources/Kokoros"] | ||
| path = backend/rust/kokoros/sources/Kokoros | ||
| url = https://github.com/lucasjinreal/Kokoros |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,80 @@ | ||
| ARG BASE_IMAGE=ubuntu:24.04 | ||
|
|
||
| FROM ${BASE_IMAGE} AS builder | ||
| ARG BACKEND=kokoros | ||
| ARG BUILD_TYPE | ||
| ENV BUILD_TYPE=${BUILD_TYPE} | ||
| ARG CUDA_MAJOR_VERSION | ||
| ARG CUDA_MINOR_VERSION | ||
| ARG SKIP_DRIVERS=false | ||
| ENV CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION} | ||
| ENV CUDA_MINOR_VERSION=${CUDA_MINOR_VERSION} | ||
| ENV DEBIAN_FRONTEND=noninteractive | ||
| ARG TARGETARCH | ||
| ARG TARGETVARIANT | ||
| ARG UBUNTU_VERSION=2404 | ||
|
|
||
| RUN apt-get update && \ | ||
| apt-get install -y --no-install-recommends \ | ||
| build-essential \ | ||
| git ccache \ | ||
| ca-certificates \ | ||
| make cmake wget \ | ||
| curl unzip \ | ||
| clang \ | ||
| pkg-config \ | ||
| libssl-dev \ | ||
| espeak-ng libespeak-ng-dev \ | ||
| libsonic-dev libpcaudio-dev \ | ||
| libopus-dev \ | ||
| protobuf-compiler && \ | ||
| apt-get clean && \ | ||
| rm -rf /var/lib/apt/lists/* | ||
|
|
||
| # Cuda | ||
| ENV PATH=/usr/local/cuda/bin:${PATH} | ||
|
|
||
| # CuBLAS requirements | ||
| RUN <<EOT bash | ||
| if ( [ "${BUILD_TYPE}" = "cublas" ] || [ "${BUILD_TYPE}" = "l4t" ] ) && [ "${SKIP_DRIVERS}" = "false" ]; then | ||
| apt-get update && \ | ||
| apt-get install -y --no-install-recommends \ | ||
| software-properties-common pciutils | ||
| if [ "amd64" = "$TARGETARCH" ]; then | ||
| curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/x86_64/cuda-keyring_1.1-1_all.deb | ||
| fi | ||
| if [ "arm64" = "$TARGETARCH" ]; then | ||
| if [ "${CUDA_MAJOR_VERSION}" = "13" ]; then | ||
| curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/sbsa/cuda-keyring_1.1-1_all.deb | ||
| else | ||
| curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/arm64/cuda-keyring_1.1-1_all.deb | ||
| fi | ||
| fi | ||
| dpkg -i cuda-keyring_1.1-1_all.deb && \ | ||
| rm -f cuda-keyring_1.1-1_all.deb && \ | ||
| apt-get update && \ | ||
| apt-get install -y --no-install-recommends \ | ||
| cuda-nvcc-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ | ||
| libcufft-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ | ||
| libcurand-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ | ||
| libcublas-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ | ||
| libcusparse-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ | ||
| libcusolver-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} | ||
| apt-get clean && rm -rf /var/lib/apt/lists/* | ||
| fi | ||
| EOT | ||
|
|
||
| # Install Rust | ||
| RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y | ||
| ENV PATH="/root/.cargo/bin:${PATH}" | ||
|
|
||
| COPY . /LocalAI | ||
|
|
||
| RUN git config --global --add safe.directory /LocalAI | ||
|
|
||
| RUN make -C /LocalAI/backend/rust/${BACKEND} build | ||
|
|
||
| FROM scratch | ||
| ARG BACKEND=kokoros | ||
|
|
||
| COPY --from=builder /LocalAI/backend/rust/${BACKEND}/package/. ./ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| [package] | ||
| name = "kokoros-grpc" | ||
| version = "0.1.0" | ||
| edition = "2021" | ||
|
|
||
| [[bin]] | ||
| name = "kokoros-grpc" | ||
| path = "src/main.rs" | ||
|
|
||
| [dependencies] | ||
| kokoros = { path = "sources/Kokoros/kokoros" } | ||
|
|
||
| tonic = "0.13" | ||
| prost = "0.13" | ||
| tokio = { version = "1", features = ["full"] } | ||
| tokio-stream = "0.1" | ||
| clap = { version = "4", features = ["derive"] } | ||
| tracing = "0.1" | ||
| tracing-subscriber = { version = "0.3", features = ["env-filter"] } | ||
|
|
||
| [build-dependencies] | ||
| tonic-build = "0.13" | ||
|
|
||
| [features] | ||
| default = ["cpu"] | ||
| cpu = ["kokoros/cpu"] | ||
| cuda = ["kokoros/cuda"] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| CURRENT_DIR=$(abspath ./) | ||
|
|
||
| # Enable CUDA feature when building with cublas/l4t | ||
| CARGO_FEATURES := | ||
| ifneq (,$(filter cublas l4t,$(BUILD_TYPE))) | ||
| CARGO_FEATURES := --features cuda --no-default-features | ||
| endif | ||
|
|
||
| .PHONY: kokoros-grpc | ||
| kokoros-grpc: | ||
| mkdir -p $(CURRENT_DIR)/proto | ||
| cp $(CURRENT_DIR)/../../backend.proto $(CURRENT_DIR)/proto/backend.proto | ||
| cd $(CURRENT_DIR) && \ | ||
| BACKEND_PROTO_PATH=$(CURRENT_DIR)/proto/backend.proto \ | ||
| cargo build --release $(CARGO_FEATURES) | ||
|
|
||
| .PHONY: package | ||
| package: | ||
| bash package.sh | ||
|
|
||
| .PHONY: build | ||
| build: kokoros-grpc package | ||
|
|
||
| .PHONY: clean | ||
| clean: | ||
| cargo clean | ||
| rm -rf package proto |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,15 @@ | ||
| fn main() -> Result<(), Box<dyn std::error::Error>> { | ||
| let proto_path = std::env::var("BACKEND_PROTO_PATH") | ||
| .unwrap_or_else(|_| "proto/backend.proto".to_string()); | ||
|
|
||
| let proto_dir = std::path::Path::new(&proto_path) | ||
| .parent() | ||
| .unwrap_or(std::path::Path::new(".")); | ||
|
|
||
| tonic_build::configure() | ||
| .build_server(true) | ||
| .build_client(false) | ||
| .compile_protos(&[&proto_path], &[proto_dir])?; | ||
|
|
||
| Ok(()) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| #!/bin/bash | ||
| set -e | ||
|
|
||
| CURDIR=$(dirname "$(realpath $0)") | ||
| mkdir -p $CURDIR/package/lib | ||
|
|
||
| # Copy the binary | ||
| cp -avf $CURDIR/target/release/kokoros-grpc $CURDIR/package/ | ||
|
|
||
| # Copy the run script | ||
| cp -rfv $CURDIR/run.sh $CURDIR/package/ | ||
| chmod +x $CURDIR/package/run.sh | ||
|
|
||
| # Copy ONNX Runtime shared libraries from ort build artifacts | ||
| ORT_LIB_DIR=$(find $CURDIR/target/release/build -name "libonnxruntime*.so*" -path "*/ort-sys-*/out/*" -exec dirname {} \; 2>/dev/null | head -1) | ||
| if [ -n "$ORT_LIB_DIR" ]; then | ||
| cp -avfL $ORT_LIB_DIR/libonnxruntime*.so* $CURDIR/package/lib/ 2>/dev/null || true | ||
| fi | ||
|
|
||
| # Copy espeak-ng data | ||
| if [ -d "/usr/share/espeak-ng-data" ]; then | ||
| cp -rf /usr/share/espeak-ng-data $CURDIR/package/ | ||
| elif [ -d "/usr/lib/x86_64-linux-gnu/espeak-ng-data" ]; then | ||
| cp -rf /usr/lib/x86_64-linux-gnu/espeak-ng-data $CURDIR/package/ | ||
| fi | ||
|
|
||
| # Copy ALL dynamic library dependencies of the binary | ||
| echo "Bundling dynamic library dependencies..." | ||
| ldd $CURDIR/target/release/kokoros-grpc | grep "=>" | awk '{print $3}' | while read lib; do | ||
| if [ -n "$lib" ] && [ -f "$lib" ]; then | ||
| cp -avfL "$lib" $CURDIR/package/lib/ 2>/dev/null || true | ||
| fi | ||
| done | ||
|
|
||
| # Copy CA certificates for HTTPS (needed for model auto-download) | ||
| if [ -d "/etc/ssl/certs" ]; then | ||
| mkdir -p $CURDIR/package/etc/ssl | ||
| cp -rf /etc/ssl/certs $CURDIR/package/etc/ssl/ | ||
| fi | ||
|
|
||
| # Copy the dynamic linker | ||
| if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then | ||
| cp -arfLv /lib64/ld-linux-x86-64.so.2 $CURDIR/package/lib/ld.so | ||
| elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then | ||
| cp -arfLv /lib/ld-linux-aarch64.so.1 $CURDIR/package/lib/ld.so | ||
| fi | ||
|
|
||
| echo "Packaging completed successfully" | ||
| ls -liah $CURDIR/package/ | ||
| ls -liah $CURDIR/package/lib/ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,23 @@ | ||
| #!/bin/bash | ||
| set -ex | ||
|
|
||
| CURDIR=$(dirname "$(realpath $0)") | ||
|
|
||
| export LD_LIBRARY_PATH=$CURDIR/lib:${LD_LIBRARY_PATH:-} | ||
|
|
||
| # SSL certificates for model auto-download | ||
| if [ -d "$CURDIR/etc/ssl/certs" ]; then | ||
| export SSL_CERT_DIR=$CURDIR/etc/ssl/certs | ||
| fi | ||
|
|
||
| # espeak-ng data directory | ||
| if [ -d "$CURDIR/espeak-ng-data" ]; then | ||
| export ESPEAK_NG_DATA=$CURDIR/espeak-ng-data | ||
| fi | ||
|
|
||
| # Use bundled ld.so if present (portability) | ||
| if [ -f $CURDIR/lib/ld.so ]; then | ||
| exec $CURDIR/lib/ld.so $CURDIR/kokoros-grpc "$@" | ||
| fi | ||
|
|
||
| exec $CURDIR/kokoros-grpc "$@" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| use tonic::{Request, Status}; | ||
|
|
||
| /// Returns an interceptor function if LOCALAI_GRPC_AUTH_TOKEN is set. | ||
| pub fn make_auth_interceptor( | ||
| ) -> Option<impl Fn(Request<()>) -> Result<Request<()>, Status> + Clone> { | ||
| let token = std::env::var("LOCALAI_GRPC_AUTH_TOKEN").ok()?; | ||
| if token.is_empty() { | ||
| return None; | ||
| } | ||
| let expected = format!("Bearer {}", token); | ||
| Some( | ||
| move |req: Request<()>| -> Result<Request<()>, Status> { | ||
| let meta = req.metadata(); | ||
| match meta.get("authorization") { | ||
| Some(val) => { | ||
| if val.as_bytes() == expected.as_bytes() { | ||
| Ok(req) | ||
| } else { | ||
| Err(Status::unauthenticated("invalid token")) | ||
| } | ||
| } | ||
| None => Err(Status::unauthenticated("missing authorization")), | ||
| } | ||
| }, | ||
| ) | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm ok to add it as submodule, but we need also a CI workflow to update submodules regularly