From 283f0bc4734360a52e629bd0015119734db4c38f Mon Sep 17 00:00:00 2001 From: Dave Rigby Date: Mon, 26 Feb 2024 16:02:15 +0000 Subject: [PATCH] Add --pinecone-throughput-per-user option This option specifies how many requests per second each User instance should issue. Fractional values are permitted - e.g. a value of 0.1 will issue one request per user every 10 seconds. Default value is 0 meaning 'unlimited', hence out-of-the-box behaviour is unchanged. --- locustfile.py | 12 +++++++++++- tests/integration/test_requests.py | 9 +++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/locustfile.py b/locustfile.py index 7e55c3e..275d790 100644 --- a/locustfile.py +++ b/locustfile.py @@ -7,7 +7,7 @@ from functools import wraps import gevent import grpc.experimental.gevent as grpc_gevent -from locust import FastHttpUser, User, events, tag, task +from locust import FastHttpUser, User, constant_throughput, events, tag, task from locust.env import Environment from locust.exception import StopUser from locust.runners import Runner, WorkerRunner @@ -87,6 +87,10 @@ def _(parser): "(default: %(default)s).") pc_options.add_argument("--pinecone-dataset-cache", type=str, default=".dataset_cache", help="Path to directory to cache downloaded datasets (default: %(default)s).") + pc_options.add_argument("--pinecone-throughput-per-user", type=float, default=0, + help="How many requests per second each user should issue (default: %(default)s). " + "Setting to zero will make each user issue requests as fast as possible " + "(next request sent as soon as previous one completes).") # iterations option included from locust-plugins parser.add_argument( @@ -254,11 +258,17 @@ def __init__(self, environment): self.client = PineconeSdk(self.environment, use_grpc=True) else: raise Exception(f"Invalid pinecone_mode {self.mode}") + self.target_throughput = environment.parsed_options.pinecone_throughput_per_user if isinstance(self.environment.runner, WorkerRunner): # Wait until the datset has been loaded for this environment (Runner) environment.setup_dataset_greenlet.join() + def wait_time(self): + if self.target_throughput > 0: + return constant_throughput(self.target_throughput)(self) + return 0 + @tag('query') @task def vectorQuery(self): diff --git a/tests/integration/test_requests.py b/tests/integration/test_requests.py index f3fac94..9a014da 100644 --- a/tests/integration/test_requests.py +++ b/tests/integration/test_requests.py @@ -155,3 +155,12 @@ def test_pinecone_fetch(self, index_host, mode): def test_pinecone_delete(self, index_host, mode): self.do_request(index_host, mode, 'delete', 'Delete') + + def test_pinecone_limit_throughput(self, index_host, mode): + self.do_request(index_host, mode, 'query', 'Vector (Query only)', + extra_args=["--pinecone-throughput-per-user=1"]) + + def test_pinecone_unlimited_throughput(self, index_host, mode): + # Test that a limit of 0 meaning "unlimited" is handled correctly. + self.do_request(index_host, mode, 'query', 'Vector (Query only)', + extra_args=["--pinecone-throughput-per-user=0"])