forked from elastic/elasticsearch
-
Notifications
You must be signed in to change notification settings - Fork 3
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[ML] Make Inference Services pluggable (elastic#99886)
Creates an InferenceServicePlugins interface for inference services to implement and adds a test implementation to mock an inference service.
- Loading branch information
Showing
97 changed files
with
705 additions
and
392 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
17 changes: 9 additions & 8 deletions
17
...l/inference/results/InferenceResults.java → ...ticsearch/inference/InferenceResults.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
69 changes: 69 additions & 0 deletions
69
server/src/main/java/org/elasticsearch/inference/InferenceServiceRegistry.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,69 @@ | ||
/* | ||
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
* or more contributor license agreements. Licensed under the Elastic License | ||
* 2.0 and the Server Side Public License, v 1; you may not use this file except | ||
* in compliance with, at your election, the Elastic License 2.0 or the Server | ||
* Side Public License, v 1. | ||
*/ | ||
|
||
package org.elasticsearch.inference; | ||
|
||
import org.elasticsearch.common.component.AbstractLifecycleComponent; | ||
import org.elasticsearch.common.io.stream.NamedWriteableRegistry; | ||
import org.elasticsearch.plugins.InferenceServicePlugin; | ||
|
||
import java.io.IOException; | ||
import java.util.ArrayList; | ||
import java.util.List; | ||
import java.util.Map; | ||
import java.util.Optional; | ||
import java.util.function.Function; | ||
import java.util.stream.Collectors; | ||
|
||
public class InferenceServiceRegistry extends AbstractLifecycleComponent { | ||
|
||
private final Map<String, InferenceService> services; | ||
private final List<NamedWriteableRegistry.Entry> namedWriteables = new ArrayList<>(); | ||
|
||
public InferenceServiceRegistry( | ||
List<InferenceServicePlugin> inferenceServicePlugins, | ||
InferenceServicePlugin.InferenceServiceFactoryContext factoryContext | ||
) { | ||
// TODO check names are unique | ||
services = inferenceServicePlugins.stream() | ||
.flatMap(r -> r.getInferenceServiceFactories().stream()) | ||
.map(factory -> factory.create(factoryContext)) | ||
.collect(Collectors.toMap(InferenceService::name, Function.identity())); | ||
|
||
for (var plugin : inferenceServicePlugins) { | ||
namedWriteables.addAll(plugin.getInferenceServiceNamedWriteables()); | ||
} | ||
} | ||
|
||
public Map<String, InferenceService> getServices() { | ||
return services; | ||
} | ||
|
||
public Optional<InferenceService> getService(String serviceName) { | ||
return Optional.ofNullable(services.get(serviceName)); | ||
} | ||
|
||
public List<NamedWriteableRegistry.Entry> getNamedWriteables() { | ||
return namedWriteables; | ||
} | ||
|
||
@Override | ||
protected void doStart() { | ||
|
||
} | ||
|
||
@Override | ||
protected void doStop() { | ||
|
||
} | ||
|
||
@Override | ||
protected void doClose() throws IOException { | ||
|
||
} | ||
} |
7 changes: 4 additions & 3 deletions
7
.../elasticsearch/xpack/inference/Model.java → ...va/org/elasticsearch/inference/Model.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
7 changes: 4 additions & 3 deletions
7
...arch/xpack/inference/ServiceSettings.java → ...sticsearch/inference/ServiceSettings.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
7 changes: 4 additions & 3 deletions
7
...csearch/xpack/inference/TaskSettings.java → ...elasticsearch/inference/TaskSettings.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
43 changes: 43 additions & 0 deletions
43
server/src/main/java/org/elasticsearch/plugins/InferenceServicePlugin.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,43 @@ | ||
/* | ||
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
* or more contributor license agreements. Licensed under the Elastic License | ||
* 2.0 and the Server Side Public License, v 1; you may not use this file except | ||
* in compliance with, at your election, the Elastic License 2.0 or the Server | ||
* Side Public License, v 1. | ||
*/ | ||
|
||
package org.elasticsearch.plugins; | ||
|
||
import org.elasticsearch.client.internal.Client; | ||
import org.elasticsearch.common.io.stream.NamedWriteableRegistry; | ||
import org.elasticsearch.inference.InferenceService; | ||
|
||
import java.util.List; | ||
|
||
/** | ||
* InferenceServicePlugins implement an inference service | ||
*/ | ||
public interface InferenceServicePlugin { | ||
|
||
List<Factory> getInferenceServiceFactories(); | ||
|
||
record InferenceServiceFactoryContext(Client client) {} | ||
|
||
interface Factory { | ||
/** | ||
* InferenceServices are created from the factory context | ||
*/ | ||
InferenceService create(InferenceServiceFactoryContext context); | ||
} | ||
|
||
/** | ||
* The named writables defined and used by each of the implemented | ||
* InferenceServices. Each service should define named writables for | ||
* - {@link org.elasticsearch.inference.TaskSettings} | ||
* - {@link org.elasticsearch.inference.ServiceSettings} | ||
* And optionally for {@link org.elasticsearch.inference.InferenceResults} | ||
* if the service uses a new type of result. | ||
* @return All named writables defined by the services | ||
*/ | ||
List<NamedWriteableRegistry.Entry> getInferenceServiceNamedWriteables(); | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.