Merge branch 'main' into gh-oblt/add-pull-requests-write-permission-docs-builder

Author: shainaraskas

libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java

@@ -237,20 +237,21 @@ public static void subtract(float[] v1, float[] v2, float[] result) {
     }
 
     /**
-     * calculates the spill-over score for a vector and a centroid, given its residual with
-     * its actually nearest centroid
+     * calculates the soar distance for a vector and a centroid
      * @param v1 the vector
      * @param centroid the centroid
      * @param originalResidual the residual with the actually nearest centroid
-     * @return the spill-over score (soar)
+     * @param soarLambda the lambda parameter
+     * @param rnorm distance to the nearest centroid
+     * @return the soar distance
      */
-    public static float soarResidual(float[] v1, float[] centroid, float[] originalResidual) {
+    public static float soarDistance(float[] v1, float[] centroid, float[] originalResidual, float soarLambda, float rnorm) {
         if (v1.length != centroid.length) {
             throw new IllegalArgumentException("vector dimensions differ: " + v1.length + "!=" + centroid.length);
         }
         if (originalResidual.length != v1.length) {
             throw new IllegalArgumentException("vector dimensions differ: " + originalResidual.length + "!=" + v1.length);
         }
-        return IMPL.soarResidual(v1, centroid, originalResidual);
+        return IMPL.soarDistance(v1, centroid, originalResidual, soarLambda, rnorm);
     }
 }

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java

@@ -11,6 +11,7 @@
 
 import org.apache.lucene.util.BitUtil;
 import org.apache.lucene.util.Constants;
+import org.apache.lucene.util.VectorUtil;
 
 final class DefaultESVectorUtilSupport implements ESVectorUtilSupport {
 
@@ -139,15 +140,16 @@ public void centerAndCalculateOSQStatsDp(float[] target, float[] centroid, float
     }
 
     @Override
-    public float soarResidual(float[] v1, float[] centroid, float[] originalResidual) {
+    public float soarDistance(float[] v1, float[] centroid, float[] originalResidual, float soarLambda, float rnorm) {
         assert v1.length == centroid.length;
         assert v1.length == originalResidual.length;
+        float dsq = VectorUtil.squareDistance(v1, centroid);
         float proj = 0;
         for (int i = 0; i < v1.length; i++) {
             float djk = v1[i] - centroid[i];
             proj = fma(djk, originalResidual[i], proj);
         }
-        return proj;
+        return dsq + soarLambda * proj * proj / rnorm;
     }
 
     public static int ipByteBitImpl(byte[] q, byte[] d) {

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java

@@ -37,6 +37,6 @@ public interface ESVectorUtilSupport {
 
     void centerAndCalculateOSQStatsDp(float[] target, float[] centroid, float[] centered, float[] stats);
 
-    float soarResidual(float[] v1, float[] centroid, float[] originalResidual);
+    float soarDistance(float[] v1, float[] centroid, float[] originalResidual, float soarLambda, float rnorm);
 
 }

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java

@@ -368,14 +368,17 @@ public float calculateOSQLoss(float[] target, float[] interval, float step, floa
     }
 
     @Override
-    public float soarResidual(float[] v1, float[] centroid, float[] originalResidual) {
+    public float soarDistance(float[] v1, float[] centroid, float[] originalResidual, float soarLambda, float rnorm) {
         assert v1.length == centroid.length;
         assert v1.length == originalResidual.length;
         float proj = 0;
+        float dsq = 0;
         int i = 0;
         if (v1.length > 2 * FLOAT_SPECIES.length()) {
             FloatVector projVec1 = FloatVector.zero(FLOAT_SPECIES);
             FloatVector projVec2 = FloatVector.zero(FLOAT_SPECIES);
+            FloatVector acc1 = FloatVector.zero(FLOAT_SPECIES);
+            FloatVector acc2 = FloatVector.zero(FLOAT_SPECIES);
             int unrolledLimit = FLOAT_SPECIES.loopBound(v1.length) - FLOAT_SPECIES.length();
             for (; i < unrolledLimit; i += 2 * FLOAT_SPECIES.length()) {
                 // one
@@ -384,13 +387,15 @@ public float soarResidual(float[] v1, float[] centroid, float[] originalResidual
                 FloatVector originalResidualVec0 = FloatVector.fromArray(FLOAT_SPECIES, originalResidual, i);
                 FloatVector djkVec0 = v1Vec0.sub(centroidVec0);
                 projVec1 = fma(djkVec0, originalResidualVec0, projVec1);
+                acc1 = fma(djkVec0, djkVec0, acc1);
 
                 // two
                 FloatVector v1Vec1 = FloatVector.fromArray(FLOAT_SPECIES, v1, i + FLOAT_SPECIES.length());
                 FloatVector centroidVec1 = FloatVector.fromArray(FLOAT_SPECIES, centroid, i + FLOAT_SPECIES.length());
                 FloatVector originalResidualVec1 = FloatVector.fromArray(FLOAT_SPECIES, originalResidual, i + FLOAT_SPECIES.length());
                 FloatVector djkVec1 = v1Vec1.sub(centroidVec1);
                 projVec2 = fma(djkVec1, originalResidualVec1, projVec2);
+                acc2 = fma(djkVec1, djkVec1, acc2);
             }
             // vector tail
             for (; i < FLOAT_SPECIES.loopBound(v1.length); i += FLOAT_SPECIES.length()) {
@@ -399,15 +404,18 @@ public float soarResidual(float[] v1, float[] centroid, float[] originalResidual
                 FloatVector originalResidualVec = FloatVector.fromArray(FLOAT_SPECIES, originalResidual, i);
                 FloatVector djkVec = v1Vec.sub(centroidVec);
                 projVec1 = fma(djkVec, originalResidualVec, projVec1);
+                acc1 = fma(djkVec, djkVec, acc1);
             }
             proj += projVec1.add(projVec2).reduceLanes(ADD);
+            dsq += acc1.add(acc2).reduceLanes(ADD);
         }
         // tail
         for (; i < v1.length; i++) {
             float djk = v1[i] - centroid[i];
             proj = fma(djk, originalResidual[i], proj);
+            dsq = fma(djk, djk, dsq);
         }
-        return proj;
+        return dsq + soarLambda * proj * proj / rnorm;
     }
 
     private static final VectorSpecies<Byte> BYTE_SPECIES_128 = ByteVector.SPECIES_128;

libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java

@@ -268,7 +268,7 @@ public void testOsqGridPoints() {
         }
     }
 
-    public void testSoarOverspillScore() {
+    public void testSoarDistance() {
         int size = random().nextInt(128, 512);
         float deltaEps = 1e-5f * size;
         var vector = new float[size];
@@ -279,8 +279,10 @@ public void testSoarOverspillScore() {
             centroid[i] = random().nextFloat();
             preResidual[i] = random().nextFloat();
         }
-        var expected = defaultedProvider.getVectorUtilSupport().soarResidual(vector, centroid, preResidual);
-        var result = defOrPanamaProvider.getVectorUtilSupport().soarResidual(vector, centroid, preResidual);
+        float soarLambda = random().nextFloat();
+        float rnorm = random().nextFloat();
+        var expected = defaultedProvider.getVectorUtilSupport().soarDistance(vector, centroid, preResidual, soarLambda, rnorm);
+        var result = defOrPanamaProvider.getVectorUtilSupport().soarDistance(vector, centroid, preResidual, soarLambda, rnorm);
         assertEquals(expected, result, deltaEps);
     }
 

muted-tests.yml

@@ -2,9 +2,6 @@ tests:
 - class: "org.elasticsearch.client.RestClientSingleHostIntegTests"
   issue: "https://github.com/elastic/elasticsearch/issues/102717"
   method: "testRequestResetAndAbort"
-- class: org.elasticsearch.smoketest.WatcherYamlRestIT
-  method: test {p0=watcher/usage/10_basic/Test watcher usage stats output}
-  issue: https://github.com/elastic/elasticsearch/issues/112189
 - class: org.elasticsearch.packaging.test.PackagesSecurityAutoConfigurationTests
   method: test20SecurityNotAutoConfiguredOnReInstallation
   issue: https://github.com/elastic/elasticsearch/issues/112635

server/src/main/java/module-info.java

@@ -475,6 +475,7 @@
             org.elasticsearch.serverless.apifiltering;
     exports org.elasticsearch.lucene.spatial;
     exports org.elasticsearch.inference.configuration;
+    exports org.elasticsearch.inference.validation;
     exports org.elasticsearch.monitor.metrics;
     exports org.elasticsearch.plugins.internal.rewriter to org.elasticsearch.inference;
     exports org.elasticsearch.lucene.util.automaton;

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -200,6 +200,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_REMOVE_ERROR_PARSING_8_19 = def(8_841_0_52);
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_EMBEDDING_BATCH_SIZE_8_19 = def(8_841_0_53);
     public static final TransportVersion STREAMS_LOGS_SUPPORT_8_19 = def(8_841_0_54);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_INPUT_TYPE_8_19 = def(8_841_0_55);
 
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
@@ -308,6 +309,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_REMOVE_ERROR_PARSING = def(9_102_0_00);
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_EMBEDDING_BATCH_SIZE = def(9_103_0_00);
     public static final TransportVersion STREAMS_LOGS_SUPPORT = def(9_104_0_00);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_INPUT_TYPE = def(9_105_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/KMeansLocal.java

@@ -202,7 +202,7 @@ private int[] assignSpilled(FloatVectorValues vectors, List<int[]> neighborhoods
                     continue;
                 }
                 float[] neighborCentroid = centroids[neighbor];
-                float soar = distanceSoar(diffs, vector, neighborCentroid, vectorCentroidDist);
+                float soar = ESVectorUtil.soarDistance(vector, neighborCentroid, diffs, soarLambda, vectorCentroidDist);
                 if (soar < minSoar) {
                     bestAssignment = neighbor;
                     minSoar = soar;
@@ -215,13 +215,6 @@ private int[] assignSpilled(FloatVectorValues vectors, List<int[]> neighborhoods
         return spilledAssignments;
     }
 
-    private float distanceSoar(float[] residual, float[] vector, float[] centroid, float rnorm) {
-        // TODO: combine these to be more efficient
-        float dsq = VectorUtil.squareDistance(vector, centroid);
-        float rproj = ESVectorUtil.soarResidual(vector, centroid, residual);
-        return dsq + soarLambda * rproj * rproj / rnorm;
-    }
-
     /**
      * cluster using a lloyd k-means algorithm that is not neighbor aware
      *

server/src/main/java/org/elasticsearch/inference/InferenceService.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.validation.ServiceIntegrationValidator;
 
 import java.io.Closeable;
 import java.util.EnumSet;
@@ -248,4 +249,14 @@ default void updateModelsWithDynamicFields(List<Model> model, ActionListener<Lis
      * after ensuring the node's internals are set up (for example if this ensures the internal ES client is ready for use).
      */
     default void onNodeStarted() {}
+
+    /**
+     * Get the service integration validator for the given task type.
+     * This allows services to provide custom validation logic.
+     * @param taskType The task type
+     * @return The service integration validator or null if the default should be used
+     */
+    default ServiceIntegrationValidator getServiceIntegrationValidator(TaskType taskType) {
+        return null;
+    }
 }