Skip to content

Commit 36ff4cd

Browse files
Jan-Kazlouski-elasticJan-Kazlouski-elastic
committed
Merge remote-tracking branch 'origin/main' into feature/llama-embeding-completion
# Conflicts: # server/src/main/java/org/elasticsearch/TransportVersions.java
2 parents 9573a48 + e2bc52e commit 36ff4cd

File tree

481 files changed

+12027
-5931
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

481 files changed

+12027
-5931
lines changed

‎benchmarks/src/main/java/org/elasticsearch/benchmark/vector/VectorScorerBenchmark.java‎ renamed to ‎benchmarks/src/main/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmark.java‎

Lines changed: 5 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -55,24 +55,23 @@
5555
/**
5656
* Benchmark that compares various scalar quantized vector similarity function
5757
* implementations;: scalar, lucene's panama-ized, and Elasticsearch's native.
58-
* Run with ./gradlew -p benchmarks run --args 'VectorScorerBenchmark'
58+
* Run with ./gradlew -p benchmarks run --args 'Int7uScorerBenchmark'
5959
*/
60-
public class VectorScorerBenchmark {
60+
public class Int7uScorerBenchmark {
6161

6262
static {
6363
LogConfigurator.configureESLogging(); // native access requires logging to be initialized
6464
}
6565

6666
@Param({ "96", "768", "1024" })
67-
int dims;
68-
int size = 2; // there are only two vectors to compare
67+
public int dims;
68+
final int size = 2; // there are only two vectors to compare
6969

7070
Directory dir;
7171
IndexInput in;
7272
VectorScorerFactory factory;
7373

74-
byte[] vec1;
75-
byte[] vec2;
74+
byte[] vec1, vec2;
7675
float vec1Offset;
7776
float vec2Offset;
7877
float scoreCorrectionConstant;
@@ -139,39 +138,6 @@ public void setup() throws IOException {
139138
nativeDotScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.DOT_PRODUCT, values, queryVec).get();
140139
luceneSqrScorerQuery = luceneScorer(values, VectorSimilarityFunction.EUCLIDEAN, queryVec);
141140
nativeSqrScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.EUCLIDEAN, values, queryVec).get();
142-
143-
// sanity
144-
var f1 = dotProductLucene();
145-
var f2 = dotProductNative();
146-
var f3 = dotProductScalar();
147-
if (f1 != f2) {
148-
throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
149-
}
150-
if (f1 != f3) {
151-
throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
152-
}
153-
// square distance
154-
f1 = squareDistanceLucene();
155-
f2 = squareDistanceNative();
156-
f3 = squareDistanceScalar();
157-
if (f1 != f2) {
158-
throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
159-
}
160-
if (f1 != f3) {
161-
throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
162-
}
163-
164-
var q1 = dotProductLuceneQuery();
165-
var q2 = dotProductNativeQuery();
166-
if (q1 != q2) {
167-
throw new AssertionError("query: lucene[" + q1 + "] != " + "native[" + q2 + "]");
168-
}
169-
170-
var sqr1 = squareDistanceLuceneQuery();
171-
var sqr2 = squareDistanceNativeQuery();
172-
if (sqr1 != sqr2) {
173-
throw new AssertionError("query: lucene[" + q1 + "] != " + "native[" + q2 + "]");
174-
}
175141
}
176142

177143
@TearDown

‎benchmarks/src/main/java/org/elasticsearch/benchmark/vector/JDKVectorInt7uBenchmark.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public class JDKVectorInt7uBenchmark {
5252

5353
Arena arena;
5454

55-
@Param({ "1", "128", "207", "256", "300", "512", "702", "1024" })
55+
@Param({ "1", "128", "207", "256", "300", "512", "702", "1024", "1536", "2048" })
5656
public int size;
5757

5858
@Setup(Level.Iteration)

‎benchmarks/src/main/java/org/elasticsearch/benchmark/vector/OptimizedScalarQuantizerBenchmark.java‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ public class OptimizedScalarQuantizerBenchmark {
4343

4444
float[] vector;
4545
float[] centroid;
46-
byte[] destination;
46+
byte[] legacyDestination;
47+
int[] destination;
4748

4849
@Param({ "1", "4", "7" })
4950
byte bits;
@@ -54,7 +55,8 @@ public class OptimizedScalarQuantizerBenchmark {
5455
public void init() {
5556
ThreadLocalRandom random = ThreadLocalRandom.current();
5657
// random byte arrays for binary methods
57-
destination = new byte[dims];
58+
legacyDestination = new byte[dims];
59+
destination = new int[dims];
5860
vector = new float[dims];
5961
centroid = new float[dims];
6062
for (int i = 0; i < dims; ++i) {
@@ -65,13 +67,20 @@ public void init() {
6567

6668
@Benchmark
6769
public byte[] scalar() {
68-
osq.scalarQuantize(vector, destination, bits, centroid);
69-
return destination;
70+
osq.legacyScalarQuantize(vector, legacyDestination, bits, centroid);
71+
return legacyDestination;
72+
}
73+
74+
@Benchmark
75+
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
76+
public byte[] legacyVector() {
77+
osq.legacyScalarQuantize(vector, legacyDestination, bits, centroid);
78+
return legacyDestination;
7079
}
7180

7281
@Benchmark
7382
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
74-
public byte[] vector() {
83+
public int[] vector() {
7584
osq.scalarQuantize(vector, destination, bits, centroid);
7685
return destination;
7786
}

‎benchmarks/src/test/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmarkTests.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.benchmark.vector;
11+
12+
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
13+
14+
import org.apache.lucene.util.Constants;
15+
import org.elasticsearch.test.ESTestCase;
16+
import org.junit.BeforeClass;
17+
import org.openjdk.jmh.annotations.Param;
18+
19+
import java.util.Arrays;
20+
21+
public class Int7uScorerBenchmarkTests extends ESTestCase {
22+
23+
final double delta = 1e-3;
24+
final int dims;
25+
26+
public Int7uScorerBenchmarkTests(int dims) {
27+
this.dims = dims;
28+
}
29+
30+
@BeforeClass
31+
public static void skipWindows() {
32+
assumeFalse("doesn't work on windows yet", Constants.WINDOWS);
33+
}
34+
35+
public void testDotProduct() throws Exception {
36+
for (int i = 0; i < 100; i++) {
37+
var bench = new Int7uScorerBenchmark();
38+
bench.dims = dims;
39+
bench.setup();
40+
try {
41+
float expected = bench.dotProductScalar();
42+
assertEquals(expected, bench.dotProductLucene(), delta);
43+
assertEquals(expected, bench.dotProductNative(), delta);
44+
45+
expected = bench.dotProductLuceneQuery();
46+
assertEquals(expected, bench.dotProductNativeQuery(), delta);
47+
} finally {
48+
bench.teardown();
49+
}
50+
}
51+
}
52+
53+
public void testSquareDistance() throws Exception {
54+
for (int i = 0; i < 100; i++) {
55+
var bench = new Int7uScorerBenchmark();
56+
bench.dims = dims;
57+
bench.setup();
58+
try {
59+
float expected = bench.squareDistanceScalar();
60+
assertEquals(expected, bench.squareDistanceLucene(), delta);
61+
assertEquals(expected, bench.squareDistanceNative(), delta);
62+
63+
expected = bench.squareDistanceLuceneQuery();
64+
assertEquals(expected, bench.squareDistanceNativeQuery(), delta);
65+
} finally {
66+
bench.teardown();
67+
}
68+
}
69+
}
70+
71+
@ParametersFactory
72+
public static Iterable<Object[]> parametersFactory() {
73+
try {
74+
var params = Int7uScorerBenchmark.class.getField("dims").getAnnotationsByType(Param.class)[0].value();
75+
return () -> Arrays.stream(params).map(Integer::parseInt).map(i -> new Object[] { i }).iterator();
76+
} catch (NoSuchFieldException e) {
77+
throw new AssertionError(e);
78+
}
79+
}
80+
}

‎build-conventions/src/main/java/org/elasticsearch/gradle/internal/conventions/PublishPlugin.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,8 @@ private void addNameAndDescriptionToPom(Project project, NamedDomainObjectSet<Ma
175175
private static void configureWithShadowPlugin(Project project, MavenPublication publication) {
176176
var shadow = project.getExtensions().getByType(ShadowExtension.class);
177177
shadow.component(publication);
178+
publication.artifact(project.getTasks().named("javadocJar"));
179+
publication.artifact(project.getTasks().named("sourcesJar"));
178180
}
179181

180182
private static void addScmInfo(XmlProvider xml, GitInfo gitInfo) {

‎build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/ElasticsearchTestBasePlugin.java‎

Lines changed: 104 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,11 @@
3434
import java.io.File;
3535
import java.util.List;
3636
import java.util.Map;
37+
import java.util.stream.Stream;
3738

3839
import javax.inject.Inject;
3940

41+
import static java.util.stream.Collectors.joining;
4042
import static org.elasticsearch.gradle.internal.util.ParamsUtils.loadBuildParams;
4143
import static org.elasticsearch.gradle.util.FileUtils.mkdirs;
4244
import static org.elasticsearch.gradle.util.GradleUtils.maybeConfigure;
@@ -173,6 +175,16 @@ public void execute(Task t) {
173175
// we use 'temp' relative to CWD since this is per JVM and tests are forbidden from writing to CWD
174176
nonInputProperties.systemProperty("java.io.tmpdir", test.getWorkingDir().toPath().resolve("temp"));
175177

178+
SourceSetContainer sourceSets = project.getExtensions().getByType(SourceSetContainer.class);
179+
SourceSet mainSourceSet = sourceSets.findByName(SourceSet.MAIN_SOURCE_SET_NAME);
180+
SourceSet testSourceSet = sourceSets.findByName(SourceSet.TEST_SOURCE_SET_NAME);
181+
if ("test".equals(test.getName()) && mainSourceSet != null && testSourceSet != null) {
182+
FileCollection mainRuntime = mainSourceSet.getRuntimeClasspath();
183+
FileCollection testRuntime = testSourceSet.getRuntimeClasspath();
184+
FileCollection testOnlyFiles = testRuntime.minus(mainRuntime);
185+
test.doFirst(task -> test.environment("es.entitlement.testOnlyPath", testOnlyFiles.getAsPath()));
186+
}
187+
176188
test.systemProperties(getProviderFactory().systemPropertiesPrefixedBy("tests.").get());
177189
test.systemProperties(getProviderFactory().systemPropertiesPrefixedBy("es.").get());
178190

@@ -205,46 +217,122 @@ public void execute(Task t) {
205217
}
206218

207219
/*
208-
* If this project builds a shadow JAR than any unit tests should test against that artifact instead of
220+
* If this project builds a shadow JAR then any unit tests should test against that artifact instead of
209221
* compiled class output and dependency jars. This better emulates the runtime environment of consumers.
210222
*/
211223
project.getPluginManager().withPlugin("com.gradleup.shadow", p -> {
212224
if (test.getName().equals(JavaPlugin.TEST_TASK_NAME)) {
213225
// Remove output class files and any other dependencies from the test classpath, since the shadow JAR includes these
214-
SourceSetContainer sourceSets = project.getExtensions().getByType(SourceSetContainer.class);
215-
FileCollection mainRuntime = sourceSets.getByName(SourceSet.MAIN_SOURCE_SET_NAME).getRuntimeClasspath();
216226
// Add any "shadow" dependencies. These are dependencies that are *not* bundled into the shadow JAR
217227
Configuration shadowConfig = project.getConfigurations().getByName(ShadowBasePlugin.CONFIGURATION_NAME);
218228
// Add the shadow JAR artifact itself
219229
FileCollection shadowJar = project.files(project.getTasks().named("shadowJar"));
220-
FileCollection testRuntime = sourceSets.getByName(SourceSet.TEST_SOURCE_SET_NAME).getRuntimeClasspath();
230+
FileCollection mainRuntime = mainSourceSet.getRuntimeClasspath();
231+
FileCollection testRuntime = testSourceSet.getRuntimeClasspath();
221232
test.setClasspath(testRuntime.minus(mainRuntime).plus(shadowConfig).plus(shadowJar));
222233
}
223234
});
224235
});
225-
configureImmutableCollectionsPatch(project);
236+
configureJavaBaseModuleOptions(project);
237+
configureEntitlements(project);
238+
}
239+
240+
/**
241+
* Computes and sets the {@code --patch-module=java.base} and {@code --add-opens=java.base} JVM command line options.
242+
*/
243+
private void configureJavaBaseModuleOptions(Project project) {
244+
project.getTasks().withType(Test.class).matching(task -> task.getName().equals("test")).configureEach(test -> {
245+
FileCollection patchedImmutableCollections = patchedImmutableCollections(project);
246+
if (patchedImmutableCollections != null) {
247+
test.getInputs().files(patchedImmutableCollections);
248+
test.systemProperty("tests.hackImmutableCollections", "true");
249+
}
250+
251+
FileCollection entitlementBridge = entitlementBridge(project);
252+
if (entitlementBridge != null) {
253+
test.getInputs().files(entitlementBridge);
254+
}
255+
256+
test.getJvmArgumentProviders().add(() -> {
257+
String javaBasePatch = Stream.concat(
258+
singleFilePath(patchedImmutableCollections).map(str -> str + "/java.base"),
259+
singleFilePath(entitlementBridge)
260+
).collect(joining(File.pathSeparator));
261+
262+
return javaBasePatch.isEmpty()
263+
? List.of()
264+
: List.of("--patch-module=java.base=" + javaBasePatch, "--add-opens=java.base/java.util=ALL-UNNAMED");
265+
});
266+
});
226267
}
227268

228-
private void configureImmutableCollectionsPatch(Project project) {
269+
private Stream<String> singleFilePath(FileCollection collection) {
270+
return Stream.ofNullable(collection).filter(fc -> fc.isEmpty() == false).map(FileCollection::getSingleFile).map(File::toString);
271+
}
272+
273+
private static FileCollection patchedImmutableCollections(Project project) {
229274
String patchProject = ":test:immutable-collections-patch";
230275
if (project.findProject(patchProject) == null) {
231-
return; // build tests may not have this project, just skip
276+
return null; // build tests may not have this project, just skip
232277
}
233278
String configurationName = "immutableCollectionsPatch";
234279
FileCollection patchedFileCollection = project.getConfigurations()
235280
.create(configurationName, config -> config.setCanBeConsumed(false));
236281
var deps = project.getDependencies();
237282
deps.add(configurationName, deps.project(Map.of("path", patchProject, "configuration", "patch")));
238-
project.getTasks().withType(Test.class).matching(task -> task.getName().equals("test")).configureEach(test -> {
239-
test.getInputs().files(patchedFileCollection);
240-
test.systemProperty("tests.hackImmutableCollections", "true");
241-
test.getJvmArgumentProviders()
242-
.add(
243-
() -> List.of(
244-
"--patch-module=java.base=" + patchedFileCollection.getSingleFile() + "/java.base",
245-
"--add-opens=java.base/java.util=ALL-UNNAMED"
246-
)
283+
return patchedFileCollection;
284+
}
285+
286+
private static FileCollection entitlementBridge(Project project) {
287+
return project.getConfigurations().findByName("entitlementBridge");
288+
}
289+
290+
/**
291+
* Sets the required JVM options and system properties to enable entitlement enforcement on tests.
292+
* <p>
293+
* One command line option is set in {@link #configureJavaBaseModuleOptions} out of necessity,
294+
* since the command line can have only one {@code --patch-module} option for a given module.
295+
*/
296+
private static void configureEntitlements(Project project) {
297+
Configuration agentConfig = project.getConfigurations().create("entitlementAgent");
298+
Project agent = project.findProject(":libs:entitlement:agent");
299+
if (agent != null) {
300+
agentConfig.defaultDependencies(
301+
deps -> { deps.add(project.getDependencies().project(Map.of("path", ":libs:entitlement:agent"))); }
302+
);
303+
}
304+
FileCollection agentFiles = agentConfig;
305+
306+
Configuration bridgeConfig = project.getConfigurations().create("entitlementBridge");
307+
Project bridge = project.findProject(":libs:entitlement:bridge");
308+
if (bridge != null) {
309+
bridgeConfig.defaultDependencies(
310+
deps -> { deps.add(project.getDependencies().project(Map.of("path", ":libs:entitlement:bridge"))); }
311+
);
312+
}
313+
FileCollection bridgeFiles = bridgeConfig;
314+
315+
project.getTasks().withType(Test.class).configureEach(test -> {
316+
// See also SystemJvmOptions.maybeAttachEntitlementAgent.
317+
318+
// Agent
319+
if (agentFiles.isEmpty() == false) {
320+
test.getInputs().files(agentFiles);
321+
test.systemProperty("es.entitlement.agentJar", agentFiles.getAsPath());
322+
test.systemProperty("jdk.attach.allowAttachSelf", true);
323+
}
324+
325+
// Bridge
326+
if (bridgeFiles.isEmpty() == false) {
327+
String modulesContainingEntitlementInstrumentation = "java.logging,java.net.http,java.naming,jdk.net";
328+
test.getInputs().files(bridgeFiles);
329+
// Tests may not be modular, but the JDK still is
330+
test.jvmArgs(
331+
"--add-exports=java.base/org.elasticsearch.entitlement.bridge=ALL-UNNAMED,"
332+
+ modulesContainingEntitlementInstrumentation
247333
);
334+
}
248335
});
249336
}
337+
250338
}

‎build-tools-internal/src/test/java/org/elasticsearch/gradle/internal/release/ReleaseNotesGeneratorTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ private ChangelogEntry makeHighlightsEntry(int pr, boolean notable) {
188188
}
189189

190190
private String getResource(String name) throws Exception {
191-
return Files.readString(Paths.get(Objects.requireNonNull(this.getClass().getResource(name)).toURI()), StandardCharsets.UTF_8);
191+
return Files.readString(Paths.get(Objects.requireNonNull(this.getClass().getResource(name)).toURI()), StandardCharsets.UTF_8).replace("\r", "");
192192
}
193193

194194
private void writeResource(String name, String contents) throws Exception {

0 commit comments

Comments
 (0)