Turbo quant rest experiments

.github/workflows/build-hf-dataset.yaml

@@ -1,24 +1,81 @@
 name: Build HF Dataset
 
-# Stub on master.
+# Spins up a single Hetzner machine, downloads a HuggingFace embeddings dataset,
+# splits train/test, computes brute-force kNN ground truth, packages
+# vectors.npy + tests.jsonl into a tarball, and uploads it to GCS.
 
 on:
   workflow_dispatch:
     inputs:
       hf_dataset:
         description: "HF dataset id, e.g. Qdrant/dbpedia-entities-openai3-text-embedding-3-large-1536-100K"
         required: true
+        default: "Qdrant/gte-multilingual-ads-1M"
       output_name:
         description: "Output tarball basename"
         required: true
+        default: "ads-gte-multilingual-1M-768-angular"
       vector_column:
         description: "Embedding column name"
-        default: "embedding"
+        default: "gte"
+
+concurrency:
+  group: hetzner-machines
+
+env:
+  HCLOUD_TOKEN: ${{ secrets.HCLOUD_TOKEN }}
+  SERVER_NAME: build-hf-dataset-${{ github.run_id }}
+  GCS_PATH: gs://ann-filtered-benchmark/datasets/
 
 jobs:
-  stub:
+  build:
     runs-on: ubuntu-latest
     steps:
-      - run: |
-          echo "This is the master stub for the Build HF Dataset workflow."
-          exit 0
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+      - uses: webfactory/ssh-agent@d4b9b8ff72958532804b70bbe600ad43b36d5f2e # v0.8.0
+        with:
+          ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Setup CI
+        run: bash -x tools/setup_ci.sh
+
+      - name: Create server
+        uses: ./.github/workflows/actions/create-server-with-retry
+        with:
+          server_name: ${{ env.SERVER_NAME }}
+          server_type: ccx43
+          region: fsn1
+
+      - name: Build dataset on remote
+        env:
+          GCS_KEY: ${{ secrets.GCS_KEY }}
+          GCS_SECRET: ${{ secrets.GCS_SECRET }}
+        run: |
+          set -euo pipefail
+          source tools/ssh.sh
+          IP=$(bash tools/hetzner/get_public_ip.sh "$SERVER_NAME")
+
+          scp_with_retry -o StrictHostKeyChecking=no \
+            scripts/build_hf_dataset.py "root@$IP:/root/"
+
+          ssh_with_retry -o StrictHostKeyChecking=no \
+            -o ServerAliveInterval=30 -o ServerAliveCountMax=20 "root@$IP" \
+            GCS_KEY="$GCS_KEY" GCS_SECRET="$GCS_SECRET" bash -s <<EOF
+          set -euxo pipefail
+          export DEBIAN_FRONTEND=noninteractive
+          apt-get update
+          apt-get install -y python3-pip
+
+          pip3 install --no-cache-dir datasets numpy faiss-cpu boto3
+
+          python3 /root/build_hf_dataset.py \
+            --hf-dataset "${{ inputs.hf_dataset }}" \
+            --output-name "${{ inputs.output_name }}" \
+            --vector-column "${{ inputs.vector_column }}" \
+            --gcs-uri "$GCS_PATH"
+          EOF
+
+      - name: Teardown server
+        if: always()
+        continue-on-error: true
+        run: bash -x tools/hetzner/remove_server.sh "$SERVER_NAME"

.github/workflows/manual-benchmarks-cascade.yaml

@@ -76,7 +76,7 @@ jobs:
           CONFIG_FILE="benchmark_cascade/benchmark-configs.json"
           BENCHMARK_SET="${{ inputs.benchmark_set }}"
           CURRENT_BATCH=${{ inputs.current_batch || 0 }}
-          BATCH_SIZE=10
+          BATCH_SIZE=12
 
           echo "benchmark_set=$BENCHMARK_SET" >> $GITHUB_OUTPUT
 
@@ -202,7 +202,7 @@ jobs:
             MACHINES_FIRST=false
 
             SERVER_NAME="benchmark-cascade-server-${i}"
-            CLIENT_NAME="benchmark-cascade-client-${i}"
+            CLIENT_NAME="benchmark-cascade-server-${i}"
             MACHINES_INFO="${MACHINES_INFO}{\"server_name\":\"$SERVER_NAME\",\"client_name\":\"$CLIENT_NAME\"}"
           done
 
@@ -252,7 +252,7 @@ jobs:
     if: inputs.process_results == false
     runs-on: ubuntu-latest
     strategy:
-      max-parallel: 10
+      max-parallel: 12
       fail-fast: false
       matrix:
         config: ${{ fromJSON(needs.prepareBenchmarks.outputs.matrix) }}
@@ -291,7 +291,7 @@ jobs:
           max_retries: 5
 
       - name: Create Client
-        if: ${{ inputs.machines_per_bench == true || inputs.current_batch == 0 }}
+        if: ${{ (inputs.machines_per_bench == true || inputs.current_batch == 0) && steps.extract_names.outputs.client_name != steps.extract_names.outputs.server_name }}
         uses: ./.github/workflows/actions/create-server-with-retry
         with:
           server_name: ${{ steps.extract_names.outputs.client_name }}
@@ -307,7 +307,7 @@ jobs:
       contents: read
       packages: read
     strategy:
-      max-parallel: 10
+      max-parallel: 12
       fail-fast: false
       matrix:
         config: ${{ fromJSON(needs.prepareBenchmarks.outputs.matrix) }}