Throw validation exception on train reject

Signed-off-by: John Mazanec <[email protected]>
opensearch-project · Oct 31, 2021 · 5b62017 · 5b62017
1 parent b3027a1
commit 5b62017
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/knn/plugin/transport/TrainingJobRouterTransportAction.java b/src/main/java/org/opensearch/knn/plugin/transport/TrainingJobRouterTransportAction.java
@@ -20,6 +20,7 @@
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Strings;
+import org.opensearch.common.ValidationException;
 import org.opensearch.common.collect.ImmutableOpenMap;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.search.builder.SearchSourceBuilder;
@@ -70,7 +71,9 @@ protected void routeRequest(TrainingModelRequest request, ActionListener<Trainin
                     DiscoveryNode node = selectNode(request.getPreferredNodeId(), response);
 
                     if (node == null) {
-                        listener.onFailure(new RejectedExecutionException("Cluster does not have capacity to train"));
+                        ValidationException exception = new ValidationException();
+                        exception.addValidationError("Cluster does not have capacity to train");
+                        listener.onFailure(exception);
                         return;
                     }
 

diff --git a/src/main/java/org/opensearch/knn/training/TrainingJobRunner.java b/src/main/java/org/opensearch/knn/training/TrainingJobRunner.java
@@ -15,6 +15,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.action.ActionListener;
 import org.opensearch.action.index.IndexResponse;
+import org.opensearch.common.ValidationException;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
 import org.opensearch.knn.indices.ModelState;
@@ -82,7 +83,9 @@ public void execute(TrainingJob trainingJob, ActionListener<IndexResponse> liste
         // the number of training jobs that enter this function. Although the training threadpool size will also prevent
         // this, we want to prevent this before we perform any serialization.
         if (!semaphore.tryAcquire()) {
-            throw new IllegalStateException("Unable to run training job: No training capacity on node.");
+            ValidationException exception = new ValidationException();
+            exception.addValidationError("Unable to run training job: No training capacity on node.");
+            throw exception;
         }
 
         jobCount.incrementAndGet();