Skip to content

Commit dd134b7

Browse files
ioanatiaioanatia
committed
Draft for semantic_text telemetry
1 parent aab83a8 commit dd134b7

File tree

2 files changed

+127
-6
lines changed

2 files changed

+127
-6
lines changed

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java‎

Lines changed: 66 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,30 +102,92 @@ public int hashCode() {
102102
}
103103
}
104104

105-
public static final InferenceFeatureSetUsage EMPTY = new InferenceFeatureSetUsage(List.of());
105+
public static class SemanticTextStats implements ToXContentObject, Writeable {
106+
private final Long totalFieldCount;
107+
private final Long indexCount;
108+
private final Long denseFieldCount;
109+
private final Long sparseFieldCount;
110+
private final Long denseInferenceIdCount;
111+
private final Long sparseInferenceIdCount;
112+
113+
public SemanticTextStats(
114+
Long totalFieldCount,
115+
Long indexCount,
116+
Long sparseFieldCount,
117+
Long denseFieldCount,
118+
Long denseInferenceIdCount,
119+
Long sparseInferenceIdCount
120+
) {
121+
this.totalFieldCount = totalFieldCount;
122+
this.indexCount = indexCount;
123+
this.sparseFieldCount = sparseFieldCount;
124+
this.denseFieldCount = denseFieldCount;
125+
this.denseInferenceIdCount = denseInferenceIdCount;
126+
this.sparseInferenceIdCount = sparseInferenceIdCount;
127+
}
128+
129+
public SemanticTextStats(StreamInput in) throws IOException {
130+
this.totalFieldCount = in.readLong();
131+
this.indexCount = in.readLong();
132+
this.denseFieldCount = in.readLong();
133+
this.denseInferenceIdCount = in.readLong();
134+
this.sparseInferenceIdCount = in.readLong();
135+
this.sparseFieldCount = in.readLong();
136+
}
137+
138+
@Override
139+
public void writeTo(StreamOutput out) throws IOException {
140+
out.writeLong(totalFieldCount);
141+
out.writeLong(indexCount);
142+
out.writeLong(denseFieldCount);
143+
out.writeLong(denseInferenceIdCount);
144+
out.writeLong(sparseInferenceIdCount);
145+
out.writeLong(sparseFieldCount);
146+
}
147+
148+
@Override
149+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
150+
builder.startObject();
151+
builder.field("total_fields", totalFieldCount);
152+
builder.field("indices", indexCount);
153+
builder.field("dense_fields", denseFieldCount);
154+
builder.field("dense_inference_ids", denseInferenceIdCount);
155+
builder.field("sparse_fields", sparseFieldCount);
156+
builder.field("sparse_inference_ids", sparseInferenceIdCount);
157+
builder.endObject();
158+
return builder;
159+
}
160+
}
161+
162+
public static final InferenceFeatureSetUsage EMPTY = new InferenceFeatureSetUsage(List.of(), null);
106163

107164
private final Collection<ModelStats> modelStats;
165+
private final SemanticTextStats semanticTextStats;
108166

109-
public InferenceFeatureSetUsage(Collection<ModelStats> modelStats) {
167+
public InferenceFeatureSetUsage(Collection<ModelStats> modelStats, SemanticTextStats semanticTextStats) {
110168
super(XPackField.INFERENCE, true, true);
111169
this.modelStats = modelStats;
170+
this.semanticTextStats = semanticTextStats;
112171
}
113172

114173
public InferenceFeatureSetUsage(StreamInput in) throws IOException {
115174
super(in);
116175
this.modelStats = in.readCollectionAsList(ModelStats::new);
176+
this.semanticTextStats = new SemanticTextStats(in);
117177
}
118178

119179
@Override
120180
protected void innerXContent(XContentBuilder builder, Params params) throws IOException {
121181
super.innerXContent(builder, params);
122182
builder.xContentList("models", modelStats);
183+
builder.field("semantic_text", semanticTextStats);
123184
}
124185

125186
@Override
126187
public void writeTo(StreamOutput out) throws IOException {
127188
super.writeTo(out);
128189
out.writeCollection(modelStats);
190+
semanticTextStats.writeTo(out);
129191
}
130192

131193
@Override
@@ -137,11 +199,11 @@ public TransportVersion getMinimalSupportedVersion() {
137199
public boolean equals(Object o) {
138200
if (o == null || getClass() != o.getClass()) return false;
139201
InferenceFeatureSetUsage that = (InferenceFeatureSetUsage) o;
140-
return Objects.equals(modelStats, that.modelStats);
202+
return Objects.equals(modelStats, that.modelStats) && Objects.equals(semanticTextStats, that.semanticTextStats);
141203
}
142204

143205
@Override
144206
public int hashCode() {
145-
return Objects.hashCode(modelStats);
207+
return Objects.hash(modelStats, semanticTextStats);
146208
}
147209
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java‎

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
import org.elasticsearch.client.internal.Client;
1515
import org.elasticsearch.client.internal.OriginSettingClient;
1616
import org.elasticsearch.cluster.ClusterState;
17+
import org.elasticsearch.cluster.metadata.IndexMetadata;
18+
import org.elasticsearch.cluster.metadata.InferenceFieldMetadata;
1719
import org.elasticsearch.cluster.service.ClusterService;
1820
import org.elasticsearch.common.Strings;
1921
import org.elasticsearch.inference.ModelConfigurations;
@@ -29,6 +31,8 @@
2931
import org.elasticsearch.xpack.core.inference.InferenceFeatureSetUsage;
3032
import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
3133

34+
import java.util.HashMap;
35+
import java.util.List;
3236
import java.util.Map;
3337
import java.util.TreeMap;
3438

@@ -62,19 +66,74 @@ protected void localClusterStateOperation(
6266
GetInferenceModelAction.Request getInferenceModelAction = new GetInferenceModelAction.Request("_all", TaskType.ANY, false);
6367
client.execute(GetInferenceModelAction.INSTANCE, getInferenceModelAction, ActionListener.wrap(response -> {
6468
Map<String, InferenceFeatureSetUsage.ModelStats> stats = new TreeMap<>();
65-
for (ModelConfigurations model : response.getEndpoints()) {
69+
List<ModelConfigurations> endpoints = response.getEndpoints();
70+
for (ModelConfigurations model : endpoints) {
6671
String statKey = model.getService() + ":" + model.getTaskType().name();
6772
InferenceFeatureSetUsage.ModelStats stat = stats.computeIfAbsent(
6873
statKey,
6974
key -> new InferenceFeatureSetUsage.ModelStats(model.getService(), model.getTaskType())
7075
);
7176
stat.add();
7277
}
73-
InferenceFeatureSetUsage usage = new InferenceFeatureSetUsage(stats.values());
78+
79+
InferenceFeatureSetUsage usage = new InferenceFeatureSetUsage(
80+
stats.values(),
81+
getSemanticTextStats(state.getMetadata().indicesAllProjects(), endpoints)
82+
);
7483
listener.onResponse(new XPackUsageFeatureResponse(usage));
7584
}, e -> {
7685
logger.warn(Strings.format("Retrieving inference usage failed with error: %s", e.getMessage()), e);
7786
listener.onResponse(new XPackUsageFeatureResponse(InferenceFeatureSetUsage.EMPTY));
7887
}));
7988
}
89+
90+
private static InferenceFeatureSetUsage.SemanticTextStats getSemanticTextStats(
91+
Iterable<IndexMetadata> indicesMetadata,
92+
List<ModelConfigurations> modelConfigurations
93+
) {
94+
long fieldCount = 0;
95+
long indexCount = 0;
96+
97+
Map<String, Long> inferenceIdsCounts = new HashMap<>();
98+
99+
for (IndexMetadata indexMetadata : indicesMetadata) {
100+
Map<String, InferenceFieldMetadata> inferenceFields = indexMetadata.getInferenceFields();
101+
102+
fieldCount += inferenceFields.size();
103+
indexCount += inferenceFields.isEmpty() ? 0 : 1;
104+
105+
inferenceFields.forEach((fieldName, inferenceFieldMetadata) -> {
106+
String inferenceId = inferenceFieldMetadata.getInferenceId();
107+
inferenceIdsCounts.compute(inferenceId, (k, v) -> v == null ? 1 : v + 1);
108+
});
109+
}
110+
111+
long sparseFieldsCount = 0;
112+
long denseFieldsCount = 0;
113+
long denseInferenceIdCount = 0;
114+
long sparseInferenceIdCount = 0;
115+
for (ModelConfigurations model : modelConfigurations) {
116+
String inferenceId = model.getInferenceEntityId();
117+
118+
if (inferenceIdsCounts.containsKey(inferenceId) == false) {
119+
continue;
120+
}
121+
if (model.getTaskType() == TaskType.SPARSE_EMBEDDING) {
122+
sparseFieldsCount += inferenceIdsCounts.get(inferenceId);
123+
sparseInferenceIdCount += 1;
124+
} else {
125+
denseFieldsCount += inferenceIdsCounts.get(inferenceId);
126+
denseInferenceIdCount += 1;
127+
}
128+
}
129+
130+
return new InferenceFeatureSetUsage.SemanticTextStats(
131+
fieldCount,
132+
indexCount,
133+
sparseFieldsCount,
134+
denseFieldsCount,
135+
denseInferenceIdCount,
136+
sparseInferenceIdCount
137+
);
138+
}
80139
}

0 commit comments

Comments
 (0)