diff options
author | Bjørn Christian Seime <bjorncs@vespa.ai> | 2023-10-19 12:34:47 +0200 |
---|---|---|
committer | Bjørn Christian Seime <bjorncs@vespa.ai> | 2023-10-19 12:50:29 +0200 |
commit | 6ec94f4b9991cc610404c49c4fdb1476877ab647 (patch) | |
tree | f3a4bdbb82b07564349b0d61d5ee6cfdf1409c09 /config-model-api | |
parent | 8489726da2b8052d928ca6ae1425d8ae0b7bb119 (diff) |
Track new metrics `malloc_peak` and `malloc_current`
Move class to config-model-api to be shared with internal config server integration.
Diffstat (limited to 'config-model-api')
3 files changed, 78 insertions, 0 deletions
diff --git a/config-model-api/abi-spec.json b/config-model-api/abi-spec.json index b28401f1873..6b663800b67 100644 --- a/config-model-api/abi-spec.json +++ b/config-model-api/abi-spec.json @@ -1420,6 +1420,31 @@ ], "fields" : [ ] }, + "com.yahoo.config.model.api.OnnxMemoryStats" : { + "superClass" : "java.lang.Record", + "interfaces" : [ ], + "attributes" : [ + "public", + "final", + "record" + ], + "methods" : [ + "public void <init>(long, long, long, long)", + "public static com.yahoo.config.model.api.OnnxMemoryStats fromJson(com.fasterxml.jackson.databind.JsonNode)", + "public static com.yahoo.config.model.api.OnnxMemoryStats fromJson(com.yahoo.config.application.api.ApplicationFile)", + "public static com.yahoo.path.Path memoryStatsFilePath(com.yahoo.path.Path)", + "public long peakMemoryUsage()", + "public com.fasterxml.jackson.databind.JsonNode toJson()", + "public final java.lang.String toString()", + "public final int hashCode()", + "public final boolean equals(java.lang.Object)", + "public long vmSize()", + "public long vmRss()", + "public long mallocPeak()", + "public long mallocCurrent()" + ], + "fields" : [ ] + }, "com.yahoo.config.model.api.OnnxModelCost$Calculator" : { "superClass" : "java.lang.Object", "interfaces" : [ ], diff --git a/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxMemoryStats.java b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxMemoryStats.java new file mode 100644 index 00000000000..4e660c6fe73 --- /dev/null +++ b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxMemoryStats.java @@ -0,0 +1,49 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +package com.yahoo.config.model.api; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.yahoo.config.application.api.ApplicationFile; +import com.yahoo.config.application.api.ApplicationPackage; +import com.yahoo.path.Path; + +import java.io.IOException; +import java.util.Optional; + +/** + * Memory statistics as reported by vespa-analyze-onnx-model. + * + * @author bjorncs + */ +public record OnnxMemoryStats(long vmSize, long vmRss, long mallocPeak, long mallocCurrent) { + private static final String VM_SIZE_FIELD = "vm_size", VM_RSS_FIELD = "vm_rss", + MALLOC_PEAK_FIELD = "malloc_peak", MALLOC_CURRENT_FIELD = "malloc_current"; + private static final ObjectMapper jsonParser = new ObjectMapper(); + + /** Parse output from `vespa-analyze-onnx-model --probe-types` */ + public static OnnxMemoryStats fromJson(JsonNode json) { + return new OnnxMemoryStats(json.get(VM_SIZE_FIELD).asLong(), json.get(VM_RSS_FIELD).asLong(), + // Temporarily allow missing fields until old config model versions are gone + Optional.ofNullable(json.get(MALLOC_PEAK_FIELD)).map(JsonNode::asLong).orElse(0L), + Optional.ofNullable(json.get(MALLOC_CURRENT_FIELD)).map(JsonNode::asLong).orElse(0L)); + } + + /** @see #fromJson(JsonNode) */ + public static OnnxMemoryStats fromJson(ApplicationFile file) throws IOException { + return fromJson(jsonParser.readTree(file.createReader())); + } + + public static Path memoryStatsFilePath(Path modelPath) { + var fileName = modelPath.getRelative().replaceAll("[^\\w\\d\\$@_]", "_") + ".memory_stats"; + return ApplicationPackage.MODELS_GENERATED_REPLICATED_DIR.append(fileName); + } + + public long peakMemoryUsage() { return Long.max(vmSize, Long.max(vmRss, Long.max(mallocPeak, mallocCurrent))); } + + public JsonNode toJson() { + return jsonParser.createObjectNode().put(VM_SIZE_FIELD, vmSize).put(VM_RSS_FIELD, vmRss) + .put(MALLOC_PEAK_FIELD, mallocPeak).put(MALLOC_CURRENT_FIELD, mallocCurrent); + } +} + diff --git a/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java index e6fe3ce18b5..44f3d63f8af 100644 --- a/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java +++ b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java @@ -2,6 +2,8 @@ package com.yahoo.config.model.api; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import com.yahoo.config.ModelReference; import com.yahoo.config.application.api.ApplicationFile; import com.yahoo.config.application.api.ApplicationPackage; @@ -31,4 +33,6 @@ public interface OnnxModelCost { @Override public void registerModel(URI uri) {} }; } + + } |