Add offheap memory tracking to Bench YAML reports (#671)

* corrected offheap calculation and add to reported statistics

Author: MarkWolters

jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/QueryTester.java

@@ -98,17 +98,23 @@ public List<Metric> run(
             diagnostics.capturePostPhaseSnapshot("Query");
 
             // Add memory and disk metrics to results
-            var systemSnapshot = diagnostics.getLatestSystemSnapshot();
+            // Get both pre and post snapshots to calculate maximum memory usage
+            var preSnapshot = diagnostics.getPrePhaseSystemSnapshot();
+            var postSnapshot = diagnostics.getLatestSystemSnapshot();
             var diskSnapshot = diagnostics.getLatestDiskSnapshot();
 
-            if (systemSnapshot != null) {
-                // Max heap usage in MB
+            if (preSnapshot != null && postSnapshot != null) {
+                // Calculate max heap usage across pre and post snapshots
+                long maxHeapUsed = Math.max(preSnapshot.memoryStats.heapUsed,
+                                           postSnapshot.memoryStats.heapUsed);
                 results.add(Metric.of("search.system.max_heap_mb", "Max heap usage (MB)", ".1f",
-                        systemSnapshot.memoryStats.heapUsed / (1024.0 * 1024.0)));
+                        maxHeapUsed / (1024.0 * 1024.0)));
 
-                // Max off-heap usage (direct + mapped) in MB
+                // Calculate max off-heap usage (direct + mapped) across pre and post snapshots
+                long maxOffHeapUsed = Math.max(preSnapshot.memoryStats.getTotalOffHeapMemory(),
+                                               postSnapshot.memoryStats.getTotalOffHeapMemory());
                 results.add(Metric.of("search.system.max_offheap_mb", "Max offheap usage (MB)", ".1f",
-                        systemSnapshot.memoryStats.getTotalOffHeapMemory() / (1024.0 * 1024.0)));
+                        maxOffHeapUsed / (1024.0 * 1024.0)));
             }
 
             if (diskSnapshot != null) {

jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/diagnostics/BenchmarkDiagnostics.java

@@ -233,6 +233,20 @@ public SystemMonitor.SystemSnapshot getLatestSystemSnapshot() {
         return snapshots.isEmpty() ? null : snapshots.get(snapshots.size() - 1);
     }
 
+    /**
+     * Gets the pre-phase system snapshot (second to last), or null if not available
+     */
+    public SystemMonitor.SystemSnapshot getPrePhaseSystemSnapshot() {
+        // The pre-phase snapshot is the second to last entry in the snapshots list, since the last one is the post-phase snapshot
+        var numSnapshots = snapshots.size();
+        if (numSnapshots < 2) {
+            return null;
+        } else {
+            SystemMonitor.SystemSnapshot prePhaseSnapshot = snapshots.get(numSnapshots - 2);
+            return prePhaseSnapshot;
+        }
+    }
+
     /**
      * Gets the latest disk usage snapshot, or null if none captured
      */

jvector-examples/yaml-configs/run-config.yml

@@ -17,7 +17,7 @@ console:
     count:      [ visited ]  # [visited, expanded, expanded base layer]
     accuracy:   [ recall ]   # [recall, MAP]
   metrics:
-    system: [ max_heap_mb ] # [ max_heap_mb, max_offheap_mb ]
+    system: [ max_heap_mb, max_offheap_mb ] # [ max_heap_mb, max_offheap_mb ]
 #    disk: total_file_size_mb # [ total_file_size_mb, file_count ]
 #    construction: index_build_time_s # [ index_build_time_s, index_quant_time_s, search_quant_time_s ]