use of org.opensearch.ml.indices.MLIndicesHandler.ML_TASK_INDEX in project ml-commons by opensearch-project.
the class GetTaskTransportAction method doExecute.
@Override
protected void doExecute(Task task, ActionRequest request, ActionListener<MLTaskGetResponse> actionListener) {
MLTaskGetRequest mlTaskGetRequest = MLTaskGetRequest.fromActionRequest(request);
String taskId = mlTaskGetRequest.getTaskId();
GetRequest getRequest = new GetRequest(ML_TASK_INDEX).id(taskId);
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
client.get(getRequest, ActionListener.wrap(r -> {
log.info("Completed Get Task Request, id:{}", taskId);
if (r != null && r.isExists()) {
try (XContentParser parser = createXContentParserFromRegistry(xContentRegistry, r.getSourceAsBytesRef())) {
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
MLTask mlTask = MLTask.parse(parser);
actionListener.onResponse(MLTaskGetResponse.builder().mlTask(mlTask).build());
} catch (Exception e) {
log.error("Failed to parse ml task" + r.getId(), e);
actionListener.onFailure(e);
}
} else {
actionListener.onFailure(new MLResourceNotFoundException("Fail to find task"));
}
}, e -> {
if (e instanceof IndexNotFoundException) {
actionListener.onFailure(new MLResourceNotFoundException("Fail to find task"));
} else {
log.error("Failed to get ML task " + taskId, e);
actionListener.onFailure(e);
}
}));
} catch (Exception e) {
log.error("Failed to get ML task " + taskId, e);
actionListener.onFailure(e);
}
}
use of org.opensearch.ml.indices.MLIndicesHandler.ML_TASK_INDEX in project ml-commons by opensearch-project.
the class MLTaskManager method createMLTask.
/**
* Create ML task. Will init ML task index first if absent.
* @param mlTask ML task
* @param listener action listener
*/
public void createMLTask(MLTask mlTask, ActionListener<IndexResponse> listener) {
mlIndicesHandler.initMLTaskIndex(ActionListener.wrap(indexCreated -> {
if (!indexCreated) {
listener.onFailure(new RuntimeException("No response to create ML task index"));
return;
}
IndexRequest request = new IndexRequest(ML_TASK_INDEX);
try (XContentBuilder builder = XContentFactory.jsonBuilder();
ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
request.source(mlTask.toXContent(builder, ToXContent.EMPTY_PARAMS)).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
client.index(request, ActionListener.runBefore(listener, () -> context.restore()));
} catch (Exception e) {
log.error("Failed to create AD task for " + mlTask.getFunctionName() + ", " + mlTask.getTaskType(), e);
listener.onFailure(e);
}
}, e -> {
log.error("Failed to create ML index", e);
listener.onFailure(e);
}));
}
Aggregations