chore(ci): add firefox support for wasm tests and benchmarks

.github/workflows/benchmark_wasm_client.yml

@@ -78,6 +78,10 @@ jobs:
     needs: setup-instance
     if: needs.setup-instance.result != 'skipped'
     runs-on: ${{ needs.setup-instance.outputs.runner-name }}
+    strategy:
+      max-parallel: 1
+      matrix:
+        browser: [ chrome, firefox ]
     steps:
       - name: Checkout tfhe-rs repo with tags
         uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938
@@ -106,12 +110,12 @@ jobs:
       - name: Install web resources
         run: |
           make install_node
-          make install_chrome_browser
-          make install_chrome_web_driver
+          make install_${{ matrix.browser }}_browser
+          make install_${{ matrix.browser }}_driver
 
       - name: Run benchmarks
         run: |
-          make bench_web_js_api_parallel_chrome_ci
+          make bench_web_js_api_parallel_${{ matrix.browser }}_ci
 
       - name: Parse results
         run: |
@@ -124,12 +128,16 @@ jobs:
           --commit-date "${{ env.COMMIT_DATE }}" \
           --bench-date "${{ env.BENCH_DATE }}" \
           --key-gen
+          rm tfhe/wasm_pk_gen.csv
 
+      # Run these benchmarks only once
       - name: Measure public key and ciphertext sizes in HL Api
+        if:  matrix.browser == 'chrome'
         run: |
           make measure_hlapi_compact_pk_ct_sizes
 
       - name: Parse key and ciphertext sizes results
+        if:  matrix.browser == 'chrome'
         run: |
           python3 ./ci/benchmark_parser.py tfhe/hlapi_cpk_and_cctl_sizes.csv ${{ env.RESULTS_FILENAME }} \
           --key-gen \
@@ -138,7 +146,7 @@ jobs:
       - name: Upload parsed results artifact
         uses: actions/upload-artifact@50769540e7f4bd5e21e526ee35c689e35e0d6874
         with:
-          name: ${{ github.sha }}_wasm
+          name: ${{ github.sha }}_wasm_${{ matrix.browser }}
           path: ${{ env.RESULTS_FILENAME }}
 
       - name: Checkout Slab repo
@@ -160,7 +168,7 @@ jobs:
         uses: rtCamp/action-slack-notify@4e5fb42d249be6a45a298f3c9543b111b02f7907
         env:
           SLACK_COLOR: ${{ job.status }}
-          SLACK_MESSAGE: "WASM benchmarks finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"
+          SLACK_MESSAGE: "WASM benchmarks (${{ matrix.browser }}) finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"
 
   teardown-instance:
     name: Teardown instance (wasm-client-benchmarks)

Makefile

@@ -167,12 +167,13 @@ install_web_resource:
 	echo "$(checksum) $(filename)" > checksum && \
 	sha256sum -c checksum && \
 	rm checksum && \
-	unzip $(filename)
+	$(decompress_cmd) $(filename)
 
 install_chrome_browser: url = "https://storage.googleapis.com/chrome-for-testing-public/128.0.6613.137/linux64/chrome-linux64.zip"
 install_chrome_browser: checksum = "c5d7da679f3a353ae4e4420ab113de06d4bd459152f5b17558390c02d9520566"
 install_chrome_browser: dest = "$(WEB_RUNNER_DIR)/chrome"
 install_chrome_browser: filename = "chrome-linux64.zip"
+install_chrome_browser: decompress_cmd = unzip
 
 .PHONY: install_chrome_browser # Install Chrome browser for Linux
 install_chrome_browser: install_web_resource
@@ -181,10 +182,29 @@ install_chrome_web_driver: url = "https://storage.googleapis.com/chrome-for-test
 install_chrome_web_driver: checksum = "f041092f403fb7455a6da2871070b6587c32814a3e3c2b0a794d3d4aa4739151"
 install_chrome_web_driver: dest = "$(WEB_RUNNER_DIR)/chrome"
 install_chrome_web_driver: filename = "chromedriver-linux64.zip"
+install_chrome_web_driver: decompress_cmd = unzip
 
 .PHONY: install_chrome_web_driver # Install Chrome web driver for Linux
 install_chrome_web_driver: install_web_resource
 
+install_firefox_browser: url = "https://download-installer.cdn.mozilla.net/pub/firefox/releases/131.0/linux-x86_64/en-US/firefox-131.0.tar.bz2"
+install_firefox_browser: checksum = "4ca8504a62a31472ecb8c3a769d4301dd4ac692d4cc5d51b8fe2cf41e7b11106"
+install_firefox_browser: dest = "$(WEB_RUNNER_DIR)/firefox"
+install_firefox_browser: filename = "firefox-131.0.tar.bz2"
+install_firefox_browser: decompress_cmd = tar -xvf
+
+.PHONY: install_firefox_browser # Install firefox browser for Linux
+install_firefox_browser: install_web_resource
+
+install_firefox_web_driver: url = "https://github.com/mozilla/geckodriver/releases/download/v0.35.0/geckodriver-v0.35.0-linux64.tar.gz"
+install_firefox_web_driver: checksum = "ac26e9ba8f3b8ce0fbf7339b9c9020192f6dcfcbf04a2bcd2af80dfe6bb24260"
+install_firefox_web_driver: dest = "$(WEB_RUNNER_DIR)/firefox"
+install_firefox_web_driver: filename = "geckodriver-v0.35.0-linux64.tar.gz"
+install_firefox_web_driver: decompress_cmd = tar -xvf
+
+.PHONY: install_firefox_web_driver # Install firefox web driver for Linux
+install_firefox_web_driver: install_web_resource
+
 .PHONY: check_linelint_installed # Check if linelint newline linter is installed
 check_linelint_installed:
 	@printf "\n" | linelint - > /dev/null 2>&1 || \
@@ -922,16 +942,31 @@ test_web_js_api_parallel_chrome: driver_path = "$(WEB_RUNNER_DIR)/chrome/chromed
 test_web_js_api_parallel_chrome: browser_kind = chrome
 test_web_js_api_parallel_chrome: filter = Test
 
-.PHONY: test_web_js_api_parallel_chrome # Run tests for the web wasm api
+.PHONY: test_web_js_api_parallel_chrome # Run tests for the web wasm api on Chrome
 test_web_js_api_parallel_chrome: run_web_js_api_parallel
 
-.PHONY: test_web_js_api_parallel_chrome_ci # Run tests for the web wasm api
+.PHONY: test_web_js_api_parallel_chrome_ci # Run tests for the web wasm api on Chrome
 test_web_js_api_parallel_chrome_ci: setup_venv
 	source ~/.nvm/nvm.sh && \
 	nvm install $(NODE_VERSION) && \
 	nvm use $(NODE_VERSION) && \
 	$(MAKE) test_web_js_api_parallel_chrome
 
+test_web_js_api_parallel_firefox: browser_path = "$(WEB_RUNNER_DIR)/firefox/firefox/firefox"
+test_web_js_api_parallel_firefox: driver_path = "$(WEB_RUNNER_DIR)/firefox/geckodriver"
+test_web_js_api_parallel_firefox: browser_kind = firefox
+test_web_js_api_parallel_firefox: filter = Test
+
+.PHONY: test_web_js_api_parallel_firefox # Run tests for the web wasm api on Firefox
+test_web_js_api_parallel_firefox: run_web_js_api_parallel
+
+.PHONY: test_web_js_api_parallel_firefox_ci # Run tests for the web wasm api on Firefox
+test_web_js_api_parallel_firefox_ci: setup_venv
+	source ~/.nvm/nvm.sh && \
+	nvm install $(NODE_VERSION) && \
+	nvm use $(NODE_VERSION) && \
+	$(MAKE) test_web_js_api_parallel_firefox
+
 .PHONY: no_tfhe_typo # Check we did not invert the h and f in tfhe
 no_tfhe_typo:
 	@./scripts/no_tfhe_typo.sh
@@ -1099,6 +1134,21 @@ bench_web_js_api_parallel_chrome_ci: setup_venv
 	nvm use $(NODE_VERSION) && \
 	$(MAKE) bench_web_js_api_parallel_chrome
 
+bench_web_js_api_parallel_firefox: browser_path = "$(WEB_RUNNER_DIR)/firefox/firefox/firefox"
+bench_web_js_api_parallel_firefox: driver_path = "$(WEB_RUNNER_DIR)/firefox/geckodriver"
+bench_web_js_api_parallel_firefox: browser_kind = firefox
+bench_web_js_api_parallel_firefox: filter = Bench
+
+.PHONY: bench_web_js_api_parallel_firefox # Run benchmarks for the web wasm api
+bench_web_js_api_parallel_firefox: run_web_js_api_parallel
+
+.PHONY: bench_web_js_api_parallel_firefox_ci # Run benchmarks for the web wasm api
+bench_web_js_api_parallel_firefox_ci: setup_venv
+	source ~/.nvm/nvm.sh && \
+	nvm install $(NODE_VERSION) && \
+	nvm use $(NODE_VERSION) && \
+	$(MAKE) bench_web_js_api_parallel_firefox
+
 #
 # Utility tools
 #

ci/webdriver.py

@@ -22,9 +22,11 @@ import time
 from bs4 import BeautifulSoup
 from selenium import webdriver
 from selenium.common.exceptions import TimeoutException
-from selenium.webdriver.chrome.options import Options
-from selenium.webdriver.chrome.service import Service
+from selenium.webdriver.chrome.options import Options as ChromeOptions
+from selenium.webdriver.chrome.service import Service as ChromeService
 from selenium.webdriver.common.by import By
+from selenium.webdriver.firefox.options import Options as FirefoxOptions
+from selenium.webdriver.firefox.service import Service as FirefoxService
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.ui import WebDriverWait
 
@@ -107,7 +109,7 @@ class BrowserKind(enum.Enum):
     """
 
     chrome = 1
-    # firefox = 2
+    firefox = 2
 
 
 class Driver:
@@ -125,37 +127,52 @@ class Driver:
         self.browser_path = browser_path
         self.driver_path = driver_path
 
+        self._is_threaded_logs = threaded_logs
+        self._log_thread = None
+
         self.browser_kind = browser_kind
 
-        self.options = Options()
+        match self.browser_kind:
+            case BrowserKind.chrome:
+                self.options = ChromeOptions()
+            case BrowserKind.firefox:
+                self.options = FirefoxOptions()
+
         self.options.binary_location = self.browser_path
         self.options.add_argument("--headless")
         if os.getuid() == 0:
             # If user ID is root then driver needs to run in no-sandbox mode.
-            print("Script is running as root, running browser with --no-sandbox for compatibility")
+            print(
+                "Script is running as root, running browser with --no-sandbox for compatibility"
+            )
             self.options.add_argument("--no-sandbox")
 
         self._driver = None
 
         self.shutting_down = False
 
-        self._log_thread = None
-        if threaded_logs:
-            self._log_thread = threading.Thread(target=self._threaded_logs)
-
-    def set_capability(self, name, value):
-        self.options.set_capability(name, value)
-
     def get_driver(self):
         if self._driver is None:
-            driver_service = Service(self.driver_path)
 
             match self.browser_kind:
                 case BrowserKind.chrome:
+                    driver_service = ChromeService(self.driver_path)
+                    self.options.set_capability("goog:loggingPrefs", {"browser": "ALL"})
                     self._driver = webdriver.Chrome(
                         service=driver_service, options=self.options
                     )
-                # TODO: Add Firefox support
+                    if self._is_threaded_logs:
+                        self._log_thread = threading.Thread(target=self._threaded_logs)
+                case BrowserKind.firefox:
+                    driver_service = FirefoxService(self.driver_path)
+                    self.options.log.level = "trace"
+                    self.options.enable_bidi = True
+                    self._driver = webdriver.Firefox(
+                        service=driver_service, options=self.options
+                    )
+                    self._driver.script.add_console_message_handler(
+                        self._on_console_logs
+                    )
                 case _:
                     print(
                         f"{self.browser_kind.name.capitalize()} browser driver is not supported"
@@ -189,6 +206,16 @@ class Driver:
     def find_element(self, element_id):
         return self.get_driver().find_element(By.ID, element_id)
 
+    def _on_console_logs(self, log):
+        """
+        Callback used for retrieving console log using BiDi protocol reling on websocket
+        """
+        # Filter out useless message
+        if "using deprecated parameters" in log.text:
+            return
+
+        print(f"{log.level.upper()}: {log.text}")
+
     def print_log(self, log_type):
         logs = self.get_driver().get_log(log_type)
         for log in logs:
@@ -325,14 +352,18 @@ def run_case(driver, case):
     return json.loads(benchmark_results) if benchmark_results else None
 
 
-def dump_benchmark_results(results):
+def dump_benchmark_results(results, browser_kind):
     """
     Dump as JSON benchmark results into a file.
     If `results` is an empty dict then this function is a no-op.
 
     :param results: benchmark results as :class:`dict`
+    :param browser_kind: browser as :class:`BrowserKind`
     """
     if results:
+        results = {
+            "_".join((key, browser_kind.name)): val for key, val in results.items()
+        }
         pathlib.Path("tfhe/wasm_benchmark_results.json").write_text(json.dumps(results))
 
 
@@ -433,12 +464,6 @@ def main():
     driver = Driver(
         args.browser_path, args.driver_path, browser_kind, threaded_logs=True
     )
-    match browser_kind:
-        case BrowserKind.chrome:
-            driver.set_capability("goog:loggingPrefs", {"browser": "ALL"})
-        case _:
-            # A no-op for browser that are not supported
-            pass
 
     driver.get_page(f"http://{args.address}:{args.port}", timeout_seconds=10)
 
@@ -463,7 +488,7 @@ def main():
             else:
                 failures.append(case.id)
 
-    dump_benchmark_results(benchmark_results)
+    dump_benchmark_results(benchmark_results, browser_kind)
 
     # Close the browser
     driver.quit()