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 | |
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.
4 files changed, 82 insertions, 18 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) {} }; } + + } diff --git a/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java b/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java index 5649cd51c95..0f89a839a26 100644 --- a/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java +++ b/config-model/src/main/java/com/yahoo/vespa/model/ml/OnnxModelProbe.java @@ -8,6 +8,7 @@ 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.config.model.api.OnnxMemoryStats; import com.yahoo.io.IOUtils; import com.yahoo.path.Path; import com.yahoo.tensor.TensorType; @@ -45,7 +46,7 @@ public class OnnxModelProbe { String jsonInput = createJsonInput(app.getFileReference(modelPath).getAbsolutePath(), inputTypes); var jsonOutput = callVespaAnalyzeOnnxModel(jsonInput); outputType = outputTypeFromJson(jsonOutput, outputName); - writeMemoryStats(app, modelPath, MemoryStats.fromJson(jsonOutput)); + writeMemoryStats(app, modelPath, OnnxMemoryStats.fromJson(jsonOutput)); if ( ! outputType.equals(TensorType.empty)) { writeProbedOutputType(app, modelPath, contextKey, outputType); } @@ -56,16 +57,11 @@ public class OnnxModelProbe { return outputType; } - private static void writeMemoryStats(ApplicationPackage app, Path modelPath, MemoryStats memoryStats) throws IOException { - String path = app.getFileReference(memoryStatsPath(modelPath)).getAbsolutePath(); + private static void writeMemoryStats(ApplicationPackage app, Path modelPath, OnnxMemoryStats memoryStats) throws IOException { + String path = app.getFileReference(OnnxMemoryStats.memoryStatsFilePath(modelPath)).getAbsolutePath(); IOUtils.writeFile(path, memoryStats.toJson().toPrettyString(), false); } - private static Path memoryStatsPath(Path modelPath) { - var fileName = OnnxModelInfo.asValidIdentifier(modelPath.getRelative()) + ".memory_stats"; - return ApplicationPackage.MODELS_GENERATED_REPLICATED_DIR.append(fileName); - } - private static String createContextKey(String onnxName, Map<String, TensorType> inputTypes) { StringBuilder key = new StringBuilder().append(onnxName).append(":"); inputTypes.entrySet().stream().sorted(Map.Entry.comparingByKey()) @@ -161,14 +157,4 @@ public class OnnxModelProbe { } return jsonParser.readTree(output.toString()); } - - public record MemoryStats(long vmSize, long vmRss) { - static MemoryStats fromJson(JsonNode json) { - return new MemoryStats(json.get("vm_size").asLong(), json.get("vm_rss").asLong()); - } - JsonNode toJson() { - return jsonParser.createObjectNode().put("vm_size", vmSize).put("vm_rss", vmRss); - } - } - } |