From 96f8c27694057ad02dbc912e2d920e07bbf1b9de Mon Sep 17 00:00:00 2001 From: Perry Krug <299362+krugster@users.noreply.github.com> Date: Wed, 16 Apr 2025 09:56:15 +0100 Subject: [PATCH 1/4] Update README.md --- README.md | 63 +++++++++++++++++++------------------------------------ 1 file changed, 21 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index 54d0ee9..ea5301a 100644 --- a/README.md +++ b/README.md @@ -110,7 +110,6 @@ The following databases are currently supported by VSB: * [Pinecone](vsb/databases/pinecone/README.md) * [pgvector](vsb/databases/pgvector/README.md) -* [OpenSearch](vsb/databases/opensearch/README.md) > [!TIP] > You can also display the list of supported databases using the following command: @@ -158,6 +157,7 @@ Common parameters include the following: * `--users=`: Specify the number of concurrent users (connections) to simulate. * `--skip_populate`: Skip populating the database with data and immediately perform the Run phase. +* `--namespace_name`: Which namespace to write into (defaults to "") ## Use cases @@ -174,28 +174,12 @@ range of vector databases. It can be used to perform a range of tasks, including ### Synthetic Workloads Sometimes the workload you want to model doesn't exist in any provided dataset, or you want -to test a specific aspect of a database's performance. In these cases, you can use -a _synthetic_ workload to generate custom workloads with specific -characteristics. - -There are three modes of synthetic workloads, however the most common is the -`synthetic-proportional` workload: - -* `synthetic-proportional` workloads populate the database with an initial set of -records, then run a series of assorted request operations (inserts, queries, deletes, updates) in proportion to the given ratios. - -* `synthetic` workloads generate a fixed number of records and queries with the given -distribution, then runs population and query phases in sequence. - -* `synthetic-runbook` workloads generate a fixed number of records and queries, and -splits them into a series of multiple 'populate -> run' steps. This is useful - for testing how a database performs when data is loaded incrementally. - -Some important parameters for synthetic workloads include: +to test a specific aspect of a database's performance. In these cases, you can use the `synthetic`, `synthetic-runbook`, +and `synthetic-proportional` workloads to generate custom workloads with specific characteristics. Some +important parameters for synthetic workloads include: * `--synthetic_records`: The number of records to generate for the synthetic workload. -* `--synthetic_requests`: The number of requests to generate for the run phase of - the synthetic workload. +* `--synthetic_requests`: The number of requests to generate for the synthetic workload. * `--synthetic_dimensions`: The dimensionality of generated vectors. * `--synthetic_query_distribution`: The distribution of query/fetch IDs for synthetic proportional workloads. * `--synthetic_record_ratio`: The distribution of record vectors in space for synthetic proportional workloads. @@ -203,39 +187,34 @@ Some important parameters for synthetic workloads include: * `--synthetic_query_ratio`: The proportion of query operations for synthetic proportional workloads. * `--synthetic_metadata`: The metadata schema to use for each record. -#### **Defining Synthetic Metadata** +Metadata is specified by providing multiple `--synthetic_metadata` flags describing a metadata field +with a key and supported value. Values can be `<# digits>n` for a random integer, `<# chars>s` for a +random string, `<# chars>s<# strings>l` for a random list of strings, or `b` for a random boolean. -Metadata can optionally be generated for each record in a synthetic workload. -Metadata is specified using one or more `--synthetic_metadata` flags. Each flag -defines a *metadata field* with a *name* and a *format specification*. +For example, `--synthetic_metadata=id:10n` generates a metadata field `id` with a random 10-digit integer. +`--synthetic_metadata=tags:5s10l` generates a metadata field `tags` with a list of 10 ranodm strings, each +5 characters long. `--synthetic_metadata=active:b` generates a metadata field `active` with a random boolean. -##### **Supported Metadata Types** -| Type | Format Spec | Example | -|------|------------------------|---------| -| Random integer with `#` digits | `<# digits>n` | `id:10n` → `{"id": 1234567890}` | -| Random alphanumeric string of `#` characters | `<# chars>s` | `username:8s` → `{"username": "aZb3Xy91"}` | -| List of `<# items>` strings, each of length `<# chars>` | `<# chars>s<# items>l` | `tags:5s10l` → `{"tags": ["apple", "delta", "omega", ...]}` | -| Random boolean (`true` or `false`) | `b` | `active:b` → `{"active": true}` | +You can see the full list of parameters by running `vsb --help`. -##### **Example Metadata Usage** -- `--synthetic_metadata=id:10n` → Generates a numeric `id` with 10 random digits* -- `--synthetic_metadata=tags:5s10l` → Generates a `tags` field containing a list of 10 random words, each 5 characters long. -- `--synthetic_metadata=username:8s` → Generates a random username with 8 characters. -- `--synthetic_metadata=active:b` → Generates an active/inactive flag as `true` or `false`. +`synthetic` workloads generate a fixed number of records and queries with the given distribution, then +runs population and query phases in sequence. -You can see the full list of parameters by running `vsb --help`. +`synthetic-runbook` workloads generate a fixed number of records and queries, and splits them into a series of +'populate -> run' steps. This is useful for testing how a database performs when data is loaded incrementally. +`synthetic-proportional` workloads populate the database with an initial set of records, then run a series of +assorted request operations (inserts, queries, deletes, updates) in proportion to the given ratios. **Example** -The following command runs a synthetic workload against Pinecone, with 1,000 initial records, -then performs 100 requests in a zipfian query distribution, made up of 30% inserts, -50% queries, 10% deletes, and 10% updates: +The following command runs a synthetic workload against Pinecone, with 10,000 initial records, +a zipfian query distribution, and 30% inserts, 50% queries, 10% deletes, and 10% updates: ```shell vsb --database=pinecone --pinecone_api_key= \ --workload=synthetic-proportional \ - --synthetic_records=1000 --synthetic_requests=100 \ + --synthetic_records=10000 --synthetic_queries=1000000 \ --synthetic_insert_ratio=0.3 --synthetic_query_ratio=0.5 \ --synthetic_delete_ratio=0.1 --synthetic_update_ratio=0.1 \ --synthetic_dimensions=768 --synthetic_query_distribution=zipfian \ From 270cfb9c6af523a6690173d1b00f6933340283bf Mon Sep 17 00:00:00 2001 From: Perry Krug <299362+krugster@users.noreply.github.com> Date: Wed, 16 Apr 2025 09:57:03 +0100 Subject: [PATCH 2/4] Update cmdline_args.py --- vsb/cmdline_args.py | 92 +++------------------------------------------ 1 file changed, 6 insertions(+), 86 deletions(-) diff --git a/vsb/cmdline_args.py b/vsb/cmdline_args.py index a41f801..352fe0d 100644 --- a/vsb/cmdline_args.py +++ b/vsb/cmdline_args.py @@ -109,6 +109,12 @@ def add_vsb_cmdline_args( "then the target will be distributed across all users. " "Specify 0 for unlimited. Default is %(default)s.", ) + general_group.add_argument( + "--namespace_name", + type=str, + default="", + help="Name of namespace to use. One will be created if it does not exist. Default is vsb-.", + ) if include_locust_args: general_group.add_argument( @@ -278,7 +284,6 @@ def add_vsb_cmdline_args( help="Name of Pinecone index to connect to. One will be created if it does not exist. Default is vsb-.", env_var="VSB__PINECONE_INDEX_NAME", ) - pinecone_group.add_argument( "--pinecone_index_spec", type=json_to_pinecone_spec, @@ -286,89 +291,6 @@ def add_vsb_cmdline_args( help="JSON spec of Pinecone index to create (if it does not exist). Default is %(default)s.", ) - opensearch_group = parser.add_argument_group( - "Options specific to OpenSearch database" - ) - opensearch_group.add_argument( - "--opensearch_index_name", - type=str, - default=None, - help="Name of OpenSearch index to connect to the OpenSearch Collection. One will be created if it does not exist. Default is vsb-.", - env_var="VSB__OPENSEARCH_INDEX_NAME", - ) - opensearch_group.add_argument( - "--opensearch_host", - type=str, - default="localhost", - help="opensearch host to connect to the OpenSearch Collection. Default is %(default)s.", - env_var="VSB__OPENSEARCH_HOST", - ) - opensearch_group.add_argument( - "--opensearch_port", - type=int, - default=9200, - help="opensearch port to connect to the OpenSearch Collection. Default is %(default)s.", - env_var="VSB__OPENSEARCH_PORT", - ) - opensearch_group.add_argument( - "--opensearch_region", - type=str, - default="us-east-1", - help="opensearch region to connect to the OpenSearch Collection. Default is %(default)s.", - env_var="VSB__OPENSEARCH_REGION", - ) - opensearch_group.add_argument( - "--opensearch_service", - type=str, - default="aoss", - help="opensearch service to connect. Default is %(default)s. which is for Amazon OpenSearch Serverless. For Amazon OpenSearch Managed cluster, use 'es'.", - env_var="VSB__OPENSEARCH_SERVICE", - ) - opensearch_group.add_argument( - "--aws_access_key", - type=str, - default=None, - help="AWS access key to connect to the OpenSearch Collection. Default is %(default)s.", - env_var="VSB__AWS_ACCESS_KEY", - ) - opensearch_group.add_argument( - "--aws_secret_key", - type=str, - default=None, - help="AWS secret access key to connect to the OpenSearch Collection. Default is %(default)s.", - env_var="VSB__AWS_SECRET_KEY", - ) - opensearch_group.add_argument( - "--aws_session_token", - type=str, - default=None, - help="AWS session token to connect to the OpenSearch Collection. Default is %(default)s.", - env_var="VSB__AWS_SESSION_TOKEN", - ) - opensearch_group.add_argument( - "--opensearch_username", - type=str, - default=None, - help="Opensearch username to use. If specified then must also specify " - "--opensearch_password. Default is %(default)s.", - env_var="VSB__OPENSEARCH_USERNAME", - ) - opensearch_group.add_argument( - "--opensearch_password", - type=str, - default=None, - help="Opensearch password to use. If specified then must also specify " - "--opensearch_username. Default is %(default)s.", - env_var="VSB__OPENSEARCH_PASSWORD", - ) - opensearch_group.add_argument( - "--opensearch_use_tls", - action=argparse.BooleanOptionalAction, - default=True, - help="Should Opensearch use SSL / TLS to connect to the server. Default is %(default)s.", - env_var="VSB__OPENSEARCH_USE_TLS", - ) - pgvector_group = parser.add_argument_group("Options specific to pgvector database") pgvector_group.add_argument( "--pgvector_host", @@ -482,8 +404,6 @@ def validate_parsed_args( "The following arguments must be specified when --database is " "'pinecone'" + formatter.format_help(), ) - case "opensearch": - pass case "pgvector": pass case _: From 4b49e9406a66733684511d8574e8cce6b19c5b95 Mon Sep 17 00:00:00 2001 From: Perry Krug <299362+krugster@users.noreply.github.com> Date: Wed, 16 Apr 2025 09:58:04 +0100 Subject: [PATCH 3/4] Update users.py --- vsb/users.py | 65 +++++++++++++++------------------------------------- 1 file changed, 19 insertions(+), 46 deletions(-) diff --git a/vsb/users.py b/vsb/users.py index 24ee12b..56b8a31 100644 --- a/vsb/users.py +++ b/vsb/users.py @@ -15,7 +15,6 @@ from vsb import metrics, metrics_tracker from vsb.databases import DB import vsb.metrics_tracker -from vsb.randomized_pacer import RandomizedPacer from vsb.vsb_types import ( RecordList, SearchRequest, @@ -75,20 +74,6 @@ def setup(self): # Nothing more to do, but sleep briefly here to prevent # us busy-looping in this state. gevent.sleep(0.1) - except StopUser: - # Explicit stop, assume callee already logged the reason; exit with - # non-zero: - logger.debug("Stopping user due to StopUser exception") - self.environment.runner.send_message( - "update_progress", - { - "user": 0, - "phase": "setup", - "error": "Stopping user due to StopUser exception", - }, - ) - self.state = self.State.Done - except Exception as e: traceback.print_exception(e) locust.log.unhandled_greenlet_exception = True @@ -244,10 +229,9 @@ def __init__(self, environment): self.workload = environment.workload_sequence[iteration] self.state = RunUser.State.Active opts = environment.parsed_options - target_throughput = opts.requests_per_sec / float(opts.num_users) - self.pacer = RandomizedPacer(target_throughput) + self.target_throughput = opts.requests_per_sec / float(opts.num_users) logger.debug( - f"Initialising RunUser id:{self.user_id}, target request/sec:{target_throughput}" + f"Initialising RunUser id:{self.user_id}, target request/sec:{self.target_throughput}" ) self.query_iter = None @@ -265,7 +249,9 @@ def wait_time(self): """Method called by locust to control how long this task should wait between executions. """ - return self.pacer.wait_time() + if self.target_throughput > 0: + return constant_throughput(self.target_throughput)(self) + return 0 def do_run(self): if not self.query_iter: @@ -288,18 +274,19 @@ def do_run(self): return try: index = self.database.get_namespace(tenant) + namespace = self.environment.parsed_options.namespace_name or "" start = time.perf_counter() match request: case SearchRequest(): - results = index.search(request) + results = index.search(request,namespace) case InsertRequest(): - results = index.insert_batch(request.records) + results = index.insert_batch(request.records,namespace) case UpdateRequest(): - results = index.update_batch(request.records) + results = index.update_batch(request.records,namespace) case FetchRequest(): - results = index.fetch_batch(request.ids) + results = index.fetch_batch(request.ids,namespace) case DeleteRequest(): - results = index.delete_batch(request.ids) + results = index.delete_batch(request.ids,namespace) case _: raise ValueError(f"Unknown request type:{request}") stop = time.perf_counter() @@ -389,9 +376,6 @@ class Phase(Enum): """Issue requests (queries) to the database and recording the results.""" Done = auto() """Final phase when all Run Users have completed""" - DoneFailed = auto() - """Final phase when all a fatal error occurred and we want to stop down, - signalling an unsuccessful run""" def __init__(self): super().__init__() @@ -486,10 +470,6 @@ def tick(self): return self.num_users, self.num_users, [RunUser] case LoadShape.Phase.Done: return None - case LoadShape.Phase.DoneFailed: - # Set exit status to 2 to indicate a failure - locust.log.unhandled_greenlet_exception = True - return None case _: raise ValueError(f"Invalid phase:{self.phase}") @@ -507,11 +487,9 @@ def _transition_phase(self, new: Phase): ] if vsb.progress is not None: self._update_progress_bar(mark_completed=True) - if new == LoadShape.Phase.DoneFailed: - msg = f"✘ {self.phase.name} cancelled" - else: - msg = f"✔ {self.phase.name} complete" - vsb.progress.update(self.progress_task_id, description=msg) + vsb.progress.update( + self.progress_task_id, description=f"✔ {self.phase.name} complete" + ) vsb.progress.stop() vsb.progress = None if hasattr(self.runner.environment, "workload_sequence"): @@ -537,10 +515,7 @@ def _transition_phase(self, new: Phase): else: phase_display_name = self.phase.name if self.phase in tracked_phases: - # Skip reporting phase end on a fatal error; it adds noise after - # the important fatal error message. - if new != LoadShape.Phase.DoneFailed: - metrics_tracker.record_phase_end(phase_display_name) + metrics_tracker.record_phase_end(phase_display_name) self.phase = new if hasattr(self.runner.environment, "workload_sequence"): workload = self.runner.environment.workload_sequence[ @@ -565,7 +540,9 @@ def _transition_phase(self, new: Phase): def on_update_progress(self, msg, **kwargs): # Fired when VSBLoadShape (running on the master) receives an # "update_progress" message. - logger.debug(f"VSBLoadShape.on_update_progress() - {msg.data}") + logger.debug( + f"VSBLoadShape.update_progress() - user:{msg.data['user']}, phase:{msg.data['phase']}" + ) match self.phase: case LoadShape.Phase.Setup: assert msg.data["phase"] == "setup" @@ -581,11 +558,7 @@ def on_update_progress(self, msg, **kwargs): f"{self.request_count} - " f"moving to TransitionFromSetup phase" ) - if "error" in msg.data: - # Setup signalled an error; cancel the benchmark - self._transition_phase(LoadShape.Phase.DoneFailed) - else: - self._transition_phase(LoadShape.Phase.TransitionFromSetup) + self._transition_phase(LoadShape.Phase.TransitionFromSetup) case LoadShape.Phase.TransitionFromSetup: logger.error( f"VSBLoadShape.update_progress() - Unexpected progress update in " From fd644f55291fe5116c928e8b0558374870bc2923 Mon Sep 17 00:00:00 2001 From: Perry Krug <299362+krugster@users.noreply.github.com> Date: Wed, 16 Apr 2025 09:58:52 +0100 Subject: [PATCH 4/4] Update pinecone.py --- vsb/databases/pinecone/pinecone.py | 66 ++++++++++++------------------ 1 file changed, 26 insertions(+), 40 deletions(-) diff --git a/vsb/databases/pinecone/pinecone.py b/vsb/databases/pinecone/pinecone.py index e8c1158..d9d268b 100644 --- a/vsb/databases/pinecone/pinecone.py +++ b/vsb/databases/pinecone/pinecone.py @@ -4,7 +4,7 @@ import vsb from vsb import logger -from pinecone import PineconeException, NotFoundException, UnauthorizedException +from pinecone import PineconeException, NotFoundException from pinecone.grpc import PineconeGRPC, GRPCIndex from tenacity import retry, stop_after_attempt, wait_exponential_jitter, after_log import grpc.experimental.gevent as grpc_gevent @@ -24,17 +24,18 @@ class PineconeNamespace(Namespace): def __init__(self, index: GRPCIndex, namespace: str): # TODO: Support multiple namespaces self.index = index + self.namespace = Namespace - def insert_batch(self, batch: RecordList): + def insert_batch(self, batch: RecordList,namespace): # Pinecone expects a list of dicts (or tuples). dicts = [dict(rec) for rec in batch] - self.index.upsert(dicts) + self.index.upsert(dicts,namespace=namespace) - def update_batch(self, batch: list[Record]): + def update_batch(self, batch: list[Record],namespace): # Pinecone treats insert and update as the same operation. - self.insert_batch(batch) + self.insert_batch(batch,namespace=namespace) - def search(self, request: SearchRequest) -> list[str]: + def search(self, request: SearchRequest,namespace) -> list[str]: @retry( wait=wait_exponential_jitter(initial=0.1, jitter=0.1), stop=stop_after_attempt(5), @@ -42,18 +43,18 @@ def search(self, request: SearchRequest) -> list[str]: ) def do_query_with_retry(): return self.index.query( - vector=request.values, top_k=request.top_k, filter=request.filter + vector=request.values, top_k=request.top_k, filter=request.filter,namespace=namespace ) result = do_query_with_retry() matches = [m["id"] for m in result["matches"]] return matches - def fetch_batch(self, request: list[str]) -> list[Record]: - return self.index.fetch(request).vectors.values + def fetch_batch(self, request: list[str],namespace) -> list[Record]: + return self.index.fetch(request,namespace=namespace).vectors.values - def delete_batch(self, request: list[str]): - self.index.delete(request) + def delete_batch(self, request: list[str],namespace): + self.index.delete(request,namespace=namespace) class PineconeDB(DB): @@ -68,52 +69,37 @@ def __init__( self.pc = PineconeGRPC(config["pinecone_api_key"]) self.skip_populate = config["skip_populate"] self.overwrite = config["overwrite"] - self.index_name = config["pinecone_index_name"] - if self.index_name is None: + index_name = config["pinecone_index_name"] + if index_name is None: # None specified, default to "vsb-" - self.index_name = f"vsb-{name}" + index_name = f"vsb-{name}" spec = config["pinecone_index_spec"] try: - self.index = self.pc.Index(name=self.index_name) + self.index = self.pc.Index(name=index_name) self.created_index = False - except UnauthorizedException: - api_key = config["pinecone_api_key"] - masked_api_key = api_key[:4] + "*" * (len(api_key) - 8) + api_key[-4:] - logger.critical( - f"PineconeDB: Got UnauthorizedException when attempting to connect " - f"to index '{self.index_name}' using API key '{masked_api_key}' - check " - f"your API key and permissions" - ) - raise StopUser() except NotFoundException: logger.info( - f"PineconeDB: Specified index '{self.index_name}' was not found, or the " - f"specified API key cannot access it. Creating new index '{self.index_name}'." + f"PineconeDB: Specified index '{index_name}' was not found, or the " + f"specified API key cannot access it. Creating new index '{index_name}'." ) self.pc.create_index( - name=self.index_name, - dimension=dimensions, - metric=metric.value, - spec=spec, + name=index_name, dimension=dimensions, metric=metric.value, spec=spec ) - self.index = self.pc.Index(name=self.index_name) + self.index = self.pc.Index(name=index_name) self.created_index = True - info = self.pc.describe_index(self.index_name) + info = self.pc.describe_index(index_name) index_dims = info["dimension"] if dimensions != index_dims: raise ValueError( - f"PineconeDB index '{self.index_name}' has incorrect dimensions - expected:{dimensions}, found:{index_dims}" + f"PineconeDB index '{index_name}' has incorrect dimensions - expected:{dimensions}, found:{index_dims}" ) index_metric = info["metric"] if metric.value != index_metric: raise ValueError( - f"PineconeDB index '{self.index_name}' has incorrect metric - expected:{metric.value}, found:{index_metric}" + f"PineconeDB index '{index_name}' has incorrect metric - expected:{metric.value}, found:{index_metric}" ) - def close(self): - self.index.close() - def get_batch_size(self, sample_record: Record) -> int: # Return the largest batch size possible, based on the following # constraints: @@ -149,15 +135,15 @@ def initialize_population(self): return if not self.created_index and not self.overwrite: msg = ( - f"PineconeDB: Index '{self.index_name}' already exists - cowardly " + f"PineconeDB: Index '{self.index.name}' already exists - cowardly " f"refusing to overwrite existing data. Specify --overwrite to " - f"delete it, or specify --skip_populate to skip population phase." + f"delete it, or specify --skip-populate to skip population phase." ) logger.critical(msg) raise StopUser() try: logger.info( - f"PineconeDB: Deleting existing index '{self.index_name}' before " + f"PineconeDB: Deleting existing index '{self.index.name}' before " f"population (--overwrite=True)" ) self.index.delete(delete_all=True)