diff --git a/.github/workflows/rust-vector-tests.yml b/.github/workflows/rust-vector-tests.yml index e00dc2f12b..a22817ab24 100644 --- a/.github/workflows/rust-vector-tests.yml +++ b/.github/workflows/rust-vector-tests.yml @@ -76,7 +76,7 @@ jobs: vectors-e2e: - name: S3 Vectors Integration (${{ matrix.volume }} volume${{ matrix.hardened && ', hardened' || '' }}) + name: S3 Vectors Integration (${{ matrix.volume }} volume${{ matrix.hardened && ', hardened' || '' }}${{ matrix.worker && ', worker' || '' }}) runs-on: ubuntu-24.04 timeout-minutes: 45 strategy: @@ -87,6 +87,16 @@ jobs: # volume, filer, gateway, engine) under mutual TLS, read JWTs # required, and every chunk encrypted by the gateway. hardened: ["", "1"] + worker: [""] + # Worker: the admin server and the Lance maintenance worker join the + # cluster, for the tests that build vector indexes through a job. + include: + - volume: go + hardened: "" + worker: "1" + - volume: rust + hardened: "" + worker: "1" env: CARGO_PROFILE_DEV_DEBUG: "0" CARGO_BUILD_JOBS: "2" @@ -96,6 +106,7 @@ jobs: VECTOR_TLS: ${{ matrix.hardened }} VECTOR_READ_JWT: ${{ matrix.hardened }} VECTOR_ENCRYPT_VOLUMES: ${{ matrix.hardened }} + VECTOR_WORKER: ${{ matrix.worker }} steps: - uses: actions/checkout@v7 with: @@ -123,6 +134,18 @@ jobs: run: | go test ./weed/s3api/s3vectors -count=1 go test ./weed/command -run 'Vector|Vectors' -count=1 + # The worker workspace has no protoc of its own for lance's build + # scripts; use the vendored one its lock names, as its own workflow does. + - name: Use the vendored protoc for the worker + if: matrix.worker == '1' + run: | + cd seaweed-worker + cargo fetch --locked + version=$(awk '/^name = "protoc-bin-vendored-linux-x86_64"$/{found=1; next} found && /^version = /{gsub(/"/,"",$3); print $3; exit}' Cargo.lock) + test -n "$version" || { echo "protoc-bin-vendored-linux-x86_64 is not in Cargo.lock" >&2; exit 1; } + protoc=$(find ~/.cargo/registry/src -path "*protoc-bin-vendored-linux-x86_64-$version/bin/protoc" | head -1) + test -x "$protoc" || { echo "no vendored protoc $version in the registry" >&2; exit 1; } + echo "PROTOC=$protoc" >> "$GITHUB_ENV" - name: Build integration binaries working-directory: test/s3vectors # dev profile: correctness suite, no release codegen needed. weed-volume @@ -131,6 +154,7 @@ jobs: set -e make ./weed ./weed-vector CARGO_PROFILE=dev if [ "${{ matrix.volume }}" = "rust" ]; then make ./weed-volume CARGO_PROFILE=dev; fi + if [ "${{ matrix.worker }}" = "1" ]; then make ./weed-vector-cli ./weed-worker CARGO_PROFILE=dev; fi - uses: actions/setup-python@v7 with: python-version: "3.12" @@ -154,7 +178,7 @@ jobs: if: always() uses: actions/upload-artifact@v7 with: - name: vectors-e2e-diagnostics-${{ matrix.volume }}${{ matrix.hardened && '-hardened' || '' }} + name: vectors-e2e-diagnostics-${{ matrix.volume }}${{ matrix.hardened && '-hardened' || '' }}${{ matrix.worker && '-worker' || '' }} path: | ${{ runner.temp }}/vector-tests ${{ runner.temp }}/vector-results.xml diff --git a/other/metrics/grafana_seaweedfs.json b/other/metrics/grafana_seaweedfs.json index 0c5f93d4e4..b9e7e64813 100644 --- a/other/metrics/grafana_seaweedfs.json +++ b/other/metrics/grafana_seaweedfs.json @@ -11708,6 +11708,101 @@ } ], "pluginVersion": "10.3.1" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "title": "Vector Indexes Built", + "description": "S3 Vectors tables the Lance worker indexed after they crossed the row threshold. Each table is built once; later rows are folded in by index optimization.", + "type": "timeseries", + "id": 245, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 64 + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "drawStyle": "line", + "fillOpacity": 10, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 4, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "unit": "short", + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + } + ] + } + }, + "overrides": [] + }, + "options": { + "legend": { + "calcs": [ + "lastNotNull", + "max" + ], + "displayMode": "table", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum(increase(SeaweedFS_worker_lance_vector_indexes_built_total{cluster=~\"$cluster\"}[$__rate_interval]))", + "range": true, + "refId": "A", + "legendFormat": "indexes built" + } + ] } ] }, diff --git a/seaweed-worker/crates/lance/src/jobs/mod.rs b/seaweed-worker/crates/lance/src/jobs/mod.rs index a9c3bcfb51..b86c8b0ab5 100644 --- a/seaweed-worker/crates/lance/src/jobs/mod.rs +++ b/seaweed-worker/crates/lance/src/jobs/mod.rs @@ -10,6 +10,7 @@ pub mod cleanup; pub mod compact; pub mod indices; pub mod sort; +pub mod vector_index; use std::collections::HashMap; use std::sync::Arc; @@ -58,6 +59,12 @@ pub fn handlers( .with_signing_key(signing_key.clone()) .with_metrics(metrics.clone()), ), + Arc::new( + vector_index::VectorIndexHandler::new(namespace_url.clone()) + .with_fallback(fallback.clone()) + .with_signing_key(signing_key.clone()) + .with_metrics(metrics.clone()), + ), Arc::new( sort::SortHandler::new(namespace_url.clone()) .with_fallback(fallback.clone()) diff --git a/seaweed-worker/crates/lance/src/jobs/vector_index.rs b/seaweed-worker/crates/lance/src/jobs/vector_index.rs new file mode 100644 index 0000000000..f07b43d084 --- /dev/null +++ b/seaweed-worker/crates/lance/src/jobs/vector_index.rs @@ -0,0 +1,354 @@ +use std::collections::HashMap; + +use anyhow::{Result, anyhow, bail}; +use async_trait::async_trait; +use seaweed_worker_core::config_form::{form, int_or, int_value, number_field}; +use seaweed_worker_core::pb::{ + AdminRuntimeDefaults, ConfigValue, DetectionComplete, DetectionProposals, ExecuteJobRequest, + JobCompleted, JobProgressUpdate, JobProposal, JobResult, JobTypeCapability, JobTypeDescriptor, + RunDetectionRequest, +}; +use seaweed_worker_core::{DetectionSender, ExecutionSender, JobHandler}; +use tracing::warn; + +use crate::catalog::{NamespaceClient, SigningKey, parse_id}; +use crate::dataset::{self, OpenTable}; +use crate::jobs::{clamp, string_list, table_id}; +use crate::vector_index::{INDEX_NAME, build, index_names, partitions_for}; + +pub const JOB_TYPE: &str = "lance_vector_index"; + +const DEFAULT_MIN_ROWS: i64 = 50_000; +const MIN_ROWS_FLOOR: i64 = 1; +const MIN_ROWS_CEILING: i64 = 1_000_000_000; +/// The namespace an S3 Vectors bucket keeps its indexes in. +const VECTOR_NAMESPACE: &str = "vector"; +/// The column S3 Vectors writes its vectors to. +const VECTOR_COLUMN: &str = "vector"; +/// A build trains on the whole table, so the admin's 90 seconds are not enough. +/// The run window is as long, or it would cut an attempt short first. +const EXECUTION_TIMEOUT_SECONDS: i32 = 3600; + +/// Builds the vector index of an S3 Vectors table once it holds enough rows +/// for an exact scan to be the slow way. Tables below the threshold, and rows +/// written after a build, are searched exactly; `lance_optimize_indices` then +/// keeps the index covering. +pub struct VectorIndexHandler { + namespace_url: String, + fallback: dataset::FallbackOptions, + signing_key: Option, + metrics: Option, +} + +impl VectorIndexHandler { + pub fn new(namespace_url: String) -> Self { + Self { + namespace_url, + fallback: dataset::FallbackOptions::new(), + signing_key: None, + metrics: None, + } + } + + pub fn with_metrics(mut self, metrics: Option) -> Self { + self.metrics = metrics; + self + } + + pub fn with_fallback(mut self, fallback: dataset::FallbackOptions) -> Self { + self.fallback = fallback; + self + } + + pub fn with_signing_key(mut self, signing_key: Option) -> Self { + self.signing_key = signing_key; + self + } + + fn client(&self) -> NamespaceClient { + NamespaceClient::new(self.namespace_url.clone()) + .with_fallback(&self.fallback) + .with_signing_key(self.signing_key.as_ref()) + } +} + +/// `$vector$`: a table the S3 Vectors API created. +pub fn is_vector_table(id: &[String]) -> bool { + id.len() == 3 && id[1] == VECTOR_NAMESPACE +} + +fn has_vector_column(table: &OpenTable) -> bool { + table.dataset.schema().field(VECTOR_COLUMN).is_some() +} + +async fn is_indexed(table: &OpenTable) -> Result { + Ok(!index_names(&table.dataset, VECTOR_COLUMN).await?.is_empty()) +} + +/// Why an execution leaves the table as it is, if it does: the table may have +/// changed since it was proposed. +fn skip_reason(indexed: bool, rows: usize) -> Option<&'static str> { + if indexed { + Some("already indexed") + } else if rows == 0 { + Some("no rows to index") + } else { + None + } +} + +#[async_trait] +impl JobHandler for VectorIndexHandler { + fn capability(&self) -> JobTypeCapability { + JobTypeCapability { + job_type: JOB_TYPE.to_string(), + can_detect: true, + can_execute: true, + max_detection_concurrency: 1, + max_execution_concurrency: 1, + display_name: "Lance Vector Index".to_string(), + description: + "Build the vector index of S3 Vectors tables that have grown past the threshold" + .to_string(), + weight: 40, + } + } + + fn descriptor(&self) -> JobTypeDescriptor { + let mut defaults: HashMap = HashMap::new(); + defaults.insert("min_rows".to_string(), int_value(DEFAULT_MIN_ROWS)); + + JobTypeDescriptor { + job_type: JOB_TYPE.to_string(), + display_name: "Lance Vector Index".to_string(), + description: + "Build an IVF index for S3 Vectors tables so queries stop scanning every row" + .to_string(), + icon: "fas fa-project-diagram".to_string(), + descriptor_version: 1, + worker_config_form: Some(form( + "lance-vector-index-worker", + "Vector index", + vec![number_field( + "min_rows", + "Row threshold", + "Build the index once a table holds at least this many rows", + MIN_ROWS_FLOOR, + MIN_ROWS_CEILING, + )], + defaults.clone(), + )), + worker_default_values: defaults, + admin_runtime_defaults: Some(AdminRuntimeDefaults { + job_type_max_runtime_seconds: EXECUTION_TIMEOUT_SECONDS, + execution_timeout_seconds: EXECUTION_TIMEOUT_SECONDS, + ..Default::default() + }), + ..Default::default() + } + } + + async fn detect( + &self, + request: &RunDetectionRequest, + sender: &dyn DetectionSender, + ) -> Result<()> { + let min_rows = clamp( + int_or(&request.worker_config_values, "min_rows", DEFAULT_MIN_ROWS), + MIN_ROWS_FLOOR, + MIN_ROWS_CEILING, + ) as usize; + let max_results = usize::try_from(request.max_results).unwrap_or(0); + let client = self.client(); + let tables = client.list_all_tables().await?; + + let mut proposals = Vec::new(); + let mut has_more = false; + for encoded in &tables { + let id = parse_id(encoded); + if !is_vector_table(&id) { + continue; + } + let table = match dataset::open(&client, &id, &self.fallback).await { + Ok(table) => table, + Err(err) => { + if let Some(counters) = &self.metrics { + counters.worker.object_skipped(JOB_TYPE, "open"); + } + warn!("skipping {encoded}: {err:#}"); + continue; + } + }; + if let Some(counters) = &self.metrics { + counters.worker.object_seen(JOB_TYPE); + } + if !has_vector_column(&table) { + continue; + } + match is_indexed(&table).await { + Ok(false) => {} + Ok(true) => continue, + Err(err) => { + if let Some(counters) = &self.metrics { + counters.worker.object_skipped(JOB_TYPE, "indices"); + } + warn!("skipping {encoded}: reading its indices failed: {err:#}"); + continue; + } + } + let rows = match table.dataset.count_rows(None).await { + Ok(rows) => rows, + Err(err) => { + if let Some(counters) = &self.metrics { + counters.worker.object_skipped(JOB_TYPE, "rows"); + } + warn!("skipping {encoded}: counting its rows failed: {err:#}"); + continue; + } + }; + if rows < min_rows { + continue; + } + if max_results > 0 && proposals.len() >= max_results { + has_more = true; + break; + } + let mut parameters: HashMap = HashMap::new(); + parameters.insert("table_id".to_string(), string_list(&id)); + proposals.push(JobProposal { + proposal_id: format!("{JOB_TYPE}:{encoded}"), + dedupe_key: format!("{JOB_TYPE}:{encoded}"), + job_type: JOB_TYPE.to_string(), + summary: format!("Index {encoded} ({rows} rows)"), + detail: format!( + "{rows} rows are searched by an exact scan on every query; the policy \ + indexes tables from {min_rows} rows" + ), + parameters, + ..Default::default() + }); + } + + let total = proposals.len() as i32; + sender.send_proposals(DetectionProposals { + request_id: request.request_id.clone(), + job_type: JOB_TYPE.to_string(), + proposals, + has_more, + })?; + sender.send_complete(DetectionComplete { + request_id: request.request_id.clone(), + job_type: JOB_TYPE.to_string(), + success: true, + error_message: String::new(), + total_proposals: total, + skipped_reason: String::new(), + eligible_after: None, + })?; + Ok(()) + } + + async fn execute( + &self, + request: &ExecuteJobRequest, + sender: &dyn ExecutionSender, + ) -> Result<()> { + let job = request + .job + .as_ref() + .ok_or_else(|| anyhow!("execute request carried no job"))?; + let id = table_id(&job.parameters) + .ok_or_else(|| anyhow!("job {} carried no table_id", job.job_id))?; + + let client = self.client(); + let mut table = dataset::open(&client, &id, &self.fallback).await?; + let rows = table.dataset.count_rows(None).await?; + let partitions = partitions_for(rows); + let mut output: HashMap = HashMap::new(); + output.insert("rows".to_string(), int_value(rows as i64)); + output.insert("partitions".to_string(), int_value(partitions as i64)); + + if !has_vector_column(&table) { + bail!("the table has no {VECTOR_COLUMN} column to index"); + } + if let Some(reason) = skip_reason(is_indexed(&table).await?, rows) { + sender.send_completed(JobCompleted { + request_id: request.request_id.clone(), + job_id: job.job_id.clone(), + job_type: JOB_TYPE.to_string(), + success: true, + result: Some(JobResult { + output_values: output, + summary: reason.to_string(), + ..Default::default() + }), + ..Default::default() + })?; + return Ok(()); + } + + sender.send_progress(JobProgressUpdate { + request_id: request.request_id.clone(), + job_id: job.job_id.clone(), + job_type: JOB_TYPE.to_string(), + progress_percent: 10.0, + stage: format!("indexing {rows} rows into {partitions} partitions"), + ..Default::default() + })?; + + build( + &mut table.dataset, + VECTOR_COLUMN, + INDEX_NAME, + partitions, + false, + ) + .await?; + if let Some(counters) = &self.metrics { + counters.indexes_built.inc(); + } + + sender.send_completed(JobCompleted { + request_id: request.request_id.clone(), + job_id: job.job_id.clone(), + job_type: JOB_TYPE.to_string(), + success: true, + result: Some(JobResult { + output_values: output, + summary: format!("{rows} rows indexed into {partitions} partitions"), + ..Default::default() + }), + ..Default::default() + })?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn only_vector_bucket_tables_qualify() { + assert!(is_vector_table(&parse_id("media$vector$movies"))); + assert!(!is_vector_table(&parse_id("media$ml$events"))); + assert!(!is_vector_table(&parse_id("vector$movies"))); + } + + #[test] + fn a_table_indexed_or_emptied_since_detection_is_left_alone() { + assert_eq!(skip_reason(true, 60_000), Some("already indexed")); + assert_eq!(skip_reason(true, 0), Some("already indexed")); + assert_eq!(skip_reason(false, 0), Some("no rows to index")); + assert_eq!(skip_reason(false, 60_000), None); + } + + #[test] + fn a_build_gets_longer_than_the_default_execution_timeout() { + let descriptor = VectorIndexHandler::new(String::new()).descriptor(); + let defaults = descriptor.admin_runtime_defaults.unwrap(); + assert_eq!(defaults.execution_timeout_seconds, 3600); + assert_eq!(defaults.job_type_max_runtime_seconds, 3600); + assert!(!defaults.enabled); + } +} diff --git a/seaweed-worker/crates/lance/src/metrics.rs b/seaweed-worker/crates/lance/src/metrics.rs index e414a45386..c4af41ebc4 100644 --- a/seaweed-worker/crates/lance/src/metrics.rs +++ b/seaweed-worker/crates/lance/src/metrics.rs @@ -18,6 +18,7 @@ pub struct LanceMetrics { pub worker: Metrics, pub fragments_removed: IntCounter, pub rows_indexed: IntCounter, + pub indexes_built: IntCounter, pub rows_sorted: IntCounter, pub versions_removed: IntCounter, pub bytes_reclaimed: IntCounter, @@ -35,6 +36,10 @@ impl LanceMetrics { "lance_rows_indexed_total", "Rows brought under an index that did not cover them.", )?, + indexes_built: metrics.counter( + "lance_vector_indexes_built_total", + "Vector indexes built for S3 Vectors tables.", + )?, rows_sorted: metrics.counter( "lance_rows_sorted_total", "Rows rewritten into their table's declared order.", diff --git a/test/s3vectors/Makefile b/test/s3vectors/Makefile index 77048ada54..4ac15246a9 100644 --- a/test/s3vectors/Makefile +++ b/test/s3vectors/Makefile @@ -6,6 +6,7 @@ ifneq ($(wildcard $(WORKER_MANIFEST)),) WORKER := ./weed-vector-worker endif VOLUME := ./weed-volume +LANCE_WORKER := ./weed-worker CARGO_PROFILE ?= release CARGO_TARGET_DIR ?= $(abspath ../../seaweed-vector/target) override CARGO_TARGET_DIR := $(abspath $(CARGO_TARGET_DIR)) @@ -27,6 +28,10 @@ $(CLI): FORCE cd ../../seaweed-vector && cargo build --locked --profile $(CARGO_PROFILE) -p weed-vector-core --bin weed-vector-cli cp $(CARGO_OUTPUT)/weed-vector-cli $(CLI) +$(LANCE_WORKER): FORCE + cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../seaweed-worker/Cargo.toml -p weed-lance-worker --bin weed-worker + cp $(CARGO_OUTPUT)/weed-worker $(LANCE_WORKER) + ifneq ($(WORKER),) $(WORKER): $(VEC) FORCE cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../rust/Cargo.toml -p sw-vector-distributed --bin weed-vector @@ -40,6 +45,9 @@ $(VOLUME): $(VEC) $(WORKER) FORCE FORCE: build: $(WEED) $(VEC) $(WORKER) $(VOLUME) +ifneq ($(VECTOR_WORKER),) +build: $(CLI) $(LANCE_WORKER) +endif test: build ./run_test.sh @@ -48,4 +56,4 @@ perf: build $(CLI) ./run_perf.sh clean: - rm -f $(WEED) $(VEC) $(CLI) ./weed-vector-worker $(VOLUME) + rm -f $(WEED) $(VEC) $(CLI) $(LANCE_WORKER) ./weed-vector-worker $(VOLUME) diff --git a/test/s3vectors/conftest.py b/test/s3vectors/conftest.py index e54816b8a4..1279d60233 100644 --- a/test/s3vectors/conftest.py +++ b/test/s3vectors/conftest.py @@ -77,6 +77,20 @@ def stop(process): process.wait(timeout=10) +def pytest_configure(config): + config.addinivalue_line("markers", "worker: needs the admin server and the Lance maintenance worker") + + +def pytest_collection_modifyitems(config, items): + # Deselected rather than skipped: CI rejects a leg that reports skips. + if os.environ.get("VECTOR_WORKER"): + return + skipped = [item for item in items if item.get_closest_marker("worker")] + if skipped: + config.hook.pytest_deselected(items=skipped) + items[:] = [item for item in items if item not in skipped] + + @pytest.fixture(scope="module", params=combinations(), ids=lambda p: "-".join(p)) def cluster(request, tmp_path_factory): yield from run_cluster(request.param, tmp_path_factory, identities=True) @@ -225,17 +239,21 @@ def run_cluster(param, tmp_path_factory, identities): print(f"\nVector integration logs: {directory}") weed = binary("WEED_BINARY", "weed") engine = binary("VECTOR_BINARY", "weed-vector") + # VECTOR_WORKER=1 adds the admin server and the Lance maintenance worker, + # for the tests that build vector indexes through a plugin job. + lance_worker = binary("LANCE_WORKER_BINARY", "weed-worker") if os.environ.get("VECTOR_WORKER") else None worker = binary("VECTOR_WORKER_BINARY", "weed-vector-worker") if mode == "distributed" else None volume_binary = binary("RUST_VOLUME_SERVER_BINARY", "weed-volume") if volume == "rust" else weed processes = {} with contextlib.ExitStack() as stack: with contextlib.ExitStack() as reservations: ports = [] - for _ in range(12): + for _ in range(14): sock = reservations.enter_context(socket.socket()) sock.bind(("127.0.0.1", 0)) ports.append(sock.getsockname()[1]) - master, master_grpc, vol, vol_grpc, filer, filer_grpc, s3, s3_grpc, lance, vector, worker1, worker2 = ports + (master, master_grpc, vol, vol_grpc, filer, filer_grpc, s3, s3_grpc, lance, vector, worker1, worker2, + admin, admin_grpc) = ports master_addr = f"127.0.0.1:{master}.{master_grpc}" filer_addr = f"127.0.0.1:{filer}.{filer_grpc}" @@ -330,8 +348,22 @@ def run_cluster(param, tmp_path_factory, identities): processes, "gateway did not discover both search workers", ) + admin_url = None + if lance_worker: + (directory / "admin").mkdir() + admin_url = f"http://127.0.0.1:{admin}" + start("admin", [weed, "admin", "-ip=127.0.0.1", f"-port={admin}", f"-port.grpc={admin_grpc}", + "-port.legacy=0", f"-master={master_addr}", f"-dataDir={directory / 'admin'}"]) + wait_for(lambda: http_ready(f"{admin_url}/api/plugin/job-types"), processes, "admin not ready") + start("lance-worker", [lance_worker, "--admin", f"127.0.0.1:{admin}.{admin_grpc}", + "--namespace", f"http://127.0.0.1:{lance}", "--id", "vectors-test-worker", + "--access-key", "test", "--secret-key", "test"]) + wait_for(lambda: "lance_vector_index" in + [jt["job_type"] for jt in json.load(urllib.request.urlopen(f"{admin_url}/api/plugin/job-types", timeout=2))], + processes, "the Lance worker never registered with admin") + yield {"endpoint": endpoint, "mode": mode, "processes": processes, "directory": directory, - "catalog": f"http://127.0.0.1:{lance}"} + "catalog": f"http://127.0.0.1:{lance}", "admin": admin_url} @pytest.fixture(scope="module") diff --git a/test/s3vectors/s3vectors_test.py b/test/s3vectors/s3vectors_test.py index 375e2d6242..bc23076a11 100644 --- a/test/s3vectors/s3vectors_test.py +++ b/test/s3vectors/s3vectors_test.py @@ -492,6 +492,84 @@ def test_tag_conditions(vectors, vector_clients, index): vectors.delete_index(**other) +def exact_neighbours(rows, query, k): + import math + qn = math.sqrt(sum(x * x for x in query)) + scored = sorted(((sum(a * b for a, b in zip(v, query)) / (math.sqrt(sum(a * a for a in v)) * qn), key) + for key, v in rows.items()), reverse=True) + return {key for _, key in scored[:k]} + + +@pytest.mark.worker +def test_vector_index_job(vectors, cluster): + """The maintenance worker indexes a table past the row threshold through the + admin server, queries then run through the index, and rows written after + the build are still found.""" + import os + import random + import subprocess + import urllib.request + from pathlib import Path + + bucket, index, dim = "vectors-indexed", "idx", 16 + vectors.create_vector_bucket(vectorBucketName=bucket) + vectors.create_index(vectorBucketName=bucket, indexName=index, dataType="float32", + dimension=dim, distanceMetric="cosine") + # Embedding-like data: points around shared centres. Independent Gaussian + # vectors have no locality for an IVF index to use, so recall through one + # would only track the fraction of partitions probed. + rng = random.Random(7) + centres = [[rng.gauss(0, 1) for _ in range(dim)] for _ in range(30)] + + def point(): + centre = rng.choice(centres) + return [c + 0.3 * rng.gauss(0, 1) for c in centre] + + rows = {f"v{i}": point() for i in range(1200)} + keys = list(rows) + for offset in range(0, len(keys), 400): + vectors.put_vectors(vectorBucketName=bucket, indexName=index, + vectors=[{"key": key, "data": {"float32": rows[key]}} for key in keys[offset:offset + 400]]) + + admin = cluster["admin"] + job = "lance_vector_index" + + def call(path, method="GET", body=None): + data = json.dumps(body).encode() if body is not None else None + request = urllib.request.Request(f"{admin}/api/plugin/job-types/{job}{path}", data=data, method=method, + headers={"Content-Type": "application/json"} if data else {}) + with urllib.request.urlopen(request, timeout=120) as response: + return json.load(response) + + call("/config", "PUT", {"admin_runtime": {"enabled": True}, + "worker_config_values": {"min_rows": {"int64_value": "1000"}}}) + run = call("/run", "POST", {"timeout_seconds": 110}) + assert run["success_count"] == 1 and run["error_count"] == 0, run + + cli = Path(os.environ.get("VECTOR_CLI_BINARY", Path(__file__).resolve().parent / "weed-vector-cli")) + coverage = json.loads(subprocess.run( + [str(cli), "--namespace", cluster["catalog"], "--access-key", "test", "--secret-key", "test", + "coverage", f"{bucket}$vector${index}"], check=True, capture_output=True, text=True).stdout) + assert len(coverage) == 1 and coverage[0]["segments"] == 1 and coverage[0]["unindexed_rows"] == 0, coverage + + found = 0 + for _ in range(20): + query = point() + hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index, + queryVector={"float32": query}, topK=10)["vectors"] + assert len(hits) == 10 + found += len({h["key"] for h in hits} & exact_neighbours(rows, query, 10)) + assert found >= 180, f"recall {found}/200" + + late = point() + vectors.put_vectors(vectorBucketName=bucket, indexName=index, vectors=[{"key": "late", "data": {"float32": late}}]) + hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index, queryVector={"float32": late}, topK=1) + assert [h["key"] for h in hits["vectors"]] == ["late"] + + again = call("/run", "POST", {"timeout_seconds": 60}) + assert again["detected_count"] == 0, again + + def test_invalid_credentials(cluster): client = boto3.client("s3vectors", endpoint_url=cluster["endpoint"], aws_access_key_id="unknown", aws_secret_access_key="unknown", region_name="us-east-1") diff --git a/weed/admin/plugin/scheduler_lane.go b/weed/admin/plugin/scheduler_lane.go index ec960e4a14..2b13c1a59b 100644 --- a/weed/admin/plugin/scheduler_lane.go +++ b/weed/admin/plugin/scheduler_lane.go @@ -98,6 +98,7 @@ var jobTypeLaneMap = map[string]SchedulerLane{ "lance_compact": LaneLance, "lance_optimize_indices": LaneLance, "lance_cleanup_versions": LaneLance, + "lance_vector_index": LaneLance, } // JobTypeLane returns the scheduler lane for the given job type.