add PRODUCT_VARIANT

codelibs · Aug 17, 2024 · 355a083 · 355a083
1 parent 7757ad3
commit 355a083
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 4 deletions.
diff --git a/.github/workflows/run-elasticsearch-int8-linux.yml b/.github/workflows/run-elasticsearch-int8-linux.yml
@@ -0,0 +1,56 @@
+name: Run Elasticsearch (int8) on Linux
+
+on:
+  push:
+    branches:
+    - main
+    paths:
+    - run-elasticsearch.ipynb
+  pull_request:
+    branches:
+    - main
+  workflow_dispatch:
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Check Hardware Spec
+      run: /bin/bash ./scripts/get_hardware_info.sh
+    - name: Cache dataset
+      uses: actions/cache@v4
+      with:
+        path: |
+          dataset
+        key: ${{ runner.os }}-dataset
+    - name: Download data
+      run: /bin/bash ./scripts/setup.sh gha
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+    - name: Cache pip modules
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/pip
+        key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
+        restore-keys: |
+          ${{ runner.os }}-pip-
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    - name: Run ANN search
+      run: papermill run-elasticsearch.ipynb output-elasticsearch-linux.ipynb --log-output
+      env:
+        PRODUCT_VARIANT: "int8"
+        SETTING_QUANTIZATION: "int8"
+    - name: 'Upload Results'
+      uses: actions/upload-artifact@v4
+      with:
+        name: elasticsearch-int8-results
+        path: results.json
+
diff --git a/run-elasticsearch.ipynb b/run-elasticsearch.ipynb
@@ -71,7 +71,7 @@
     "            \"hnsw_ef_construction\": 200,\n",
     "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
-    "            \"quantization\": \"int8\",\n",
+    "            \"quantization\": \"none\",\n",
     "        },\n",
     "        \"1m-768-m48-efc200-ef100-ip\": {\n",
     "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
@@ -87,7 +87,7 @@
     "            \"hnsw_ef_construction\": 200,\n",
     "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
-    "            \"quantization\": \"int8\",\n",
+    "            \"quantization\": \"none\",\n",
     "        },\n",
     "        \"5m-768-m48-efc200-ef100-ip\": {\n",
     "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
@@ -103,10 +103,12 @@
     "            \"hnsw_ef_construction\": 200,\n",
     "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
-    "            \"quantization\": \"int8\",\n",
+    "            \"quantization\": \"none\",\n",
     "        },\n",
     "    }\n",
-    "    return DataSetConfig(**setting.get(target_name))\n"
+    "    target_setting = setting.get(target_name)\n",
+    "    target_setting[\"quantization\"] = os.getenv(\"SETTING_QUANTIZATION\", target_setting[\"quantization\"])\n",
+    "    return DataSetConfig(**target_setting)\n"
    ]
   },
   {
@@ -826,6 +828,7 @@
     "def save_results(target_config, config, results):\n",
     "    with open(\"results.json\", \"wt\", encoding=\"utf-8\") as f:\n",
     "        json.dump({\n",
+    "            \"variant\": os.getenv(\"PRODUCT_VARIANT\", \"\"),\n",
     "            \"target\": target_config,\n",
     "            \"version\": config.elasticsearch_version,\n",
     "            \"settings\": asdict(config),\n",