Merge pull request #19 from knights-analytics/gpu-support

Gpu support

Author: riccardopinosio

.github/workflows/release.yaml

@@ -11,9 +11,11 @@ permissions:
 
 jobs:
   test:
-      name: Run test suite
-      runs-on: ubuntu-latest
-      steps:
+    name: Run test suite
+    runs-on: ubuntu-latest
+    env:
+      CI: true
+    steps:
       - name: Set up Go
         uses: actions/setup-go@v2
         with:
@@ -37,12 +39,16 @@ jobs:
           path: ./artifacts/libtokenizers.a
       - uses: actions/upload-artifact@v4
         with:
-          name: onnxruntime.so
-          path: ./artifacts/onnxruntime.so
+          name: onnxruntime-linux-x64.so
+          path: ./artifacts/onnxruntime-linux-x64.so
+      - uses: actions/upload-artifact@v4
+        with:
+          name: onnxruntime-linux-x64-gpu
+          path: ./artifacts/onnxruntime-linux-x64-gpu
       - uses: actions/upload-artifact@v4
         with:
-          name: hugot-cli-linux-amd64 
-          path: ./artifacts/hugot-cli-linux-amd64
+          name: hugot-cli-linux-x64
+          path: ./artifacts/hugot-cli-linux-x64
       - name: Publish Test Report
         uses: mikepenz/action-junit-report@v4
         if: success() || failure()
@@ -63,16 +69,20 @@ jobs:
           path: .
       - uses: actions/download-artifact@v4
         with:
-          name: onnxruntime.so
+          name: onnxruntime-linux-x64.so
+          path: .
+      - uses: actions/download-artifact@v4
+        with:
+          name: hugot-cli-linux-x64
           path: .
       - uses: actions/download-artifact@v4
         with:
-          name: hugot-cli-linux-amd64
+          name: onnxruntime-linux-x64-gpu
           path: .
       - name: Display structure of downloaded files
         run: ls -R
       - uses: ncipollo/release-action@v1
         with:
-          artifacts: "libtokenizers.a, onnxruntime.so, hugot-cli-linux-amd64"
+          artifacts: "libtokenizers.a, onnxruntime-linux-x64.so, onnxruntime-linux-x64-gpu, hugot-cli-linux-x64"
           generateReleaseNotes: true
           skipIfReleaseExists: true

.github/workflows/test.yaml

@@ -10,6 +10,8 @@ jobs:
   test:
       name: Run test suite
       runs-on: ubuntu-latest
+      env:
+        CI: true
       steps:
       - name: Set up Go
         uses: actions/setup-go@v2
@@ -34,12 +36,16 @@ jobs:
           path: ./artifacts/libtokenizers.a
       - uses: actions/upload-artifact@v4
         with:
-          name: onnxruntime.so
-          path: ./artifacts/onnxruntime.so
+          name: onnxruntime-linux-x64.so
+          path: ./artifacts/onnxruntime-linux-x64.so
       - uses: actions/upload-artifact@v4
         with:
-          name: hugot-cli-linux-amd64 
-          path: ./artifacts/hugot-cli-linux-amd64
+          name: onnxruntime-linux-x64-gpu
+          path: ./artifacts/onnxruntime-linux-x64-gpu
+      - uses: actions/upload-artifact@v4
+        with:
+          name: hugot-cli-linux-x64
+          path: ./artifacts/hugot-cli-linux-x64
       - name: Publish Test Report
         uses: mikepenz/action-junit-report@v4
         if: success() || failure()

Dockerfile

@@ -1,6 +1,7 @@
 ARG GO_VERSION=1.22.2
 ARG RUST_VERSION=1.77
-ARG ONNXRUNTIME_VERSION=1.17.1
+ARG ONNXRUNTIME_VERSION=1.17.3
+ARG CUDA_VERSION=12.4
 
 #--- rust build of tokenizer ---
 
@@ -12,40 +13,54 @@ RUN git clone https://github.com/knights-analytics/tokenizers -b main && \
 
 #--- build and test layer ---
 
-FROM public.ecr.aws/amazonlinux/amazonlinux:2023 AS building
+FROM public.ecr.aws/amazonlinux/amazonlinux:2023 AS hugot-build
 ARG GO_VERSION
 ARG ONNXRUNTIME_VERSION
 
 RUN dnf -y install gcc jq bash tar xz gzip glibc-static libstdc++ wget zip git && \
     ln -s /usr/lib64/libstdc++.so.6 /usr/lib64/libstdc++.so && \
     dnf install -y 'dnf-command(config-manager)' && \
-    dnf config-manager \
-    --add-repo https://download.fedoraproject.org/pub/fedora/linux/releases/39/Everything/x86_64/os/ && \
+    dnf config-manager --add-repo https://download.fedoraproject.org/pub/fedora/linux/releases/39/Everything/x86_64/os/ && \
+    # from fedora
+    dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/fedora39/x86_64/cuda-fedora39.repo && \
+    dnf install -y cuda-cudart-12-4 libcublas-12-4 libcurand-12-4 libcufft-12-4 && \
+    # from rhel
+    dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel9/x86_64/cuda-rhel9.repo && \
+    dnf install -y libcudnn8 && \
     dnf clean all
 
+# go
 RUN curl -LO https://golang.org/dl/go${GO_VERSION}.linux-amd64.tar.gz && \
     tar -C /usr/local -xzf go${GO_VERSION}.linux-amd64.tar.gz && \
     rm go${GO_VERSION}.linux-amd64.tar.gz
 ENV PATH="$PATH:/usr/local/go/bin"
 
+# tokenizer
 COPY --from=tokenizer /tokenizers/target/release/libtokenizers.a /usr/lib/libtokenizers.a
+
+# onnxruntime cpu and gpu
 RUN curl -LO https://github.com/microsoft/onnxruntime/releases/download/v${ONNXRUNTIME_VERSION}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz && \
    tar -xzf onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz && \
-   mv ./onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}/lib/libonnxruntime.so.${ONNXRUNTIME_VERSION} /usr/lib64/onnxruntime.so
+   mv ./onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}/lib/libonnxruntime.so.${ONNXRUNTIME_VERSION} /usr/lib64/onnxruntime.so && \
+   curl -LO https://github.com/microsoft/onnxruntime/releases/download/v${ONNXRUNTIME_VERSION}/onnxruntime-linux-x64-gpu-cuda12-${ONNXRUNTIME_VERSION}.tgz && \
+   tar -xzf onnxruntime-linux-x64-gpu-cuda12-${ONNXRUNTIME_VERSION}.tgz && \
+   mv ./onnxruntime-linux-x64-gpu-${ONNXRUNTIME_VERSION}/lib /usr/lib64/onnxruntime-gpu
 
+# build gotestsum and test2json
 RUN GOOS=linux GOARCH=amd64 CGO_ENABLED=0 go build -o test2json -ldflags="-s -w" cmd/test2json && mv test2json /usr/local/bin/test2json && \
     curl -LO https://github.com/gotestyourself/gotestsum/releases/download/v1.11.0/gotestsum_1.11.0_linux_amd64.tar.gz && \
     tar -xzf gotestsum_1.11.0_linux_amd64.tar.gz --directory /usr/local/bin
 
-# build cli
+# build cli binary
 COPY . /build
 WORKDIR /build
 RUN cd ./cmd && CGO_ENABLED=1 GOOS=linux GOARCH=amd64 go build -a -o ./target main.go
 
 # NON-PRIVILEDGED USER
 # create non-priviledged testuser with id: 1000
-RUN dnf install --disablerepo=* --enablerepo=amazonlinux --allowerasing -y dirmngr && dnf clean all
-RUN useradd -u 1000 -m testuser && chown -R testuser:testuser /build
+RUN dnf install --disablerepo=* --enablerepo=amazonlinux --allowerasing -y dirmngr sudo which && dnf clean all
+RUN useradd -u 1000 -m testuser && chown -R testuser:testuser /build && usermod -a -G wheel testuser
+RUN echo "testuser ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers.d/testuser
 
 # ENTRYPOINT
 COPY ./scripts/entrypoint.sh /entrypoint.sh
@@ -54,10 +69,10 @@ RUN sed -i 's/\r//g' /entrypoint.sh
 RUN chmod +x /entrypoint.sh
 ENTRYPOINT ["/entrypoint.sh"]
 
-#--- artifacts layer
-
+# artifacts layer
 FROM scratch AS artifacts
 
-COPY --from=building /usr/lib64/onnxruntime.so onnxruntime.so
-COPY --from=building /usr/lib/libtokenizers.a libtokenizers.a
-COPY --from=building /build/cmd/target /hugot-cli-linux-amd64
+COPY --from=hugot-build /usr/lib64/onnxruntime.so onnxruntime-linux-x64.so
+COPY --from=hugot-build /usr/lib64/onnxruntime-gpu onnxruntime-linux-x64-gpu
+COPY --from=hugot-build /usr/lib/libtokenizers.a libtokenizers.a
+COPY --from=hugot-build /build/cmd/target /hugot-cli-linux-x64

README.md

@@ -12,6 +12,8 @@ The goal of this library is to provide an easy, scalable, and hassle-free way to
 2. Hassle-free and performant production use: we exclusively support onnx exports of huggingface models. Pytorch transformer models that don't have an onnx version can be easily exported to onnx via [huggingface optimum](https://huggingface.co/docs/optimum/index), and used with the library
 3. Run on your hardware: this library is for those who want to run transformer models tightly coupled with their go applications, without the performance drawbacks of having to hit a rest API, or the hassle of setting up and maintaining e.g. a python RPC service that talks to go.
 
+We support inference on CPU and on all accelerators supported by ONNXRuntime. Note, however, that currently only CPU and GPU inference on nvidia GPU (with cuda) are tested (see below).
+
 ## Why
 
 Developing and fine-tuning transformer models with the huggingface python library is a great experience, but if your production stack is golang-based being able to reliably deploy and scale the resulting pytorch models can be challenging and require quite some setup. This library aims to allow you to just lift-and-shift your python model and use the same huggingface pipelines you use for development for inference in a go application.
@@ -20,7 +22,7 @@ Developing and fine-tuning transformer models with the huggingface python librar
 
 For the golang developer or ML engineer who wants to run transformer piplines on their own hardware, tightly coupled with their own application.
 
-## What is already there
+## Implemented pipelines
 
 Currently, we have implementations for the following transfomer pipelines:
 
@@ -32,11 +34,34 @@ Implementations for additional pipelines will follow. We also very gladly accept
 
 Hugot can be used both as a library and as a command-line application. See below for usage instructions.
 
+## Hardware acceleration 🚀
+
+Hugot now also supports the following accelerator backends for your inference:
+ - CUDA (tested). See below for setup instructions.
+ - TensorRT (untested)
+ - DirectML (untested)
+ - CoreML (untested)
+ - OpenVINO (untested)
+
+Please help us out by testing the untested options above and providing feedback, good or bad!
+
+To use Hugot with nvidia gpu acceleration, you need to have the following:
+
+- The cuda gpu version of onnxruntime on the machine/docker container. You can see how we get that by looking at the [Dockerfile](./Dockerfile). You can also get the onnxruntime libraries that we use for testing from the release. Just download the gpu .so libraries and put them in /usr/lib64.
+- the nvidia driver for your graphics card
+- the required cuda libraries installed on your system that are compatible with the onnxruntime gpu version you use. See [here](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html). For instance, for onnxruntime-gpu 17.3, we need CUDA 12.x (any minor version should be compatible) and cuDNN 8.9.2.26.
+
+On the last point above, you can install CUDA 12.x by installing the full cuda toolkit, but that's quite a big package. In our testing on awslinux/fedora, we have been able to limit the libraries needed to run hugot with nvidia gpu acceleration to just these:
+
+- cuda-cudart-12-4 libcublas-12-4 libcurand-12-4 libcufft-12-4 (from fedora repo)
+- libcudnn8 (from RHEL repo, for cuDNN)
+
+On different distros (e.g. Ubuntu), you should be able to install the equivalent packages and gpu inference should work.
+
 ## Limitations
 
 Apart from the fact that only the aforementioned pipelines are currently implemented, the current limitations are:
-    - the library and cli are only tested on amd64-linux
-    - only CPU inference is supported
+    - the library and cli are only built/tested on amd64-linux
 
 Pipelines are also tested on specifically NLP use cases. In particular, we use the following models for testing:
 - feature extraction: all-MiniLM-L6-v2
@@ -196,6 +221,8 @@ session, err := hugot.NewSession(
 
 InterOpNumThreads and IntraOpNumThreads constricts each goroutine's call to a single core, greatly reducing locking and cache penalties. Disabling CpuMemArena and MemPattern skips pre-allocation of some memory structures, increasing latency, but also throughput efficiency.
 
+For GPU the config above also applies. We are still testing the optimum GPU configuration, whether it is better to run in parallel or with a single thread, and what size of input batch is fastest.
+
 ## Contributing
 
 ### Development environment

cmd/main.go

@@ -13,11 +13,12 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/mattn/go-isatty"
+	"github.com/urfave/cli/v2"
+
 	"github.com/knights-analytics/hugot"
 	"github.com/knights-analytics/hugot/pipelines"
 	util "github.com/knights-analytics/hugot/utils"
-	"github.com/mattn/go-isatty"
-	"github.com/urfave/cli/v2"
 )
 
 var modelPath string

cmd/main_test.go

@@ -2,15 +2,15 @@ package main
 
 import (
 	"context"
+	_ "embed"
 	"fmt"
 	"os"
 	"path"
 	"testing"
 
-	_ "embed"
+	"github.com/urfave/cli/v2"
 
 	util "github.com/knights-analytics/hugot/utils"
-	"github.com/urfave/cli/v2"
 )
 
 //go:embed testData/textClassification.jsonl

compose-dev.yaml

@@ -4,7 +4,14 @@ services:
     build:
       context: .
       dockerfile: ./Dockerfile
-      target: building
+      target: hugot-build
     volumes:
       - $src_dir:/home/testuser/repositories/hugot
-    entrypoint: ["tail", "-f", "/dev/null"]
+    entrypoint: ["tail", "-f", "/dev/null"]
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

compose-test.yaml

@@ -5,10 +5,18 @@ services:
     build:
       context: .
       dockerfile: ./Dockerfile
-      target: building
+      target: hugot-build
     volumes:
       - $test_folder:/test
       - $test_folder/../scripts/run-unit-tests-container.sh:/run-unit-tests-container.sh
     environment:
       - HOST_UID=$host_uid
+      - CI=$CI
     command: /run-unit-tests-container.sh
+    # deploy:
+    #   resources:
+    #     reservations:
+    #       devices:
+    #         - driver: nvidia
+    #           count: 1
+    #           capabilities: [gpu]

downloader.go

@@ -14,6 +14,8 @@ import (
 	"time"
 
 	hfd "github.com/bodaay/HuggingFaceModelDownloader/hfdownloader"
+
+	util "github.com/knights-analytics/hugot/utils"
 )
 
 // DownloadOptions is a struct of options that can be passed to DownloadModel
@@ -138,3 +140,26 @@ func checkURL(client *http.Client, url string, authToken string) (bool, bool, er
 	}
 	return tokenizerFound, onnxFound, nil
 }
+
+func downloadModelIfNotExists(session *Session, modelName string, destination string) string {
+
+	modelNameFS := modelName
+	if strings.Contains(modelNameFS, ":") {
+		modelNameFS = strings.Split(modelName, ":")[0]
+	}
+	modelNameFS = path.Join(destination, strings.Replace(modelNameFS, "/", "_", -1))
+
+	fullModelPath := path.Join(destination, modelNameFS)
+	exists, err := util.FileSystem.Exists(context.Background(), fullModelPath)
+	if err != nil {
+		panic(err)
+	}
+	if exists {
+		return fullModelPath
+	}
+	fullModelPath, err = session.DownloadModel(modelName, destination, NewDownloadOptions())
+	if err != nil {
+		panic(err)
+	}
+	return fullModelPath
+}

go.mod

@@ -2,28 +2,27 @@ module github.com/knights-analytics/hugot
 
 go 1.22
 
+replace github.com/viant/afsc => github.com/knights-analytics/afsc v0.0.0-20240419150436-7b7def176954
+
 require (
 	github.com/bodaay/HuggingFaceModelDownloader v0.0.0-20240307153905-2f38356a6d6c
 	github.com/json-iterator/go v1.1.12
-	github.com/knights-analytics/tokenizers v0.10.0
+	github.com/knights-analytics/tokenizers v0.12.0
 	github.com/mattn/go-isatty v0.0.20
 	github.com/stretchr/testify v1.9.0
 	github.com/urfave/cli/v2 v2.27.1
-	github.com/viant/afs v1.25.0
-	github.com/viant/afsc v1.9.1
+	github.com/viant/afs v1.25.1
+	github.com/viant/afsc v1.9.2
 	github.com/yalue/onnxruntime_go v1.9.0
 )
 
 require (
-	cloud.google.com/go v0.112.2 // indirect
-	cloud.google.com/go/storage v1.40.0 // indirect
-	github.com/aws/aws-sdk-go v1.51.21 // indirect
+	github.com/aws/aws-sdk-go v1.51.28 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.4 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fatih/color v1.16.0 // indirect
 	github.com/go-errors/errors v1.5.1 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
-	github.com/kr/text v0.2.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
@@ -32,14 +31,6 @@ require (
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/xrash/smetrics v0.0.0-20240312152122-5f08fbb34913 // indirect
 	golang.org/x/crypto v0.22.0 // indirect
-	golang.org/x/net v0.24.0 // indirect
-	golang.org/x/oauth2 v0.19.0 // indirect
-	golang.org/x/sync v0.7.0 // indirect
 	golang.org/x/sys v0.19.0 // indirect
-	google.golang.org/api v0.172.0 // indirect
-	google.golang.org/genproto v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
-	google.golang.org/grpc v1.63.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )