| from __future__ import annotations |
|
|
| import json |
| import os |
| import socket |
| import statistics |
| import subprocess |
| import sys |
| import time |
| import unittest |
| from datetime import datetime, timezone |
| from urllib.request import Request, urlopen |
| from tests.real_server import RealUvicornServer |
|
|
| IMPORT_ERROR = None |
| try: |
| from backend import app |
| except ModuleNotFoundError as exc: |
| app = None |
| IMPORT_ERROR = exc |
|
|
| RESULTS_PATH = os.path.join(os.path.dirname(__file__), "benchmark-results.json") |
|
|
|
|
| def _git_sha(): |
| try: |
| repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| return subprocess.run( |
| ["git", "-c", f"safe.directory={repo}", "-C", repo, |
| "log", "-1", "--format=%H"], |
| capture_output=True, text=True, timeout=5, |
| ).stdout.strip() |
| except Exception: |
| return "" |
|
|
|
|
| def _save_results(project, queries): |
| existing = {"results": []} |
| if os.path.isfile(RESULTS_PATH): |
| try: |
| with open(RESULTS_PATH) as f: |
| existing = json.load(f) |
| except (json.JSONDecodeError, OSError): |
| pass |
| existing.setdefault("results", []).append({ |
| "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), |
| "project": project, |
| "git_sha": _git_sha(), |
| "queries": queries, |
| }) |
| existing["results"] = existing["results"][-50:] |
| with open(RESULTS_PATH, "w") as f: |
| json.dump(existing, f, indent=2, ensure_ascii=False) |
|
|
|
|
| def _print_and_collect(results_list, label, iterations, times): |
| median = statistics.median(times) |
| mn = min(times) |
| mx = max(times) |
| print(f" {label:<46s} median={median:>8.3f}ms min={mn:>8.3f}ms max={mx:>8.3f}ms (n={iterations})") |
| results_list.append({ |
| "label": label.strip(), |
| "iterations": iterations, |
| "median_ms": round(median, 3), |
| "min_ms": round(mn, 3), |
| "max_ms": round(mx, 3), |
| }) |
|
|
|
|
| def free_port(): |
| with socket.socket() as sock: |
| sock.bind(("127.0.0.1", 0)) |
| return sock.getsockname()[1] |
|
|
|
|
| |
| _server_proc = None |
| _server_port = None |
| _server = None |
|
|
|
|
| def start_server(): |
| global _server_proc, _server_port, _server |
| if _server_proc is not None: |
| return |
| _server = RealUvicornServer().start() |
| _server_proc, _server_port = _server.process, _server.port |
|
|
|
|
| def stop_server(): |
| global _server_proc, _server_port, _server |
| if _server: |
| _server.stop() |
| _server.stderr_file.close() |
| _server = _server_proc = _server_port = None |
|
|
|
|
| |
|
|
| @unittest.skipIf(app is None, f"app dependencies unavailable: {IMPORT_ERROR}") |
| class LocalSearchBenchmarkTests(unittest.TestCase): |
| _results: list = [] |
|
|
| @classmethod |
| def setUpClass(cls): |
| if app and not app.records: |
| app.load_data() |
|
|
| def _bench(self, label, fn, iterations=10): |
| times = [] |
| for _ in range(iterations): |
| start = time.perf_counter() |
| fn() |
| elapsed = (time.perf_counter() - start) * 1000 |
| times.append(elapsed) |
| _print_and_collect(self._results, label, iterations, times) |
|
|
| def test_benchmarks(self): |
| self._bench('"ζζΊ" normal', lambda: app.search("ζζΊ")) |
| self._bench('"ζζΊ" exact', lambda: app.search("ζζΊ", exact=True)) |
| self._bench('"ζει©ε½" exact', lambda: app.search("ζει©ε½", exact=True)) |
| self._bench('"ζε ι©ε½" normal', lambda: app.search("ζε ι©ε½")) |
| self._bench('"ζγζΊ" exact', lambda: app.search("ζγζΊ", exact=True)) |
| self._bench('"ζ*ζΊ" exact', lambda: app.search("ζ*ζΊ", exact=True)) |
| self._bench('"ζ?ζΊ" exact', lambda: app.search("ζ?ζΊ", exact=True)) |
| self._bench('"ζ" exact', lambda: app.search("ζ", exact=True)) |
| self._bench('"ABC" normal', lambda: app.search("ABC")) |
| self._bench('"ζζΊ" search_folders=False', |
| lambda: app.search("ζζΊ", search_folders=False)) |
| self._bench('sort by name', lambda: app.search(page=1, page_size=100, sort="name")) |
| self._bench('sort by size', lambda: app.search(page=1, page_size=100, sort="size")) |
|
|
| @classmethod |
| def tearDownClass(cls): |
| if cls._results: |
| _save_results("huggingface-Search-local", cls._results) |
|
|
|
|
| |
|
|
| class ApiSearchBenchmarkTests(unittest.TestCase): |
| _results: list = [] |
|
|
| def _bench(self, label, path, payload=None, iterations=5): |
| body = json.dumps(payload).encode() if payload else None |
| req = Request(f"http://127.0.0.1:{_server_port}{path}", data=body, |
| headers={"Content-Type": "application/json"} if body else {}) |
| times = [] |
| for i in range(iterations + 1): |
| start = time.perf_counter() |
| with urlopen(req, timeout=60) as r: |
| r.read() |
| elapsed = (time.perf_counter() - start) * 1000 |
| if i > 0: |
| times.append(elapsed) |
| _print_and_collect(self._results, label, iterations, times) |
|
|
| def test_search_api_benchmarks(self): |
| self._bench('"ζζΊ" normal', "/api/search", {"q": "ζζΊ"}) |
| self._bench('"ζζΊ" exact', "/api/search", {"q": "ζζΊ", "exact": True}) |
| self._bench('"ζει©ε½" exact', "/api/search", {"q": "ζει©ε½", "exact": True}) |
| self._bench('"ζε ι©ε½" normal', "/api/search", {"q": "ζε ι©ε½"}) |
| self._bench('"ζγζΊ" exact', "/api/search", {"q": "ζγζΊ", "exact": True}) |
| self._bench('"ζ*ζΊ" exact', "/api/search", {"q": "ζ*ζΊ", "exact": True}) |
| self._bench('"ζ?ζΊ" exact', "/api/search", {"q": "ζ?ζΊ", "exact": True}) |
| self._bench('"ζ" exact', "/api/search", {"q": "ζ", "exact": True}) |
| self._bench('"ABC" normal', "/api/search", {"q": "ABC"}) |
| self._bench('repositories', "/api/repos") |
|
|
| @classmethod |
| def tearDownClass(cls): |
| if cls._results: |
| _save_results("huggingface-Search-api", cls._results) |
|
|
|
|
| |
|
|
| try: |
| from playwright.sync_api import sync_playwright |
| except ImportError: |
| sync_playwright = None |
|
|
|
|
| @unittest.skipUnless(sync_playwright, "install requirements-test.txt, run: playwright install chromium") |
| class BrowserSearchBenchmarkTests(unittest.TestCase): |
| _results: list = [] |
|
|
| @classmethod |
| def setUpClass(cls): |
| cls.playwright = sync_playwright().start() |
| cls.browser = cls.playwright.chromium.launch( |
| headless=True, args=["--no-sandbox", "--disable-gpu", "--disable-dev-shm-usage"]) |
|
|
| @classmethod |
| def tearDownClass(cls): |
| if cls._results: |
| _save_results("huggingface-Search-browser", cls._results) |
| cls.browser.close() |
| cls.playwright.stop() |
|
|
| def _bench_search(self, label, query, iterations=3): |
| ctx = self.browser.new_context(viewport={"width": 1440, "height": 900}) |
| page = ctx.new_page() |
| times = [] |
| for i in range(iterations + 1): |
| page.goto(f"http://127.0.0.1:{_server_port}/") |
| page.wait_for_selector("#results-list .result-item", timeout=30000) |
| inp = page.locator("#search-input") |
| start = time.perf_counter() |
| inp.fill(query) |
| page.wait_for_function("""query => STATE.query === query && !STATE.isLoading && |
| displayedSearchView?.key === getSearchViewKey() |
| """, arg=query, timeout=30000) |
| elapsed = (time.perf_counter() - start) * 1000 |
| if i > 0: |
| times.append(elapsed) |
| ctx.close() |
| _print_and_collect(self._results, label, iterations, times) |
|
|
| def test_browser_search_benchmarks(self): |
| self._bench_search('"ζζΊ" normal', "ζζΊ") |
| self._bench_search('"ζει©ε½" exact', "ζει©ε½") |
| self._bench_search('"ζ*ζΊ" exact', "ζ*ζΊ") |
|
|
|
|
| def setUpModule(): |
| if os.environ.get("HF_SEARCH_RUN_BENCHMARK") != "1": |
| raise unittest.SkipTest("set HF_SEARCH_RUN_BENCHMARK=1 to run benchmarks") |
| start_server() |
|
|
|
|
| def tearDownModule(): |
| stop_server() |
|
|
