seaweedfs/test/s3vectors/perf_test.py

"""Measure the S3 Vectors data plane end to end against the cluster fixture.

Opt-in through run_perf.sh or `make perf`; run_test.sh does not collect it.
Sizes come from the environment so the same module runs a quick smoke pass or
a larger profile:

    PERF_VECTORS      vectors to ingest (10000)
    PERF_DIM          dimension (128)
    PERF_BATCH        vectors per PutVectors call, at most 500 (500)
    PERF_QUERIES      timed queries per scenario (100)
    PERF_TOPK         topK (10)
    PERF_CONCURRENCY  client threads for the throughput scenario ("1,4,8")
    PERF_RECALL       queries checked against an exact search (20)
    PERF_REPORT       JSON report path (perf-results.json in the cluster dir)

Each timed scenario starts with one untimed call. Timings only report; the
assertions check correctness (recall against an exact client-side search,
complete listings), never speed.
"""

import concurrent.futures
import json
import math
import os
from pathlib import Path
import platform
import random
import re
import statistics
import time
import uuid

import boto3
from botocore.config import Config
import pytest

VECTORS = int(os.environ.get("PERF_VECTORS", "10000"))
DIM = int(os.environ.get("PERF_DIM", "128"))
BATCH = int(os.environ.get("PERF_BATCH", "500"))
QUERIES = int(os.environ.get("PERF_QUERIES", "100"))
TOPK = int(os.environ.get("PERF_TOPK", "10"))
CONCURRENCY = [int(c) for c in os.environ.get("PERF_CONCURRENCY", "1,4,8").split(",")]
CATEGORIES = 10
RECALL_SAMPLE = int(os.environ.get("PERF_RECALL", "20"))
INDEX = "perf"

REPORT = {}


def summarize(seconds):
    ms = sorted(s * 1000 for s in seconds)

    def pct(p):
        return ms[min(len(ms) - 1, max(0, math.ceil(p * len(ms)) - 1))]

    return {"n": len(ms), "mean_ms": statistics.fmean(ms), "p50_ms": pct(0.50),
            "p95_ms": pct(0.95), "p99_ms": pct(0.99), "max_ms": ms[-1]}


def timed(call):
    start = time.perf_counter()
    result = call()
    return time.perf_counter() - start, result


def client(endpoint):
    return boto3.client("s3vectors", endpoint_url=endpoint, aws_access_key_id="test",
                        aws_secret_access_key="test", region_name="us-east-1",
                        config=Config(connect_timeout=3, read_timeout=300, retries={"max_attempts": 0},
                                      max_pool_connections=max(CONCURRENCY) + 2))


def exact_top_k(data, query, k, keep=None):
    """Cosine top-k over the ingested vectors, the reference for recall."""
    qn = math.sqrt(sum(x * x for x in query))
    scored = []
    for i, (vector, norm) in enumerate(data):
        if keep is not None and not keep(i):
            continue
        scored.append((sum(a * b for a, b in zip(vector, query)) / (norm * qn), i))
    scored.sort(reverse=True)
    return {f"v{i}" for _, i in scored[:k]}


def list_all(s3v, args, **extra):
    keys, token = [], None
    while True:
        page = s3v.list_vectors(**args, maxResults=500, **extra, **({"nextToken": token} if token else {}))
        keys.extend(vector["key"] for vector in page["vectors"])
        token = page.get("nextToken")
        if not token:
            return keys


def assert_complete(keys):
    assert len(keys) == VECTORS and set(keys) == {f"v{i}" for i in range(VECTORS)}


def recall(expected, hits):
    return len(expected & {h["key"] for h in hits}) / len(expected)


@pytest.fixture(scope="module")
def perf(cluster):
    assert 1 <= BATCH <= 500 and VECTORS >= BATCH
    rng = random.Random(1)
    data = []
    for _ in range(VECTORS):
        vector = [round(rng.gauss(0, 1), 5) for _ in range(DIM)]
        data.append((vector, math.sqrt(sum(x * x for x in vector))))
    queries = [[round(rng.gauss(0, 1), 5) for _ in range(DIM)] for _ in range(QUERIES)]

    s3v = client(cluster["endpoint"])
    bucket = "perf-" + uuid.uuid4().hex[:12]
    s3v.create_vector_bucket(vectorBucketName=bucket)
    s3v.create_index(vectorBucketName=bucket, indexName=INDEX, dataType="float32",
                     dimension=DIM, distanceMetric="cosine")
    args = {"vectorBucketName": bucket, "indexName": INDEX}
    combo = re.sub(r"^vectors-|\d+$", "", cluster["directory"].name)
    REPORT[combo] = {"config": {"vectors": VECTORS, "dim": DIM, "batch": BATCH, "queries": QUERIES,
                                "top_k": TOPK, "concurrency": CONCURRENCY, "cpus": os.cpu_count(),
                                "load_1m": os.getloadavg()[0], "machine": platform.machine()}}
    yield {"client": s3v, "endpoint": cluster["endpoint"], "args": args, "data": data,
           "queries": queries, "report": REPORT[combo]}
    s3v.delete_index(**args)
    s3v.delete_vector_bucket(vectorBucketName=bucket)
    path = Path(os.environ.get("PERF_REPORT", cluster["directory"] / "perf-results.json"))
    path.write_text(json.dumps(REPORT, indent=2))
    print_report(combo, REPORT[combo], path)


def test_ingest(perf):
    s3v, args, data = perf["client"], perf["args"], perf["data"]
    latencies = []
    start = time.perf_counter()
    for offset in range(0, VECTORS, BATCH):
        batch = [{"key": f"v{i}", "data": {"float32": data[i][0]},
                  "metadata": {"category": f"c{i % CATEGORIES}", "n": i}}
                 for i in range(offset, min(offset + BATCH, VECTORS))]
        seconds, _ = timed(lambda: s3v.put_vectors(**args, vectors=batch))
        latencies.append(seconds)
    total = time.perf_counter() - start
    tenth = max(1, len(latencies) // 10)
    perf["report"]["ingest"] = {
        "vectors_per_s": VECTORS / total, "seconds": total, "calls": summarize(latencies),
        "first_tenth_mean_ms": statistics.fmean(latencies[:tenth]) * 1000,
        "last_tenth_mean_ms": statistics.fmean(latencies[-tenth:]) * 1000,
    }


def test_upsert_existing(perf):
    s3v, args, data = perf["client"], perf["args"], perf["data"]
    keys = random.Random(2).sample(range(VECTORS), BATCH)
    batch = [{"key": f"v{i}", "data": {"float32": data[i][0]},
              "metadata": {"category": f"c{i % CATEGORIES}", "n": i}} for i in keys]
    latencies = [timed(lambda: s3v.put_vectors(**args, vectors=batch))[0] for _ in range(5)]
    perf["report"]["upsert_existing"] = summarize(latencies)
    assert_complete(list_all(s3v, args))


def run_queries(perf, name, **extra):
    s3v, args, data, queries = perf["client"], perf["args"], perf["data"], perf["queries"]
    latencies, recalls = [], []
    keep = None
    if "filter" in extra:
        keep = lambda i: i % CATEGORIES == 3  # noqa: E731
    s3v.query_vectors(**args, queryVector={"float32": queries[0]}, topK=TOPK, **extra)
    for n, query in enumerate(queries):
        seconds, response = timed(lambda: s3v.query_vectors(
            **args, queryVector={"float32": query}, topK=TOPK, returnDistance=True, **extra))
        latencies.append(seconds)
        if n < RECALL_SAMPLE:
            recalls.append(recall(exact_top_k(data, query, TOPK, keep), response["vectors"]))
    perf["report"][name] = {**summarize(latencies), "recall_at_k": statistics.fmean(recalls) if recalls else None}
    assert not recalls or statistics.fmean(recalls) >= 0.9


def test_query(perf):
    run_queries(perf, "query")


def test_query_with_metadata(perf):
    run_queries(perf, "query_return_metadata", returnMetadata=True)


def test_query_filtered(perf):
    run_queries(perf, "query_filtered_10pct", filter={"category": {"$eq": "c3"}})


def test_query_throughput(perf):
    args, queries = perf["args"], perf["queries"]
    results = {}
    for threads in CONCURRENCY:
        clients = [client(perf["endpoint"]) for _ in range(threads)]
        for s3v in clients:
            s3v.query_vectors(**args, queryVector={"float32": queries[0]}, topK=TOPK)
        work = queries * threads

        def run(worker):
            s3v = clients[worker]
            return [timed(lambda: s3v.query_vectors(**args, queryVector={"float32": query}, topK=TOPK))[0]
                    for query in work[worker::threads]]

        start = time.perf_counter()
        with concurrent.futures.ThreadPoolExecutor(threads) as pool:
            latencies = [s for chunk in pool.map(run, range(threads)) for s in chunk]
        elapsed = time.perf_counter() - start
        results[str(threads)] = {"qps": len(work) / elapsed, **summarize(latencies)}
    perf["report"]["query_throughput"] = results


def test_get_vectors(perf):
    s3v, args = perf["client"], perf["args"]
    rng = random.Random(3)
    size = min(100, VECTORS)
    s3v.get_vectors(**args, keys=["v0"])
    latencies = []
    for _ in range(max(1, QUERIES // 2)):
        keys = [f"v{i}" for i in rng.sample(range(VECTORS), size)]
        seconds, response = timed(lambda: s3v.get_vectors(
            **args, keys=keys, returnData=True, returnMetadata=True))
        latencies.append(seconds)
        assert len(response["vectors"]) == size
    perf["report"]["get_vectors"] = {**summarize(latencies), "keys": size}


def test_list_vectors(perf):
    args = perf["args"]
    results = {}
    for segments in (1, 4, 16):
        clients = [client(perf["endpoint"]) for _ in range(segments)]
        for s3v in clients:
            s3v.list_vectors(**args, maxResults=1)

        def scan(segment):
            extra = {"segmentCount": segments, "segmentIndex": segment} if segments > 1 else {}
            return list_all(clients[segment], args, **extra)

        start = time.perf_counter()
        with concurrent.futures.ThreadPoolExecutor(segments) as pool:
            listed = [key for keys in pool.map(scan, range(segments)) for key in keys]
        elapsed = time.perf_counter() - start
        assert_complete(listed)
        results[str(segments)] = {"seconds": elapsed, "vectors_per_s": VECTORS / elapsed}
    perf["report"]["list_vectors_segments"] = results


def test_delete_vectors(perf):
    s3v, args = perf["client"], perf["args"]
    size = min(100, VECTORS)
    keys = [f"v{i}" for i in random.Random(4).sample(range(VECTORS), size * min(5, VECTORS // size))]
    latencies = [timed(lambda: s3v.delete_vectors(**args, keys=keys[n:n + size]))[0]
                 for n in range(0, len(keys), size)]
    perf["report"]["delete_vectors"] = {**summarize(latencies), "keys": size}


def print_report(combo, report, path):
    cfg = report["config"]
    lines = [f"\n== S3 Vectors perf [{combo}]: {cfg['vectors']} x {cfg['dim']}d cosine, "
             f"topK {cfg['top_k']}, {cfg['cpus']} CPUs, load {cfg['load_1m']:.1f}"]
    ingest = report.get("ingest")
    if ingest:
        lines.append(f"ingest            {ingest['vectors_per_s']:9.0f} vectors/s  "
                     f"batch p50 {ingest['calls']['p50_ms']:.0f} ms, first tenth "
                     f"{ingest['first_tenth_mean_ms']:.0f} ms -> last tenth {ingest['last_tenth_mean_ms']:.0f} ms")
    for name in ("upsert_existing", "query", "query_return_metadata", "query_filtered_10pct",
                 "get_vectors", "delete_vectors"):
        if name in report:
            r = report[name]
            label = f"{name} ({r['keys']} keys)" if "keys" in r else name
            extra = f"  recall@k {r['recall_at_k']:.3f}" if r.get("recall_at_k") is not None else ""
            lines.append(f"{label:<26} p50 {r['p50_ms']:7.1f} ms  p95 {r['p95_ms']:7.1f} ms  "
                         f"p99 {r['p99_ms']:7.1f} ms{extra}")
    for threads, r in report.get("query_throughput", {}).items():
        lines.append(f"query x{threads:<3} threads    {r['qps']:7.1f} qps   p50 {r['p50_ms']:7.1f} ms  "
                     f"p95 {r['p95_ms']:7.1f} ms")
    for segments, r in report.get("list_vectors_segments", {}).items():
        lines.append(f"list {segments:>2} segments     {r['vectors_per_s']:9.0f} vectors/s  "
                     f"({r['seconds']:.2f} s)")
    lines.append(f"report: {path}")
    print("\n".join(lines))