Skip to content

Commit

Permalink
add get config api to retrieve root agent id
Browse files Browse the repository at this point in the history
Signed-off-by: Bhavana Ramaram <[email protected]>
  • Loading branch information
rbhavna committed Feb 2, 2024
1 parent 939bbee commit e8dc391
Show file tree
Hide file tree
Showing 15 changed files with 1,029 additions and 2 deletions.
10 changes: 10 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/CommonValue.java
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public class CommonValue {

public static final String MASTER_KEY = "master_key";
public static final String CREATE_TIME_FIELD = "create_time";
public static final String LAST_UPDATE_TIME_FIELD = "last_update_time";

public static final String BOX_TYPE_KEY = "box_type";
// hot node
Expand Down Expand Up @@ -359,7 +360,16 @@ public class CommonValue {
+ MASTER_KEY
+ "\": {\"type\": \"keyword\"},\n"
+ " \""
+ MLConfig.TYPE_FIELD
+ "\" : {\"type\":\"keyword\"},\n"
+ " \""
+ MLConfig.CONFIGURATION_FIELD
+ "\" : {\"type\": \"flat_object\"},\n"
+ " \""
+ CREATE_TIME_FIELD
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
+ " \""
+ LAST_UPDATE_TIME_FIELD
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"}\n"
+ " }\n"
+ "}";
Expand Down
83 changes: 83 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/Configuration.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common;

import lombok.Builder;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.Setter;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.xcontent.ToXContentObject;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;

import java.io.IOException;

import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;

@Getter
@EqualsAndHashCode
public class Configuration implements ToXContentObject, Writeable {

public static final String ROOT_AGENT_ID = "agent_id";

@Setter
private String agentId;

@Builder(toBuilder = true)
public Configuration(
String agentId
) {
this.agentId = agentId;
}

public Configuration(StreamInput input) throws IOException {
this.agentId = input.readOptionalString();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeOptionalString(agentId);
}

@Override
public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
XContentBuilder builder = xContentBuilder.startObject();
if (agentId != null) {
builder.field(ROOT_AGENT_ID, agentId);
}
return builder.endObject();
}

public static Configuration fromStream(StreamInput in) throws IOException {
Configuration configuration = new Configuration(in);
return configuration;
}

public static Configuration parse(XContentParser parser) throws IOException {
String agentId = null;

ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
String fieldName = parser.currentName();
parser.nextToken();

switch (fieldName) {
case ROOT_AGENT_ID:
agentId = parser.text();
break;
default:
parser.skipChildren();
break;
}
}
return Configuration.builder()
.agentId(agentId)
.build();
}
}
136 changes: 136 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/MLConfig.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common;

import lombok.Builder;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.Setter;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.xcontent.ToXContentObject;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;

import java.io.IOException;
import java.time.Instant;

import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;

@Getter
@EqualsAndHashCode
public class MLConfig implements ToXContentObject, Writeable {

public static final String TYPE_FIELD = "type";

public static final String CONFIGURATION_FIELD = "configuration";

public static final String CREATE_TIME_FIELD = "create_time";
public static final String LAST_UPDATE_TIME_FIELD = "last_update_time";

@Setter
private String type;

private Configuration configuration;
private final Instant createTime;
private Instant lastUpdateTime;

@Builder(toBuilder = true)
public MLConfig(
String type,
Configuration configuration,
Instant createTime,
Instant lastUpdateTime
) {
this.type = type;
this.configuration = configuration;
this.createTime = createTime;
this.lastUpdateTime = lastUpdateTime;
}

public MLConfig(StreamInput input) throws IOException {
this.type = input.readOptionalString();
if (input.readBoolean()) {
configuration = new Configuration(input);
}
createTime = input.readOptionalInstant();
lastUpdateTime = input.readOptionalInstant();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeOptionalString(type);
if (configuration != null) {
out.writeBoolean(true);
configuration.writeTo(out);
} else {
out.writeBoolean(false);
}
out.writeOptionalInstant(createTime);
out.writeOptionalInstant(lastUpdateTime);
}

@Override
public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
XContentBuilder builder = xContentBuilder.startObject();
if (type != null) {
builder.field(TYPE_FIELD, type);
}
if (configuration != null) {
builder.field(CONFIGURATION_FIELD, configuration);
}
if (createTime != null) {
builder.field(CREATE_TIME_FIELD, createTime.toEpochMilli());
}
if (lastUpdateTime != null) {
builder.field(LAST_UPDATE_TIME_FIELD, lastUpdateTime.toEpochMilli());
}
return builder.endObject();
}

public static MLConfig fromStream(StreamInput in) throws IOException {
MLConfig mlConfig = new MLConfig(in);
return mlConfig;
}

public static MLConfig parse(XContentParser parser) throws IOException {
String type = null;
Configuration configuration = null;
Instant createTime = null;
Instant lastUpdateTime = null;

ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
String fieldName = parser.currentName();
parser.nextToken();

switch (fieldName) {
case TYPE_FIELD:
type = parser.text();
break;
case CONFIGURATION_FIELD:
configuration = Configuration.parse(parser);
break;
case CREATE_TIME_FIELD:
createTime = Instant.ofEpochMilli(parser.longValue());
break;
case LAST_UPDATE_TIME_FIELD:
lastUpdateTime = Instant.ofEpochMilli(parser.longValue());
break;
default:
parser.skipChildren();
break;
}
}
return MLConfig.builder()
.type(type)
.configuration(configuration)
.createTime(createTime)
.lastUpdateTime(lastUpdateTime)
.build();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common.transport.config;

import org.opensearch.action.ActionType;

public class MLConfigGetAction extends ActionType<MLConfigGetResponse> {
public static final MLConfigGetAction INSTANCE = new MLConfigGetAction();
public static final String NAME = "cluster:admin/opensearch/ml/config/get";

private MLConfigGetAction() { super(NAME, MLConfigGetResponse::new);}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common.transport.config;

import lombok.Builder;
import lombok.Getter;
import org.opensearch.action.ActionRequest;
import org.opensearch.action.ActionRequestValidationException;
import org.opensearch.core.common.io.stream.InputStreamStreamInput;
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 java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UncheckedIOException;

import static org.opensearch.action.ValidateActions.addValidationError;

@Getter
public class MLConfigGetRequest extends ActionRequest {

String configId;

@Builder
public MLConfigGetRequest(String configId) {
this.configId = configId;
}

public MLConfigGetRequest(StreamInput in) throws IOException {
super(in);
this.configId = in.readString();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
super.writeTo(out);
out.writeString(this.configId);
}

@Override
public ActionRequestValidationException validate() {
ActionRequestValidationException exception = null;

if (this.configId == null) {
exception = addValidationError("ML config id can't be null", exception);
}

return exception;
}

public static MLConfigGetRequest fromActionRequest(ActionRequest actionRequest) {
if (actionRequest instanceof MLConfigGetRequest) {
return (MLConfigGetRequest) actionRequest;
}

try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
actionRequest.writeTo(osso);
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
return new MLConfigGetRequest(input);
}
} catch (IOException e) {
throw new UncheckedIOException("failed to parse ActionRequest into MLConfigGetRequest", e);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common.transport.config;

import lombok.Builder;
import org.opensearch.core.action.ActionResponse;
import org.opensearch.core.common.io.stream.InputStreamStreamInput;
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.core.xcontent.ToXContentObject;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.ml.common.MLConfig;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UncheckedIOException;

public class MLConfigGetResponse extends ActionResponse implements ToXContentObject {
MLConfig mlConfig;

@Builder
public MLConfigGetResponse(MLConfig mlConfig) {
this.mlConfig = mlConfig;
}

public MLConfigGetResponse(StreamInput in) throws IOException {
super(in);
mlConfig = MLConfig.fromStream(in);
}

@Override
public void writeTo(StreamOutput out) throws IOException{
mlConfig.writeTo(out);
}

@Override
public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
return mlConfig.toXContent(xContentBuilder, params);
}

public static MLConfigGetResponse fromActionResponse(ActionResponse actionResponse) {
if (actionResponse instanceof MLConfigGetResponse) {
return (MLConfigGetResponse) actionResponse;
}

try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
actionResponse.writeTo(osso);
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
return new MLConfigGetResponse(input);
}
} catch (IOException e) {
throw new UncheckedIOException("failed to parse ActionResponse into MLConfigGetResponse", e);
}
}

}
Loading

0 comments on commit e8dc391

Please sign in to comment.