Add micro-benchmarks module in k-NN plugin for benchmark streaming ve…

…ctors to JNI layer functionality.

Signed-off-by: Navneet Verma <[email protected]>

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Features
 ### Enhancements
 * Make the HitQueue size more appropriate for exact search [#1549](https://github.com/opensearch-project/k-NN/pull/1549)
+* Add micro-benchmark module in k-NN plugin for benchmark streaming vectors to JNI layer functionality. [#1583](https://github.com/opensearch-project/k-NN/pull/1583)
 ### Bug Fixes
 ### Infrastructure
 ### Documentation

jni/include/org_opensearch_knn_jni_FaissService.h

@@ -122,6 +122,14 @@ JNIEXPORT jbyteArray JNICALL Java_org_opensearch_knn_jni_FaissService_trainIndex
 JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectors
   (JNIEnv *, jclass, jlong, jobjectArray);
 
+/*
+ * Class:     org_opensearch_knn_jni_FaissService
+ * Method:    transferVectorsV2
+ * Signature: (J[[F)J
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectorsV2
+        (JNIEnv *, jclass, jlong, jobjectArray);
+
 /*
  * Class:     org_opensearch_knn_jni_FaissService
  * Method:    freeVectors

jni/src/org_opensearch_knn_jni_FaissService.cpp

@@ -191,6 +191,24 @@ JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectors
     return (jlong) vect;
 }
 
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectorsV2(JNIEnv * env, jclass cls,
+jlong vectorsPointerJ,
+        jobjectArray vectorsJ)
+{
+    std::vector<float> *vect;
+    if ((long) vectorsPointerJ == 0) {
+        vect = new std::vector<float>;
+    } else {
+        vect = reinterpret_cast<std::vector<float>*>(vectorsPointerJ);
+    }
+
+    int dim = jniUtil.GetInnerDimensionOf2dJavaFloatArray(env, vectorsJ);
+    auto dataset = jniUtil.Convert2dJavaObjectArrayToCppFloatVector(env, vectorsJ, dim);
+    vect->insert(vect->end(), dataset.begin(), dataset.end());
+
+    return (jlong) vect;
+}
+
 JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_freeVectors(JNIEnv * env, jclass cls,
                                                                             jlong vectorsPointerJ)
 {

micro-benchmarks/README.md

@@ -0,0 +1,97 @@
+# OpenSearch K-NN Microbenchmark Suite
+
+This directory contains the microbenchmark suite of Opensearch K-NN Plugin. It relies on [JMH](http://openjdk.java.net/projects/code-tools/jmh/).
+
+This module draws a lot of inspiration from [Opensearch benchmarks](https://github.com/opensearch-project/OpenSearch/tree/main/benchmarks). 
+
+## Purpose
+
+Micro benchmarks are intended to spot performance regressions in performance-critical components.
+
+The microbenchmark suite is also handy for ad-hoc micro benchmarks but please remove them again before merging your PR.
+
+## Getting Started
+
+Just run `gradlew -p micro-benchmarks run` from the project root
+directory. It will build all microbenchmarks, execute them and print
+the result.
+
+## Running Microbenchmarks
+
+Running via an IDE is not supported as the results are meaningless
+because we have no control over the JVM running the benchmarks.
+
+If you want to run a specific benchmark class like, say,
+`TransferVectorsBenchmarks`, you can use `--args`:
+
+```
+gradlew -p micro-benchmarks run --args ' TransferVectorsBenchmarks'
+```
+
+Setting Heap while running the benchmarks
+```
+./gradlew -p micro-benchmarks run --args ' -gc true ' -Djvm.heap.size=4g
+```
+
+Everything in the `'` gets sent on the command line to JMH. The leading ` `
+inside the `'`s is important. Without it parameters are sometimes sent to
+gradle.
+
+## Adding Microbenchmarks
+
+Before adding a new microbenchmark, make yourself familiar with the JMH API. You can check our existing microbenchmarks and also the
+[JMH samples](http://hg.openjdk.java.net/code-tools/jmh/file/tip/jmh-samples/src/main/java/org/openjdk/jmh/samples/).
+
+In contrast to tests, the actual name of the benchmark class is not relevant to JMH. However, stick to the naming convention and
+end the class name of a benchmark with `Benchmark`. To have JMH execute a benchmark, annotate the respective methods with `@Benchmark`.
+
+## Tips and Best Practices
+
+To get realistic results, you should exercise care when running benchmarks. Here are a few tips:
+
+### Do
+
+* Ensure that the system executing your microbenchmarks has as little load as possible. Shutdown every process that can cause unnecessary
+  runtime jitter. Watch the `Error` column in the benchmark results to see the run-to-run variance.
+* Ensure to run enough warmup iterations to get the benchmark into a stable state. If you are unsure, don't change the defaults.
+* Avoid CPU migrations by pinning your benchmarks to specific CPU cores. On Linux you can use `taskset`.
+* Fix the CPU frequency to avoid Turbo Boost from kicking in and skewing your results. On Linux you can use `cpufreq-set` and the
+  `performance` CPU governor.
+* Vary the problem input size with `@Param`.
+* Use the integrated profilers in JMH to dig deeper if benchmark results to not match your hypotheses:
+    * Add `-prof gc` to the options to check whether the garbage collector runs during a microbenchmarks and skews
+      your results. If so, try to force a GC between runs (`-gc true`) but watch out for the caveats.
+    * Add `-prof perf` or `-prof perfasm` (both only available on Linux) to see hotspots.
+* Have your benchmarks peer-reviewed.
+
+### Don't
+
+* Blindly believe the numbers that your microbenchmark produces but verify them by measuring e.g. with `-prof perfasm`.
+* Run more threads than your number of CPU cores (in case you run multi-threaded microbenchmark).
+* Look only at the `Score` column and ignore `Error`. Instead, take countermeasures to keep `Error` low / variance explainable.
+
+## Disassembling
+
+Disassembling is fun! Maybe not always useful, but always fun! Generally, you'll want to install `perf` and FCML's `hsdis`.
+`perf` is generally available via `apg-get install perf` or `pacman -S perf`. FCML is a little more involved. This worked
+on 2020-08-01:
+
+```
+wget https://github.com/swojtasiak/fcml-lib/releases/download/v1.2.2/fcml-1.2.2.tar.gz
+tar xf fcml*
+cd fcml*
+./configure
+make
+cd example/hsdis
+make
+sudo cp .libs/libhsdis.so.0.0.0 /usr/lib/jvm/java-14-adoptopenjdk/lib/hsdis-amd64.so
+```
+
+If you want to disassemble a single method do something like this:
+
+```
+gradlew -p micro-benchmarks run --args ' MemoryStatsBenchmark -jvmArgs "-XX:+UnlockDiagnosticVMOptions -XX:CompileCommand=print,*.yourMethodName -XX:PrintAssemblyOptions=intel"
+```
+
+
+If you want `perf` to find the hot methods for you then do add `-prof:perfasm`.

micro-benchmarks/build.gradle

@@ -0,0 +1,68 @@
+/*
+ *  Copyright OpenSearch Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ */
+import org.opensearch.gradle.info.BuildParams
+
+apply plugin: 'opensearch.build'
+apply plugin: 'application'
+apply plugin: 'java'
+apply plugin: 'io.freefair.lombok'
+
+assemble.enabled = false
+
+application {
+    mainClass = 'org.openjdk.jmh.Main'
+}
+
+test.enabled = false
+
+repositories {
+    mavenLocal()
+    maven { url "https://aws.oss.sonatype.org/content/repositories/snapshots" }
+    mavenCentral()
+    maven { url "https://plugins.gradle.org/m2/" }
+}
+
+dependencies {
+    // This will take root project as the dependency
+    api(project(':'))
+    api "org.openjdk.jmh:jmh-core:$versions.jmh"
+    annotationProcessor "org.openjdk.jmh:jmh-generator-annprocess:$versions.jmh"
+    // Dependencies of JMH
+    runtimeOnly 'net.sf.jopt-simple:jopt-simple:5.0.4'
+    runtimeOnly 'org.apache.commons:commons-math3:3.6.1'
+}
+
+// enable the JMH's BenchmarkProcessor to generate the final benchmark classes
+// needs to be added separately otherwise Gradle will quote it and javac will fail
+compileJava.options.compilerArgs.addAll(["-processor", "org.openjdk.jmh.generators.BenchmarkProcessor"])
+
+
+run {
+    // This is required for C++ code
+    systemProperty "java.library.path", "$rootDir/jni/release"
+    executable = "${BuildParams.runtimeJavaHome}/bin/java"
+    var jvmHeapSize = System.getProperty("jvm.heap.size", "6g")
+    jvmArgs("-Xms" + jvmHeapSize, "-Xmx"  + jvmHeapSize)
+}
+
+
+// No licenses for our benchmark deps (we don't ship benchmarks)
+tasks.named("dependencyLicenses").configure { it.enabled = false }
+dependenciesInfo.enabled = false
+
+thirdPartyAudit.ignoreViolations(
+        // these classes intentionally use JDK internal API (and this is ok since the project is maintained by Oracle employees)
+        'org.openjdk.jmh.util.Utils'
+)
+
+spotless {
+    java {
+        // IDEs can sometimes run annotation processors that leave files in
+        // here, causing Spotless to complain. Even though this path ought not
+        // to exist, exclude it anyway in order to avoid spurious failures.
+        targetExclude 'src/main/generated/**/*.java'
+    }
+}
+

micro-benchmarks/src/main/java/org/opensearch/knn/TransferVectorsBenchmarks.java

@@ -0,0 +1,87 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.knn;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.opensearch.knn.jni.JNIService;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * The class provides runs some benchmarks and provide the performance data around how much time it will take to
+ * transfer vectors from java to jni layer for different configuration.
+ */
+@Warmup(iterations = 1, timeUnit = TimeUnit.SECONDS, time = 300)
+@Measurement(iterations = 1, timeUnit = TimeUnit.SECONDS, time = 300)
+@Fork(3)
+@BenchmarkMode(Mode.SingleShotTime)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@State(Scope.Benchmark)
+public class TransferVectorsBenchmarks {
+    private static final Random random = new Random(1212121212);
+    private static final int TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED = 1000000;
+
+    @Param({ "128", "256", "384", "512" })
+    private int dimension;
+
+    @Param({ "100000", "500000", "1000000" })
+    private int vectorsPerTransfer;
+
+    private List<float[]> vectorList;
+
+    @Setup(Level.Invocation)
+    public void setup() {
+        vectorList = new ArrayList<>();
+        for (int i = 0; i < TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED; i++) {
+            vectorList.add(generateRandomVector(dimension));
+        }
+    }
+
+    @Benchmark
+    public void transferVectors() {
+        long vectorsAddress = 0;
+        List<float[]> vectorToTransfer = new ArrayList<>();
+        for (float[] floats : vectorList) {
+            if (vectorToTransfer.size() == vectorsPerTransfer) {
+                vectorsAddress = JNIService.transferVectorsV2(vectorsAddress, vectorToTransfer.toArray(new float[][] {}));
+                vectorToTransfer = new ArrayList<>();
+            }
+            vectorToTransfer.add(floats);
+        }
+        if (!vectorToTransfer.isEmpty()) {
+            vectorsAddress = JNIService.transferVectorsV2(vectorsAddress, vectorToTransfer.toArray(new float[][] {}));
+        }
+        JNIService.freeVectors(vectorsAddress);
+    }
+
+    private float[] generateRandomVector(int dimensions) {
+        float[] vector = new float[dimensions];
+        for (int i = 0; i < dimensions; i++) {
+            vector[i] = -500 + (float) random.nextGaussian() * (1000);
+        }
+        return vector;
+    }
+}

micro-benchmarks/src/main/resources/log4j2.properties

@@ -0,0 +1,19 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+#
+# Modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+#
+
+appender.console.type = Console
+appender.console.name = console
+appender.console.layout.type = PatternLayout
+appender.console.layout.pattern = [%d{ISO8601}][%-5p][%-25c] %marker %m%n
+
+# Do not log at all if it is not really critical - we're in a benchmark
+rootLogger.level = error
+rootLogger.appenderRef.console.ref = console

settings.gradle

@@ -8,4 +8,5 @@ rootProject.name = 'opensearch-knn'
 include ":qa"
 include ":qa:rolling-upgrade"
 include ":qa:restart-upgrade"
+include ":micro-benchmarks"
 

src/main/java/org/opensearch/knn/index/KNNSettings.java

@@ -5,6 +5,7 @@
 
 package org.opensearch.knn.index;
 
+import lombok.extern.log4j.Log4j2;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.OpenSearchParseException;
@@ -48,6 +49,7 @@
  * 2. KNN settings to enable/disable plugin, circuit breaker settings
  * 3. KNN settings to manage graphs loaded in native memory
  */
+@Log4j2
 public class KNNSettings {
 
     private static final Logger logger = LogManager.getLogger(KNNSettings.class);
@@ -81,6 +83,7 @@ public class KNNSettings {
     /**
      * Default setting values
      */
+    public static final boolean KNN_DEFAULT_FAISS_AVX2_DISABLED_VALUE = false;
     public static final String INDEX_KNN_DEFAULT_SPACE_TYPE = "l2";
     public static final Integer INDEX_KNN_DEFAULT_ALGO_PARAM_M = 16;
     public static final Integer INDEX_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH = 100;
@@ -232,7 +235,11 @@ public class KNNSettings {
         Dynamic
     );
 
-    public static final Setting<Boolean> KNN_FAISS_AVX2_DISABLED_SETTING = Setting.boolSetting(KNN_FAISS_AVX2_DISABLED, false, NodeScope);
+    public static final Setting<Boolean> KNN_FAISS_AVX2_DISABLED_SETTING = Setting.boolSetting(
+        KNN_FAISS_AVX2_DISABLED,
+        KNN_DEFAULT_FAISS_AVX2_DISABLED_VALUE,
+        NodeScope
+    );
 
     /**
      * Dynamic settings
@@ -386,7 +393,19 @@ public static double getCircuitBreakerUnsetPercentage() {
     }
 
     public static boolean isFaissAVX2Disabled() {
-        return KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX2_DISABLED);
+        try {
+            return KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX2_DISABLED);
+        } catch (Exception e) {
+            // In some UTs we identified that cluster setting is not set properly an leads to NPE. This check will avoid
+            // those cases and will still return the default value.
+            log.warn(
+                "Unable to get setting value {} from cluster settings. Using default value as {}",
+                KNN_FAISS_AVX2_DISABLED,
+                KNN_DEFAULT_FAISS_AVX2_DISABLED_VALUE,
+                e
+            );
+            return KNN_DEFAULT_FAISS_AVX2_DISABLED_VALUE;
+        }
     }
 
     public static Integer getFilteredExactSearchThreshold(final String indexName) {