Change model APIs for multi tenant

Signed-off-by: zane-neo <zaniu@amazon.com>
zane-neo · Apr 22, 2024 · 1212972 · 1212972
1 parent 06773b4
commit 1212972
Show file tree

Hide file tree

Showing 16 changed files with 389 additions and 218 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/MLModel.java b/common/src/main/java/org/opensearch/ml/common/MLModel.java
@@ -40,6 +40,8 @@
 public class MLModel implements ToXContentObject {
     @Deprecated
     public static final String ALGORITHM_FIELD = "algorithm";
+
+    public static final String TENANT_ID_FIELD = "tenant_id";
     public static final String FUNCTION_NAME_FIELD = "function_name";
     public static final String MODEL_NAME_FIELD = "name";
     public static final String MODEL_GROUP_ID_FIELD = "model_group_id";
@@ -133,6 +135,7 @@ public class MLModel implements ToXContentObject {
     private Connector connector;
     private String connectorId;
     private Guardrails guardrails;
+    private String tenantId;
 
     @Builder(toBuilder = true)
     public MLModel(String name,
@@ -166,7 +169,8 @@ public MLModel(String name,
             Boolean isHidden,
             Connector connector,
             String connectorId,
-            Guardrails guardrails) {
+            Guardrails guardrails,
+            String tenantId) {
         this.name = name;
         this.modelGroupId = modelGroupId;
         this.algorithm = algorithm;
@@ -200,6 +204,7 @@ public MLModel(String name,
         this.connector = connector;
         this.connectorId = connectorId;
         this.guardrails = guardrails;
+        this.tenantId = tenantId;
     }
 
     public MLModel(StreamInput input) throws IOException {
@@ -442,6 +447,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
         if (guardrails != null) {
             builder.field(GUARDRAILS_FIELD, guardrails);
         }
+        if (tenantId != null) {
+            builder.field(TENANT_ID_FIELD, tenantId);
+        }
         builder.endObject();
         return builder;
     }
@@ -486,6 +494,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
         Connector connector = null;
         String connectorId = null;
         Guardrails guardrails = null;
+        String tenantId = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -617,6 +626,9 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
                 case GUARDRAILS_FIELD:
                     guardrails = Guardrails.parse(parser);
                     break;
+                case TENANT_ID_FIELD:
+                    tenantId = parser.text();
+                    break;
                 default:
                     parser.skipChildren();
                     break;
@@ -656,6 +668,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
                 .connector(connector)
                 .connectorId(connectorId)
                 .guardrails(guardrails)
+                .tenantId(tenantId)
                 .build();
     }
 

diff --git a/common/src/main/java/org/opensearch/ml/common/MLModelGroup.java b/common/src/main/java/org/opensearch/ml/common/MLModelGroup.java
@@ -35,13 +35,16 @@ public class MLModelGroup implements ToXContentObject {
 
     public static final String ACCESS = "access"; //assigned to public, private, or null when model group created
     public static final String MODEL_GROUP_ID_FIELD = "model_group_id"; //unique ID assigned to each model group
+
+    public static final String TENANT_ID_FIELD = "tenant_id";
     public static final String CREATED_TIME_FIELD = "created_time"; //model group created time stamp
     public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time"; //updated whenever a new model version is created
 
 
     @Setter
     private String name;
     private String description;
+    @Setter
     private int latestVersion;
     private List<String> backendRoles;
     private User owner;
@@ -50,14 +53,18 @@ public class MLModelGroup implements ToXContentObject {
 
     private String modelGroupId;
 
+    private String tenantId;
+
     private Instant createdTime;
+    @Setter
     private Instant lastUpdatedTime;
 
 
     @Builder(toBuilder = true)
     public MLModelGroup(String name, String description, int latestVersion,
                         List<String> backendRoles, User owner, String access,
                         String modelGroupId,
+                        String tenantId,
                         Instant createdTime,
                         Instant lastUpdatedTime) {
         this.name = Objects.requireNonNull(name, "model group name must not be null");
@@ -67,6 +74,7 @@ public MLModelGroup(String name, String description, int latestVersion,
         this.owner = owner;
         this.access = access;
         this.modelGroupId = modelGroupId;
+        this.tenantId = tenantId;
         this.createdTime = createdTime;
         this.lastUpdatedTime = lastUpdatedTime;
     }
@@ -132,6 +140,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (modelGroupId != null) {
             builder.field(MODEL_GROUP_ID_FIELD, modelGroupId);
         }
+        if (tenantId != null) {
+            builder.field(TENANT_ID_FIELD, tenantId);
+        }
         if (createdTime != null) {
             builder.field(CREATED_TIME_FIELD, createdTime.toEpochMilli());
         }
@@ -150,6 +161,7 @@ public static MLModelGroup parse(XContentParser parser) throws IOException {
         User owner = null;
         String access = null;
         String modelGroupId = null;
+        String tenantId = null;
         Instant createdTime = null;
         Instant lastUpdateTime = null;
 
@@ -184,6 +196,9 @@ public static MLModelGroup parse(XContentParser parser) throws IOException {
                 case MODEL_GROUP_ID_FIELD:
                     modelGroupId = parser.text();
                     break;
+                case TENANT_ID_FIELD:
+                    tenantId = parser.text();
+                    break;
                 case CREATED_TIME_FIELD:
                     createdTime = Instant.ofEpochMilli(parser.longValue());
                     break;
@@ -203,6 +218,7 @@ public static MLModelGroup parse(XContentParser parser) throws IOException {
                 .owner(owner)
                 .access(access)
                 .modelGroupId(modelGroupId)
+                .tenantId(tenantId)
                 .createdTime(createdTime)
                 .lastUpdatedTime(lastUpdateTime)
                 .build();

diff --git a/common/src/main/java/org/opensearch/ml/common/transport/AbstractGetRequest.java b/common/src/main/java/org/opensearch/ml/common/transport/AbstractGetRequest.java
@@ -0,0 +1,29 @@
+/*
+ *
+ *  * Copyright OpenSearch Contributors
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.opensearch.ml.common.transport;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.Setter;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.core.common.io.stream.StreamInput;
+
+import java.io.IOException;
+
+@NoArgsConstructor
+@AllArgsConstructor
+public abstract class AbstractGetRequest extends ActionRequest {
+    @Setter
+    @Getter
+    private String tenantId;
+
+    public AbstractGetRequest(StreamInput in) throws IOException {
+        super(in);
+    }
+}
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/connector/MLConnectorGetRequest.java b/common/src/main/java/org/opensearch/ml/common/transport/connector/MLConnectorGetRequest.java
@@ -14,6 +14,7 @@
 import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.ml.common.transport.AbstractGetRequest;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -24,11 +25,10 @@
 
 @Getter
 @Setter
-public class MLConnectorGetRequest extends ActionRequest {
+public class MLConnectorGetRequest extends AbstractGetRequest {
 
     String connectorId;
 
-    String tenantId;
     boolean returnContent;
 
     @Builder

diff --git a/common/src/main/java/org/opensearch/ml/common/transport/model/MLModelGetRequest.java b/common/src/main/java/org/opensearch/ml/common/transport/model/MLModelGetRequest.java
@@ -16,6 +16,7 @@
 import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.ml.common.transport.AbstractGetRequest;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -27,7 +28,7 @@
 @Getter
 @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
 @ToString
-public class MLModelGetRequest extends ActionRequest {
+public class MLModelGetRequest extends AbstractGetRequest {
 
     String modelId;
     boolean returnContent;

diff --git a/common/src/main/java/org/opensearch/ml/common/transport/register/MLRegisterModelInput.java b/common/src/main/java/org/opensearch/ml/common/transport/register/MLRegisterModelInput.java
@@ -7,6 +7,7 @@
 
 import lombok.Builder;
 import lombok.Data;
+import lombok.Setter;
 import org.opensearch.Version;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -96,6 +97,9 @@ public class MLRegisterModelInput implements ToXContentObject, Writeable {
     private Boolean isHidden;
     private Guardrails guardrails;
 
+    @Setter
+    private String tenantId;
+
     @Builder(toBuilder = true)
     public MLRegisterModelInput(FunctionName functionName,
             String modelName,

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java
@@ -58,6 +58,12 @@ public void initModelIndexIfAbsent(ActionListener<Boolean> listener) {
         initMLIndexIfAbsent(MLIndex.MODEL, listener);
     }
 
+    public Boolean initModelIndexIfAbsent() throws ExecutionException, InterruptedException {
+        PlainActionFuture<Boolean> actionFuture = PlainActionFuture.newFuture();
+        initModelIndexIfAbsent(actionFuture);
+        return actionFuture.get();
+    }
+
     public void initMLTaskIndex(ActionListener<Boolean> listener) {
         initMLIndexIfAbsent(MLIndex.TASK, listener);
     }