Skip to content

Commit

Permalink
load class from config file; tune execute part
Browse files Browse the repository at this point in the history
Signed-off-by: Yaliang Wu <ylwu@amazon.com>
  • Loading branch information
ylwu-amzn committed Nov 23, 2021
1 parent ecb9417 commit db67aa2
Show file tree
Hide file tree
Showing 38 changed files with 757 additions and 485 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,10 @@
import org.opensearch.action.ActionFuture;
import org.opensearch.action.ActionListener;
import org.opensearch.action.support.PlainActionFuture;
import org.opensearch.ml.common.parameter.Input;
import org.opensearch.ml.common.parameter.MLInput;
import org.opensearch.ml.common.parameter.MLOutput;
import org.opensearch.ml.common.parameter.Output;

/**
* A client to provide interfaces for machine learning jobs. This will be used by other plugins.
Expand Down Expand Up @@ -64,19 +66,19 @@ default ActionFuture<MLOutput> train(MLInput mlInput) {

/**
* Execute ML algorithm.
* @param mlInput
* @return
* @param input input data
* @return output
*/
default ActionFuture<MLOutput> execute(MLInput mlInput) {
PlainActionFuture<MLOutput> actionFuture = PlainActionFuture.newFuture();
execute(mlInput, actionFuture);
default ActionFuture<Output> execute(Input input) {
PlainActionFuture<Output> actionFuture = PlainActionFuture.newFuture();
execute(input, actionFuture);
return actionFuture;
}

/**
* Execute ML algorithm
* @param mlInput
* @param listener
* @param input input data
* @param listener action listener
*/
void execute(MLInput mlInput, ActionListener<MLOutput> listener);
void execute(Input input, ActionListener<Output> listener);
}
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,10 @@
import lombok.experimental.FieldDefaults;
import org.opensearch.action.ActionListener;
import org.opensearch.client.node.NodeClient;
import org.opensearch.ml.common.parameter.Input;
import org.opensearch.ml.common.parameter.MLInput;
import org.opensearch.ml.common.parameter.MLOutput;
import org.opensearch.ml.common.parameter.Output;
import org.opensearch.ml.common.transport.execute.MLExecuteTaskAction;
import org.opensearch.ml.common.transport.execute.MLExecuteTaskRequest;
import org.opensearch.ml.common.transport.execute.MLExecuteTaskResponse;
Expand Down Expand Up @@ -65,10 +67,9 @@ public void train(MLInput mlInput, ActionListener<MLOutput> listener) {
}

@Override
public void execute(MLInput mlInput, ActionListener<MLOutput> listener) {
validateMLInput(mlInput, false);
public void execute(Input input, ActionListener<Output> listener) {
MLExecuteTaskRequest executeTaskRequest = MLExecuteTaskRequest.builder()
.mlInput(mlInput)
.input(input)
.build();

client.execute(MLExecuteTaskAction.INSTANCE, executeTaskRequest, ActionListener.wrap(response -> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,17 @@
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.opensearch.action.ActionListener;
import org.opensearch.common.io.stream.StreamOutput;
import org.opensearch.common.xcontent.XContentBuilder;
import org.opensearch.ml.common.dataframe.DataFrame;
import org.opensearch.ml.common.parameter.Input;
import org.opensearch.ml.common.parameter.MLAlgoName;
import org.opensearch.ml.common.parameter.MLAlgoParams;
import org.opensearch.ml.common.parameter.MLInput;
import org.opensearch.ml.common.parameter.MLOutput;
import org.opensearch.ml.common.parameter.MLOutputType;
import org.opensearch.ml.common.parameter.MLTrainingOutput;
import org.opensearch.ml.common.parameter.Output;

import java.io.IOException;

Expand Down Expand Up @@ -67,8 +71,13 @@ public void train(MLInput mlInput, ActionListener<MLOutput> listener) {
}

@Override
public void execute(MLInput mlInput, ActionListener<MLOutput> listener) {
listener.onResponse(new MLOutput() {
public void execute(Input input, ActionListener<Output> listener) {
listener.onResponse(new Output() {
@Override
public void writeTo(StreamOutput out) {

}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*
* Modifications Copyright OpenSearch Contributors. See
* GitHub history for details.
*/

package org.opensearch.ml.common;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.common.xcontent.DeprecationHandler;
import org.opensearch.common.xcontent.NamedXContentRegistry;
import org.opensearch.common.xcontent.XContentFactory;
import org.opensearch.common.xcontent.XContentParser;
import org.opensearch.common.xcontent.XContentType;
import org.opensearch.ml.common.dataset.MLInputDataType;
import org.opensearch.ml.common.parameter.Input;
import org.opensearch.ml.common.parameter.MLAlgoName;
import org.opensearch.ml.common.parameter.MLAlgoParams;
import org.opensearch.ml.common.parameter.MLOutputType;

import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Constructor;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;


public class MLCommonsClassLoader {

private static final Logger logger = LogManager.getLogger(MLCommonsClassLoader.class);
private static Map<Enum<?>, Class<?>> parameterClassMap = new HashMap<>();

private static Map<Enum<?>, Class<?>> mlAlgoClassMap = new HashMap<>();

static {
loadClassMapping(MLCommonsClassLoader.class, "/ml-commons-config.yml");
}

public static void loadClassMapping(Class<?> resource, String configFile) {
InputStream inputStream = resource.getResourceAsStream(configFile);
try (
XContentParser parser = XContentFactory.xContent(XContentType.YAML)
.createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, inputStream)
) {
if (parser.currentToken() == null) {
parser.nextToken();
}
while(parser.nextToken() != XContentParser.Token.END_OBJECT) {
XContentParser.Token currentToken = parser.currentToken();
if (currentToken == XContentParser.Token.FIELD_NAME) {
String key = parser.currentName();
if ("ml_algo_param_class".equals(key)) {
parseMLAlgoParams(parser, parameterClassMap, k -> MLAlgoName.fromString(k));
} else if ("ml_input_data_set_class".equals(key)) {
parseMLAlgoParams(parser, parameterClassMap, k -> MLInputDataType.fromString(k));
} else if ("ml_output_class".equals(key)) {
parseMLAlgoParams(parser, parameterClassMap, k -> MLOutputType.fromString(k));
} else if ("ml_algo_class".equals(key)) {
parseMLAlgoParams(parser, mlAlgoClassMap, k -> MLAlgoName.fromString(k));
} else if ("executable_function_class".equals(key)) {
parseMLAlgoParams(parser, mlAlgoClassMap, k -> MLAlgoName.fromString(k));
}
} else {
parser.nextToken();
}
}
} catch (Exception e) {
e.printStackTrace();
}
}

@SuppressWarnings("unchecked")
public static <T extends Enum<T>, S, I extends Object> S initInstance(T type, I in, Class<?> constructorParamClass) {
Class<?> clazz = constructorParamClass == MLAlgoParams.class || constructorParamClass == Input.class ? mlAlgoClassMap.get(type) : parameterClassMap.get(type);
if (clazz == null) {
throw new IllegalArgumentException("Can't find class for type " + type);
}
try {
Constructor<?> constructor = clazz.getConstructor(constructorParamClass);
return (S) constructor.newInstance(in);
} catch (Exception e) {
return null;
}
}

private static <T extends Enum<?>> void parseMLAlgoParams(XContentParser parser, Map<T, Class<?>> map, Function<String, T> enumParser) throws IOException, ClassNotFoundException {
String key = null;
String value = null;
while(parser.nextToken() != XContentParser.Token.END_OBJECT) {
XContentParser.Token currentToken = parser.currentToken();
if (currentToken == XContentParser.Token.FIELD_NAME) {
key = parser.currentName();
} else if (currentToken == XContentParser.Token.VALUE_STRING) {
value = parser.text();
if (key != null) {
try {
map.put(enumParser.apply(key), Class.forName(value));
} catch (Exception e) {
logger.error("Failed to load class mapping for " + key, e);
}
}
key = null;
}
}
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import lombok.Getter;
import lombok.NonNull;
import lombok.experimental.FieldDefaults;
import org.opensearch.ml.common.dataframe.DataFrameType;
import org.opensearch.ml.common.dataframe.DefaultDataFrame;

/**
Expand All @@ -38,9 +39,17 @@ public DataFrameInputDataset(@NonNull DataFrame dataFrame) {
this.dataFrame = dataFrame;
}

public DataFrameInputDataset(StreamInput streaminput) throws IOException {
public DataFrameInputDataset(StreamInput in) throws IOException {
super(MLInputDataType.DATA_FRAME);
this.dataFrame = new DefaultDataFrame(streaminput);
DataFrameType dataFrameType = in.readEnum(DataFrameType.class);
switch (dataFrameType) {
case DEFAULT:
this.dataFrame = new DefaultDataFrame(in);
break;
default:
this.dataFrame = null;
break;
}
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,28 @@

package org.opensearch.ml.common.dataset;

import lombok.Getter;

public enum MLInputDataType {
SEARCH_QUERY,
DATA_FRAME;
SEARCH_QUERY("search_query"),
DATA_FRAME("data_frame");


@Getter
private final String name;

MLInputDataType(String name) {
this.name = name;
}

public String toString() {
return name;
}

public static MLInputDataType fromString(String name){
for(MLInputDataType e : MLInputDataType.values()){
if(e.name.equals(name)) return e;
}
return null;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,6 @@ public SearchQueryInputDataset(@NonNull List<String> indices, @NonNull SearchSou
public SearchQueryInputDataset(StreamInput streaminput) throws IOException {
super(MLInputDataType.SEARCH_QUERY);
String searchString = streaminput.readString();
// SearchModule searchModule = new SearchModule(Settings.EMPTY, false, Collections.emptyList());
// XContentParser parser = XContentType.JSON.xContent().createParser(new NamedXContentRegistry(searchModule.getNamedXContents()), LoggingDeprecationHandler.INSTANCE, searchString);
XContentParser parser = XContentType.JSON.xContent().createParser(xContentRegistry, LoggingDeprecationHandler.INSTANCE, searchString);
this.searchSourceBuilder = SearchSourceBuilder.fromXContent(parser);
this.indices = streaminput.readStringList();
Expand Down
20 changes: 20 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/parameter/Input.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*
* Modifications Copyright OpenSearch Contributors. See
* GitHub history for details.
*/

package org.opensearch.ml.common.parameter;

import org.opensearch.common.io.stream.Writeable;
import org.opensearch.common.xcontent.ToXContentObject;

public interface Input extends ToXContentObject, Writeable {

MLAlgoName getFunctionName();
}
Loading

0 comments on commit db67aa2

Please sign in to comment.