diff --git a/dependencies.yaml b/dependencies.yaml index bedafbb05f..45d027a116 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -249,6 +249,14 @@ files: key: opensearch includes: - bench_python_opensearch + py_postgres_cuvs_bench: + output: pyproject + pyproject_dir: python/cuvs_bench + extras: + table: project.optional-dependencies + key: postgres + includes: + - bench_python_postgres py_synthesize_dataset_cuvs_bench: output: pyproject pyproject_dir: python/cuvs_bench @@ -703,6 +711,12 @@ dependencies: - output_types: [conda, pyproject] packages: - opensearch-py>=2.4.0 + bench_python_postgres: + common: + - output_types: pyproject + packages: + - psycopg[binary]>=3.1 + - pgvector>=0.3 bench_python_synthesize_dataset: common: - output_types: pyproject diff --git a/fern/docs.yml b/fern/docs.yml index 77ee0dae2a..0b23d2d378 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -221,6 +221,8 @@ navigation: path: "./pages/cuvs_bench/synthesize_dataset.md" - page: "Backends" path: "./pages/cuvs_bench/pluggable_backend.md" + - page: "PostgreSQL Backend" + path: "./pages/cuvs_bench/postgres.md" - page: "cuVS Bench Parameter Tuning Guide" hidden: true path: "./pages/cuvs_bench/param_tuning.md" diff --git a/fern/pages/cuvs_bench/postgres.md b/fern/pages/cuvs_bench/postgres.md new file mode 100644 index 0000000000..037adb4ea4 --- /dev/null +++ b/fern/pages/cuvs_bench/postgres.md @@ -0,0 +1,256 @@ +# PostgreSQL: pgvector and pg_cuvs + +This optional backend runs vector searches through PostgreSQL using the same +corpus, zero-based row IDs, ground truth, and cuvs-bench recall calculation. + +| Algorithm | Build | Search setting | Search execution | +| --- | --- | --- | --- | +| `pgvector_hnsw` | pgvector HNSW | `ef_search` | CPU | +| `pgvector_ivfflat` | pgvector IVFFlat | `probes` | CPU | +| `pgcuvs_cagra` | pg_cuvs CAGRA | `candidates` (`cuvs.k`) | GPU sidecar | +| `pgcuvs_hnsw_import` | CAGRA plus `USING pg_cuvs_hnsw` export | `ef_search` | pgvector CPU | + +Only Euclidean (`<->`, `vector_l2_ops`) and single-client latency mode are +implemented. `batch_size` is retained in the orchestration/output interface; +it does not create SQL batching, concurrency, or GPU saturation. QPS is the +inverse of mean serial query latency. Compare PostgreSQL rows to each other +at matched recall. Native C++ cuVS rows measure a different boundary and must +not be presented as equivalent SQL application timings. + +## Install + +Use Python 3.12 or newer. From this cuVS checkout: + +```bash +python -m venv .venv +. .venv/bin/activate +python -m pip install -e './python/cuvs_bench[postgres]' -Crapidsai.disable-cuda=true +``` + +The optional extra installs psycopg 3 and pgvector's Python adapter. It does +not install a PostgreSQL server, the pgvector server extension, pg_cuvs, or +its daemon. The RAPIDS build option above avoids requiring `nvcc` on the +SQL benchmark client; GPU work runs in the PostgreSQL server's daemon. +CPU pgvector algorithms do not require pg_cuvs or a GPU on the server. + +On a PostgreSQL server with pgvector installed, create a dedicated database: + +```bash +createdb cuvs_bench +psql -X -v ON_ERROR_STOP=1 -d cuvs_bench -c 'CREATE EXTENSION vector;' +export PGDATABASE=cuvs_bench +# Remote server: set PGHOST/PGPORT/PGUSER; use .pgpass or PGSERVICE for credentials. +``` + +The role needs CREATE privilege on this database. CAGRA measurements also need +SELECT privilege on `pg_stat_gpu_search` for the VRAM and GPU-dispatch checks. +An administrator must grant this to an ordinary benchmark role, for example +`GRANT SELECT ON pg_stat_gpu_search TO benchmark_role;` (replace the role name). +Database CREATE privilege alone does not grant access to this view. The GPU +verification below used the bootstrap's administrative role. +Every build loads its own +corpus into a fresh `cuvs_bench_` schema. Normal cleanup drops only that +schema, including its indexes. It does not replace public tables or existing +indexes. Abrupt client termination can leave a schema behind; inspect those +schemas before removing abandoned benchmark data. Run build and search in the +same orchestrator call. Separate-process search-only and persistent index +reuse are not supported; `force` does not change that lifetime. +The optional `index_dir` connection override sets a PostgreSQL `PGC_SUSET` +parameter and needs the corresponding administrative privilege. Omit it to +use the server's configured directory when connecting as an ordinary role. + +## GPU server setup and build + +The pgcuvs algorithms require PostgreSQL 16, pgvector >= 0.7.0, pg_cuvs >= 0.8.0, +CUDA 12.4 and cuVS 26.04. PostgreSQL must preload pg_cuvs and its GPU daemon +must be running. The client may connect remotely and needs no GPU. + +On a fresh Ubuntu 22.04 A100 host with a working NVIDIA driver, systemd and +passwordless sudo, the public pg_cuvs bootstrap installs the toolchain, builds +the extension and daemon, configures PostgreSQL 16 and downloads wiki_all_1M. +It changes system packages and PostgreSQL configuration; use a dedicated host. +The fresh-host setup below was exercised at this pinned public source revision; +the verification environment and GPU test scope are recorded below: + +```bash +git clone https://github.com/pg-cuvs/pg_cuvs.git "$HOME/pg_cuvs" +git -C "$HOME/pg_cuvs" checkout 8b73d65fb8a3d6f66cdfe69eec4969c1f2fc753c +bash "$HOME/pg_cuvs/infra/brev/bootstrap.sh" +# Keep the server build toolchain out of the benchmark client's Python PATH. +( +. "$HOME/.pg_cuvs_env" +cd "$HOME/pg_cuvs" +# Explicit rebuild commands; bootstrap already performs these steps. +make -j2 +sudo -E env PATH="$PATH" CONDA_PREFIX="$CONDA_PREFIX" make install install-server +) +systemctl is-active postgresql@16-main pg-cuvs-server +createdb cuvs_bench +psql -X -v ON_ERROR_STOP=1 -d cuvs_bench \ + -c 'CREATE EXTENSION IF NOT EXISTS vector;' \ + -c 'CREATE EXTENSION pg_cuvs;' +psql -X -d cuvs_bench -c 'SHOW shared_preload_libraries;' +nvidia-smi +``` + +The bootstrap's [public source and assumptions](https://github.com/pg-cuvs/pg_cuvs/blob/8b73d65fb8a3d6f66cdfe69eec4969c1f2fc753c/infra/brev/bootstrap.sh) +are the authority for server setup. Use the same `cuvs.index_dir` as the daemon +(`index_dir` in backend options overrides it for the connection). Never add a +Conda environment's entire library directory to system ldconfig. This backend +does not change system configuration or restart the daemon. + +## Obtain data and run + +From the cuVS checkout, download the standard million-row binary dataset: + +```bash +. .venv/bin/activate +mkdir -p datasets/wiki_all_1M +curl -fL https://data.rapids.ai/raft/datasets/wiki_all_1M/wiki_all_1M.tar \ + -o datasets/wiki_all_1M/wiki_all_1M.tar +tar -xf datasets/wiki_all_1M/wiki_all_1M.tar -C datasets/wiki_all_1M +export PGDATABASE=cuvs_bench +``` + +The supplied ground truth describes the full corpus. Using `subset_size` +without recomputing ground truth invalidates recall. Run the small parameter +group first (it still loads the complete corpus): + +```python +from cuvs_bench.orchestrator import BenchmarkOrchestrator +from cuvs_bench.run.data_export import write_results_to_csv + +results = BenchmarkOrchestrator(backend_type="postgres").run_benchmark( + dataset="wiki_all_1M", + dataset_path="datasets", + algorithms="pgvector_hnsw,pgvector_ivfflat", + groups="test", + count=10, + search_mode="latency", +) +for result in results: + if not result.success: + raise RuntimeError(result.error_message) +write_results_to_csv(results, "wiki_all_1M", "datasets", 10, 10000) +``` + +After GPU server setup, set `algorithms` to +`"pgvector_hnsw,pgvector_ivfflat,pgcuvs_cagra,pgcuvs_hnsw_import"`. +Set `groups="base"` for the larger build/search parameter grids. For a custom +sweep, copy YAML from `python/cuvs_bench/cuvs_bench/config/algos/`, edit the +`build` and `search` lists, and pass +`algorithm_configuration="/absolute/path/to/custom.yaml"`. +One index is built per build cell and reused for every search point in that +cell. Keep `ef_search` or `candidates` >= `count`. The HNSW export uses +`USING pg_cuvs_hnsw`, rather than deprecated `pg_cuvs_build_hnsw()`. + +CSV files appear under `datasets/wiki_all_1M/result/{build,search}/`. Standard +cuvs-bench plotting tools can consume them. Retain every build cell and search +point when selecting matched-recall Pareto points. Apply a recall threshold +before selecting the fastest point. + +### A bounded wiki_all_1M CLI validation sweep + +The following grid uses one build cell and two search settings per algorithm. +It searches all 10,000 supplied queries, with `k=10`. The resource settings +below were used on a 14-vCPU, 100-GiB RAM A100 host. They are session settings, +not changes to the PostgreSQL configuration file. + +```bash +mkdir -p pg-validation-grids +python - <<'PY' +from pathlib import Path +import yaml + +graph = {"graph_degree": [32], "intermediate_graph_degree": [128], + "build_algo": ["ivf_pq"]} +grids = { + "pgcuvs_cagra": (graph, {"candidates": [64, 200]}), + "pgcuvs_hnsw_import": ({**graph, "mode": ["nsw"]}, {"ef_search": [64, 200]}), + "pgvector_hnsw": ({"m": [16], "ef_construction": [64]}, {"ef_search": [64, 200]}), + "pgvector_ivfflat": ({"lists": [1000]}, {"probes": [4, 16]}), +} +for algo, (build, search) in grids.items(): + Path(f"pg-validation-grids/{algo}.yaml").write_text(yaml.safe_dump( + {"name": algo, "groups": {"validation": {"build": build, "search": search}}} + )) +Path("pg-backend.yaml").write_text("backend: postgres\nwarmup_queries: 10\n") +PY +export PGDATABASE=cuvs_bench +export PGOPTIONS='-c maintenance_work_mem=8GB -c max_parallel_maintenance_workers=7' +for algo in pgcuvs_cagra pgcuvs_hnsw_import pgvector_hnsw pgvector_ivfflat; do + python -m cuvs_bench.run --backend-config pg-backend.yaml \ + --dataset wiki_all_1M --dataset-path datasets \ + --configuration pg-validation-grids --algorithms "$algo" \ + --groups validation -k 10 --batch-size 10000 --search-mode latency --force +done +``` + +Check for one build row and two raw search rows per algorithm. CAGRA search +rows must include `observed_search_mode=gpu_cagra` and `gpu_queries=10000`. +Build CSVs preserve `index_size_bytes` and `index_size_source`. These runs +validate execution and serialization; they do not establish a speedup over +another engine, machine, or measurement boundary. +To exercise the GPU build axis as well, set `graph_degree` to `[32, 64]` in +the two GPU grids. Each then produces two build rows and four raw search rows. + +## Measurement contract + +- File reading, binary COPY and ANALYZE are outside the build timer. +- Build time covers CREATE INDEX. HNSW export includes the CAGRA build and + conversion. YAML records the source graph and export parameters. +- Query time includes client adaptation, SQL planning/execution, IPC/GPU work + for CAGRA, heap fetch and result transfer. Warmup and EXPLAIN are excluded. + The backend checks for an Index Scan before and after each search point; a + sequential-scan fallback becomes a failed result. CAGRA connections explicitly + disable CPU HNSW fallback. Outside the timer, daemon observations must report + resident `gpu_cagra` mode and exactly one additional GPU search per measured + query. Missing/reset counters, CPU mode or an unavailable daemon fail the + point instead of producing a GPU-labeled measurement. +- IDs are corpus offsets. Missing neighbors stay `-1` with infinite distance. + The orchestrator calculates recall using actual IDs and ground truth. +- CPU index size is `pg_relation_size`. CAGRA size comes from the daemon's + `pg_stat_gpu_search.vram_bytes`; missing accounting fails the build. Disk + bytes and daemon VRAM bytes are different resources, labeled in metadata. + +## Validate + +CPU contracts require no database. Live tests need a disposable database with +the vector extension installed: + +```bash +python -m pip install pytest +python -m pytest python/cuvs_bench/cuvs_bench/tests/test_postgres.py -m 'not postgres' +CUVS_BENCH_POSTGRES_DSN='dbname=cuvs_bench' \ + python -m pytest python/cuvs_bench/cuvs_bench/tests/test_postgres.py +CUVS_BENCH_POSTGRES_GPU_DSN='dbname=cuvs_bench' \ + python -m pytest python/cuvs_bench/cuvs_bench/tests/test_postgres.py -m postgres_gpu +``` + +Set the GPU DSN only for a server built with real CUDA/cuVS, not a CPU +reference implementation. The opt-in tests build CAGRA and export HNSW from +a 4096-by-32 corpus, check real GPU mode/counters, and check that CPU HNSW +search does not increment its source CAGRA's search counter. + +Fresh-host verification on 2026-10-16 used Ubuntu 22.04.5, A100-SXM4-80GB, +driver 580.126.09, PostgreSQL 16.15, pgvector 0.8.7, pg_cuvs 0.8.0 at the source +revision above, and CUDA 12.4/cuVS 26.04. The client was installed from this +cuvs-bench source with Python 3.12.15. Both real GPU cases passed, along with +the 28 CPU and telemetry-contract cases. CPU-only test success by itself is +not evidence that GPU kernels ran. + +That initial GPU run predates the client-side connection-error, group-selection, +thread-option, and test-marker fixes. Those later changes were checked with CPU +integration and wheel tests; GPU execution of the final revision has not been +repeated. + +Always provide ground truth for recall comparisons. Without it, the current +result interface leaves recall at zero rather than marking it unmeasured. +Python API dry-run search results must not be passed to the generic CSV +exporter as measurements: it currently retains zero-valued skipped rows. +The CLI skips export for dry runs. + +This backend adapts the public +[pg_cuvs prototype](https://github.com/pg-cuvs/pg_cuvs/tree/53a8ca0ff82f3bb2073bcd4c7ee8a7d8dc49cdd0/bench/cuvs_bench_backend) +to current cuvs-bench interfaces. It changes no GPU kernels and makes no +performance improvement claim. diff --git a/python/cuvs_bench/cuvs_bench/backends/postgres.py b/python/cuvs_bench/cuvs_bench/backends/postgres.py new file mode 100644 index 0000000000..23c92e8d4a --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/postgres.py @@ -0,0 +1,544 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026, JAESOL SHIN. + +"""PostgreSQL pgvector and pg_cuvs benchmarks with end-to-end SQL timing. + +The server and extensions are managed externally. Each build owns a fresh +schema, loads zero-based corpus IDs, and builds one index before sweeping all +search parameters. cleanup() removes only that schema, including GPU indexes. +""" + +import json +import time +import uuid +from pathlib import Path + +import numpy as np + +from .base import BenchmarkBackend, BuildResult, SearchResult +from ..orchestrator.config_loaders import ( + BenchmarkConfig, + ConfigLoader, + IndexConfig, +) + +_BUILD_KEYS = { + "pgvector_hnsw": {"m", "ef_construction"}, + "pgvector_ivfflat": {"lists"}, + "pgcuvs_cagra": { + "graph_degree", + "intermediate_graph_degree", + "build_algo", + }, + "pgcuvs_hnsw_import": { + "graph_degree", + "intermediate_graph_degree", + "build_algo", + "mode", + }, +} +_SEARCH_GUC = { + "pgvector_hnsw": ("ef_search", "hnsw.ef_search"), + "pgvector_ivfflat": ("probes", "ivfflat.probes"), + "pgcuvs_cagra": ("candidates", "cuvs.k"), + "pgcuvs_hnsw_import": ("ef_search", "hnsw.ef_search"), +} + + +def _positive_int(value): + if isinstance(value, bool) or not isinstance(value, (int, np.integer)): + raise ValueError("Index parameters must be positive integers") + if value <= 0: + raise ValueError("Index parameters must be positive integers") + return int(value) + + +def _build_params(algo, params): + if algo not in _BUILD_KEYS: + raise ValueError(f"Unsupported PostgreSQL algorithm: {algo}") + if set(params) - _BUILD_KEYS[algo]: + raise ValueError(f"Unsupported build parameters for {algo}") + validated = {} + for key, value in params.items(): + if key == "build_algo": + if value not in ("ivf_pq", "nn_descent", "auto"): + raise ValueError("Unsupported CAGRA build_algo") + elif key == "mode": + if value not in ("nsw", "hnsw", "hnswlib", "hnswlib_file"): + raise ValueError("Unsupported pg_cuvs HNSW mode") + else: + value = _positive_int(value) + validated[key] = value + if validated.get("intermediate_graph_degree", 128) < validated.get( + "graph_degree", 64 + ): + raise ValueError("intermediate_graph_degree must be >= graph_degree") + return validated + + +class PostgresConfigLoader(ConfigLoader): + """Use standard dataset YAML and Cartesian build/search parameter grids.""" + + def __init__(self, config_path=None): + self.config_path = str( + config_path or Path(__file__).parent.parent / "config" + ) + + @property + def backend_type(self): + return "postgres" + + def _discover_algo_groups( + self, dataset_conf, dataset, dataset_path, **kwargs + ): + algorithms = kwargs.get("algorithms") + allowed = ( + set(algorithms.split(",")) if algorithms else set(_BUILD_KEYS) + ) + if allowed - _BUILD_KEYS.keys(): + raise ValueError(f"Unsupported PostgreSQL algorithm: {algorithms}") + groups = kwargs.get("groups") + selected_groups = set(groups.split(",")) if groups else None + algo_groups = kwargs.get("algo_groups") + explicit_groups = ( + {tuple(item.split(".")) for item in algo_groups.split(",")} + if algo_groups + else set() + ) + if any(len(pair) != 2 or not all(pair) for pair in explicit_groups): + raise ValueError("Expected algo_groups as algorithm.group pairs") + result = {} + for path in self.gather_algorithm_configs( + self.config_path, kwargs.get("algorithm_configuration") + ): + cfg = self.load_yaml_file(path) + algo = cfg.get("name") + if algo not in allowed: + continue + for group, params in cfg.get("groups", {}).items(): + if ( + selected_groups is None + or group in selected_groups + or (algo, group) in explicit_groups + ): + result[algo, group] = (algo, group, params, {}) + if explicit_groups - result.keys(): + raise ValueError("Requested PostgreSQL algo_groups not found") + if not result: + raise ValueError("No PostgreSQL algorithm configuration selected") + return list(result.values()) + + def _build_benchmark_configs( + self, + dataset_config, + dataset_conf, + dataset, + dataset_path, + expanded_groups, + **kwargs, + ): + configs = [] + for algo, group, _, builds, searches, _ in expanded_groups: + for build in builds: + build = _build_params(algo, build) + label = ".".join( + [algo, group] + + [f"{k}{v}" for k, v in sorted(build.items())] + ) + idx = IndexConfig( + name=label, + algo=algo, + build_param=build, + search_params=searches, + file=str(Path(dataset_path) / dataset / "index" / label), + ) + conn_cfg = { + key: kwargs[key] + for key in ("conninfo", "index_dir", "warmup_queries") + if key in kwargs + } + configs.append( + BenchmarkConfig( + indexes=[idx], + backend_config={ + "name": label, + "algo": algo, + "group": group, + **conn_cfg, + }, + ) + ) + return configs + + +class PostgresBackend(BenchmarkBackend): + """Run single-client SQL latency sweeps in a backend-owned schema. + + Requires a database with vector installed; pgcuvs_* also need pg_cuvs + preloaded and its daemon running. Standard libpq PG* variables and service + files are supported, so credentials need not appear in benchmark YAML. + """ + + def __init__(self, config): + super().__init__(config) + self._conn = None + self._schema = None + self._index = None + + @property + def algo(self): + return self.config["algo"] + + def _connect(self): + if self._conn is not None: + return self._conn + try: + import psycopg + from pgvector.psycopg import register_vector + except ImportError as exc: + raise ImportError( + "PostgreSQL backend needs pip install 'cuvs-bench[postgres]'" + ) from exc + try: + self._conn = psycopg.connect( + self.config.get("conninfo", ""), + autocommit=True, + prepare_threshold=None, + ) + except Exception: + # libpq parsing errors can echo malformed credential tokens. + raise RuntimeError( + "PostgreSQL connection failed; check libpq connection settings" + ) from None + installed = { + row[0] + for row in self._conn.execute( + "SELECT extname FROM pg_extension" + ).fetchall() + } + required = {"vector"} + if self.algo.startswith("pgcuvs_"): + required.add("pg_cuvs") + if required - installed: + raise RuntimeError( + f"Install database extensions: {required - installed}" + ) + register_vector(self._conn) + if "pg_cuvs" in installed: + self._conn.execute( + "SELECT set_config('enable_cuvs', %s, false)", + ("on" if self.algo == "pgcuvs_cagra" else "off",), + ) + if self.algo == "pgcuvs_cagra": + self._conn.execute("SET cuvs.search_mode = 'cagra'") + self._conn.execute("SET cuvs.cpu_hnsw_fallback = off") + if "index_dir" in self.config: + self._conn.execute( + "SELECT set_config('cuvs.index_dir', %s, false)", + (self.config["index_dir"],), + ) + return self._conn + + def _metadata(self, idx): + return { + "group": self.config.get("group", "base"), + "index_name": idx.file, + "measurement_boundary": "end_to_end_sql", + "axis": "single_client_latency", + "build_params": json.dumps(idx.build_param, sort_keys=True), + } + + def _create_index(self, name, method, params): + from psycopg import sql + + statement = sql.SQL( + "CREATE INDEX {} ON {}.vectors USING {} (embedding vector_l2_ops)" + ).format( + sql.Identifier(name), + sql.Identifier(self._schema), + sql.Identifier(method), + ) + if params: + statement += sql.SQL(" WITH ({})").format( + sql.SQL(", ").join( + sql.SQL("{} = {}").format( + sql.Identifier(k), sql.Literal(v) + ) + for k, v in params.items() + ) + ) + self._conn.execute(statement) + + def build(self, dataset, indexes, force=False, dry_run=False): + """COPY the corpus, then time only CREATE INDEX (including GPU export).""" + idx = indexes[0] + result = BuildResult( + index_path=idx.file, + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params=dict(idx.build_param), + metadata=self._metadata(idx), + ) + try: + if len(indexes) != 1 or idx.algo != self.algo: + raise ValueError( + "Expected one index matching the backend algorithm" + ) + params = _build_params(self.algo, idx.build_param) + if dataset.distance_metric not in ("euclidean", "l2"): + raise ValueError( + "PostgreSQL backend currently supports euclidean only" + ) + if dry_run: + result.metadata.update(dry_run=True, skipped=True) + return result + if self._schema is not None: + raise RuntimeError( + "Use a fresh backend for each build configuration" + ) + conn = self._connect() + from psycopg import sql + + vectors = dataset.training_vectors + if vectors.ndim != 2 or not len(vectors) or not vectors.shape[1]: + raise ValueError("Expected a nonempty two-dimensional corpus") + schema = "cuvs_bench_" + uuid.uuid4().hex + conn.execute( + sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)) + ) + self._schema = schema + conn.execute( + sql.SQL( + "CREATE TABLE {}.vectors (id bigint, embedding vector({}))" + ).format(sql.Identifier(schema), sql.Literal(vectors.shape[1])) + ) + with conn.cursor().copy( + sql.SQL( + "COPY {}.vectors (id, embedding) FROM STDIN (FORMAT BINARY)" + ).format(sql.Identifier(schema)) + ) as copy: + copy.set_types(["int8", "vector"]) + for i, vector in enumerate(vectors): + copy.write_row((i, vector)) + conn.execute( + sql.SQL("ANALYZE {}.vectors").format(sql.Identifier(schema)) + ) + started = time.perf_counter() + if self.algo in ("pgvector_hnsw", "pgvector_ivfflat"): + self._create_index( + "ann", self.algo.removeprefix("pgvector_"), params + ) + else: + mode = params.pop("mode", "nsw") + source = ( + "ann" if self.algo == "pgcuvs_cagra" else "source_graph" + ) + self._create_index(source, "cagra", params) + if self.algo == "pgcuvs_hnsw_import": + self._create_index( + "ann", + "pg_cuvs_hnsw", + { + "source": f"{schema}.source_graph", + "mode": mode, + }, + ) + result.build_time_seconds = time.perf_counter() - started + self._index = idx + relation = f"{schema}.ann" + if self.algo == "pgcuvs_cagra": + row = conn.execute( + "SELECT vram_bytes FROM pg_stat_gpu_search " + "WHERE index_oid = %s::regclass", + (relation,), + ).fetchone() + if row is None or row[0] is None: + raise RuntimeError( + "No resident CAGRA VRAM accounting available" + ) + result.metadata["index_size_source"] = "daemon_vram_bytes" + else: + row = conn.execute( + "SELECT pg_relation_size(%s::regclass)", (relation,) + ).fetchone() + result.metadata["index_size_source"] = "pg_relation_size" + result.index_size_bytes = int(row[0]) + result.metadata["index_size_bytes"] = result.index_size_bytes + except Exception as exc: + result.success = False + result.error_message = str(exc) + return result + + def _query(self, vector, k): + from psycopg import sql + + return sql.SQL( + "SELECT id, embedding <-> %s AS distance FROM {}.vectors " + "ORDER BY distance LIMIT %s" + ).format(sql.Identifier(self._schema)), (vector, k) + + def _check_plan(self, statement, params): + from psycopg import sql + + plan = self._conn.execute( + sql.SQL("EXPLAIN (FORMAT JSON) ") + statement, params + ).fetchone()[0][0]["Plan"] + + def uses_index(node): + return ( + node.get("Node Type") == "Index Scan" + and node.get("Index Name") == "ann" + ) or any(uses_index(child) for child in node.get("Plans", [])) + + if not uses_index(plan): + raise RuntimeError( + "Expected Index Scan on ann; refusing fallback timing" + ) + + def _gpu_search_count(self): + row = self._conn.execute( + "SELECT resident, search_mode, search_count " + "FROM pg_stat_gpu_search WHERE index_oid = %s::regclass", + (f"{self._schema}.ann",), + ).fetchone() + if row is None or not row[0] or row[1] != "gpu_cagra": + raise RuntimeError("Expected resident GPU CAGRA execution") + return int(row[2]) + + def search( + self, + dataset, + indexes, + k, + batch_size=10000, + mode="latency", + force=False, + search_threads=None, + dry_run=False, + ): + """Return one result per search setting; the orchestrator computes recall.""" + idx = indexes[0] + results = [] + for search in idx.search_params: + result = SearchResult( + neighbors=np.empty((0, 0), dtype=np.int64), + distances=np.empty((0, 0), dtype=np.float32), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=[dict(search)], + metadata=self._metadata(idx), + ) + results.append(result) + try: + if mode != "latency" or search_threads not in (None, 1, "1"): + raise ValueError( + "PostgreSQL supports single-client latency only" + ) + key, guc = _SEARCH_GUC[self.algo] + if set(search) != {key}: + raise ValueError(f"Expected search parameter {key}") + value, k = _positive_int(search[key]), _positive_int(k) + if key != "probes" and value < k: + raise ValueError(f"{key} must be >= k") + if dry_run: + result.metadata.update(dry_run=True, skipped=True) + continue + if self._index != idx: + raise RuntimeError("Build this index before searching") + queries = dataset.query_vectors + if not len(queries): + raise ValueError("Expected at least one query") + conn = self._conn + conn.execute("SET enable_seqscan = off") + conn.execute( + "SELECT set_config(%s, %s, false)", (guc, str(value)) + ) + self._check_plan(*self._query(queries[0], k)) + for query in queries[ + : int(self.config.get("warmup_queries", 10)) + ]: + conn.execute(*self._query(query, k)).fetchall() + gpu_before = ( + self._gpu_search_count() + if self.algo == "pgcuvs_cagra" + else None + ) + ids = np.full((len(queries), k), -1, dtype=np.int64) + distances = np.full( + (len(queries), k), np.inf, dtype=np.float32 + ) + latencies = [] + for i, query in enumerate(queries): + statement, bindings = self._query(query, k) + started = time.perf_counter() + rows = conn.execute(statement, bindings).fetchall() + latencies.append(time.perf_counter() - started) + for j, (neighbor, distance) in enumerate(rows): + ids[i, j], distances[i, j] = neighbor, distance + self._check_plan(*self._query(queries[0], k)) + if gpu_before is not None: + gpu_queries = self._gpu_search_count() - gpu_before + if gpu_queries != len(queries): + raise RuntimeError( + "GPU query count does not match measured queries" + ) + result.metadata.update( + observed_search_mode="gpu_cagra", + gpu_queries=gpu_queries, + ) + total = sum(latencies) + result.neighbors, result.distances = ids, distances + result.search_time_ms = total * 1000 + result.queries_per_second = len(queries) / total + result.latency_percentiles = dict( + zip( + ("p50_ms", "p95_ms", "p99_ms"), + np.percentile(latencies, [50, 95, 99]) * 1000, + ) + ) + result.metadata.update( + n_queries=len(queries), + latency_seconds=total / len(queries), + ) + except Exception as exc: + result.success = False + result.error_message = str(exc) + return results + + def cleanup(self): + """Remove only this backend's schema and always close its connection.""" + if self._conn is None: + return + try: + if self._schema is not None: + from psycopg import sql + + self._conn.execute( + sql.SQL("DROP SCHEMA {} CASCADE").format( + sql.Identifier(self._schema) + ) + ) + finally: + self._conn.close() + self._conn = None + self._schema = None + self._index = None + + +def register(): + """Register the optional backend without importing database dependencies.""" + from .registry import ( + _CONFIG_LOADER_REGISTRY, + get_registry, + register_backend, + register_config_loader, + ) + + if not get_registry().is_registered("postgres"): + register_backend("postgres", PostgresBackend) + if "postgres" not in _CONFIG_LOADER_REGISTRY: + register_config_loader("postgres", PostgresConfigLoader) diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pgcuvs_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pgcuvs_cagra.yaml new file mode 100644 index 0000000000..e2395d4165 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pgcuvs_cagra.yaml @@ -0,0 +1,17 @@ +# PostgreSQL pg_cuvs GPU build and resident GPU search. +name: pgcuvs_cagra +groups: + base: + build: + graph_degree: [32, 64] + intermediate_graph_degree: [128] + build_algo: [ivf_pq] + search: + candidates: [16, 32, 64, 100, 200, 400] + test: + build: + graph_degree: [32] + intermediate_graph_degree: [128] + build_algo: [ivf_pq] + search: + candidates: [16, 64] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pgcuvs_hnsw_import.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pgcuvs_hnsw_import.yaml new file mode 100644 index 0000000000..a8bf17230c --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pgcuvs_hnsw_import.yaml @@ -0,0 +1,19 @@ +# PostgreSQL pg_cuvs GPU CAGRA build plus export, then pgvector CPU search. +name: pgcuvs_hnsw_import +groups: + base: + build: + graph_degree: [32, 64] + intermediate_graph_degree: [128] + build_algo: [ivf_pq] + mode: [nsw] + search: + ef_search: [16, 32, 64, 128, 256, 512] + test: + build: + graph_degree: [32] + intermediate_graph_degree: [128] + build_algo: [ivf_pq] + mode: [nsw] + search: + ef_search: [16, 64] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pgvector_hnsw.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pgvector_hnsw.yaml new file mode 100644 index 0000000000..c898d67d92 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pgvector_hnsw.yaml @@ -0,0 +1,15 @@ +# PostgreSQL pgvector CPU build and CPU search. +name: pgvector_hnsw +groups: + base: + build: + m: [16, 32] + ef_construction: [64, 128] + search: + ef_search: [10, 20, 40, 80, 120, 200, 400] + test: + build: + m: [16] + ef_construction: [64] + search: + ef_search: [10, 40] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pgvector_ivfflat.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pgvector_ivfflat.yaml new file mode 100644 index 0000000000..0c7bfdb249 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pgvector_ivfflat.yaml @@ -0,0 +1,13 @@ +# PostgreSQL pgvector CPU build and CPU search. +name: pgvector_ivfflat +groups: + base: + build: + lists: [100, 1000, 4000] + search: + probes: [1, 4, 8, 16, 32, 64, 100] + test: + build: + lists: [4] + search: + probes: [1, 4] diff --git a/python/cuvs_bench/cuvs_bench/tests/test_postgres.py b/python/cuvs_bench/cuvs_bench/tests/test_postgres.py new file mode 100644 index 0000000000..a1fc418dcf --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_postgres.py @@ -0,0 +1,633 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026, JAESOL SHIN. + +"""PostgreSQL backend contracts and opt-in live database tests.""" + +import os +import csv +import socket +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import yaml + +from cuvs_bench.backends.base import Dataset, SearchResult +from cuvs_bench.backends.postgres import ( + PostgresBackend, + PostgresConfigLoader, + register, +) +from cuvs_bench.backends.registry import get_backend_class, get_config_loader +from cuvs_bench.orchestrator.config_loaders import BenchmarkConfig, IndexConfig +from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator + + +@pytest.fixture +def config_dir(tmp_path): + (tmp_path / "datasets").mkdir() + (tmp_path / "datasets" / "datasets.yaml").write_text( + yaml.safe_dump([{"name": "small", "dims": 4}]) + ) + (tmp_path / "algos").mkdir() + (tmp_path / "algos" / "pgvector_hnsw.yaml").write_text( + yaml.safe_dump( + { + "name": "pgvector_hnsw", + "groups": { + "test": { + "build": {"m": [8, 16], "ef_construction": [64]}, + "search": {"ef_search": [10, 40]}, + } + }, + } + ) + ) + return tmp_path + + +def dataset(): + rng = np.random.default_rng(7) + base = rng.normal(size=(128, 4)).astype(np.float32) + queries = base[:5].copy() + dists = ((queries[:, None] - base[None]) ** 2).sum(axis=2) + return Dataset( + "small", + training_vectors=base, + query_vectors=queries, + groundtruth_neighbors=np.argsort(dists, axis=1)[:, :5], + ) + + +def index(algo="pgvector_hnsw", build=None, search=None): + return IndexConfig( + name=algo, + algo=algo, + build_param=build or {"m": 8, "ef_construction": 64}, + search_params=search or [{"ef_search": 10}, {"ef_search": 40}], + file="unused-index-label", + ) + + +def test_registration_is_idempotent(): + register() + register() + assert get_backend_class("postgres") is PostgresBackend + assert get_config_loader("postgres") is PostgresConfigLoader + + +def test_loader_keeps_search_sweep_on_one_build(config_dir): + loader = PostgresConfigLoader(config_dir) + ds, configs = loader.load( + "small", str(config_dir), algorithms="pgvector_hnsw", groups="test" + ) + assert ds.dims == 4 + assert len(configs) == 2 + assert len({cfg.index_name for cfg in configs}) == 2 + assert {cfg.build_params["m"] for cfg in configs} == {8, 16} + assert all( + cfg.search_params_list == [{"ef_search": 10}, {"ef_search": 40}] + for cfg in configs + ) + + +def test_loader_rejects_unknown_algorithm(config_dir): + with pytest.raises(ValueError, match="algorithm"): + PostgresConfigLoader(config_dir).load( + "small", str(config_dir), algorithms="pgvector_typo" + ) + + +def test_custom_group_replaces_default_search_grid(config_dir, tmp_path): + override = tmp_path / "override.yaml" + override.write_text( + yaml.safe_dump( + { + "name": "pgvector_hnsw", + "groups": { + "test": { + "build": {"m": [8], "ef_construction": [64]}, + "search": {"ef_search": [80]}, + } + }, + } + ) + ) + _, configs = PostgresConfigLoader(config_dir).load( + "small", + str(config_dir), + algorithms="pgvector_hnsw", + groups="test", + algorithm_configuration=str(override), + ) + assert len(configs) == 1 + assert configs[0].build_params == {"m": 8, "ef_construction": 64} + assert configs[0].search_params_list == [{"ef_search": 80}] + + +def test_loader_rejects_unknown_group(config_dir): + with pytest.raises(ValueError, match="configuration"): + PostgresConfigLoader(config_dir).load( + "small", str(config_dir), groups="missing" + ) + + +def test_explicit_algo_group_adds_to_selected_groups(config_dir): + path = config_dir / "algos" / "pgvector_hnsw.yaml" + config = yaml.safe_load(path.read_text()) + config["groups"]["base"] = { + "build": {"m": [32], "ef_construction": [64]}, + "search": {"ef_search": [80]}, + } + path.write_text(yaml.safe_dump(config)) + _, configs = PostgresConfigLoader(config_dir).load( + "small", + str(config_dir), + algorithms="pgvector_hnsw", + groups="base", + algo_groups="pgvector_hnsw.test,pgvector_hnsw.test", + ) + assert len(configs) == 3 + assert {cfg.build_params["m"] for cfg in configs} == {8, 16, 32} + assert sum(cfg.backend_config["group"] == "test" for cfg in configs) == 2 + + +def test_explicit_algo_group_selects_only_its_algorithm(config_dir): + (config_dir / "algos" / "pgvector_ivfflat.yaml").write_text( + yaml.safe_dump( + { + "name": "pgvector_ivfflat", + "groups": { + "base": { + "build": {"lists": [16]}, + "search": {"probes": [4]}, + }, + "test": { + "build": {"lists": [4]}, + "search": {"probes": [1]}, + }, + }, + } + ) + ) + _, configs = PostgresConfigLoader(config_dir).load( + "small", + str(config_dir), + algorithms="pgvector_hnsw,pgvector_ivfflat", + groups="base", + algo_groups="pgvector_hnsw.test", + ) + assert len(configs) == 3 + assert { + (cfg.backend_config["algo"], cfg.backend_config["group"]) + for cfg in configs + } == {("pgvector_hnsw", "test"), ("pgvector_ivfflat", "base")} + + +@pytest.mark.parametrize( + "selector", + ["pgvector_hnsw.missing", "pgvector_typo.test", "pgvector_hnsw", ".test"], +) +def test_loader_rejects_unmatched_explicit_groups(config_dir, selector): + with pytest.raises(ValueError, match="algo_groups"): + PostgresConfigLoader(config_dir).load( + "small", + str(config_dir), + algorithms="pgvector_hnsw", + groups="test", + algo_groups=selector, + ) + + +def test_dry_run_does_not_connect(monkeypatch): + backend = PostgresBackend({"algo": "pgvector_hnsw"}) + monkeypatch.setattr( + backend, "_connect", lambda: pytest.fail("dry run connected") + ) + ds, idx = dataset(), index() + assert backend.build(ds, [idx], dry_run=True).success + results = backend.search(ds, [idx], 5, dry_run=True) + assert len(results) == 2 + assert all(result.metadata["dry_run"] for result in results) + backend.cleanup() + + +@pytest.mark.parametrize("threads", [None, 1, "1"]) +def test_single_client_thread_setting_accepts_cli_string(threads): + backend = PostgresBackend({"algo": "pgvector_hnsw"}) + results = backend.search( + dataset(), [index()], 5, search_threads=threads, dry_run=True + ) + assert all(result.success for result in results) + + +@pytest.mark.parametrize("threads", [2, "2", "1,2"]) +def test_multi_client_thread_setting_is_rejected(threads): + backend = PostgresBackend({"algo": "pgvector_hnsw"}) + results = backend.search( + dataset(), [index()], 5, search_threads=threads, dry_run=True + ) + assert all(not result.success for result in results) + + +@pytest.mark.parametrize("distance", ["cosine", "inner_product"]) +def test_build_rejects_unimplemented_metric(distance): + ds = dataset() + ds.distance_metric = distance + result = PostgresBackend({"algo": "pgvector_hnsw"}).build(ds, [index()]) + assert not result.success + assert "euclidean" in result.error_message + + +@pytest.mark.parametrize( + "params", + [{"m": "8); DROP SCHEMA public CASCADE; --"}, {"m": 0}, {"typo": 8}], +) +def test_build_rejects_invalid_params_before_connect(params, monkeypatch): + backend = PostgresBackend({"algo": "pgvector_hnsw"}) + monkeypatch.setattr( + backend, "_connect", lambda: pytest.fail("invalid params connected") + ) + result = backend.build(dataset(), [index(build=params)]) + assert not result.success + + +def test_backend_does_not_export_connection_string(): + backend = PostgresBackend( + {"algo": "pgvector_hnsw", "conninfo": "password=private"} + ) + result = backend.build(dataset(), [index()], dry_run=True) + assert "private" not in str(result.to_json()) + + +@pytest.mark.parametrize("missing_module", ["psycopg", "pgvector.psycopg"]) +def test_build_reports_optional_extra_installation( + missing_module, monkeypatch +): + monkeypatch.setitem(sys.modules, missing_module, None) + result = PostgresBackend({"algo": "pgvector_hnsw"}).build( + dataset(), [index()] + ) + assert not result.success + assert "cuvs-bench[postgres]" in result.error_message + + +def test_cpu_marker_selection_excludes_gpu_with_dsns_set(monkeypatch): + monkeypatch.setenv("CUVS_BENCH_POSTGRES_DSN", "not-a-live-database") + monkeypatch.setenv("CUVS_BENCH_POSTGRES_GPU_DSN", "not-a-live-database") + collected = subprocess.run( + [ + sys.executable, + "-m", + "pytest", + str(Path(__file__).resolve()), + "--collect-only", + "-q", + "-m", + "not postgres", + "-p", + "no:cacheprovider", + "-o", + "markers=postgres: live PostgreSQL tests\n" + "postgres_gpu: real GPU PostgreSQL tests", + ], + capture_output=True, + text=True, + check=True, + ) + assert "test_dry_run_does_not_connect" in collected.stdout + assert "test_real_gpu_build_search_and_hnsw_export" not in collected.stdout + + +def test_malformed_conninfo_does_not_leak_to_results_or_stdout( + monkeypatch, capsys +): + psycopg = pytest.importorskip("psycopg") + password = "synthetic" + chr(37) + "password" + conninfo = f"postgresql://benchmark_role:{password}@localhost/cuvs_bench" + dns = Mock(side_effect=AssertionError("unexpected DNS lookup")) + connect = Mock( + side_effect=AssertionError("unexpected database connection") + ) + monkeypatch.setattr(socket, "getaddrinfo", dns) + monkeypatch.setattr(psycopg.Connection, "_connect_gen", connect) + idx = index() + config = BenchmarkConfig( + indexes=[idx], + backend_config={"algo": "pgvector_hnsw", "conninfo": conninfo}, + ) + orchestrator = BenchmarkOrchestrator(backend_type="postgres") + orchestrator.config_loader = SimpleNamespace( + load=lambda **kwargs: (None, [config]) + ) + monkeypatch.setattr( + orchestrator, "_create_dataset", lambda config: dataset() + ) + results = orchestrator._run_sweep( + build=True, + search=False, + force=False, + dry_run=False, + count=5, + batch_size=10000, + search_mode="latency", + search_threads=None, + ) + assert len(results) == 1 and not results[0].success + assert password not in results[0].error_message + assert password not in str(results[0].to_json()) + captured = capsys.readouterr() + assert password not in captured.out + captured.err + dns.assert_not_called() + connect.assert_not_called() + + +@pytest.mark.parametrize( + "row", + [ + None, + (False, "gpu_cagra", 1), + (True, "cpu_hnsw", 1), + (True, "cpu_fallback", 1), + (True, "brute_force", 1), + ], +) +def test_cagra_rejects_missing_or_non_gpu_observations(row): + backend = PostgresBackend({"algo": "pgcuvs_cagra"}) + backend._schema = "isolated_test" + backend._conn = SimpleNamespace( + execute=lambda *args: SimpleNamespace(fetchone=lambda: row) + ) + with pytest.raises(RuntimeError, match="resident GPU CAGRA"): + backend._gpu_search_count() + + +def test_cagra_reads_observed_query_count(): + backend = PostgresBackend({"algo": "pgcuvs_cagra"}) + backend._schema = "isolated_test" + backend._conn = SimpleNamespace( + execute=lambda *args: SimpleNamespace( + fetchone=lambda: (True, "gpu_cagra", 42) + ) + ) + assert backend._gpu_search_count() == 42 + + +@pytest.fixture +def live_backend(): + conninfo = os.environ.get("CUVS_BENCH_POSTGRES_DSN") + if not conninfo: + pytest.skip("Set CUVS_BENCH_POSTGRES_DSN to a disposable database") + pytest.importorskip("psycopg") + backend = PostgresBackend( + {"algo": "pgvector_hnsw", "conninfo": conninfo, "warmup_queries": 1} + ) + try: + yield backend + finally: + backend.cleanup() + + +@pytest.mark.postgres +@pytest.mark.parametrize( + "algo,build,search", + [ + ( + "pgvector_hnsw", + {"m": 8, "ef_construction": 64}, + [{"ef_search": 40}], + ), + ("pgvector_ivfflat", {"lists": 4}, [{"probes": 4}]), + ], +) +def test_live_build_search_and_cleanup(live_backend, algo, build, search): + ds, idx = dataset(), index(algo, build, search) + live_backend.config["algo"] = algo + built = live_backend.build(ds, [idx]) + assert built.success, built.error_message + assert built.index_size_bytes > 0 + results = live_backend.search(ds, [idx], 5) + assert len(results) == 1 + result = results[0] + assert result.success, result.error_message + assert result.neighbors.shape == (5, 5) + assert result.distances.shape == (5, 5) + BenchmarkOrchestrator._finalize_search_result(result, ds, 5) + assert result.recall == 1.0 + assert result.queries_per_second > 0 + assert result.latency_percentiles["p50_ms"] > 0 + assert result.metadata["measurement_boundary"] == "end_to_end_sql" + schema = live_backend._schema + conn = live_backend._conn + live_backend.cleanup() + assert conn.closed + import psycopg + + with psycopg.connect(os.environ["CUVS_BENCH_POSTGRES_DSN"]) as check: + assert ( + check.execute( + "SELECT 1 FROM pg_namespace WHERE nspname = %s", (schema,) + ).fetchone() + is None + ) + + +@pytest.mark.postgres +def test_live_rejects_seqscan(live_backend): + from psycopg import sql + + ds, idx = dataset(), index() + assert live_backend.build(ds, [idx]).success + live_backend._conn.execute( + sql.SQL("DROP INDEX {}.ann").format( + sql.Identifier(live_backend._schema) + ) + ) + result = live_backend.search(ds, [idx], 5)[0] + assert not result.success + assert "Index Scan" in result.error_message + + +@pytest.mark.postgres +def test_live_rejects_index_loss_after_initial_plan(live_backend, monkeypatch): + from psycopg import sql + + ds, idx = dataset(), index(search=[{"ef_search": 40}]) + assert live_backend.build(ds, [idx]).success + original = live_backend._check_plan + calls = 0 + + def drop_after_first_check(statement, params): + nonlocal calls + original(statement, params) + if calls == 0: + live_backend._conn.execute( + sql.SQL("DROP INDEX {}.ann").format( + sql.Identifier(live_backend._schema) + ) + ) + calls += 1 + + monkeypatch.setattr(live_backend, "_check_plan", drop_after_first_check) + result = live_backend.search(ds, [idx], 5)[0] + assert not result.success + assert "Index Scan" in result.error_message + + +@pytest.mark.postgres +@pytest.mark.parametrize( + "after,success", [(0, False), (104, False), (105, True)] +) +def test_cagra_counter_guard_with_cpu_fixture( + live_backend, monkeypatch, after, success +): + # Only the telemetry contract is injected; this is not a GPU execution test. + ds, idx = dataset(), index(search=[{"ef_search": 40}]) + assert live_backend.build(ds, [idx]).success + live_backend.config["algo"] = "pgcuvs_cagra" + idx.search_params = [{"candidates": 40}] + counts = iter([100, after]) + monkeypatch.setattr( + live_backend, "_gpu_search_count", lambda: next(counts) + ) + result = live_backend.search(ds, [idx], 5)[0] + assert result.success is success + if success: + assert result.metadata["gpu_queries"] == len(ds.query_vectors) + assert result.metadata["observed_search_mode"] == "gpu_cagra" + else: + assert "GPU query count" in result.error_message + + +@pytest.mark.postgres +def test_live_missing_neighbors_count_as_misses(live_backend): + ds = dataset() + ds.training_vectors = ds.training_vectors[:3] + idx = index() + assert live_backend.build(ds, [idx]).success + result = live_backend.search(ds, [idx], 5)[0] + assert result.success, result.error_message + assert np.all(result.neighbors[:, 3:] == -1) + assert np.all(np.isinf(result.distances[:, 3:])) + + +@pytest.mark.postgres +def test_live_index_size_survives_csv_export(live_backend, tmp_path): + from cuvs_bench.run.data_export import write_results_to_csv + + ds, idx = dataset(), index() + built = live_backend.build(ds, [idx]) + assert built.success, built.error_message + write_results_to_csv([built], ds.name, str(tmp_path), 5, 10000) + path = tmp_path / ds.name / "result" / "build" / "pgvector_hnsw,base.csv" + with path.open() as handle: + row = next(csv.DictReader(handle)) + assert int(row["index_size_bytes"]) == built.index_size_bytes > 0 + assert row["index_size_source"] == "pg_relation_size" + + +def test_orchestrator_recomputes_recall_from_neighbor_ids(): + ds = dataset() + result = SearchResult( + neighbors=ds.groundtruth_neighbors.copy(), + distances=np.zeros((5, 5)), + search_time_ms=1, + queries_per_second=5000, + recall=0, + algorithm="pgvector_hnsw", + search_params=[{}], + ) + BenchmarkOrchestrator._finalize_search_result(result, ds, 5) + assert result.recall == 1 + + +@pytest.mark.postgres +@pytest.mark.postgres_gpu +@pytest.mark.parametrize("algo", ["pgcuvs_cagra", "pgcuvs_hnsw_import"]) +def test_real_gpu_build_search_and_hnsw_export(algo): + conninfo = os.environ.get("CUVS_BENCH_POSTGRES_GPU_DSN") + if not conninfo: + pytest.skip("Set CUVS_BENCH_POSTGRES_GPU_DSN on a real GPU server") + pytest.importorskip("psycopg") + from psycopg.conninfo import make_conninfo + + # Exercise the connection override with a valid conflicting startup value. + if algo == "pgcuvs_cagra": + conninfo = make_conninfo( + conninfo, options="-c cuvs.cpu_hnsw_fallback=on" + ) + rng = np.random.default_rng(20261016) + base = rng.normal(size=(4096, 32)).astype(np.float32) + queries = base[:16].copy() + distances = ((queries[:, None] - base[None]) ** 2).sum(axis=2) + ds = Dataset( + "gpu-smoke", + training_vectors=base, + query_vectors=queries, + groundtruth_neighbors=np.argsort(distances, axis=1)[:, :10], + ) + build = { + "graph_degree": 32, + "intermediate_graph_degree": 128, + "build_algo": "ivf_pq", + } + if algo == "pgcuvs_hnsw_import": + build["mode"] = "nsw" + search = ( + [{"candidates": 128}] + if algo == "pgcuvs_cagra" + else [{"ef_search": 128}] + ) + idx = index(algo, build, search) + backend = PostgresBackend( + { + "algo": algo, + "conninfo": conninfo, + "warmup_queries": 4, + } + ) + try: + built = backend.build(ds, [idx]) + assert built.success, built.error_message + assert built.index_size_bytes > 0 + source = f"{backend._schema}.source_graph" + if algo == "pgcuvs_cagra": + assert ( + backend._conn.execute( + "SHOW cuvs.cpu_hnsw_fallback" + ).fetchone()[0] + == "off" + ) + else: + before = backend._conn.execute( + "SELECT search_count FROM pg_stat_gpu_search " + "WHERE index_oid=%s::regclass", + (source,), + ).fetchone()[0] + result = backend.search(ds, [idx], 10)[0] + assert result.success, result.error_message + BenchmarkOrchestrator._finalize_search_result(result, ds, 10) + assert result.neighbors.shape == (16, 10) + # A broad sanity bound catches broken IDs, not production ANN quality. + assert result.recall > 0.8 + if algo == "pgcuvs_cagra": + assert result.metadata["observed_search_mode"] == "gpu_cagra" + assert result.metadata["gpu_queries"] == len(queries) + else: + after = backend._conn.execute( + "SELECT search_count FROM pg_stat_gpu_search " + "WHERE index_oid=%s::regclass", + (source,), + ).fetchone()[0] + assert after == before + assert ( + backend._conn.execute("SHOW enable_cuvs").fetchone()[0] + == "off" + ) + finally: + backend.cleanup() diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 87c3e83bb9..1710951f62 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -50,15 +50,21 @@ synthesize_dataset = [ "numpy>=2.0,<3.0", "tqdm", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +postgres = [ + "pgvector>=0.3", + "psycopg[binary]>=3.1", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. [project.urls] Homepage = "https://github.com/nvidia/cuvs" [project.entry-points."cuvs_bench.backends"] elastic = "cuvs_bench.backends.elasticsearch:register" +postgres = "cuvs_bench.backends.postgres:register" [project.entry-points."cuvs_bench.config_loaders"] elastic = "cuvs_bench.backends.elasticsearch:register" +postgres = "cuvs_bench.backends.postgres:register" [tool.setuptools.package-data] "*" = ["*.*", "VERSION"] @@ -66,6 +72,8 @@ elastic = "cuvs_bench.backends.elasticsearch:register" [tool.pytest.ini_options] markers = [ "opensearch: tests that require a live OpenSearch node (run with '-m opensearch')", + "postgres: tests that require a live PostgreSQL database (CPU or GPU)", + "postgres_gpu: tests that require CUVS_BENCH_POSTGRES_GPU_DSN, pg_cuvs and a real GPU daemon", ] [tool.isort]