From dd134b7b71e9f2c485675f0c79beacb062466f23 Mon Sep 17 00:00:00 2001 From: Ioana Tagirta Date: Thu, 28 Aug 2025 21:33:34 +0200 Subject: [PATCH] Draft for semantic_text telemetry --- .../inference/InferenceFeatureSetUsage.java | 70 +++++++++++++++++-- .../action/TransportInferenceUsageAction.java | 63 ++++++++++++++++- 2 files changed, 127 insertions(+), 6 deletions(-) diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java index 1a75d079adec2..38d20b55546d7 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java @@ -102,30 +102,92 @@ public int hashCode() { } } - public static final InferenceFeatureSetUsage EMPTY = new InferenceFeatureSetUsage(List.of()); + public static class SemanticTextStats implements ToXContentObject, Writeable { + private final Long totalFieldCount; + private final Long indexCount; + private final Long denseFieldCount; + private final Long sparseFieldCount; + private final Long denseInferenceIdCount; + private final Long sparseInferenceIdCount; + + public SemanticTextStats( + Long totalFieldCount, + Long indexCount, + Long sparseFieldCount, + Long denseFieldCount, + Long denseInferenceIdCount, + Long sparseInferenceIdCount + ) { + this.totalFieldCount = totalFieldCount; + this.indexCount = indexCount; + this.sparseFieldCount = sparseFieldCount; + this.denseFieldCount = denseFieldCount; + this.denseInferenceIdCount = denseInferenceIdCount; + this.sparseInferenceIdCount = sparseInferenceIdCount; + } + + public SemanticTextStats(StreamInput in) throws IOException { + this.totalFieldCount = in.readLong(); + this.indexCount = in.readLong(); + this.denseFieldCount = in.readLong(); + this.denseInferenceIdCount = in.readLong(); + this.sparseInferenceIdCount = in.readLong(); + this.sparseFieldCount = in.readLong(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeLong(totalFieldCount); + out.writeLong(indexCount); + out.writeLong(denseFieldCount); + out.writeLong(denseInferenceIdCount); + out.writeLong(sparseInferenceIdCount); + out.writeLong(sparseFieldCount); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field("total_fields", totalFieldCount); + builder.field("indices", indexCount); + builder.field("dense_fields", denseFieldCount); + builder.field("dense_inference_ids", denseInferenceIdCount); + builder.field("sparse_fields", sparseFieldCount); + builder.field("sparse_inference_ids", sparseInferenceIdCount); + builder.endObject(); + return builder; + } + } + + public static final InferenceFeatureSetUsage EMPTY = new InferenceFeatureSetUsage(List.of(), null); private final Collection modelStats; + private final SemanticTextStats semanticTextStats; - public InferenceFeatureSetUsage(Collection modelStats) { + public InferenceFeatureSetUsage(Collection modelStats, SemanticTextStats semanticTextStats) { super(XPackField.INFERENCE, true, true); this.modelStats = modelStats; + this.semanticTextStats = semanticTextStats; } public InferenceFeatureSetUsage(StreamInput in) throws IOException { super(in); this.modelStats = in.readCollectionAsList(ModelStats::new); + this.semanticTextStats = new SemanticTextStats(in); } @Override protected void innerXContent(XContentBuilder builder, Params params) throws IOException { super.innerXContent(builder, params); builder.xContentList("models", modelStats); + builder.field("semantic_text", semanticTextStats); } @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeCollection(modelStats); + semanticTextStats.writeTo(out); } @Override @@ -137,11 +199,11 @@ public TransportVersion getMinimalSupportedVersion() { public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; InferenceFeatureSetUsage that = (InferenceFeatureSetUsage) o; - return Objects.equals(modelStats, that.modelStats); + return Objects.equals(modelStats, that.modelStats) && Objects.equals(semanticTextStats, that.semanticTextStats); } @Override public int hashCode() { - return Objects.hashCode(modelStats); + return Objects.hash(modelStats, semanticTextStats); } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java index 7973373b021bb..9d5e93ef0acd7 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java @@ -14,6 +14,8 @@ import org.elasticsearch.client.internal.Client; import org.elasticsearch.client.internal.OriginSettingClient; import org.elasticsearch.cluster.ClusterState; +import org.elasticsearch.cluster.metadata.IndexMetadata; +import org.elasticsearch.cluster.metadata.InferenceFieldMetadata; import org.elasticsearch.cluster.service.ClusterService; import org.elasticsearch.common.Strings; import org.elasticsearch.inference.ModelConfigurations; @@ -29,6 +31,8 @@ import org.elasticsearch.xpack.core.inference.InferenceFeatureSetUsage; import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction; +import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.TreeMap; @@ -62,7 +66,8 @@ protected void localClusterStateOperation( GetInferenceModelAction.Request getInferenceModelAction = new GetInferenceModelAction.Request("_all", TaskType.ANY, false); client.execute(GetInferenceModelAction.INSTANCE, getInferenceModelAction, ActionListener.wrap(response -> { Map stats = new TreeMap<>(); - for (ModelConfigurations model : response.getEndpoints()) { + List endpoints = response.getEndpoints(); + for (ModelConfigurations model : endpoints) { String statKey = model.getService() + ":" + model.getTaskType().name(); InferenceFeatureSetUsage.ModelStats stat = stats.computeIfAbsent( statKey, @@ -70,11 +75,65 @@ protected void localClusterStateOperation( ); stat.add(); } - InferenceFeatureSetUsage usage = new InferenceFeatureSetUsage(stats.values()); + + InferenceFeatureSetUsage usage = new InferenceFeatureSetUsage( + stats.values(), + getSemanticTextStats(state.getMetadata().indicesAllProjects(), endpoints) + ); listener.onResponse(new XPackUsageFeatureResponse(usage)); }, e -> { logger.warn(Strings.format("Retrieving inference usage failed with error: %s", e.getMessage()), e); listener.onResponse(new XPackUsageFeatureResponse(InferenceFeatureSetUsage.EMPTY)); })); } + + private static InferenceFeatureSetUsage.SemanticTextStats getSemanticTextStats( + Iterable indicesMetadata, + List modelConfigurations + ) { + long fieldCount = 0; + long indexCount = 0; + + Map inferenceIdsCounts = new HashMap<>(); + + for (IndexMetadata indexMetadata : indicesMetadata) { + Map inferenceFields = indexMetadata.getInferenceFields(); + + fieldCount += inferenceFields.size(); + indexCount += inferenceFields.isEmpty() ? 0 : 1; + + inferenceFields.forEach((fieldName, inferenceFieldMetadata) -> { + String inferenceId = inferenceFieldMetadata.getInferenceId(); + inferenceIdsCounts.compute(inferenceId, (k, v) -> v == null ? 1 : v + 1); + }); + } + + long sparseFieldsCount = 0; + long denseFieldsCount = 0; + long denseInferenceIdCount = 0; + long sparseInferenceIdCount = 0; + for (ModelConfigurations model : modelConfigurations) { + String inferenceId = model.getInferenceEntityId(); + + if (inferenceIdsCounts.containsKey(inferenceId) == false) { + continue; + } + if (model.getTaskType() == TaskType.SPARSE_EMBEDDING) { + sparseFieldsCount += inferenceIdsCounts.get(inferenceId); + sparseInferenceIdCount += 1; + } else { + denseFieldsCount += inferenceIdsCounts.get(inferenceId); + denseInferenceIdCount += 1; + } + } + + return new InferenceFeatureSetUsage.SemanticTextStats( + fieldCount, + indexCount, + sparseFieldsCount, + denseFieldsCount, + denseInferenceIdCount, + sparseInferenceIdCount + ); + } }