mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-10-30 08:50:48 +01:00
s3vectors: fix shared API behavior and test Go/Rust volume parity (#12216)
* s3vectors: reuse shared authenticated identity context * s3vectors: encode empty results as arrays * test: cover S3 Vectors with Go and Rust volume servers * ci: require both volume implementations in vector integration tests * s3vectors: keep identityOf as a shim over the shared identity context Concurrent branches add identityOf call sites; deleting the helper would break the build when they merge in either order. Route it through GetIdentityNameFromContext so every caller resolves the same identity. * test: wait for the vector engine to register before running The gateway resolves weed-vector lazily through the master cluster list, so a listening port does not mean the engine is reachable yet; the first vector write or query could race KeepConnected and fail. Restore the cluster.ps registration wait the previous harness had, and ignore the pytest cache directory the suite creates. * test: wait for master registration across OSS and enterprise Use the exact engine registration event after AddClusterNode completes. The current enterprise cluster.ps does not list vector nodes, so waiting on its output prevents the shared fixture from starting there. * s3vectors: keep contextWithIdentity as a shim alongside identityOf In-flight callers still write the caller name through it; route the write through SetIdentityNameInContext so both helpers agree on the shared key whichever lands first.
This commit is contained in:
1 parent
88565df885
commit
4ef123d764
12 files changed
+431
-118
No files matched your search
+23
-9
@@ -1,28 +1,42 @@
|
||||
# End-to-end S3 Vectors conformance run:
|
||||
# weed mini (cluster + lance catalog) + weed s3 (s3vectors port) +
|
||||
# weed-vector (VectorService gRPC), then boto3 s3vectors calls.
|
||||
WEED := ./weed
|
||||
VEC := ./weed-vector
|
||||
WORKER_MANIFEST := ../../rust/sw-vector-distributed/Cargo.toml
|
||||
ifneq ($(wildcard $(WORKER_MANIFEST)),)
|
||||
WORKER := ./weed-vector-worker
|
||||
endif
|
||||
VOLUME := ./weed-volume
|
||||
CARGO_PROFILE ?= release
|
||||
CARGO_TARGET_DIR ?= $(abspath ../../seaweed-vector/target)
|
||||
CARGO_OUTPUT := $(CARGO_TARGET_DIR)/$(if $(filter dev,$(CARGO_PROFILE)),debug,$(CARGO_PROFILE))
|
||||
export CARGO_TARGET_DIR
|
||||
|
||||
.PHONY: all test build clean FORCE
|
||||
|
||||
all: test
|
||||
|
||||
# No prerequisites: go/cargo decide for themselves whether a rebuild is needed.
|
||||
$(WEED): FORCE
|
||||
go build -o $(WEED) ../../weed
|
||||
|
||||
$(VEC): FORCE
|
||||
cargo build --release --manifest-path ../../seaweed-vector/Cargo.toml -p weed-vector-server
|
||||
cp ../../seaweed-vector/target/release/weed-vector $(VEC)
|
||||
cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../seaweed-vector/Cargo.toml -p weed-vector-server
|
||||
cp $(CARGO_OUTPUT)/weed-vector $(VEC)
|
||||
|
||||
ifneq ($(WORKER),)
|
||||
$(WORKER): $(VEC) FORCE
|
||||
cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../rust/Cargo.toml -p sw-vector-distributed --bin weed-vector
|
||||
cp $(CARGO_OUTPUT)/weed-vector $(WORKER)
|
||||
endif
|
||||
|
||||
$(VOLUME): $(VEC) $(WORKER) FORCE
|
||||
cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../seaweed-volume/Cargo.toml --bin weed-volume
|
||||
cp $(CARGO_OUTPUT)/weed-volume $(VOLUME)
|
||||
|
||||
FORCE:
|
||||
|
||||
build: $(WEED) $(VEC)
|
||||
build: $(WEED) $(VEC) $(WORKER) $(VOLUME)
|
||||
|
||||
test: build
|
||||
./run_test.sh
|
||||
|
||||
clean:
|
||||
rm -f $(WEED) $(VEC)
|
||||
rm -rf mini-data
|
||||
rm -f $(WEED) $(VEC) ./weed-vector-worker $(VOLUME)
|
||||
@@ -0,0 +1,23 @@
|
||||
# S3 Vectors integration tests
|
||||
|
||||
Run `python3 -m pip install -r requirements.txt`, then `make test`.
|
||||
|
||||
The suite starts a master, volume, filer, authenticated S3/Lance gateway, and
|
||||
vector engine on temporary ports. The same 12 API tests run with both Go and
|
||||
Rust volume servers. Coverage includes bucket/index lifecycle, bucket policies,
|
||||
empty results, ARN addressing, metadata, filters, cosine distance, atomic upserts,
|
||||
pagination, segmented listing, authentication, and writer loss.
|
||||
|
||||
Server logs and data stay in pytest's temporary directory for diagnosis. Every
|
||||
child process stops on exit. To select a volume server, set
|
||||
`VOLUME_SERVER_IMPL=go` or `rust`.
|
||||
|
||||
In an enterprise checkout containing `rust/sw-vector-distributed`, the build and
|
||||
default matrix also include distributed queries with two search workers. Those
|
||||
queries must still succeed after the writer stops. `VECTOR_QUERY_MODE=single`
|
||||
or `distributed` selects one query mode.
|
||||
|
||||
To use prebuilt binaries, set `WEED_BINARY`, `VECTOR_BINARY`,
|
||||
`RUST_VOLUME_SERVER_BINARY`, and (for distributed queries) `VECTOR_WORKER_BINARY`,
|
||||
then run `./run_test.sh`. `PYTHON` selects the Python interpreter. The Makefile
|
||||
accepts `CARGO_PROFILE=dev` for faster local builds.
|
||||
@@ -0,0 +1,157 @@
|
||||
import contextlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
import boto3
|
||||
from botocore.config import Config
|
||||
import pytest
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def combinations():
|
||||
volumes = os.environ.get("VOLUME_SERVER_IMPL", "go,rust").split(",")
|
||||
default_modes = "single,distributed" if (HERE / "../../rust/sw-vector-distributed/Cargo.toml").is_file() else "single"
|
||||
modes = os.environ.get("VECTOR_QUERY_MODE", default_modes).split(",")
|
||||
if not set(volumes) <= {"go", "rust"} or not set(modes) <= {"single", "distributed"}:
|
||||
raise ValueError("invalid VOLUME_SERVER_IMPL or VECTOR_QUERY_MODE")
|
||||
return [(volume, mode) for volume in volumes for mode in modes]
|
||||
|
||||
|
||||
def binary(variable, name):
|
||||
path = Path(os.environ.get(variable, str(HERE / name))).resolve()
|
||||
if not path.is_file() or not os.access(path, os.X_OK):
|
||||
pytest.fail(f"Missing executable {path}; run make build or set {variable}")
|
||||
return str(path)
|
||||
|
||||
|
||||
def wait_for(check, processes, message):
|
||||
deadline = time.monotonic() + 60
|
||||
while time.monotonic() < deadline:
|
||||
for name, process in processes.items():
|
||||
if process.poll() is not None:
|
||||
pytest.fail(f"{name} exited with {process.returncode}; {message}")
|
||||
try:
|
||||
if check():
|
||||
return
|
||||
except (OSError, urllib.error.URLError):
|
||||
pass
|
||||
time.sleep(0.2)
|
||||
pytest.fail(message)
|
||||
|
||||
|
||||
def listening(port):
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
|
||||
return True
|
||||
|
||||
|
||||
def http_ready(url):
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=0.5):
|
||||
return True
|
||||
except urllib.error.HTTPError:
|
||||
return True
|
||||
|
||||
|
||||
def stop(process):
|
||||
if process.poll() is None:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait(timeout=10)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=combinations(), ids=lambda p: "-".join(p))
|
||||
def cluster(request, tmp_path_factory):
|
||||
volume, mode = request.param
|
||||
directory = tmp_path_factory.mktemp(f"vectors-{volume}-{mode}")
|
||||
print(f"\nVector integration logs: {directory}")
|
||||
weed = binary("WEED_BINARY", "weed")
|
||||
engine = binary("VECTOR_BINARY", "weed-vector")
|
||||
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):
|
||||
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_addr = f"127.0.0.1:{master}.{master_grpc}"
|
||||
filer_addr = f"127.0.0.1:{filer}.{filer_grpc}"
|
||||
|
||||
def start(name, args):
|
||||
log = stack.enter_context((directory / f"{name}.log").open("w"))
|
||||
env = dict(os.environ, AWS_ACCESS_KEY_ID="test", AWS_SECRET_ACCESS_KEY="test",
|
||||
AWS_DEFAULT_REGION="us-east-1", AWS_EC2_METADATA_DISABLED="true")
|
||||
process = subprocess.Popen(args, cwd=directory, stdout=log, stderr=log, env=env)
|
||||
processes[name] = process
|
||||
stack.callback(stop, process)
|
||||
return process
|
||||
|
||||
for name in ("master", "volume", "filer"):
|
||||
(directory / name).mkdir()
|
||||
start("master", [weed, "-v=1", "master", "-ip=127.0.0.1", "-ip.bind=127.0.0.1",
|
||||
f"-port={master}", f"-port.grpc={master_grpc}",
|
||||
f"-mdir={directory / 'master'}", "-peers=none", "-volumeSizeLimitMB=32"])
|
||||
wait_for(lambda: http_ready(f"http://127.0.0.1:{master}/dir/status"), processes, "master not ready")
|
||||
prefix = [volume_binary] if volume == "rust" else [weed, "volume"]
|
||||
flag = "--" if volume == "rust" else "-"
|
||||
start("volume", prefix + [f"{flag}{key}={value}" for key, value in {
|
||||
"master": master_addr, "ip": "127.0.0.1", "ip.bind": "127.0.0.1",
|
||||
"port": vol, "port.grpc": vol_grpc, "dir": directory / "volume",
|
||||
"max": 16, "preStopSeconds": 0,
|
||||
}.items()])
|
||||
wait_for(lambda: listening(vol_grpc), processes, "volume not ready")
|
||||
start("filer", [weed, "filer", f"-master={master_addr}", "-ip=127.0.0.1",
|
||||
"-ip.bind=127.0.0.1", f"-port={filer}", f"-port.grpc={filer_grpc}",
|
||||
f"-defaultStoreDir={directory / 'filer'}"])
|
||||
wait_for(lambda: listening(filer_grpc), processes, "filer not ready")
|
||||
endpoint = f"http://127.0.0.1:{s3}"
|
||||
start("s3", [weed, "-v=1", "s3", f"-filer={filer_addr}", f"-config={HERE / 's3.json'}",
|
||||
"-ip=127.0.0.1", "-ip.bind=127.0.0.1", f"-port={s3}",
|
||||
f"-port.grpc={s3_grpc}", "-port.iceberg=0", f"-port.lance={lance}"])
|
||||
wait_for(lambda: http_ready(f"http://127.0.0.1:{lance}/v1/namespace/root/list"), processes, "catalog not ready")
|
||||
start("vector", [engine, "--namespace", f"http://127.0.0.1:{lance}",
|
||||
"--masters", master_addr, "--listen", f"127.0.0.1:{vector}",
|
||||
"--storage-option", "aws_access_key_id=test",
|
||||
"--storage-option", "aws_secret_access_key=test"])
|
||||
wait_for(lambda: listening(vector), processes, "vector engine not ready")
|
||||
|
||||
wait_for(
|
||||
lambda: f"Cluster: vector node 127.0.0.1:{vector} added to group ''"
|
||||
in (directory / "master.log").read_text(),
|
||||
processes,
|
||||
"vector engine never registered with the master",
|
||||
)
|
||||
|
||||
if worker:
|
||||
for i, port in enumerate((worker1, worker2)):
|
||||
start(f"worker-{i}", [worker, "--masters", master_addr, "--listen", f"127.0.0.1:{port}"])
|
||||
wait_for(lambda: listening(port), processes, "vector worker not ready")
|
||||
wait_for(
|
||||
lambda: all(
|
||||
f"S3 Vectors: discovered search worker 127.0.0.1:{port}"
|
||||
in (directory / "s3.log").read_text()
|
||||
for port in (worker1, worker2)
|
||||
),
|
||||
processes,
|
||||
"gateway did not discover both search workers",
|
||||
)
|
||||
yield {"endpoint": endpoint, "mode": mode, "processes": processes, "directory": directory}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vectors(cluster):
|
||||
return boto3.client("s3vectors", endpoint_url=cluster["endpoint"],
|
||||
aws_access_key_id="test", aws_secret_access_key="test", region_name="us-east-1",
|
||||
config=Config(connect_timeout=3, read_timeout=20, retries={"max_attempts": 0}))
|
||||
@@ -0,0 +1,2 @@
|
||||
boto3>=1.40
|
||||
pytest>=8
|
||||
@@ -1,65 +1,4 @@
|
||||
#!/usr/bin/env bash
|
||||
# Spin a mini cluster plus the vectors gateway and engine, run the boto3
|
||||
# conformance subset, tear everything down.
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
DATA=mini-data
|
||||
rm -rf "$DATA"
|
||||
PIDS=()
|
||||
cleanup() {
|
||||
kill "${PIDS[@]}" 2>/dev/null || true
|
||||
rm -rf "$DATA"
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
./weed mini -dir "$DATA" \
|
||||
-s3=false -webdav=false -admin.ui=false &
|
||||
PIDS+=($!)
|
||||
|
||||
./weed s3 \
|
||||
-filer 127.0.0.1:8888 \
|
||||
-config s3.json \
|
||||
-ip 127.0.0.1 \
|
||||
-ip.bind 127.0.0.1 \
|
||||
-port 8334 \
|
||||
-port.lance 9102 &
|
||||
PIDS+=($!)
|
||||
|
||||
./weed-vector \
|
||||
--namespace http://127.0.0.1:9102 \
|
||||
--masters 127.0.0.1:9333 \
|
||||
--storage-option aws_access_key_id=test \
|
||||
--storage-option aws_secret_access_key=test \
|
||||
--listen 127.0.0.1:9301 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Any HTTP response (even an auth rejection) means the lance catalog is up.
|
||||
for _ in $(seq 1 60); do
|
||||
code=$(curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:9102/v1/namespace/root/list || true)
|
||||
[ "$code" != "000" ] && break
|
||||
sleep 1
|
||||
done
|
||||
[ "${code:-000}" != "000" ] || { echo "lance catalog on :9102 never came up" >&2; exit 1; }
|
||||
|
||||
for _ in $(seq 1 60); do
|
||||
(exec 3<>/dev/tcp/127.0.0.1/9301) 2>/dev/null && break
|
||||
sleep 1
|
||||
done
|
||||
(exec 3<>/dev/tcp/127.0.0.1/9301) 2>/dev/null || { echo "vector server on :9301 never came up" >&2; exit 1; }
|
||||
|
||||
# A listening engine still has to land its KeepConnected registration before
|
||||
# the gateway's lazy discovery can find it — poll the cluster list for the
|
||||
# vector node rather than race the first request against it.
|
||||
vector_registered() {
|
||||
# grep -q would close the pipe on first match and SIGPIPE the shell mid-write
|
||||
# under `set -o pipefail`; a full read lets cluster.ps finish.
|
||||
printf 'cluster.ps\n' | ./weed shell -master 127.0.0.1:9333 2>/dev/null | grep "vector servers" >/dev/null
|
||||
}
|
||||
for _ in $(seq 1 60); do
|
||||
vector_registered && break
|
||||
sleep 1
|
||||
done
|
||||
vector_registered || { echo "vector server never registered with the master" >&2; exit 1; }
|
||||
|
||||
python3 -m pytest -xvs s3vectors_test.py
|
||||
exec "${PYTHON:-python3}" -m pytest -xvs s3vectors_test.py "$@"
|
||||
+154
-20
@@ -1,31 +1,17 @@
|
||||
"""Boto3 conformance subset for the S3 Vectors API gateway.
|
||||
|
||||
Run via `make test` — a mini cluster, an s3 gateway serving the vectors API on its main port, and a
|
||||
weed-vector are up by then.
|
||||
"""
|
||||
"""Run the S3 Vectors API against Go/Rust volumes and single/distributed queries."""
|
||||
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
from botocore.exceptions import ClientError
|
||||
import pytest
|
||||
|
||||
ENDPOINT = "http://127.0.0.1:8334"
|
||||
BUCKET = "vectors-bucket"
|
||||
INDEX = "idx"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vectors():
|
||||
return boto3.client(
|
||||
"s3vectors",
|
||||
endpoint_url=ENDPOINT,
|
||||
aws_access_key_id="test",
|
||||
aws_secret_access_key="test",
|
||||
region_name="us-east-1",
|
||||
)
|
||||
|
||||
|
||||
def test_vector_lifecycle(vectors):
|
||||
vectors.create_vector_bucket(vectorBucketName=BUCKET)
|
||||
buckets = [b["vectorBucketName"] for b in vectors.list_vector_buckets()["vectorBuckets"]]
|
||||
@@ -68,8 +54,7 @@ def test_vector_lifecycle(vectors):
|
||||
returnMetadata=True,
|
||||
filter={"color": {"$eq": "red"}},
|
||||
)["vectors"]
|
||||
# b is the second-nearest vector but blue, so an ignored filter would
|
||||
# return ["a", "b"] here.
|
||||
# b is closer than c but excluded by the filter.
|
||||
assert [h["key"] for h in hits] == ["a", "c"]
|
||||
distances = {h["key"]: h["distance"] for h in hits}
|
||||
assert distances["a"] == pytest.approx(0.0, abs=1e-6)
|
||||
@@ -110,3 +95,152 @@ def test_vector_bucket_policy(vectors):
|
||||
assert exc.value.response["Error"]["Code"] == "NotFoundException"
|
||||
finally:
|
||||
vectors.delete_vector_bucket(vectorBucketName=BUCKET)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def index(vectors):
|
||||
bucket = "vectors-" + uuid.uuid4().hex[:16]
|
||||
vectors.create_vector_bucket(vectorBucketName=bucket)
|
||||
vectors.create_index(vectorBucketName=bucket, indexName=INDEX, dataType="float32",
|
||||
dimension=4, distanceMetric="euclidean",
|
||||
metadataConfiguration={"nonFilterableMetadataKeys": ["private"]})
|
||||
args = {"vectorBucketName": bucket, "indexName": INDEX}
|
||||
yield args
|
||||
vectors.delete_index(**args)
|
||||
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
||||
|
||||
|
||||
def expect_error(code, operation, **kwargs):
|
||||
with pytest.raises(ClientError) as exc:
|
||||
operation(**kwargs)
|
||||
assert exc.value.response["Error"]["Code"] == code
|
||||
|
||||
|
||||
def test_empty_catalog(vectors):
|
||||
bucket = "empty-" + uuid.uuid4().hex[:16]
|
||||
assert vectors.list_vector_buckets(prefix=bucket)["vectorBuckets"] == []
|
||||
vectors.create_vector_bucket(vectorBucketName=bucket)
|
||||
try:
|
||||
assert vectors.list_indexes(vectorBucketName=bucket)["indexes"] == []
|
||||
finally:
|
||||
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
||||
|
||||
|
||||
def test_empty_index(vectors, index):
|
||||
assert vectors.get_vectors(**index, keys=["absent"])["vectors"] == []
|
||||
assert vectors.list_vectors(**index)["vectors"] == []
|
||||
assert vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=3)["vectors"] == []
|
||||
expect_error("ValidationException", vectors.query_vectors, **index,
|
||||
queryVector={"float32": [1, 0, 0]}, topK=3)
|
||||
|
||||
|
||||
def test_upsert_and_invalid_batch(vectors, index):
|
||||
vectors.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"revision": 1}}])
|
||||
vectors.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"revision": 2}}])
|
||||
expected = [{"key": "a", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"revision": 2}}]
|
||||
for invalid in (
|
||||
[{"key": "a", "data": {"float32": [1, 1, 1, 1]}}, {"key": "bad", "data": {"float32": [1, 2]}}],
|
||||
[{"key": "a", "data": {"float32": [1, 1, 1, 1]}}, {"key": "a", "data": {"float32": [2, 2, 2, 2]}}],
|
||||
):
|
||||
expect_error("ValidationException", vectors.put_vectors, **index, vectors=invalid)
|
||||
assert vectors.list_vectors(**index, returnData=True, returnMetadata=True)["vectors"] == expected
|
||||
|
||||
|
||||
def test_metadata_and_arn(vectors, index):
|
||||
description = vectors.get_index(**index)["index"]
|
||||
assert description["metadataConfiguration"] == {"nonFilterableMetadataKeys": ["private"]}
|
||||
arn = {"indexArn": description["indexArn"]}
|
||||
metadata = {"color": "red", "private": "secret", "nested": {"x": [1, True, "text"]}}
|
||||
vectors.put_vectors(**arn, vectors=[{"key": "quote'key", "data": {"float32": [1, 0, 0, 0]}, "metadata": metadata}])
|
||||
assert vectors.get_vectors(**arn, keys=["quote'key"], returnMetadata=True)["vectors"][0]["metadata"] == metadata
|
||||
hits = vectors.query_vectors(**arn, queryVector={"float32": [1, 0, 0, 0]}, topK=1, returnMetadata=True)["vectors"]
|
||||
assert hits[0]["metadata"] == metadata
|
||||
assert "distance" not in hits[0]
|
||||
expect_error("ValidationException", vectors.query_vectors, **arn, queryVector={"float32": [1, 0, 0, 0]},
|
||||
topK=1, filter={"$and": [{"color": "red"}, {"private": "secret"}]})
|
||||
vectors.delete_vectors(**arn, keys=["quote'key"])
|
||||
assert vectors.list_vectors(**arn)["vectors"] == []
|
||||
|
||||
|
||||
def test_list_pagination_and_segments(vectors, index):
|
||||
rows = [{"key": f"key-{i:02}", "data": {"float32": [i, 1, 0, 0]}, "metadata": {"row": i}} for i in range(17)]
|
||||
vectors.put_vectors(**index, vectors=rows)
|
||||
|
||||
def collect(**extra):
|
||||
result, tokens = [], set()
|
||||
while True:
|
||||
page = vectors.list_vectors(**index, maxResults=2, returnData=True, returnMetadata=True, **extra)
|
||||
result.extend(page["vectors"])
|
||||
token = page.get("nextToken")
|
||||
if not token:
|
||||
return result
|
||||
assert token not in tokens
|
||||
tokens.add(token)
|
||||
extra["nextToken"] = token
|
||||
|
||||
assert sorted(collect(), key=lambda r: r["key"]) == rows
|
||||
segmented = [row for segment in range(4) for row in collect(segmentCount=4, segmentIndex=segment)]
|
||||
assert sorted(segmented, key=lambda r: r["key"]) == rows
|
||||
expect_error("ValidationException", vectors.list_vectors, **index, nextToken="invalid")
|
||||
|
||||
|
||||
def test_filter_types(vectors, index):
|
||||
values = [1, "1", True, "true"]
|
||||
vectors.put_vectors(**index, vectors=[{"key": str(i), "data": {"float32": [1, 0, 0, 0]}, "metadata": {"value": value}}
|
||||
for i, value in enumerate(values)])
|
||||
for i, value in enumerate(values):
|
||||
hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=4,
|
||||
filter={"value": {"$eq": value}})["vectors"]
|
||||
assert [hit["key"] for hit in hits] == [str(i)]
|
||||
|
||||
|
||||
def test_cosine_metric(vectors):
|
||||
bucket = "cosine-" + uuid.uuid4().hex[:16]
|
||||
vectors.create_vector_bucket(vectorBucketName=bucket)
|
||||
args = {"vectorBucketName": bucket, "indexName": INDEX}
|
||||
vectors.create_index(**args, dataType="float32", dimension=2, distanceMetric="cosine")
|
||||
try:
|
||||
vectors.put_vectors(**args, vectors=[{"key": "far", "data": {"float32": [10, 0]}},
|
||||
{"key": "near", "data": {"float32": [1, 1]}}])
|
||||
result = vectors.query_vectors(**args, queryVector={"float32": [1, 0]}, topK=2, returnDistance=True)
|
||||
assert result["distanceMetric"] == "cosine"
|
||||
assert [v["key"] for v in result["vectors"]] == ["far", "near"]
|
||||
assert result["vectors"][0]["distance"] == pytest.approx(0, abs=1e-6)
|
||||
expect_error("ValidationException", vectors.query_vectors, **args, queryVector={"float32": [1, 0, 0]}, topK=1)
|
||||
finally:
|
||||
vectors.delete_index(**args)
|
||||
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
||||
|
||||
|
||||
def test_delete_bucket_conflict(vectors, index):
|
||||
expect_error("ConflictException", vectors.delete_vector_bucket, vectorBucketName=index["vectorBucketName"])
|
||||
assert vectors.get_index(**index)["index"]["dimension"] == 4
|
||||
|
||||
|
||||
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")
|
||||
with pytest.raises(ClientError) as exc:
|
||||
client.list_vector_buckets()
|
||||
assert exc.value.response["ResponseMetadata"]["HTTPStatusCode"] == 403
|
||||
|
||||
|
||||
def test_query_without_writer(vectors, index, cluster):
|
||||
from conftest import stop
|
||||
|
||||
vectors.put_vectors(**index, vectors=[{"key": "survivor", "data": {"float32": [1, 0, 0, 0]}}])
|
||||
stop(cluster["processes"].pop("vector"))
|
||||
args = dict(index, queryVector={"float32": [1, 0, 0, 0]}, topK=1)
|
||||
if cluster["mode"] == "single":
|
||||
expect_error("ServiceUnavailableException", vectors.query_vectors, **args)
|
||||
return
|
||||
deadline = time.monotonic() + 15
|
||||
while True:
|
||||
try:
|
||||
result = vectors.query_vectors(**args)
|
||||
break
|
||||
except ClientError as exc:
|
||||
if exc.response["Error"]["Code"] != "ServiceUnavailableException" or time.monotonic() >= deadline:
|
||||
raise
|
||||
time.sleep(0.2)
|
||||
assert [v["key"] for v in result["vectors"]] == ["survivor"]
|
||||
Reference in new issue
Block a user