Add TrainKNN Runner/Operation for Benchmarking Approximate KNN Algorithms #556

finnroblin · 2024-06-19T22:17:12Z

Description

Customers may want to benchmark approximate k-NN search algorithms that require a training step. For example, the k-NN plugin with the FAISS engine and IVF method requires a training step to cluster database vectors. Then search can be performed against a smaller number of cluster centroids instead of the entire database.

Benchmarking these algorithms requires a new operation-type to train the model. This PR adds an operation that calls the k-NN Train Model API to train a model and then polls the Get Model API until model training succeeds or fails.

Issues Resolved

This PR is a task towards resolving Issue 332 opened in opensearch-benchmarks-workloads to add a train model benchmark workload. Please see this issue for more context about why this feature is important for improving benchmarking workflow.

Testing

New functionality includes unit tests

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

VijayanB · 2024-06-19T22:56:10Z

osbenchmark/worker_coordinator/runner.py

@@ -651,6 +651,72 @@ def error_description(self, error_details):
    def __repr__(self, *args, **kwargs):
        return "bulk-index"

+class TrainKNNModel(Runner): 
+    """
+    Trains model named model_id until training is complete or timeout is exceeded. 


max retry is reached

VijayanB · 2024-06-19T22:58:38Z

osbenchmark/worker_coordinator/runner.py

+            model_response = await opensearch.transport.perform_request("GET", model_uri)
+
+            if 'state' not in model_response.keys() or current_number_retries > max_retries:
+                request_context_holder.on_client_request_end()


request_context_holder.on_client_request_start() is missing

VijayanB · 2024-06-19T22:59:46Z

osbenchmark/worker_coordinator/runner.py

+        while True: 
+            model_response = await opensearch.transport.perform_request("GET", model_uri)
+
+            if 'state' not in model_response.keys() or current_number_retries > max_retries:


i don't think you can combine both the condition. Can you please break it into two condition and raise appropriate exception message?

VijayanB · 2024-06-19T23:01:36Z

osbenchmark/worker_coordinator/runner.py

+    def inspect_model_response(model_response):
+        if model_response['state'] == 'created':
+            return 1, "models_trained"
+
+        if model_response['state'] == 'failed':
+            raise Exception("Failed to create model: {}".format(model_response))


remove this

VijayanB · 2024-06-19T23:02:08Z

osbenchmark/workload/workload.py

@@ -746,6 +747,8 @@ def from_hyphenated_string(cls, v):
            return OperationType.RegisterMlModel
        elif v == "deploy-ml-model":
            return OperationType.DeployMlModel
+        elif v == "train-k-n-n-model":


Suggested change

elif v == "train-k-n-n-model":

elif v == "train-knn-model":

I'll fix this by changing the name to TrainKnnModel. It seems like hyphens are automatically placed between every capitalized letter.

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

VijayanB · 2024-06-20T17:38:27Z

osbenchmark/worker_coordinator/runner.py

+            if model_response['state'] == 'created':
+                self.logger.info(
+                    "Training model [%s] was completed successfully.", model_id)
+                break


nit: return instead of break

VijayanB · 2024-06-20T17:40:26Z

osbenchmark/worker_coordinator/runner.py

+                break
+
+            # training failed.
+            self.logger.error(


IMO, we should add model_response['state'] == 'failed' to validate failure condition, and, raise exception if state contains unexpected value. This will help us in case knn decided to add new status

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

finnroblin · 2024-06-21T23:55:22Z

osbenchmark/worker_coordinator/test.py

Did not mean to include this file, I will discard on the next commit.

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

…xception Signed-off-by: Finn Roblin <finnrobl@amazon.com>

VijayanB · 2024-06-27T17:58:31Z

osbenchmark/worker_coordinator/runner.py

+
+    async def __call__(self, opensearch, params):
+        model_id = parse_string_parameter("model_id", params)
+        should_ignore_if_model_DNE = params.get("ignore-if-model-does-not-exist", False)


Suggested change

should_ignore_if_model_DNE = params.get("ignore-if-model-does-not-exist", False)

ignore_if_model_does_not_exist = params.get("ignore-if-model-does-not-exist", False)

VijayanB · 2024-06-27T18:01:41Z

osbenchmark/worker_coordinator/runner.py

+    NAME = "delete-knn-model"
+
+    async def __call__(self, opensearch, params):
+        model_id = parse_string_parameter("model_id", params)


did you check can we use mandatory?

Yes, these helpers enforce mandatory if they don't have a third argument. So if model_id is not provided parse_string_parameter will throw an error.

VijayanB · 2024-06-27T18:02:32Z

osbenchmark/worker_coordinator/runner.py

+            and response["status"] == 404
+            and not should_ignore_if_model_DNE
+        ):
+            self.logger.error(


may be warning

VijayanB · 2024-06-27T18:03:13Z

osbenchmark/worker_coordinator/runner.py

+
+        if (
+            "error" in response.keys()
+            and response["status"] == 404


nit: create constant for 404

VijayanB · 2024-06-27T18:05:02Z

osbenchmark/worker_coordinator/runner.py

+
+        request_context_holder.on_client_request_end()
+
+        if (


you can flip if condition with if response is success return success, and, check for error condition one by one.

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

VijayanB · 2024-07-10T19:00:17Z

@finnroblin Looks good to me.

added TrainKNN runner + tests

3d138cb

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

finnroblin requested review from IanHoang, gkamat, beaioun, cgchinmay, rishabh6788 and VijayanB as code owners June 19, 2024 22:17

finnroblin mentioned this pull request Jun 19, 2024

[FEATURE] Add Train Model KNN Workload opensearch-project/opensearch-benchmark-workloads#332

Closed

2 tasks

VijayanB reviewed Jun 19, 2024

View reviewed changes

Addressed Vijay's feedback, cleaned up tests, added TimeoutError

13da7b8

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

VijayanB reviewed Jun 20, 2024

View reviewed changes

finnroblin added 2 commits June 21, 2024 16:41

Added delete-knn-model runner

9b206c0

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

Addressed Vijay feedback opensearch-project#2

9fac862

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

finnroblin commented Jun 21, 2024

View reviewed changes

osbenchmark/worker_coordinator/test.py Outdated

Copy link

Contributor Author

finnroblin Jun 21, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did not mean to include this file, I will discard on the next commit.

finnroblin mentioned this pull request Jun 24, 2024

Add vectorsearch training workload opensearch-project/opensearch-benchmark-workloads#333

Merged

finnroblin added 3 commits June 24, 2024 16:57

Added request ignore on 404; removed local testing file

16e493b

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

Fix linting errors

b43cbc6

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

Change DeleteKnnRunner behavior to return success status instead of e…

a78351e

…xception Signed-off-by: Finn Roblin <finnrobl@amazon.com>

finnroblin requested a review from VijayanB June 26, 2024 21:30

VijayanB reviewed Jun 27, 2024

View reviewed changes

Vijay feedback & reordered conditionals

95d445e

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

Revert autoformatter changes to comment at top of runner.py

904ae3f

Signed-off-by: Finn Roblin <finnrobl@amazon.com>

finnroblin requested a review from VijayanB June 27, 2024 23:22

IanHoang approved these changes Jul 18, 2024

View reviewed changes

IanHoang merged commit 01346bb into opensearch-project:main Jul 18, 2024
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add TrainKNN Runner/Operation for Benchmarking Approximate KNN Algorithms #556

Add TrainKNN Runner/Operation for Benchmarking Approximate KNN Algorithms #556

finnroblin commented Jun 19, 2024

VijayanB Jun 19, 2024

VijayanB Jun 19, 2024

VijayanB Jun 19, 2024

VijayanB Jun 19, 2024

VijayanB Jun 19, 2024

finnroblin Jun 20, 2024

VijayanB Jun 20, 2024

VijayanB Jun 20, 2024

finnroblin Jun 21, 2024

VijayanB Jun 27, 2024

VijayanB Jun 27, 2024

finnroblin Jun 27, 2024

VijayanB Jun 27, 2024

VijayanB Jun 27, 2024

VijayanB Jun 27, 2024

VijayanB commented Jul 10, 2024 •

edited

Loading

	should_ignore_if_model_DNE = params.get("ignore-if-model-does-not-exist", False)
	ignore_if_model_does_not_exist = params.get("ignore-if-model-does-not-exist", False)

Add TrainKNN Runner/Operation for Benchmarking Approximate KNN Algorithms #556

Add TrainKNN Runner/Operation for Benchmarking Approximate KNN Algorithms #556

Conversation

finnroblin commented Jun 19, 2024

Description

Issues Resolved

Testing

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

VijayanB commented Jul 10, 2024 • edited Loading

VijayanB commented Jul 10, 2024 •

edited

Loading