From 11097765365c3baa72f2d6ab82b00cca4cf8aba6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= <boschi1997@gmail.com>
Date: Fri, 26 Jan 2024 16:20:51 +0100
Subject: [PATCH] add base workflow for running benchmarks on demand (#240)

* add base workflow for running benchmarks on demand

* add base workflow for running benchmarks on demand
---
 .github/workflows/run-benchmarks.yml       | 47 ++++++++++++++++++++++
 ragstack-e2e-tests/benchmarks/README.md    |  2 +-
 ragstack-e2e-tests/benchmarks/testcases.py |  5 ++-
 3 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 .github/workflows/run-benchmarks.yml

diff --git a/.github/workflows/run-benchmarks.yml b/.github/workflows/run-benchmarks.yml
new file mode 100644
index 000000000..6cf1f9927
--- /dev/null
+++ b/.github/workflows/run-benchmarks.yml
@@ -0,0 +1,47 @@
+name: Run benchmarks
+
+on:
+  workflow_dispatch:
+    inputs:
+      intensity:
+        description: 'Intensity of the tests (1-5)'
+        default: "2"
+
+jobs:
+  run-benchmarks:
+    name: Run benchmarks
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: Set up python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.11'
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install tox
+          
+          
+      - name: Run
+        env:
+          OPEN_AI_KEY: "${{ secrets.E2E_TESTS_OPEN_AI_KEY }}"
+          NVIDIA_API_KEY: "${{ secrets.E2E_TESTS_NVIDIA_API_KEY }}"
+        run: |
+          cd ragstack-e2e-tests
+          poetry install --no-root
+          poetry run python benchmarks/runner.py -t all -i ${{ github.event.inputs.intensity }}
+          poetry run python benchmarks/visualize.py
+
+      - name: Dump logs
+        if: always()
+        run: |
+          cd ragstack-e2e-tests
+          if [ -f "testcases.log" ]; then
+            cat testcases.log
+          fi
diff --git a/ragstack-e2e-tests/benchmarks/README.md b/ragstack-e2e-tests/benchmarks/README.md
index 646d4fb6c..76429af36 100644
--- a/ragstack-e2e-tests/benchmarks/README.md
+++ b/ragstack-e2e-tests/benchmarks/README.md
@@ -12,7 +12,7 @@ export NVIDIA_API_KEY=<your key>
 
 Then you can run the benchmark with the following command:
 ```
-poetry run python benchmarks/runner.py -t * 
+poetry run python benchmarks/runner.py -t all
 ```
 
 After the run, all the reports are stored in the `reports` folder.
diff --git a/ragstack-e2e-tests/benchmarks/testcases.py b/ragstack-e2e-tests/benchmarks/testcases.py
index b43b2c241..7532c034e 100644
--- a/ragstack-e2e-tests/benchmarks/testcases.py
+++ b/ragstack-e2e-tests/benchmarks/testcases.py
@@ -1,3 +1,4 @@
+import os
 import sys
 
 from langchain.text_splitter import CharacterTextSplitter
@@ -98,13 +99,13 @@ def embeddings_batch100_chunk512(embeddings_fn):
 
 
 def openai_ada002(batch_size):
-    return OpenAIEmbeddings(chunk_size=batch_size)
+    return OpenAIEmbeddings(chunk_size=batch_size, api_key=os.environ.get("OPEN_AI_KEY"))
 
 
 def nvidia_nvolveqa40k(batch_size):
     # 50 is the max supported batch size
     return NVIDIAEmbeddings(
-        model="nvolveqa_40k", max_batch_size=min(batch_size, 50), model_type="query"
+        model="nvolveqa_40k", max_batch_size=max(50, batch_size), model_type="query"
     )