From 6fae9aceb93209a02e728e76c2aba1754f4bbd79 Mon Sep 17 00:00:00 2001
From: David Roberts <dave.roberts@elastic.co>
Date: Wed, 20 Sep 2023 13:26:10 +0100
Subject: [PATCH] [ML] Assume ELSER v2 has the same memory requirement as ELSER
 v1

While we wait for our improved memory estimation functionality
we can continue to use the same hardcoded ELSER memory requirement
for v2 that we used for v1. The hardcoded value is likely an
overestimate, but we don't have time to do anything better for
the current release.
---
 .../StartTrainedModelDeploymentAction.java    | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/StartTrainedModelDeploymentAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/StartTrainedModelDeploymentAction.java
index b00dcfd731aee..ad9ab7088fef5 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/StartTrainedModelDeploymentAction.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/StartTrainedModelDeploymentAction.java
@@ -62,11 +62,11 @@ public class StartTrainedModelDeploymentAction extends ActionType<CreateTrainedM
     private static final ByteSizeValue MEMORY_OVERHEAD = ByteSizeValue.ofMb(240);
 
     /**
-     * The ELSER model turned out to use more memory then what we usually estimate.
-     * We overwrite the estimate with this static value for ELSER V1 for now. Soon to be
-     * replaced with a better estimate provided by the model.
+     * The ELSER model turned out to use more memory than what we usually estimate.
+     * We overwrite the estimate with this static value for ELSER v1 and v2 for now.
+     * Soon to be replaced with a better estimate provided by the model.
      */
-    private static final ByteSizeValue ELSER_1_MEMORY_USAGE = ByteSizeValue.ofMb(2004);
+    private static final ByteSizeValue ELSER_1_OR_2_MEMORY_USAGE = ByteSizeValue.ofMb(2004);
 
     public StartTrainedModelDeploymentAction() {
         super(NAME, CreateTrainedModelAssignmentAction.Response::new);
@@ -713,14 +713,14 @@ public static long estimateMemoryUsageBytes(
     ) {
         // While loading the model in the process we need twice the model size.
 
-        // 1. If ELSER v1 then 2004MB
+        // 1. If ELSER v1 or v2 then 2004MB
         // 2. If static memory and dynamic memory are not set then 240MB + 2 * model size
         // 3. Else static memory + dynamic memory * allocations + model size
 
         // The model size is still added in option 3 to account for the temporary requirement to hold the zip file in memory
-        // in `pytorch_inference`.
-        if (isElserV1Model(modelId)) {
-            return ELSER_1_MEMORY_USAGE.getBytes();
+        // in `pytorch_inference`.
+        if (isElserV1Or2Model(modelId)) {
+            return ELSER_1_OR_2_MEMORY_USAGE.getBytes();
         } else {
             long baseSize = MEMORY_OVERHEAD.getBytes() + 2 * totalDefinitionLength;
             if (perDeploymentMemoryBytes == 0 && perAllocationMemoryBytes == 0) {
@@ -734,7 +734,7 @@ public static long estimateMemoryUsageBytes(
         }
     }
 
-    private static boolean isElserV1Model(String modelId) {
-        return modelId.startsWith(".elser_model_1");
+    private static boolean isElserV1Or2Model(String modelId) {
+        return modelId.startsWith(".elser_model_1") || modelId.startsWith(".elser_model_2");
     }
 }