From aa99d80f377f8e4363664dd613d7dce396eef90f Mon Sep 17 00:00:00 2001
From: Xun Zhang <xunzh@amazon.com>
Date: Tue, 7 May 2024 17:08:15 -0700
Subject: [PATCH] add deploySetting in registering pretrained models

Signed-off-by: Xun Zhang <xunzh@amazon.com>
---
 .../src/main/java/org/opensearch/ml/engine/ModelHelper.java  | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/ModelHelper.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/ModelHelper.java
index 50c514599e..9ef104050e 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/ModelHelper.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/ModelHelper.java
@@ -27,6 +27,7 @@
 
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.FunctionName;
+import org.opensearch.ml.common.model.MLDeploySetting;
 import org.opensearch.ml.common.model.MLModelConfig;
 import org.opensearch.ml.common.model.MLModelFormat;
 import org.opensearch.ml.common.model.QuestionAnsweringModelConfig;
@@ -70,6 +71,7 @@ public void downloadPrebuiltModelConfig(
         boolean deployModel = registerModelInput.isDeployModel();
         String[] modelNodeIds = registerModelInput.getModelNodeIds();
         String modelGroupId = registerModelInput.getModelGroupId();
+        MLDeploySetting mlDeploySetting = registerModelInput.getDeploySetting();
         try {
             AccessController.doPrivileged((PrivilegedExceptionAction<Void>) () -> {
 
@@ -104,7 +106,8 @@ public void downloadPrebuiltModelConfig(
                     .modelNodeIds(modelNodeIds)
                     .isHidden(isHidden)
                     .modelGroupId(modelGroupId)
-                    .functionName(FunctionName.from((functionName)));
+                    .functionName(FunctionName.from((functionName)))
+                    .deploySetting(mlDeploySetting);
 
                 config.entrySet().forEach(entry -> {
                     switch (entry.getKey().toString()) {