Revert "Expose model registry to SemanticTextFieldMapper (elastic#126635)"

This reverts commit c906cc0.

Author: kderusso

server/src/main/java/org/elasticsearch/inference/MinimalServiceSettings.java

@@ -249,13 +249,18 @@ private static void validateFieldNotPresent(String field, Object fieldValue, Tas
         }
     }
 
+    public ModelConfigurations toModelConfigurations(String inferenceEntityId) {
+        return new ModelConfigurations(inferenceEntityId, taskType, service == null ? UNKNOWN_SERVICE : service, this);
+    }
+
     /**
      * Checks if the given {@link MinimalServiceSettings} is equivalent to the current definition.
      */
     public boolean canMergeWith(MinimalServiceSettings other) {
         return taskType == other.taskType
             && Objects.equals(dimensions, other.dimensions)
             && similarity == other.similarity
-            && elementType == other.elementType;
+            && elementType == other.elementType
+            && (service == null || service.equals(other.service));
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -197,7 +197,6 @@ public class InferencePlugin extends Plugin
     private final SetOnce<ElasticInferenceServiceComponents> elasticInferenceServiceComponents = new SetOnce<>();
     private final SetOnce<InferenceServiceRegistry> inferenceServiceRegistry = new SetOnce<>();
     private final SetOnce<ShardBulkInferenceActionFilter> shardBulkInferenceActionFilter = new SetOnce<>();
-    private final SetOnce<ModelRegistry> modelRegistry = new SetOnce<>();
     private List<InferenceServiceExtension> inferenceServiceExtensions;
 
     public InferencePlugin(Settings settings) {
@@ -261,8 +260,8 @@ public Collection<?> createComponents(PluginServices services) {
         var amazonBedrockRequestSenderFactory = new AmazonBedrockRequestSender.Factory(serviceComponents.get(), services.clusterService());
         amazonBedrockFactory.set(amazonBedrockRequestSenderFactory);
 
-        modelRegistry.set(new ModelRegistry(services.clusterService(), services.client()));
-        services.clusterService().addListener(modelRegistry.get());
+        ModelRegistry modelRegistry = new ModelRegistry(services.clusterService(), services.client());
+        services.clusterService().addListener(modelRegistry);
 
         if (inferenceServiceExtensions == null) {
             inferenceServiceExtensions = new ArrayList<>();
@@ -300,7 +299,7 @@ public Collection<?> createComponents(PluginServices services) {
                     elasicInferenceServiceFactory.get(),
                     serviceComponents.get(),
                     inferenceServiceSettings,
-                    modelRegistry.get(),
+                    modelRegistry,
                     authorizationHandler
                 )
             )
@@ -318,14 +317,14 @@ public Collection<?> createComponents(PluginServices services) {
         var serviceRegistry = new InferenceServiceRegistry(inferenceServices, factoryContext);
         serviceRegistry.init(services.client());
         for (var service : serviceRegistry.getServices().values()) {
-            service.defaultConfigIds().forEach(modelRegistry.get()::addDefaultIds);
+            service.defaultConfigIds().forEach(modelRegistry::addDefaultIds);
         }
         inferenceServiceRegistry.set(serviceRegistry);
 
         var actionFilter = new ShardBulkInferenceActionFilter(
             services.clusterService(),
             serviceRegistry,
-            modelRegistry.get(),
+            modelRegistry,
             getLicenseState(),
             services.indexingPressure()
         );
@@ -335,7 +334,7 @@ public Collection<?> createComponents(PluginServices services) {
         var inferenceStats = new PluginComponentBinding<>(InferenceStats.class, InferenceStats.create(meterRegistry));
 
         components.add(serviceRegistry);
-        components.add(modelRegistry.get());
+        components.add(modelRegistry);
         components.add(httpClientManager);
         components.add(inferenceStats);
 
@@ -499,16 +498,11 @@ public Map<String, MetadataFieldMapper.TypeParser> getMetadataMappers() {
         return Map.of(SemanticInferenceMetadataFieldsMapper.NAME, SemanticInferenceMetadataFieldsMapper.PARSER);
     }
 
-    // Overridable for tests
-    protected Supplier<ModelRegistry> getModelRegistry() {
-        return () -> modelRegistry.get();
-    }
-
     @Override
     public Map<String, Mapper.TypeParser> getMappers() {
         return Map.of(
             SemanticTextFieldMapper.CONTENT_TYPE,
-            SemanticTextFieldMapper.parser(getModelRegistry()),
+            SemanticTextFieldMapper.PARSER,
             OffsetSourceFieldMapper.CONTENT_TYPE,
             OffsetSourceFieldMapper.PARSER
         );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java

@@ -7,8 +7,6 @@
 
 package org.elasticsearch.xpack.inference.mapper;
 
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.DocIdSetIterator;
@@ -20,7 +18,6 @@
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.search.join.ScoreMode;
 import org.apache.lucene.util.BitSet;
-import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.cluster.metadata.InferenceFieldMetadata;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.BytesReference;
@@ -78,7 +75,6 @@
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
 import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryBuilder;
 import org.elasticsearch.xpack.inference.highlight.SemanticTextHighlighter;
-import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -93,7 +89,6 @@
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
-import java.util.function.Supplier;
 
 import static org.elasticsearch.inference.TaskType.SPARSE_EMBEDDING;
 import static org.elasticsearch.inference.TaskType.TEXT_EMBEDDING;
@@ -117,7 +112,6 @@
  * A {@link FieldMapper} for semantic text fields.
  */
 public class SemanticTextFieldMapper extends FieldMapper implements InferenceFieldMapper {
-    private static final Logger logger = LogManager.getLogger(SemanticTextFieldMapper.class);
     public static final NodeFeature SEMANTIC_TEXT_IN_OBJECT_FIELD_FIX = new NodeFeature("semantic_text.in_object_field_fix");
     public static final NodeFeature SEMANTIC_TEXT_SINGLE_FIELD_UPDATE_FIX = new NodeFeature("semantic_text.single_field_update_fix");
     public static final NodeFeature SEMANTIC_TEXT_DELETE_FIX = new NodeFeature("semantic_text.delete_fix");
@@ -133,12 +127,10 @@ public class SemanticTextFieldMapper extends FieldMapper implements InferenceFie
     public static final String CONTENT_TYPE = "semantic_text";
     public static final String DEFAULT_ELSER_2_INFERENCE_ID = DEFAULT_ELSER_ID;
 
-    public static final TypeParser parser(Supplier<ModelRegistry> modelRegistry) {
-        return new TypeParser(
-            (n, c) -> new Builder(n, c::bitSetProducer, c.getIndexSettings(), modelRegistry.get()),
-            List.of(validateParserContext(CONTENT_TYPE))
-        );
-    }
+    public static final TypeParser PARSER = new TypeParser(
+        (n, c) -> new Builder(n, c::bitSetProducer, c.getIndexSettings()),
+        List.of(validateParserContext(CONTENT_TYPE))
+    );
 
     public static BiConsumer<String, MappingParserContext> validateParserContext(String type) {
         return (n, c) -> {
@@ -150,7 +142,6 @@ public static BiConsumer<String, MappingParserContext> validateParserContext(Str
     }
 
     public static class Builder extends FieldMapper.Builder {
-        private final ModelRegistry modelRegistry;
         private final boolean useLegacyFormat;
 
         private final Parameter<String> inferenceId = Parameter.stringParam(
@@ -208,21 +199,14 @@ public static Builder from(SemanticTextFieldMapper mapper) {
             Builder builder = new Builder(
                 mapper.leafName(),
                 mapper.fieldType().getChunksField().bitsetProducer(),
-                mapper.fieldType().getChunksField().indexSettings(),
-                mapper.modelRegistry
+                mapper.fieldType().getChunksField().indexSettings()
             );
             builder.init(mapper);
             return builder;
         }
 
-        public Builder(
-            String name,
-            Function<Query, BitSetProducer> bitSetProducer,
-            IndexSettings indexSettings,
-            ModelRegistry modelRegistry
-        ) {
+        public Builder(String name, Function<Query, BitSetProducer> bitSetProducer, IndexSettings indexSettings) {
             super(name);
-            this.modelRegistry = modelRegistry;
             this.useLegacyFormat = InferenceMetadataFieldsMapper.isEnabled(indexSettings.getSettings()) == false;
             this.inferenceFieldBuilder = c -> createInferenceField(
                 c,
@@ -280,32 +264,9 @@ public SemanticTextFieldMapper build(MapperBuilderContext context) {
             if (useLegacyFormat && multiFieldsBuilder.hasMultiFields()) {
                 throw new IllegalArgumentException(CONTENT_TYPE + " field [" + leafName() + "] does not support multi-fields");
             }
-
-            if (modelSettings.get() == null) {
-                try {
-                    var resolvedModelSettings = modelRegistry.getMinimalServiceSettings(inferenceId.get());
-                    if (resolvedModelSettings != null) {
-                        modelSettings.setValue(resolvedModelSettings);
-                    }
-                } catch (ResourceNotFoundException exc) {
-                    // We allow the inference ID to be unregistered at this point.
-                    // This will delay the creation of sub-fields, so indexing and querying for this field won't work
-                    // until the corresponding inference endpoint is created.
-                }
-            }
-
             if (modelSettings.get() != null) {
                 validateServiceSettings(modelSettings.get());
-            } else {
-                logger.warn(
-                    "The field [{}] references an unknown inference ID [{}]. "
-                        + "Indexing and querying this field will not work correctly until the corresponding "
-                        + "inference endpoint is created.",
-                    leafName(),
-                    inferenceId.get()
-                );
             }
-
             final String fullName = context.buildFullName(leafName());
 
             if (context.isInNestedContext()) {
@@ -326,8 +287,7 @@ public SemanticTextFieldMapper build(MapperBuilderContext context) {
                     useLegacyFormat,
                     meta.getValue()
                 ),
-                builderParams(this, context),
-                modelRegistry
+                builderParams(this, context)
             );
         }
 
@@ -368,17 +328,9 @@ private SemanticTextFieldMapper copySettings(SemanticTextFieldMapper mapper, Map
         }
     }
 
-    private final ModelRegistry modelRegistry;
-
-    private SemanticTextFieldMapper(
-        String simpleName,
-        MappedFieldType mappedFieldType,
-        BuilderParams builderParams,
-        ModelRegistry modelRegistry
-    ) {
+    private SemanticTextFieldMapper(String simpleName, MappedFieldType mappedFieldType, BuilderParams builderParams) {
         super(simpleName, mappedFieldType, builderParams);
         ensureMultiFields(builderParams.multiFields().iterator());
-        this.modelRegistry = modelRegistry;
     }
 
     private void ensureMultiFields(Iterator<FieldMapper> mappers) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/registry/ModelRegistry.java

@@ -36,7 +36,6 @@
 import org.elasticsearch.cluster.ClusterStateAckListener;
 import org.elasticsearch.cluster.ClusterStateListener;
 import org.elasticsearch.cluster.SimpleBatchedAckListenerTaskExecutor;
-import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.metadata.ProjectMetadata;
 import org.elasticsearch.cluster.service.ClusterService;
@@ -141,18 +140,18 @@ public static UnparsedModel unparsedModelFromMap(ModelConfigMap modelConfigMap)
     private static final String MODEL_ID_FIELD = "model_id";
     private static final Logger logger = LogManager.getLogger(ModelRegistry.class);
 
+    private final ClusterService clusterService;
     private final OriginSettingClient client;
     private final Map<String, InferenceService.DefaultConfigId> defaultConfigIds;
 
     private final MasterServiceTaskQueue<MetadataTask> metadataTaskQueue;
     private final AtomicBoolean upgradeMetadataInProgress = new AtomicBoolean(false);
     private final Set<String> preventDeletionLock = Collections.newSetFromMap(new ConcurrentHashMap<>());
 
-    private volatile Metadata lastMetadata;
-
     public ModelRegistry(ClusterService clusterService, Client client) {
         this.client = new OriginSettingClient(client, ClientHelper.INFERENCE_ORIGIN);
         this.defaultConfigIds = new ConcurrentHashMap<>();
+        this.clusterService = clusterService;
         var executor = new SimpleBatchedAckListenerTaskExecutor<MetadataTask>() {
             @Override
             public Tuple<ClusterState, ClusterStateAckListener> executeTask(MetadataTask task, ClusterState clusterState) throws Exception {
@@ -225,18 +224,11 @@ public void clearDefaultIds() {
      * @throws ResourceNotFoundException if the specified id is guaranteed to not exist in the cluster.
      */
     public MinimalServiceSettings getMinimalServiceSettings(String inferenceEntityId) throws ResourceNotFoundException {
-        synchronized (this) {
-            assert lastMetadata != null : "initial cluster state not set yet";
-            if (lastMetadata == null) {
-                throw new IllegalStateException("initial cluster state not set yet");
-            }
-        }
         var config = defaultConfigIds.get(inferenceEntityId);
         if (config != null) {
             return config.settings();
         }
-        var project = lastMetadata.getProject(ProjectId.DEFAULT);
-        var state = ModelRegistryMetadata.fromState(project);
+        var state = ModelRegistryMetadata.fromState(clusterService.state().projectState().metadata());
         var existing = state.getMinimalServiceSettings(inferenceEntityId);
         if (state.isUpgraded() && existing == null) {
             throw new ResourceNotFoundException(inferenceEntityId + " does not exist in this cluster.");
@@ -692,14 +684,10 @@ private ActionListener<BulkResponse> getStoreIndexListener(
                 if (updateClusterState) {
                     var storeListener = getStoreMetadataListener(inferenceEntityId, listener);
                     try {
+                        var projectId = clusterService.state().projectState().projectId();
                         metadataTaskQueue.submitTask(
                             "add model [" + inferenceEntityId + "]",
-                            new AddModelMetadataTask(
-                                ProjectId.DEFAULT,
-                                inferenceEntityId,
-                                new MinimalServiceSettings(model),
-                                storeListener
-                            ),
+                            new AddModelMetadataTask(projectId, inferenceEntityId, new MinimalServiceSettings(model), storeListener),
                             timeout
                         );
                     } catch (Exception exc) {
@@ -866,9 +854,10 @@ public void onFailure(Exception exc) {
                     }
                 };
                 try {
+                    var projectId = clusterService.state().projectState().projectId();
                     metadataTaskQueue.submitTask(
                         "delete models [" + inferenceEntityIds + "]",
-                        new DeleteModelMetadataTask(ProjectId.DEFAULT, inferenceEntityIds, clusterStateListener),
+                        new DeleteModelMetadataTask(projectId, inferenceEntityIds, clusterStateListener),
                         null
                     );
                 } catch (Exception exc) {
@@ -946,13 +935,6 @@ static List<InferenceService.DefaultConfigId> taskTypeMatchedDefaults(
 
     @Override
     public void clusterChanged(ClusterChangedEvent event) {
-        if (lastMetadata == null || event.metadataChanged()) {
-            // keep track of the last applied cluster state
-            synchronized (this) {
-                lastMetadata = event.state().metadata();
-            }
-        }
-
         if (event.localNodeMaster() == false) {
             return;
         }
@@ -1002,7 +984,7 @@ public void onResponse(GetInferenceModelAction.Response response) {
                     metadataTaskQueue.submitTask(
                         "model registry auto upgrade",
                         new UpgradeModelsMetadataTask(
-                            ProjectId.DEFAULT,
+                            clusterService.state().metadata().getProject().id(),
                             map,
                             ActionListener.running(() -> upgradeMetadataInProgress.set(false))
                         ),