From fa4e950852ed8b37570806c225e87cb7bcd1fe7a Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Tue, 12 Nov 2024 07:42:23 +0100 Subject: [PATCH 01/16] Deduplicate non-empty InternalAggregation metadata when deserializing (#116589) --- .../search/aggregations/InternalAggregation.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregation.java b/server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregation.java index 38cab1761d409..b829afb0c23b0 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregation.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregation.java @@ -35,7 +35,6 @@ */ public abstract class InternalAggregation implements Aggregation, NamedWriteable { protected final String name; - protected final Map metadata; /** @@ -53,12 +52,14 @@ protected InternalAggregation(String name, Map metadata) { */ protected InternalAggregation(StreamInput in) throws IOException { final String name = in.readString(); + final Map metadata = in.readGenericMap(); if (in instanceof DelayableWriteable.Deduplicator d) { this.name = d.deduplicate(name); + this.metadata = metadata == null || metadata.isEmpty() ? metadata : d.deduplicate(metadata); } else { this.name = name; + this.metadata = metadata; } - metadata = in.readGenericMap(); } @Override From bfb30d2e72f9980a1f9d917ad6f1e3acf4bbff00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20Fred=C3=A9n?= <109296772+jfreden@users.noreply.github.com> Date: Tue, 12 Nov 2024 08:42:34 +0100 Subject: [PATCH 02/16] [DOCS] Remove tech preview from bulk create/update/delete roles (#116601) Mark bulk create/update/delete roles GA in 9.0 and 8.17 --- docs/reference/rest-api/security/bulk-create-roles.asciidoc | 1 - docs/reference/rest-api/security/bulk-delete-roles.asciidoc | 1 - 2 files changed, 2 deletions(-) diff --git a/docs/reference/rest-api/security/bulk-create-roles.asciidoc b/docs/reference/rest-api/security/bulk-create-roles.asciidoc index a198f49383907..560e8b74cdd2c 100644 --- a/docs/reference/rest-api/security/bulk-create-roles.asciidoc +++ b/docs/reference/rest-api/security/bulk-create-roles.asciidoc @@ -1,7 +1,6 @@ [role="xpack"] [[security-api-bulk-put-role]] === Bulk create or update roles API -preview::[] ++++ Bulk create or update roles API ++++ diff --git a/docs/reference/rest-api/security/bulk-delete-roles.asciidoc b/docs/reference/rest-api/security/bulk-delete-roles.asciidoc index a782b5e37fcb9..b9978c89bef3a 100644 --- a/docs/reference/rest-api/security/bulk-delete-roles.asciidoc +++ b/docs/reference/rest-api/security/bulk-delete-roles.asciidoc @@ -1,7 +1,6 @@ [role="xpack"] [[security-api-bulk-delete-role]] === Bulk delete roles API -preview::[] ++++ Bulk delete roles API ++++ From f121e09fbbfe5adee3198620dfd3840e2c792297 Mon Sep 17 00:00:00 2001 From: Liam Thompson <32779855+leemthompo@users.noreply.github.com> Date: Tue, 12 Nov 2024 10:59:20 +0100 Subject: [PATCH 03/16] [DOCS] Connectors 8.16.0 release notes (#115856) --- .../docs/connectors-release-notes.asciidoc | 10 +++- .../connectors-release-notes-8.16.0.asciidoc | 53 +++++++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) create mode 100644 docs/reference/connector/docs/release-notes/connectors-release-notes-8.16.0.asciidoc diff --git a/docs/reference/connector/docs/connectors-release-notes.asciidoc b/docs/reference/connector/docs/connectors-release-notes.asciidoc index 723671b049bf2..e1ed082365c00 100644 --- a/docs/reference/connector/docs/connectors-release-notes.asciidoc +++ b/docs/reference/connector/docs/connectors-release-notes.asciidoc @@ -4,7 +4,13 @@ Release notes ++++ -[INFO] +[NOTE] ==== -Prior to version 8.16.0, the connector release notes were published as part of the https://www.elastic.co/guide/en/enterprise-search/current/changelog.html[Enterprise Search documentation]. +Prior to version *8.16.0*, the connector release notes were published as part of the {enterprise-search-ref}/changelog.html[Enterprise Search documentation]. ==== + +*Release notes*: + +* <> + +include::release-notes/connectors-release-notes-8.16.0.asciidoc[] diff --git a/docs/reference/connector/docs/release-notes/connectors-release-notes-8.16.0.asciidoc b/docs/reference/connector/docs/release-notes/connectors-release-notes-8.16.0.asciidoc new file mode 100644 index 0000000000000..7608336073176 --- /dev/null +++ b/docs/reference/connector/docs/release-notes/connectors-release-notes-8.16.0.asciidoc @@ -0,0 +1,53 @@ +[[es-connectors-release-notes-8-16-0]] +=== 8.16.0 connectors release notes + +[discrete] +[[es-connectors-release-notes-deprecation-notice]] +==== Deprecation notices + +* *Direct index access for connectors and sync jobs* ++ +IMPORTANT: Directly accessing connector and sync job state through `.elastic-connectors*` indices is deprecated, and will be disallowed entirely in a future release. + +* Instead, the Elasticsearch Connector APIs should be used. Connectors framework code now uses the <> by default. +See https://github.com/elastic/connectors/pull/2884[*PR 2902*]. + +* *Docker `enterprise-search` namespace deprecation* ++ +IMPORTANT: The `enterprise-search` Docker namespace is deprecated and will be discontinued in a future release. ++ +Starting in `8.16.0`, Docker images are being transitioned to the new `integrations` namespace, which will become the sole location for future releases. This affects the https://github.com/elastic/connectors[Elastic Connectors] and https://github.com/elastic/data-extraction-service[Elastic Data Extraction Service]. ++ +During this transition period, images are published to both namespaces: ++ +** *Example*: ++ +Deprecated namespace:: +`docker.elastic.co/enterprise-search/elastic-connectors:v8.16.0` ++ +New namespace:: +`docker.elastic.co/integrations/elastic-connectors:v8.16.0` ++ +Users should migrate to the new `integrations` namespace as soon as possible to ensure continued access to future releases. + +[discrete] +[[es-connectors-release-notes-8-16-0-enhancements]] +==== Enhancements + +* Docker images now use Chainguard's Wolfi base image (`docker.elastic.co/wolfi/jdk:openjdk-11-dev`), replacing the previous `ubuntu:focal` base. + +* The Sharepoint Online connector now works with the `Sites.Selected` permission instead of the broader permission `Sites.Read.All`. +See https://github.com/elastic/connectors/pull/2762[*PR 2762*]. + +* Starting in 8.16.0, connectors will start using proper SEMVER, with `MAJOR.MINOR.PATCH`, which aligns with Elasticsearch/Kibana versions. This drops the previous `.BUILD` suffix, which we used to release connectors between Elastic stack releases. Going forward, these inter-stack-release releases will be suffixed instead with `+`, aligning with Elastic Agent and conforming to SEMVER. +See https://github.com/elastic/connectors/pull/2749[*PR 2749*]. + +* Connector logs now use UTC timestamps, instead of machine-local timestamps. This only impacts logging output. +See https://github.com/elastic/connectors/pull/2695[*PR 2695*]. + +[discrete] +[[es-connectors-release-notes-8-16-0-bug-fixes]] +==== Bug fixes + +* The Dropbox connector now fetches the files from team shared folders. +See https://github.com/elastic/connectors/pull/2718[*PR 2718*]. \ No newline at end of file From d34c5630cae240dafb2134441cf132d4280e1ce7 Mon Sep 17 00:00:00 2001 From: David Kyle Date: Tue, 12 Nov 2024 10:43:19 +0000 Subject: [PATCH 04/16] [ML] Avoid the .ml-stats index in post test cleanup (#116476) Fixes ml yaml rest tests failing in the post clean up with a search_phase_execution_exception against the .ml-stats index. The fix is to use another method to find reference ingest pipelines avoid the call to _ml/trained_models/_stats --- muted-tests.yml | 51 ------------------- .../integration/MlRestTestStateCleaner.java | 44 ++++++++-------- 2 files changed, 21 insertions(+), 74 deletions(-) diff --git a/muted-tests.yml b/muted-tests.yml index 3273b203b0982..ddd806d49ae5f 100644 --- a/muted-tests.yml +++ b/muted-tests.yml @@ -17,9 +17,6 @@ tests: - 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.xpack.test.rest.XPackRestIT - method: test {p0=ml/inference_processor/Test create processor with missing mandatory fields} - issue: https://github.com/elastic/elasticsearch/issues/112191 - class: org.elasticsearch.xpack.esql.action.ManyShardsIT method: testRejection issue: https://github.com/elastic/elasticsearch/issues/112406 @@ -142,9 +139,6 @@ tests: - class: org.elasticsearch.search.SearchServiceTests method: testParseSourceValidation issue: https://github.com/elastic/elasticsearch/issues/115936 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/inference_crud/Test delete given model referenced by pipeline} - issue: https://github.com/elastic/elasticsearch/issues/115970 - class: org.elasticsearch.index.reindex.ReindexNodeShutdownIT method: testReindexWithShutdown issue: https://github.com/elastic/elasticsearch/issues/115996 @@ -168,48 +162,27 @@ tests: - class: org.elasticsearch.xpack.ml.integration.DatafeedJobsRestIT method: testLookbackWithIndicesOptions issue: https://github.com/elastic/elasticsearch/issues/116127 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/inference_crud/Test delete given model with alias referenced by pipeline} - issue: https://github.com/elastic/elasticsearch/issues/116133 - class: org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT method: test {categorize.Categorize SYNC} issue: https://github.com/elastic/elasticsearch/issues/113054 - class: org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT method: test {categorize.Categorize ASYNC} issue: https://github.com/elastic/elasticsearch/issues/113055 -- class: org.elasticsearch.xpack.inference.InferenceRestIT - method: test {p0=inference/40_semantic_text_query/Query a field that uses the default ELSER 2 endpoint} - issue: https://github.com/elastic/elasticsearch/issues/114376 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/inference_crud/Test force delete given model with alias referenced by pipeline} - issue: https://github.com/elastic/elasticsearch/issues/116136 - class: org.elasticsearch.xpack.test.rest.XPackRestIT method: test {p0=transform/transforms_start_stop/Test start already started transform} issue: https://github.com/elastic/elasticsearch/issues/98802 - class: org.elasticsearch.action.search.SearchPhaseControllerTests method: testProgressListener issue: https://github.com/elastic/elasticsearch/issues/116149 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/forecast/Test forecast unknown job} - issue: https://github.com/elastic/elasticsearch/issues/116150 - class: org.elasticsearch.xpack.test.rest.XPackRestIT method: test {p0=terms_enum/10_basic/Test security} issue: https://github.com/elastic/elasticsearch/issues/116178 - class: org.elasticsearch.search.basic.SearchWithRandomDisconnectsIT method: testSearchWithRandomDisconnects issue: https://github.com/elastic/elasticsearch/issues/116175 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/start_stop_datafeed/Test start datafeed given index pattern with no matching indices} - issue: https://github.com/elastic/elasticsearch/issues/116220 - class: org.elasticsearch.search.basic.SearchWhileRelocatingIT method: testSearchAndRelocateConcurrentlyRandomReplicas issue: https://github.com/elastic/elasticsearch/issues/116145 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/filter_crud/Test update filter} - issue: https://github.com/elastic/elasticsearch/issues/116271 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/get_datafeeds/Test explicit get all datafeeds} - issue: https://github.com/elastic/elasticsearch/issues/116284 - class: org.elasticsearch.xpack.deprecation.DeprecationHttpIT method: testDeprecatedSettingsReturnWarnings issue: https://github.com/elastic/elasticsearch/issues/108628 @@ -231,24 +204,9 @@ tests: - class: org.elasticsearch.threadpool.SimpleThreadPoolIT method: testThreadPoolMetrics issue: https://github.com/elastic/elasticsearch/issues/108320 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/jobs_crud/Test put job deprecated bucket span} - issue: https://github.com/elastic/elasticsearch/issues/116419 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/explain_data_frame_analytics/Test both job id and body} - issue: https://github.com/elastic/elasticsearch/issues/116433 -- class: org.elasticsearch.smoketest.MlWithSecurityIT - method: test {yaml=ml/inference_crud/Test force delete given model with alias referenced by pipeline} - issue: https://github.com/elastic/elasticsearch/issues/116443 - class: org.elasticsearch.xpack.downsample.ILMDownsampleDisruptionIT method: testILMDownsampleRollingRestart issue: https://github.com/elastic/elasticsearch/issues/114233 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/data_frame_analytics_crud/Test put config with unknown field in outlier detection analysis} - issue: https://github.com/elastic/elasticsearch/issues/116458 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/evaluate_data_frame/Test outlier_detection with query} - issue: https://github.com/elastic/elasticsearch/issues/116484 - class: org.elasticsearch.xpack.kql.query.KqlQueryBuilderTests issue: https://github.com/elastic/elasticsearch/issues/116487 - class: org.elasticsearch.reservedstate.service.FileSettingsServiceTests @@ -263,12 +221,6 @@ tests: - class: org.elasticsearch.xpack.logsdb.qa.StandardVersusLogsIndexModeRandomDataDynamicMappingChallengeRestIT method: testMatchAllQuery issue: https://github.com/elastic/elasticsearch/issues/116536 -- class: org.elasticsearch.xpack.test.rest.XPackRestIT - method: test {p0=ml/inference_crud/Test force delete given model referenced by pipeline} - issue: https://github.com/elastic/elasticsearch/issues/116555 -- class: org.elasticsearch.smoketest.MlWithSecurityIT - method: test {yaml=ml/data_frame_analytics_crud/Test delete given stopped config} - issue: https://github.com/elastic/elasticsearch/issues/116608 - class: org.elasticsearch.xpack.esql.ccq.MultiClusterSpecIT method: test {categorize.Categorize} issue: https://github.com/elastic/elasticsearch/issues/116434 @@ -284,9 +236,6 @@ tests: - class: org.elasticsearch.packaging.test.BootstrapCheckTests method: test20RunWithBootstrapChecks issue: https://github.com/elastic/elasticsearch/issues/116620 -- class: org.elasticsearch.smoketest.MlWithSecurityIT - method: test {yaml=ml/inference_crud/Test force delete given model referenced by pipeline} - issue: https://github.com/elastic/elasticsearch/issues/116624 - class: org.elasticsearch.packaging.test.DockerTests method: test011SecurityEnabledStatus issue: https://github.com/elastic/elasticsearch/issues/116628 diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java index 6f6224d505327..25d9509ecdc7a 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java @@ -10,14 +10,15 @@ import org.elasticsearch.client.Request; import org.elasticsearch.client.Response; import org.elasticsearch.client.RestClient; -import org.elasticsearch.common.xcontent.support.XContentMapValues; import org.elasticsearch.test.rest.ESRestTestCase; import java.io.IOException; +import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Set; -import java.util.stream.Collectors; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasSize; public class MlRestTestStateCleaner { @@ -30,24 +31,29 @@ public MlRestTestStateCleaner(Logger logger, RestClient adminClient) { } public void resetFeatures() throws IOException { - waitForMlStatsIndexToInitialize(); - deleteAllTrainedModelIngestPipelines(); + deletePipelinesWithInferenceProcessors(); // This resets all features, not just ML, but they should have been getting reset between tests anyway so it shouldn't matter adminClient.performRequest(new Request("POST", "/_features/_reset")); } @SuppressWarnings("unchecked") - private void deleteAllTrainedModelIngestPipelines() throws IOException { - final Request getAllTrainedModelStats = new Request("GET", "/_ml/trained_models/_stats"); - getAllTrainedModelStats.addParameter("size", "10000"); - final Response trainedModelsStatsResponse = adminClient.performRequest(getAllTrainedModelStats); + private void deletePipelinesWithInferenceProcessors() throws IOException { + final Response pipelinesResponse = adminClient.performRequest(new Request("GET", "/_ingest/pipeline")); + final Map pipelines = ESRestTestCase.entityAsMap(pipelinesResponse); + + var pipelinesWithInferenceProcessors = new HashSet(); + for (var entry : pipelines.entrySet()) { + var pipelineDef = (Map) entry.getValue(); // each top level object is a separate pipeline + var processors = (List>) pipelineDef.get("processors"); + for (var processor : processors) { + assertThat(processor.entrySet(), hasSize(1)); + if ("inference".equals(processor.keySet().iterator().next())) { + pipelinesWithInferenceProcessors.add(entry.getKey()); + } + } + } - final List> pipelines = (List>) XContentMapValues.extractValue( - "trained_model_stats.ingest.pipelines", - ESRestTestCase.entityAsMap(trainedModelsStatsResponse) - ); - Set pipelineIds = pipelines.stream().flatMap(m -> m.keySet().stream()).collect(Collectors.toSet()); - for (String pipelineId : pipelineIds) { + for (String pipelineId : pipelinesWithInferenceProcessors) { try { adminClient.performRequest(new Request("DELETE", "/_ingest/pipeline/" + pipelineId)); } catch (Exception ex) { @@ -55,12 +61,4 @@ private void deleteAllTrainedModelIngestPipelines() throws IOException { } } } - - private void waitForMlStatsIndexToInitialize() throws IOException { - ESRestTestCase.ensureHealth(adminClient, ".ml-stats-*", (request) -> { - request.addParameter("wait_for_no_initializing_shards", "true"); - request.addParameter("level", "shards"); - request.addParameter("timeout", "30s"); - }); - } } From bcf1bd4c969eebfd1a55dbcd078060ba1522de94 Mon Sep 17 00:00:00 2001 From: Iraklis Psaroudakis Date: Tue, 12 Nov 2024 13:08:04 +0200 Subject: [PATCH 05/16] Ensure Fleet REST yaml tests work in serverless (#115869) Use new gradle plugin for the yaml tests so they can be executed in serverless CI as well. Relates ES-8275 --- .../action/search/SearchRequestBuilder.java | 8 +++++ x-pack/plugin/fleet/qa/rest/build.gradle | 35 +++++++++++++------ .../xpack/fleet/FleetRestIT.java | 24 +++++++++++-- .../test/fleet/20_wait_for_checkpoints.yml | 5 +-- .../yamlRestTest/resources}/roles.yml | 0 5 files changed, 57 insertions(+), 15 deletions(-) rename x-pack/plugin/fleet/qa/rest/{ => src/yamlRestTest/resources}/roles.yml (100%) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java b/server/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java index afbfe129c302e..2927c394da3d4 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java @@ -130,6 +130,14 @@ public SearchRequestBuilder setWaitForCheckpoints(Map waitForChe return this; } + /** + * Set the timeout for the {@link #setWaitForCheckpoints(Map)} request. + */ + public SearchRequestBuilder setWaitForCheckpointsTimeout(final TimeValue waitForCheckpointsTimeout) { + request.setWaitForCheckpointsTimeout(waitForCheckpointsTimeout); + return this; + } + /** * Specifies what type of requested indices to ignore and wildcard indices expressions. *

diff --git a/x-pack/plugin/fleet/qa/rest/build.gradle b/x-pack/plugin/fleet/qa/rest/build.gradle index fda9251c7ef34..dec624bc3cc56 100644 --- a/x-pack/plugin/fleet/qa/rest/build.gradle +++ b/x-pack/plugin/fleet/qa/rest/build.gradle @@ -1,8 +1,15 @@ -apply plugin: 'elasticsearch.legacy-yaml-rest-test' +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ -dependencies { - yamlRestTestImplementation(testArtifact(project(xpackModule('core')))) -} +import org.elasticsearch.gradle.internal.info.BuildParams + +apply plugin: 'elasticsearch.internal-yaml-rest-test' +apply plugin: 'elasticsearch.yaml-rest-compat-test' +apply plugin: 'elasticsearch.internal-test-artifact' restResources { restApi { @@ -10,11 +17,17 @@ restResources { } } -testClusters.configureEach { - testDistribution = 'DEFAULT' - setting 'xpack.security.enabled', 'true' - setting 'xpack.license.self_generated.type', 'trial' - extraConfigFile 'roles.yml', file('roles.yml') - user username: 'elastic_admin', password: 'admin-password' - user username: 'fleet_unprivileged_secrets', password: 'password', role: 'unprivileged_secrets' +artifacts { + restXpackTests(new File(projectDir, "src/yamlRestTest/resources/rest-api-spec/test")) +} + +tasks.named('yamlRestTest') { + usesDefaultDistribution() +} +tasks.named('yamlRestCompatTest') { + usesDefaultDistribution() +} +if (BuildParams.inFipsJvm){ + // This test cluster is using a BASIC license and FIPS 140 mode is not supported in BASIC + tasks.named("yamlRestTest").configure{enabled = false } } diff --git a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java index 202149abf11e1..bc49649bc1139 100644 --- a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java +++ b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java @@ -12,8 +12,12 @@ import org.elasticsearch.common.settings.SecureString; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.util.concurrent.ThreadContext; +import org.elasticsearch.test.cluster.ElasticsearchCluster; +import org.elasticsearch.test.cluster.local.distribution.DistributionType; +import org.elasticsearch.test.cluster.util.resource.Resource; import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate; import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase; +import org.junit.ClassRule; public class FleetRestIT extends ESClientYamlSuiteTestCase { @@ -21,14 +25,30 @@ public FleetRestIT(final ClientYamlTestCandidate testCandidate) { super(testCandidate); } + @ClassRule + public static ElasticsearchCluster cluster = ElasticsearchCluster.local() + .distribution(DistributionType.DEFAULT) + .setting("xpack.license.self_generated.type", "basic") + .setting("xpack.security.enabled", "true") + .rolesFile(Resource.fromClasspath("roles.yml")) + .user("elastic_admin", "admin-password", "superuser", true) + .user("fleet_unprivileged_secrets", "password", "unprivileged_secrets", true) + .build(); + + @Override + protected String getTestRestCluster() { + return cluster.getHttpAddresses(); + } + @Override protected Settings restClientSettings() { - String authentication = basicAuthHeaderValue("elastic_admin", new SecureString("admin-password".toCharArray())); - return Settings.builder().put(super.restClientSettings()).put(ThreadContext.PREFIX + ".Authorization", authentication).build(); + String token = basicAuthHeaderValue("elastic_admin", new SecureString("admin-password".toCharArray())); + return Settings.builder().put(ThreadContext.PREFIX + ".Authorization", token).build(); } @ParametersFactory public static Iterable parameters() throws Exception { return ESClientYamlSuiteTestCase.createParameters(); } + } diff --git a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml index 5610502a65d23..4c168c8feb0cd 100644 --- a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml +++ b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml @@ -105,6 +105,7 @@ setup: index: "test-after-refresh" allow_partial_search_results: false wait_for_checkpoints: 2 + wait_for_checkpoints_timeout: 1m body: { query: { match_all: {} } } --- @@ -115,7 +116,7 @@ setup: body: - { "allow_partial_search_results": false, wait_for_checkpoints: 1 } - { query: { match_all: { } } } - - { "allow_partial_search_results": false, wait_for_checkpoints: 2 } + - { "allow_partial_search_results": false, wait_for_checkpoints: 2, wait_for_checkpoints_timeout: 1m } - { query: { match_all: { } } } - match: { responses.0._shards.successful: 1 } @@ -128,7 +129,7 @@ setup: - {query: { match_all: {} } } - { "index": "test-alias", "allow_partial_search_results": false, wait_for_checkpoints: 1 } - { query: { match_all: { } } } - - {"index": "test-refresh-disabled", "allow_partial_search_results": false, wait_for_checkpoints: 2} + - { "index": "test-refresh-disabled", "allow_partial_search_results": false, wait_for_checkpoints: 2, wait_for_checkpoints_timeout: 1m } - {query: { match_all: {} } } - match: { responses.0._shards.successful: 1 } diff --git a/x-pack/plugin/fleet/qa/rest/roles.yml b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/roles.yml similarity index 100% rename from x-pack/plugin/fleet/qa/rest/roles.yml rename to x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/roles.yml From cca70d7eff1540fad2ff19081aac6f5f82cc68f5 Mon Sep 17 00:00:00 2001 From: David Kyle Date: Tue, 12 Nov 2024 11:14:06 +0000 Subject: [PATCH 06/16] [ML] Batch the chunks (#115477) Models running on an ml node have a queue of requests, when that queue is full new requests are rejected. A large document can chunk into hundreds of requests and in extreme cases a single large document can overflow the queue. Avoid this by batches of chunks keeping certain number of requests in flight. --- .../ElasticsearchInternalService.java | 103 ++++++++++++--- .../EmbeddingRequestChunkerTests.java | 13 ++ .../ElasticsearchInternalServiceTests.java | 122 ++++++++++++++++-- 3 files changed, 205 insertions(+), 33 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java index 83249266c79ab..fe83acc8574aa 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java @@ -68,6 +68,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.Stream; @@ -680,25 +681,13 @@ public void chunkedInfer( esModel.getConfigurations().getChunkingSettings() ).batchRequestsWithListeners(listener); - for (var batch : batchedRequests) { - var inferenceRequest = buildInferenceRequest( - esModel.mlNodeDeploymentId(), - EmptyConfigUpdate.INSTANCE, - batch.batch().inputs(), - inputType, - timeout - ); - - ActionListener mlResultsListener = batch.listener() - .delegateFailureAndWrap( - (l, inferenceResult) -> translateToChunkedResult(model.getTaskType(), inferenceResult.getInferenceResults(), l) - ); - - var maybeDeployListener = mlResultsListener.delegateResponse( - (l, exception) -> maybeStartDeployment(esModel, exception, inferenceRequest, mlResultsListener) - ); - - client.execute(InferModelAction.INSTANCE, inferenceRequest, maybeDeployListener); + if (batchedRequests.isEmpty()) { + listener.onResponse(List.of()); + } else { + // Avoid filling the inference queue by executing the batches in series + // Each batch contains up to EMBEDDING_MAX_BATCH_SIZE inference request + var sequentialRunner = new BatchIterator(esModel, inputType, timeout, batchedRequests); + sequentialRunner.run(); } } else { listener.onFailure(notElasticsearchModelException(model)); @@ -1018,6 +1007,82 @@ static TaskType inferenceConfigToTaskType(InferenceConfig config) { } } + /** + * Iterates over the batch executing a limited number requests at a time to avoid + * filling the ML node inference queue. + * + * First, a single request is executed, which can also trigger deploying a model + * if necessary. When this request is successfully executed, a callback executes + * N requests in parallel next. Each of these requests also has a callback that + * executes one more request, so that at all time N requests are in-flight. This + * continues until all requests are executed. + */ + class BatchIterator { + private static final int NUM_REQUESTS_INFLIGHT = 20; // * batch size = 200 + + private final AtomicInteger index = new AtomicInteger(); + private final ElasticsearchInternalModel esModel; + private final List requestAndListeners; + private final InputType inputType; + private final TimeValue timeout; + + BatchIterator( + ElasticsearchInternalModel esModel, + InputType inputType, + TimeValue timeout, + List requestAndListeners + ) { + this.esModel = esModel; + this.requestAndListeners = requestAndListeners; + this.inputType = inputType; + this.timeout = timeout; + } + + void run() { + // The first request may deploy the model, and upon completion runs + // NUM_REQUESTS_INFLIGHT in parallel. + inferenceExecutor.execute(() -> inferBatch(NUM_REQUESTS_INFLIGHT, true)); + } + + private void inferBatch(int runAfterCount, boolean maybeDeploy) { + int batchIndex = index.getAndIncrement(); + if (batchIndex >= requestAndListeners.size()) { + return; + } + executeRequest(batchIndex, maybeDeploy, () -> { + for (int i = 0; i < runAfterCount; i++) { + // Subsequent requests may not deploy the model, because the first request + // already did so. Upon completion, it runs one more request. + inferenceExecutor.execute(() -> inferBatch(1, false)); + } + }); + } + + private void executeRequest(int batchIndex, boolean maybeDeploy, Runnable runAfter) { + EmbeddingRequestChunker.BatchRequestAndListener batch = requestAndListeners.get(batchIndex); + var inferenceRequest = buildInferenceRequest( + esModel.mlNodeDeploymentId(), + EmptyConfigUpdate.INSTANCE, + batch.batch().inputs(), + inputType, + timeout + ); + logger.trace("Executing batch index={}", batchIndex); + + ActionListener listener = batch.listener() + .delegateFailureAndWrap( + (l, inferenceResult) -> translateToChunkedResult(esModel.getTaskType(), inferenceResult.getInferenceResults(), l) + ); + if (runAfter != null) { + listener = ActionListener.runAfter(listener, runAfter); + } + if (maybeDeploy) { + listener = listener.delegateResponse((l, exception) -> maybeStartDeployment(esModel, exception, inferenceRequest, l)); + } + client.execute(InferModelAction.INSTANCE, inferenceRequest, listener); + } + } + public static class Configuration { public static InferenceServiceConfiguration get() { return configuration.getOrCompute(); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java index c1be537a6b0a7..4fdf254101d3e 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java @@ -24,12 +24,25 @@ import java.util.concurrent.atomic.AtomicReference; import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.empty; import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.startsWith; public class EmbeddingRequestChunkerTests extends ESTestCase { + public void testEmptyInput() { + var embeddingType = randomFrom(EmbeddingRequestChunker.EmbeddingType.values()); + var batches = new EmbeddingRequestChunker(List.of(), 100, 100, 10, embeddingType).batchRequestsWithListeners(testListener()); + assertThat(batches, empty()); + } + + public void testBlankInput() { + var embeddingType = randomFrom(EmbeddingRequestChunker.EmbeddingType.values()); + var batches = new EmbeddingRequestChunker(List.of(""), 100, 100, 10, embeddingType).batchRequestsWithListeners(testListener()); + assertThat(batches, hasSize(1)); + } + public void testShortInputsAreSingleBatch() { String input = "one chunk"; var embeddingType = randomFrom(EmbeddingRequestChunker.EmbeddingType.values()); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java index 89a27a921cbea..9a4d0dda82238 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java @@ -12,6 +12,7 @@ import org.apache.logging.log4j.Level; import org.elasticsearch.ElasticsearchStatusException; import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.LatchedActionListener; import org.elasticsearch.action.support.PlainActionFuture; import org.elasticsearch.client.internal.Client; import org.elasticsearch.cluster.service.ClusterService; @@ -65,6 +66,7 @@ import org.elasticsearch.xpack.inference.InferencePlugin; import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsTests; import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker; +import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunkingSettings; import org.elasticsearch.xpack.inference.services.ServiceFields; import org.junit.After; import org.junit.Before; @@ -72,12 +74,14 @@ import org.mockito.Mockito; import java.util.ArrayList; +import java.util.Arrays; import java.util.EnumSet; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -832,16 +836,16 @@ public void testParsePersistedConfig() { } } - public void testChunkInfer_E5WithNullChunkingSettings() { + public void testChunkInfer_E5WithNullChunkingSettings() throws InterruptedException { testChunkInfer_e5(null); } - public void testChunkInfer_E5ChunkingSettingsSet() { + public void testChunkInfer_E5ChunkingSettingsSet() throws InterruptedException { testChunkInfer_e5(ChunkingSettingsTests.createRandomChunkingSettings()); } @SuppressWarnings("unchecked") - private void testChunkInfer_e5(ChunkingSettings chunkingSettings) { + private void testChunkInfer_e5(ChunkingSettings chunkingSettings) throws InterruptedException { var mlTrainedModelResults = new ArrayList(); mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults()); mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults()); @@ -889,6 +893,9 @@ private void testChunkInfer_e5(ChunkingSettings chunkingSettings) { gotResults.set(true); }, ESTestCase::fail); + var latch = new CountDownLatch(1); + var latchedListener = new LatchedActionListener<>(resultsListener, latch); + service.chunkedInfer( model, null, @@ -897,22 +904,23 @@ private void testChunkInfer_e5(ChunkingSettings chunkingSettings) { InputType.SEARCH, new ChunkingOptions(null, null), InferenceAction.Request.DEFAULT_TIMEOUT, - ActionListener.runAfter(resultsListener, () -> terminate(threadPool)) + latchedListener ); + latch.await(); assertTrue("Listener not called", gotResults.get()); } - public void testChunkInfer_SparseWithNullChunkingSettings() { + public void testChunkInfer_SparseWithNullChunkingSettings() throws InterruptedException { testChunkInfer_Sparse(null); } - public void testChunkInfer_SparseWithChunkingSettingsSet() { + public void testChunkInfer_SparseWithChunkingSettingsSet() throws InterruptedException { testChunkInfer_Sparse(ChunkingSettingsTests.createRandomChunkingSettings()); } @SuppressWarnings("unchecked") - private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) { + private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) throws InterruptedException { var mlTrainedModelResults = new ArrayList(); mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults()); mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults()); @@ -936,6 +944,7 @@ private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) { var service = createService(client); var gotResults = new AtomicBoolean(); + var resultsListener = ActionListener.>wrap(chunkedResponse -> { assertThat(chunkedResponse, hasSize(2)); assertThat(chunkedResponse.get(0), instanceOf(InferenceChunkedSparseEmbeddingResults.class)); @@ -955,6 +964,9 @@ private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) { gotResults.set(true); }, ESTestCase::fail); + var latch = new CountDownLatch(1); + var latchedListener = new LatchedActionListener<>(resultsListener, latch); + service.chunkedInfer( model, null, @@ -963,22 +975,23 @@ private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) { InputType.SEARCH, new ChunkingOptions(null, null), InferenceAction.Request.DEFAULT_TIMEOUT, - ActionListener.runAfter(resultsListener, () -> terminate(threadPool)) + latchedListener ); + latch.await(); assertTrue("Listener not called", gotResults.get()); } - public void testChunkInfer_ElserWithNullChunkingSettings() { + public void testChunkInfer_ElserWithNullChunkingSettings() throws InterruptedException { testChunkInfer_Elser(null); } - public void testChunkInfer_ElserWithChunkingSettingsSet() { + public void testChunkInfer_ElserWithChunkingSettingsSet() throws InterruptedException { testChunkInfer_Elser(ChunkingSettingsTests.createRandomChunkingSettings()); } @SuppressWarnings("unchecked") - private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) { + private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) throws InterruptedException { var mlTrainedModelResults = new ArrayList(); mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults()); mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults()); @@ -1022,6 +1035,9 @@ private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) { gotResults.set(true); }, ESTestCase::fail); + var latch = new CountDownLatch(1); + var latchedListener = new LatchedActionListener<>(resultsListener, latch); + service.chunkedInfer( model, null, @@ -1030,9 +1046,10 @@ private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) { InputType.SEARCH, new ChunkingOptions(null, null), InferenceAction.Request.DEFAULT_TIMEOUT, - ActionListener.runAfter(resultsListener, () -> terminate(threadPool)) + latchedListener ); + latch.await(); assertTrue("Listener not called", gotResults.get()); } @@ -1093,7 +1110,7 @@ public void testChunkInferSetsTokenization() { } @SuppressWarnings("unchecked") - public void testChunkInfer_FailsBatch() { + public void testChunkInfer_FailsBatch() throws InterruptedException { var mlTrainedModelResults = new ArrayList(); mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults()); mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults()); @@ -1129,6 +1146,9 @@ public void testChunkInfer_FailsBatch() { gotResults.set(true); }, ESTestCase::fail); + var latch = new CountDownLatch(1); + var latchedListener = new LatchedActionListener<>(resultsListener, latch); + service.chunkedInfer( model, null, @@ -1137,12 +1157,86 @@ public void testChunkInfer_FailsBatch() { InputType.SEARCH, new ChunkingOptions(null, null), InferenceAction.Request.DEFAULT_TIMEOUT, - ActionListener.runAfter(resultsListener, () -> terminate(threadPool)) + latchedListener ); + latch.await(); assertTrue("Listener not called", gotResults.get()); } + @SuppressWarnings("unchecked") + public void testChunkingLargeDocument() throws InterruptedException { + int numBatches = randomIntBetween(3, 6); + + // how many response objects to return in each batch + int[] numResponsesPerBatch = new int[numBatches]; + for (int i = 0; i < numBatches - 1; i++) { + numResponsesPerBatch[i] = ElasticsearchInternalService.EMBEDDING_MAX_BATCH_SIZE; + } + numResponsesPerBatch[numBatches - 1] = randomIntBetween(1, ElasticsearchInternalService.EMBEDDING_MAX_BATCH_SIZE); + int numChunks = Arrays.stream(numResponsesPerBatch).sum(); + + // build a doc with enough words to make numChunks of chunks + int wordsPerChunk = 10; + int numWords = numChunks * wordsPerChunk; + var input = "word ".repeat(numWords); + + Client client = mock(Client.class); + when(client.threadPool()).thenReturn(threadPool); + + // mock the inference response + doAnswer(invocationOnMock -> { + var request = (InferModelAction.Request) invocationOnMock.getArguments()[1]; + var listener = (ActionListener) invocationOnMock.getArguments()[2]; + var mlTrainedModelResults = new ArrayList(); + for (int i = 0; i < request.numberOfDocuments(); i++) { + mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults()); + } + var response = new InferModelAction.Response(mlTrainedModelResults, "foo", true); + listener.onResponse(response); + return null; + }).when(client).execute(same(InferModelAction.INSTANCE), any(InferModelAction.Request.class), any(ActionListener.class)); + + var service = createService(client); + + var gotResults = new AtomicBoolean(); + var resultsListener = ActionListener.>wrap(chunkedResponse -> { + assertThat(chunkedResponse, hasSize(1)); + assertThat(chunkedResponse.get(0), instanceOf(InferenceChunkedTextEmbeddingFloatResults.class)); + var sparseResults = (InferenceChunkedTextEmbeddingFloatResults) chunkedResponse.get(0); + assertThat(sparseResults.chunks(), hasSize(numChunks)); + + gotResults.set(true); + }, ESTestCase::fail); + + // Create model using the word boundary chunker. + var model = new MultilingualE5SmallModel( + "foo", + TaskType.TEXT_EMBEDDING, + "e5", + new MultilingualE5SmallInternalServiceSettings(1, 1, "cross-platform", null), + new WordBoundaryChunkingSettings(wordsPerChunk, 0) + ); + + var latch = new CountDownLatch(1); + var latchedListener = new LatchedActionListener<>(resultsListener, latch); + + // For the given input we know how many requests will be made + service.chunkedInfer( + model, + null, + List.of(input), + Map.of(), + InputType.SEARCH, + new ChunkingOptions(null, null), + InferenceAction.Request.DEFAULT_TIMEOUT, + latchedListener + ); + + latch.await(); + assertTrue("Listener not called with results", gotResults.get()); + } + public void testParsePersistedConfig_Rerank() { // with task settings { From 6303de34e44e32372f9d3d4be68bf6d243d9b110 Mon Sep 17 00:00:00 2001 From: Jan Kuipers <148754765+jan-elastic@users.noreply.github.com> Date: Tue, 12 Nov 2024 12:58:04 +0100 Subject: [PATCH 07/16] Fix NPE in MlMemoryAutoscalingDecider (#116650) * Fix NPE in MlMemoryAutoscalingDecider * Update docs/changelog/116650.yaml * Update 116650.yaml * Update docs/changelog/116650.yaml * better fix --- docs/changelog/116650.yaml | 5 +++++ .../xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java | 6 +++++- .../xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java | 2 +- 3 files changed, 11 insertions(+), 2 deletions(-) create mode 100644 docs/changelog/116650.yaml diff --git a/docs/changelog/116650.yaml b/docs/changelog/116650.yaml new file mode 100644 index 0000000000000..d314a918aede9 --- /dev/null +++ b/docs/changelog/116650.yaml @@ -0,0 +1,5 @@ +pr: 116650 +summary: Fix bug in ML autoscaling when some node info is unavailable +area: Machine Learning +type: bug +issues: [] diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java index bab7bb52f928f..5a06308a3c8cc 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java @@ -17,7 +17,11 @@ public static Builder builder(ByteSizeValue nodeSize, ByteSizeValue tierSize) { } public static Builder from(AutoscalingCapacity autoscalingCapacity) { - return builder(autoscalingCapacity.node().memory(), autoscalingCapacity.total().memory()); + if (autoscalingCapacity == null) { + return builder(null, null); + } else { + return builder(autoscalingCapacity.node().memory(), autoscalingCapacity.total().memory()); + } } @Override diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java index dfe0e557f749d..0ff6aece95ab1 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java @@ -809,7 +809,7 @@ static MlMemoryAutoscalingCapacity ensureScaleDown( MlMemoryAutoscalingCapacity scaleDownResult, MlMemoryAutoscalingCapacity currentCapacity ) { - if (scaleDownResult == null || currentCapacity == null) { + if (scaleDownResult == null || currentCapacity == null || currentCapacity.isUndetermined()) { return null; } MlMemoryAutoscalingCapacity newCapacity = MlMemoryAutoscalingCapacity.builder( From 85b2bab2e38409449a37adc9408902cbf79f8c8f Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Tue, 12 Nov 2024 12:17:04 +0000 Subject: [PATCH 08/16] Bump versions after 8.15.4 release --- .buildkite/pipelines/intake.yml | 2 +- .buildkite/pipelines/periodic-packaging.yml | 6 +++--- .buildkite/pipelines/periodic.yml | 10 +++++----- .ci/bwcVersions | 2 +- .ci/snapshotBwcVersions | 2 +- server/src/main/java/org/elasticsearch/Version.java | 1 + .../resources/org/elasticsearch/TransportVersions.csv | 1 + .../org/elasticsearch/index/IndexVersions.csv | 1 + 8 files changed, 14 insertions(+), 11 deletions(-) diff --git a/.buildkite/pipelines/intake.yml b/.buildkite/pipelines/intake.yml index 37ea49e3a6d95..167830d3ed8b3 100644 --- a/.buildkite/pipelines/intake.yml +++ b/.buildkite/pipelines/intake.yml @@ -56,7 +56,7 @@ steps: timeout_in_minutes: 300 matrix: setup: - BWC_VERSION: ["8.15.4", "8.16.0", "8.17.0", "9.0.0"] + BWC_VERSION: ["8.15.5", "8.16.0", "8.17.0", "9.0.0"] agents: provider: gcp image: family/elasticsearch-ubuntu-2004 diff --git a/.buildkite/pipelines/periodic-packaging.yml b/.buildkite/pipelines/periodic-packaging.yml index 788960c76e150..0f2e70addd684 100644 --- a/.buildkite/pipelines/periodic-packaging.yml +++ b/.buildkite/pipelines/periodic-packaging.yml @@ -272,8 +272,8 @@ steps: env: BWC_VERSION: 8.14.3 - - label: "{{matrix.image}} / 8.15.4 / packaging-tests-upgrade" - command: ./.ci/scripts/packaging-test.sh -Dbwc.checkout.align=true destructiveDistroUpgradeTest.v8.15.4 + - label: "{{matrix.image}} / 8.15.5 / packaging-tests-upgrade" + command: ./.ci/scripts/packaging-test.sh -Dbwc.checkout.align=true destructiveDistroUpgradeTest.v8.15.5 timeout_in_minutes: 300 matrix: setup: @@ -286,7 +286,7 @@ steps: machineType: custom-16-32768 buildDirectory: /dev/shm/bk env: - BWC_VERSION: 8.15.4 + BWC_VERSION: 8.15.5 - label: "{{matrix.image}} / 8.16.0 / packaging-tests-upgrade" command: ./.ci/scripts/packaging-test.sh -Dbwc.checkout.align=true destructiveDistroUpgradeTest.v8.16.0 diff --git a/.buildkite/pipelines/periodic.yml b/.buildkite/pipelines/periodic.yml index 7b6a6ea72fe83..f68f64332426c 100644 --- a/.buildkite/pipelines/periodic.yml +++ b/.buildkite/pipelines/periodic.yml @@ -287,8 +287,8 @@ steps: - signal_reason: agent_stop limit: 3 - - label: 8.15.4 / bwc - command: .ci/scripts/run-gradle.sh -Dbwc.checkout.align=true v8.15.4#bwcTest + - label: 8.15.5 / bwc + command: .ci/scripts/run-gradle.sh -Dbwc.checkout.align=true v8.15.5#bwcTest timeout_in_minutes: 300 agents: provider: gcp @@ -297,7 +297,7 @@ steps: buildDirectory: /dev/shm/bk preemptible: true env: - BWC_VERSION: 8.15.4 + BWC_VERSION: 8.15.5 retry: automatic: - exit_status: "-1" @@ -429,7 +429,7 @@ steps: setup: ES_RUNTIME_JAVA: - openjdk21 - BWC_VERSION: ["8.15.4", "8.16.0", "8.17.0", "9.0.0"] + BWC_VERSION: ["8.15.5", "8.16.0", "8.17.0", "9.0.0"] agents: provider: gcp image: family/elasticsearch-ubuntu-2004 @@ -471,7 +471,7 @@ steps: ES_RUNTIME_JAVA: - openjdk21 - openjdk23 - BWC_VERSION: ["8.15.4", "8.16.0", "8.17.0", "9.0.0"] + BWC_VERSION: ["8.15.5", "8.16.0", "8.17.0", "9.0.0"] agents: provider: gcp image: family/elasticsearch-ubuntu-2004 diff --git a/.ci/bwcVersions b/.ci/bwcVersions index 2e77631450825..b4a4460ff5a80 100644 --- a/.ci/bwcVersions +++ b/.ci/bwcVersions @@ -14,7 +14,7 @@ BWC_VERSION: - "8.12.2" - "8.13.4" - "8.14.3" - - "8.15.4" + - "8.15.5" - "8.16.0" - "8.17.0" - "9.0.0" diff --git a/.ci/snapshotBwcVersions b/.ci/snapshotBwcVersions index c6edc709a8ceb..7dad55b653925 100644 --- a/.ci/snapshotBwcVersions +++ b/.ci/snapshotBwcVersions @@ -1,5 +1,5 @@ BWC_VERSION: - - "8.15.4" + - "8.15.5" - "8.16.0" - "8.17.0" - "9.0.0" diff --git a/server/src/main/java/org/elasticsearch/Version.java b/server/src/main/java/org/elasticsearch/Version.java index 5e4df05c10182..909d733fd3719 100644 --- a/server/src/main/java/org/elasticsearch/Version.java +++ b/server/src/main/java/org/elasticsearch/Version.java @@ -187,6 +187,7 @@ public class Version implements VersionId, ToXContentFragment { public static final Version V_8_15_2 = new Version(8_15_02_99); public static final Version V_8_15_3 = new Version(8_15_03_99); public static final Version V_8_15_4 = new Version(8_15_04_99); + public static final Version V_8_15_5 = new Version(8_15_05_99); public static final Version V_8_16_0 = new Version(8_16_00_99); public static final Version V_8_17_0 = new Version(8_17_00_99); public static final Version V_9_0_0 = new Version(9_00_00_99); diff --git a/server/src/main/resources/org/elasticsearch/TransportVersions.csv b/server/src/main/resources/org/elasticsearch/TransportVersions.csv index b0ef5b780e775..26c518962c19a 100644 --- a/server/src/main/resources/org/elasticsearch/TransportVersions.csv +++ b/server/src/main/resources/org/elasticsearch/TransportVersions.csv @@ -131,3 +131,4 @@ 8.15.1,8702002 8.15.2,8702003 8.15.3,8702003 +8.15.4,8702003 diff --git a/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv b/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv index e3681cc975988..6cab0b513ee63 100644 --- a/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv +++ b/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv @@ -131,3 +131,4 @@ 8.15.1,8512000 8.15.2,8512000 8.15.3,8512000 +8.15.4,8512000 From 6c85934c18e22eca5109e66a66174cae339fa040 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Tue, 12 Nov 2024 12:18:50 +0000 Subject: [PATCH 09/16] Prune changelogs after 8.15.4 release --- docs/changelog/112250.yaml | 5 ----- docs/changelog/113723.yaml | 6 ------ docs/changelog/114407.yaml | 6 ------ docs/changelog/114533.yaml | 5 ----- docs/changelog/114601.yaml | 6 ------ docs/changelog/115181.yaml | 5 ----- docs/changelog/115308.yaml | 6 ------ docs/changelog/115430.yaml | 5 ----- docs/changelog/115459.yaml | 5 ----- docs/changelog/115510.yaml | 6 ------ docs/changelog/115834.yaml | 5 ----- docs/changelog/116031.yaml | 6 ------ docs/changelog/116219.yaml | 6 ------ 13 files changed, 72 deletions(-) delete mode 100644 docs/changelog/112250.yaml delete mode 100644 docs/changelog/113723.yaml delete mode 100644 docs/changelog/114407.yaml delete mode 100644 docs/changelog/114533.yaml delete mode 100644 docs/changelog/114601.yaml delete mode 100644 docs/changelog/115181.yaml delete mode 100644 docs/changelog/115308.yaml delete mode 100644 docs/changelog/115430.yaml delete mode 100644 docs/changelog/115459.yaml delete mode 100644 docs/changelog/115510.yaml delete mode 100644 docs/changelog/115834.yaml delete mode 100644 docs/changelog/116031.yaml delete mode 100644 docs/changelog/116219.yaml diff --git a/docs/changelog/112250.yaml b/docs/changelog/112250.yaml deleted file mode 100644 index edbb5667d4b9d..0000000000000 --- a/docs/changelog/112250.yaml +++ /dev/null @@ -1,5 +0,0 @@ -pr: 112250 -summary: Do not exclude empty arrays or empty objects in source filtering -area: Search -type: bug -issues: [109668] diff --git a/docs/changelog/113723.yaml b/docs/changelog/113723.yaml deleted file mode 100644 index 2cbcf49102719..0000000000000 --- a/docs/changelog/113723.yaml +++ /dev/null @@ -1,6 +0,0 @@ -pr: 113723 -summary: Fix max file size check to use `getMaxFileSize` -area: Infra/Core -type: bug -issues: - - 113705 diff --git a/docs/changelog/114407.yaml b/docs/changelog/114407.yaml deleted file mode 100644 index 4c1134a9d3834..0000000000000 --- a/docs/changelog/114407.yaml +++ /dev/null @@ -1,6 +0,0 @@ -pr: 114407 -summary: Fix synthetic source handling for `bit` type in `dense_vector` field -area: Search -type: bug -issues: - - 114402 diff --git a/docs/changelog/114533.yaml b/docs/changelog/114533.yaml deleted file mode 100644 index f45589e8de921..0000000000000 --- a/docs/changelog/114533.yaml +++ /dev/null @@ -1,5 +0,0 @@ -pr: 114533 -summary: Fix dim validation for bit `element_type` -area: Vector Search -type: bug -issues: [] diff --git a/docs/changelog/114601.yaml b/docs/changelog/114601.yaml deleted file mode 100644 index d2f563d62a639..0000000000000 --- a/docs/changelog/114601.yaml +++ /dev/null @@ -1,6 +0,0 @@ -pr: 114601 -summary: Support semantic_text in object fields -area: Vector Search -type: bug -issues: - - 114401 diff --git a/docs/changelog/115181.yaml b/docs/changelog/115181.yaml deleted file mode 100644 index 65f59d5ed0add..0000000000000 --- a/docs/changelog/115181.yaml +++ /dev/null @@ -1,5 +0,0 @@ -pr: 115181 -summary: Always check the parent breaker with zero bytes in `PreallocatedCircuitBreakerService` -area: Aggregations -type: bug -issues: [] diff --git a/docs/changelog/115308.yaml b/docs/changelog/115308.yaml deleted file mode 100644 index 163f0232a3e58..0000000000000 --- a/docs/changelog/115308.yaml +++ /dev/null @@ -1,6 +0,0 @@ -pr: 115308 -summary: "ESQL: Disable pushdown of WHERE past STATS" -area: ES|QL -type: bug -issues: - - 115281 diff --git a/docs/changelog/115430.yaml b/docs/changelog/115430.yaml deleted file mode 100644 index c2903f7751012..0000000000000 --- a/docs/changelog/115430.yaml +++ /dev/null @@ -1,5 +0,0 @@ -pr: 115430 -summary: Prevent NPE if model assignment is removed while waiting to start -area: Machine Learning -type: bug -issues: [] diff --git a/docs/changelog/115459.yaml b/docs/changelog/115459.yaml deleted file mode 100644 index b20a8f765c084..0000000000000 --- a/docs/changelog/115459.yaml +++ /dev/null @@ -1,5 +0,0 @@ -pr: 115459 -summary: Guard blob store local directory creation with `doPrivileged` -area: Infra/Core -type: bug -issues: [] diff --git a/docs/changelog/115510.yaml b/docs/changelog/115510.yaml deleted file mode 100644 index 1e71270e18f97..0000000000000 --- a/docs/changelog/115510.yaml +++ /dev/null @@ -1,6 +0,0 @@ -pr: 115510 -summary: Fix lingering license warning header in IP filter -area: License -type: bug -issues: - - 114865 diff --git a/docs/changelog/115834.yaml b/docs/changelog/115834.yaml deleted file mode 100644 index 91f9e9a4e2e41..0000000000000 --- a/docs/changelog/115834.yaml +++ /dev/null @@ -1,5 +0,0 @@ -pr: 115834 -summary: Try to simplify geometries that fail with `TopologyException` -area: Geo -type: bug -issues: [] diff --git a/docs/changelog/116031.yaml b/docs/changelog/116031.yaml deleted file mode 100644 index e30552bf3b513..0000000000000 --- a/docs/changelog/116031.yaml +++ /dev/null @@ -1,6 +0,0 @@ -pr: 116031 -summary: Resolve pipelines from template on lazy rollover write -area: Data streams -type: bug -issues: - - 112781 diff --git a/docs/changelog/116219.yaml b/docs/changelog/116219.yaml deleted file mode 100644 index aeeea68570e77..0000000000000 --- a/docs/changelog/116219.yaml +++ /dev/null @@ -1,6 +0,0 @@ -pr: 116219 -summary: "[apm-data] Apply lazy rollover on index template creation" -area: Data streams -type: bug -issues: - - 116230 From 098c8dad900bd46bfdac0156b5c10a173929e175 Mon Sep 17 00:00:00 2001 From: Jedr Blaszyk Date: Tue, 12 Nov 2024 15:27:58 +0100 Subject: [PATCH 10/16] [Docs] Fix sharepoint docs for 8.16 release (#116661) --- .../docs/_connectors-overview-table.asciidoc | 2 +- .../connector/docs/connectors-sharepoint.asciidoc | 12 +++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/docs/reference/connector/docs/_connectors-overview-table.asciidoc b/docs/reference/connector/docs/_connectors-overview-table.asciidoc index f25ea3deceeee..f5f8103349dde 100644 --- a/docs/reference/connector/docs/_connectors-overview-table.asciidoc +++ b/docs/reference/connector/docs/_connectors-overview-table.asciidoc @@ -44,7 +44,7 @@ NOTE: All connectors are available as self-managed <>|*GA*|8.12+|8.12+|8.11+|8.13+|8.13+|https://github.com/elastic/connectors/tree/main/connectors/sources/salesforce.py[View code] |<>|*GA*|8.10+|8.10+|8.11+|8.13+|8.13+|https://github.com/elastic/connectors/tree/main/connectors/sources/servicenow.py[View code] |<>|*GA*|8.9+|8.9+|8.9+|8.9+|8.9+|https://github.com/elastic/connectors/tree/main/connectors/sources/sharepoint_online.py[View code] -|<>|*Beta*|8.15+|-|8.11+|8.13+|8.14+|https://github.com/elastic/connectors/tree/main/connectors/sources/sharepoint_server.py[View code] +|<>|*Beta*|8.15+|-|8.11+|8.13+|8.15+|https://github.com/elastic/connectors/tree/main/connectors/sources/sharepoint_server.py[View code] |<>|*Preview*|8.14+|-|-|-|-|https://github.com/elastic/connectors/tree/main/connectors/sources/slack.py[View code] |<>|*Preview*|8.14+|-|-|8.13+|-|https://github.com/elastic/connectors/tree/main/connectors/sources/teams.py[View code] |<>|*Preview*|8.14+|-|8.11+|8.13+|-|https://github.com/elastic/connectors/tree/main/connectors/sources/zoom.py[View code] diff --git a/docs/reference/connector/docs/connectors-sharepoint.asciidoc b/docs/reference/connector/docs/connectors-sharepoint.asciidoc index f5590daa1e701..d7a2307a9db80 100644 --- a/docs/reference/connector/docs/connectors-sharepoint.asciidoc +++ b/docs/reference/connector/docs/connectors-sharepoint.asciidoc @@ -67,6 +67,9 @@ The following SharePoint Server versions are compatible: The following configuration fields are required to set up the connector: +`authentication`:: +Authentication mode, either *Basic* or *NTLM*. + `username`:: The username of the account for the SharePoint Server instance. @@ -133,7 +136,7 @@ The connector syncs the following SharePoint object types: [NOTE] ==== * Content from files bigger than 10 MB won't be extracted by default. Use the <> to handle larger binary files. -* Permissions are not synced. **All documents** indexed to an Elastic deployment will be visible to **all users with access** to that Elasticsearch Index. +* Permissions are not synced by default. Enable <> to sync permissions. ==== [discrete#es-connectors-sharepoint-sync-types] @@ -191,7 +194,7 @@ This connector is written in Python using the {connectors-python}[Elastic connec View the {connectors-python}/connectors/sources/sharepoint_server.py[source code for this connector^] (branch _{connectors-branch}_, compatible with Elastic _{minor-version}_). -// Closing the collapsible section +// Closing the collapsible section =============== @@ -254,6 +257,9 @@ Once connected, you'll be able to update these values in Kibana. The following configuration fields are required to set up the connector: +`authentication`:: +Authentication mode, either *Basic* or *NTLM*. + `username`:: The username of the account for the SharePoint Server instance. @@ -408,5 +414,5 @@ This connector is written in Python using the {connectors-python}[Elastic connec View the {connectors-python}/connectors/sources/sharepoint_server.py[source code for this connector^] (branch _{connectors-branch}_, compatible with Elastic _{minor-version}_). -// Closing the collapsible section +// Closing the collapsible section =============== From ade29fb8f444731a81cbf2dff2ba6c93206fccdc Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Tue, 12 Nov 2024 15:50:50 +0100 Subject: [PATCH 11/16] Deduplicate DocValueFormat objects from InternalAggregation when deserializing (#116640) --- .../elasticsearch/search/DocValueFormat.java | 21 +++++++++++++++++-- .../elasticsearch/search/SearchModule.java | 4 ++-- .../search/DocValueFormatTests.java | 4 ++-- 3 files changed, 23 insertions(+), 6 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/DocValueFormat.java b/server/src/main/java/org/elasticsearch/search/DocValueFormat.java index 51f52326907eb..a1e8eb25f4780 100644 --- a/server/src/main/java/org/elasticsearch/search/DocValueFormat.java +++ b/server/src/main/java/org/elasticsearch/search/DocValueFormat.java @@ -12,6 +12,7 @@ import org.apache.lucene.document.InetAddressPoint; import org.apache.lucene.util.BytesRef; import org.elasticsearch.TransportVersions; +import org.elasticsearch.common.io.stream.DelayableWriteable; import org.elasticsearch.common.io.stream.NamedWriteable; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; @@ -260,7 +261,7 @@ private DateTime(DateFormatter formatter, ZoneId timeZone, DateFieldMapper.Resol this.formatSortValues = formatSortValues; } - public DateTime(StreamInput in) throws IOException { + private DateTime(StreamInput in) throws IOException { String formatterPattern = in.readString(); Locale locale = in.getTransportVersion().onOrAfter(TransportVersions.DATE_TIME_DOC_VALUES_LOCALES) ? LocaleUtils.parse(in.readString()) @@ -285,6 +286,14 @@ public String getWriteableName() { return NAME; } + public static DateTime readFrom(StreamInput in) throws IOException { + final DateTime dateTime = new DateTime(in); + if (in instanceof DelayableWriteable.Deduplicator d) { + return d.deduplicate(dateTime); + } + return dateTime; + } + @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(formatter.pattern()); @@ -528,7 +537,7 @@ public Decimal(String pattern) { this.format = new DecimalFormat(pattern, SYMBOLS); } - public Decimal(StreamInput in) throws IOException { + private Decimal(StreamInput in) throws IOException { this(in.readString()); } @@ -537,6 +546,14 @@ public String getWriteableName() { return NAME; } + public static Decimal readFrom(StreamInput in) throws IOException { + final Decimal decimal = new Decimal(in); + if (in instanceof DelayableWriteable.Deduplicator d) { + return d.deduplicate(decimal); + } + return decimal; + } + @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(pattern); diff --git a/server/src/main/java/org/elasticsearch/search/SearchModule.java b/server/src/main/java/org/elasticsearch/search/SearchModule.java index fd39a95bdb75d..7a8b4e0cfe95a 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchModule.java +++ b/server/src/main/java/org/elasticsearch/search/SearchModule.java @@ -1013,8 +1013,8 @@ private void registerScoreFunction(ScoreFunctionSpec scoreFunction) { private void registerValueFormats() { registerValueFormat(DocValueFormat.BOOLEAN.getWriteableName(), in -> DocValueFormat.BOOLEAN); - registerValueFormat(DocValueFormat.DateTime.NAME, DocValueFormat.DateTime::new); - registerValueFormat(DocValueFormat.Decimal.NAME, DocValueFormat.Decimal::new); + registerValueFormat(DocValueFormat.DateTime.NAME, DocValueFormat.DateTime::readFrom); + registerValueFormat(DocValueFormat.Decimal.NAME, DocValueFormat.Decimal::readFrom); registerValueFormat(DocValueFormat.GEOHASH.getWriteableName(), in -> DocValueFormat.GEOHASH); registerValueFormat(DocValueFormat.GEOTILE.getWriteableName(), in -> DocValueFormat.GEOTILE); registerValueFormat(DocValueFormat.IP.getWriteableName(), in -> DocValueFormat.IP); diff --git a/server/src/test/java/org/elasticsearch/search/DocValueFormatTests.java b/server/src/test/java/org/elasticsearch/search/DocValueFormatTests.java index e81066a731d2e..7c9a68cbc91f1 100644 --- a/server/src/test/java/org/elasticsearch/search/DocValueFormatTests.java +++ b/server/src/test/java/org/elasticsearch/search/DocValueFormatTests.java @@ -43,8 +43,8 @@ public class DocValueFormatTests extends ESTestCase { public void testSerialization() throws Exception { List entries = new ArrayList<>(); entries.add(new Entry(DocValueFormat.class, DocValueFormat.BOOLEAN.getWriteableName(), in -> DocValueFormat.BOOLEAN)); - entries.add(new Entry(DocValueFormat.class, DocValueFormat.DateTime.NAME, DocValueFormat.DateTime::new)); - entries.add(new Entry(DocValueFormat.class, DocValueFormat.Decimal.NAME, DocValueFormat.Decimal::new)); + entries.add(new Entry(DocValueFormat.class, DocValueFormat.DateTime.NAME, DocValueFormat.DateTime::readFrom)); + entries.add(new Entry(DocValueFormat.class, DocValueFormat.Decimal.NAME, DocValueFormat.Decimal::readFrom)); entries.add(new Entry(DocValueFormat.class, DocValueFormat.GEOHASH.getWriteableName(), in -> DocValueFormat.GEOHASH)); entries.add(new Entry(DocValueFormat.class, DocValueFormat.GEOTILE.getWriteableName(), in -> DocValueFormat.GEOTILE)); entries.add(new Entry(DocValueFormat.class, DocValueFormat.IP.getWriteableName(), in -> DocValueFormat.IP)); From 94ab1a6fa7638f8905538028a330b2901c01b8b0 Mon Sep 17 00:00:00 2001 From: Pawan Kartik Date: Tue, 12 Nov 2024 15:10:33 +0000 Subject: [PATCH 12/16] Add tests for RCS1:ES|QL to verify behaviour for disconnected clusters (#116449) * Add tests for RCS1:ES|QL to verify behaviour for disconnected clusters * fix: build * Add missing assertions for ccs metadata * Address review comments --- ...ssClusterEsqlRCS1UnavailableRemotesIT.java | 286 ++++++++++++++++++ 1 file changed, 286 insertions(+) create mode 100644 x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1UnavailableRemotesIT.java diff --git a/x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1UnavailableRemotesIT.java b/x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1UnavailableRemotesIT.java new file mode 100644 index 0000000000000..b6fc43e2a6e48 --- /dev/null +++ b/x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1UnavailableRemotesIT.java @@ -0,0 +1,286 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.remotecluster; + +import org.elasticsearch.client.Request; +import org.elasticsearch.client.Response; +import org.elasticsearch.client.ResponseException; +import org.elasticsearch.common.Strings; +import org.elasticsearch.test.cluster.ElasticsearchCluster; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.json.JsonXContent; +import org.junit.Before; +import org.junit.ClassRule; +import org.junit.rules.RuleChain; +import org.junit.rules.TestRule; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.Matchers.greaterThan; + +public class CrossClusterEsqlRCS1UnavailableRemotesIT extends AbstractRemoteClusterSecurityTestCase { + private static final AtomicBoolean SSL_ENABLED_REF = new AtomicBoolean(); + + static { + fulfillingCluster = ElasticsearchCluster.local() + .name("fulfilling-cluster") + .nodes(1) + .module("x-pack-esql") + .module("x-pack-enrich") + .apply(commonClusterConfig) + .setting("remote_cluster.port", "0") + .setting("xpack.ml.enabled", "false") + .setting("xpack.security.remote_cluster_server.ssl.enabled", () -> String.valueOf(SSL_ENABLED_REF.get())) + .setting("xpack.security.remote_cluster_server.ssl.key", "remote-cluster.key") + .setting("xpack.security.remote_cluster_server.ssl.certificate", "remote-cluster.crt") + .setting("xpack.security.authc.token.enabled", "true") + .keystore("xpack.security.remote_cluster_server.ssl.secure_key_passphrase", "remote-cluster-password") + .node(0, spec -> spec.setting("remote_cluster_server.enabled", "true")) + .build(); + + queryCluster = ElasticsearchCluster.local() + .name("query-cluster") + .module("x-pack-esql") + .module("x-pack-enrich") + .apply(commonClusterConfig) + .setting("xpack.ml.enabled", "false") + .setting("xpack.security.remote_cluster_client.ssl.enabled", () -> String.valueOf(SSL_ENABLED_REF.get())) + .build(); + } + + @ClassRule + public static TestRule clusterRule = RuleChain.outerRule(fulfillingCluster).around(queryCluster); + + @Before + public void setupPreRequisites() throws IOException { + setupRolesAndPrivileges(); + loadData(); + } + + public void testEsqlRcs1UnavailableRemoteScenarios() throws Exception { + clusterShutDownWithRandomSkipUnavailable(); + remoteClusterShutdownWithSkipUnavailableTrue(); + remoteClusterShutdownWithSkipUnavailableFalse(); + } + + private void clusterShutDownWithRandomSkipUnavailable() throws Exception { + // skip_unavailable is set to a random boolean value. + // However, no clusters are stopped. Hence, we do not expect any other behaviour + // other than a 200-OK. + + configureRemoteCluster("my_remote_cluster", fulfillingCluster, true, randomBoolean(), randomBoolean()); + String query = "FROM *,my_remote_cluster:* | LIMIT 10"; + Response response = client().performRequest(esqlRequest(query)); + + Map map = responseAsMap(response); + ArrayList columns = (ArrayList) map.get("columns"); + ArrayList values = (ArrayList) map.get("values"); + Map clusters = (Map) map.get("_clusters"); + Map clusterDetails = (Map) clusters.get("details"); + Map localClusterDetails = (Map) clusterDetails.get("(local)"); + Map remoteClusterDetails = (Map) clusterDetails.get("my_remote_cluster"); + + assertOK(response); + assertThat((int) map.get("took"), greaterThan(0)); + assertThat(columns.size(), is(4)); + assertThat(values.size(), is(9)); + + assertThat((int) clusters.get("total"), is(2)); + assertThat((int) clusters.get("successful"), is(2)); + assertThat((int) clusters.get("running"), is(0)); + assertThat((int) clusters.get("skipped"), is(0)); + assertThat((int) clusters.get("partial"), is(0)); + assertThat((int) clusters.get("failed"), is(0)); + + assertThat(clusterDetails.size(), is(2)); + assertThat((int) localClusterDetails.get("took"), greaterThan(0)); + assertThat(localClusterDetails.get("status"), is("successful")); + + assertThat((int) remoteClusterDetails.get("took"), greaterThan(0)); + assertThat(remoteClusterDetails.get("status"), is("successful")); + } + + @SuppressWarnings("unchecked") + private void remoteClusterShutdownWithSkipUnavailableTrue() throws Exception { + // Remote cluster is stopped and skip unavailable is set to true. + // We expect no exception and partial results from the remaining open cluster. + + configureRemoteCluster("my_remote_cluster", fulfillingCluster, true, randomBoolean(), true); + + try { + // Stop remote cluster. + fulfillingCluster.stop(true); + + // A simple query that targets our remote cluster. + String query = "FROM *,my_remote_cluster:* | LIMIT 10"; + Response response = client().performRequest(esqlRequest(query)); + + Map map = responseAsMap(response); + ArrayList columns = (ArrayList) map.get("columns"); + ArrayList values = (ArrayList) map.get("values"); + Map clusters = (Map) map.get("_clusters"); + Map clusterDetails = (Map) clusters.get("details"); + Map localClusterDetails = (Map) clusterDetails.get("(local)"); + Map remoteClusterDetails = (Map) clusterDetails.get("my_remote_cluster"); + + // Assert results obtained from the local cluster and that remote cluster was + // skipped. + assertOK(response); + assertThat((int) map.get("took"), greaterThan(0)); + + assertThat(columns.size(), is(2)); + assertThat(values.size(), is(5)); + + assertThat((int) clusters.get("total"), is(2)); + assertThat((int) clusters.get("successful"), is(1)); + assertThat((int) clusters.get("skipped"), is(1)); + assertThat((int) clusters.get("running"), is(0)); + assertThat((int) clusters.get("partial"), is(0)); + assertThat((int) clusters.get("failed"), is(0)); + + assertThat(clusterDetails.size(), is(2)); + assertThat((int) localClusterDetails.get("took"), greaterThan(0)); + assertThat(localClusterDetails.get("status"), is("successful")); + + assertThat((int) remoteClusterDetails.get("took"), greaterThan(0)); + assertThat(remoteClusterDetails.get("status"), is("skipped")); + + } catch (ResponseException r) { + throw new AssertionError(r); + } finally { + fulfillingCluster.start(); + closeFulfillingClusterClient(); + initFulfillingClusterClient(); + } + } + + private void remoteClusterShutdownWithSkipUnavailableFalse() throws Exception { + // Remote cluster is stopped and skip_unavailable is set to false. + // Although the other cluster is open, we expect an Exception. + + configureRemoteCluster("my_remote_cluster", fulfillingCluster, true, randomBoolean(), false); + + try { + // Stop remote cluster. + fulfillingCluster.stop(true); + + // A simple query that targets our remote cluster. + String query = "FROM *,my_remote_cluster:* | LIMIT 10"; + ResponseException ex = expectThrows(ResponseException.class, () -> client().performRequest(esqlRequest(query))); + assertThat(ex.getMessage(), containsString("connect_transport_exception")); + } finally { + fulfillingCluster.start(); + closeFulfillingClusterClient(); + initFulfillingClusterClient(); + } + } + + private void setupRolesAndPrivileges() throws IOException { + var putUserRequest = new Request("PUT", "/_security/user/" + REMOTE_SEARCH_USER); + putUserRequest.setJsonEntity(""" + { + "password": "x-pack-test-password", + "roles" : ["remote_search"] + }"""); + assertOK(adminClient().performRequest(putUserRequest)); + + var putRoleOnRemoteClusterRequest = new Request("PUT", "/_security/role/" + REMOTE_SEARCH_ROLE); + putRoleOnRemoteClusterRequest.setJsonEntity(""" + { + "indices": [ + { + "names": ["points", "squares"], + "privileges": ["read", "read_cross_cluster", "create_index", "monitor"] + } + ], + "remote_indices": [ + { + "names": ["points", "squares"], + "privileges": ["read", "read_cross_cluster", "create_index", "monitor"], + "clusters": ["my_remote_cluster"] + } + ] + }"""); + assertOK(adminClient().performRequest(putRoleOnRemoteClusterRequest)); + } + + private void loadData() throws IOException { + Request createIndex = new Request("PUT", "points"); + createIndex.setJsonEntity(""" + { + "mappings": { + "properties": { + "id": { "type": "integer" }, + "score": { "type": "integer" } + } + } + } + """); + assertOK(client().performRequest(createIndex)); + + Request bulkRequest = new Request("POST", "/_bulk?refresh=true"); + bulkRequest.setJsonEntity(""" + { "index": { "_index": "points" } } + { "id": 1, "score": 75} + { "index": { "_index": "points" } } + { "id": 2, "score": 125} + { "index": { "_index": "points" } } + { "id": 3, "score": 100} + { "index": { "_index": "points" } } + { "id": 4, "score": 50} + { "index": { "_index": "points" } } + { "id": 5, "score": 150} + """); + assertOK(client().performRequest(bulkRequest)); + + createIndex = new Request("PUT", "squares"); + createIndex.setJsonEntity(""" + { + "mappings": { + "properties": { + "num": { "type": "integer" }, + "square": { "type": "integer" } + } + } + } + """); + assertOK(performRequestAgainstFulfillingCluster(createIndex)); + + bulkRequest = new Request("POST", "/_bulk?refresh=true"); + bulkRequest.setJsonEntity(""" + { "index": {"_index": "squares"}} + { "num": 1, "square": 1 } + { "index": {"_index": "squares"}} + { "num": 2, "square": 4 } + { "index": {"_index": "squares"}} + { "num": 3, "square": 9 } + { "index": {"_index": "squares"}} + { "num": 4, "square": 16 } + """); + assertOK(performRequestAgainstFulfillingCluster(bulkRequest)); + } + + private Request esqlRequest(String query) throws IOException { + XContentBuilder body = JsonXContent.contentBuilder(); + + body.startObject(); + body.field("query", query); + body.field("include_ccs_metadata", true); + body.endObject(); + + Request request = new Request("POST", "_query"); + request.setJsonEntity(Strings.toString(body)); + + return request; + } +} From b7167b73e377f7d42f56646b18908eaa7069a79f Mon Sep 17 00:00:00 2001 From: Jake Landis Date: Tue, 12 Nov 2024 09:13:37 -0600 Subject: [PATCH 13/16] Docs for monitor_stats privilege (#116533) This commit adds docs for monitor_stats and updates an example snippet to include both remote_indices and remote_cluster. --- .../security/bulk-create-roles.asciidoc | 4 +++- .../rest-api/security/create-roles.asciidoc | 22 ++++++++++++++----- .../authorization/managing-roles.asciidoc | 8 +++---- .../authorization/privileges.asciidoc | 5 +++++ 4 files changed, 27 insertions(+), 12 deletions(-) diff --git a/docs/reference/rest-api/security/bulk-create-roles.asciidoc b/docs/reference/rest-api/security/bulk-create-roles.asciidoc index 560e8b74cdd2c..37f49f2445770 100644 --- a/docs/reference/rest-api/security/bulk-create-roles.asciidoc +++ b/docs/reference/rest-api/security/bulk-create-roles.asciidoc @@ -102,7 +102,9 @@ They have no effect for remote clusters configured with the <> can be used to determine +which privileges are allowed per version. For more information, see <>. diff --git a/docs/reference/rest-api/security/create-roles.asciidoc b/docs/reference/rest-api/security/create-roles.asciidoc index a1ab892330e67..d23b9f06e2d87 100644 --- a/docs/reference/rest-api/security/create-roles.asciidoc +++ b/docs/reference/rest-api/security/create-roles.asciidoc @@ -105,7 +105,9 @@ They have no effect for remote clusters configured with the <> can be used to determine +which privileges are allowed per version. For more information, see <>. @@ -176,21 +178,29 @@ POST /_security/role/cli_or_drivers_minimal -------------------------------------------------- // end::sql-queries-permission[] -The following example configures a role with remote indices privileges on a remote cluster: +The following example configures a role with remote indices and remote cluster privileges for a remote cluster: [source,console] -------------------------------------------------- -POST /_security/role/role_with_remote_indices +POST /_security/role/only_remote_access_role { "remote_indices": [ { - "clusters": [ "my_remote" ], <1> + "clusters": ["my_remote"], <1> "names": ["logs*"], <2> "privileges": ["read", "read_cross_cluster", "view_index_metadata"] <3> } + ], + "remote_cluster": [ + { + "clusters": ["my_remote"], <1> + "privileges": ["monitor_stats"] <4> + } ] } -------------------------------------------------- -<1> The remote indices privileges apply to remote cluster with the alias `my_remote`. -<2> Privileges are granted for indices matching pattern `logs*` on the remote cluster ( `my_remote`). +<1> The remote indices and remote cluster privileges apply to remote cluster with the alias `my_remote`. +<2> Privileges are granted for indices matching pattern `logs*` on the remote cluster (`my_remote`). <3> The actual <> granted for `logs*` on `my_remote`. +<4> The actual <> granted for `my_remote`. +Note - only a subset of the cluster privileges are supported for remote clusters. diff --git a/docs/reference/security/authorization/managing-roles.asciidoc b/docs/reference/security/authorization/managing-roles.asciidoc index 535d70cbc5e9c..0c3f520605f07 100644 --- a/docs/reference/security/authorization/managing-roles.asciidoc +++ b/docs/reference/security/authorization/managing-roles.asciidoc @@ -249,12 +249,10 @@ The following describes the structure of a remote cluster permissions entry: <> and <>. This field is required. <2> The cluster level privileges for the remote cluster. The allowed values here are a subset of the -<>. This field is required. +<>. +The <> can be used to determine +which privileges are allowed here. This field is required. -The `monitor_enrich` privilege for remote clusters was introduced in version -8.15.0. Currently, this is the only privilege available for remote clusters and -is required to enable users to use the `ENRICH` keyword in ES|QL queries across -clusters. ==== Example diff --git a/docs/reference/security/authorization/privileges.asciidoc b/docs/reference/security/authorization/privileges.asciidoc index 747b1eef40441..3b69e5c1ba984 100644 --- a/docs/reference/security/authorization/privileges.asciidoc +++ b/docs/reference/security/authorization/privileges.asciidoc @@ -250,6 +250,11 @@ Privileges to list and view details on existing repositories and snapshots. + This privilege is not available in {serverless-full}. +`monitor_stats`:: +Privileges to list and view details of stats. ++ +This privilege is not available in {serverless-full}. + `monitor_text_structure`:: All read-only operations related to the <>. + From a71c132481217d2a803cc493da903d14076c9e60 Mon Sep 17 00:00:00 2001 From: Liam Thompson <32779855+leemthompo@users.noreply.github.com> Date: Tue, 12 Nov 2024 16:14:02 +0100 Subject: [PATCH 14/16] [DOCS] Update sharepoint-online connector perms (#116641) --- .../docs/connectors-sharepoint-online.asciidoc | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/reference/connector/docs/connectors-sharepoint-online.asciidoc b/docs/reference/connector/docs/connectors-sharepoint-online.asciidoc index 95ff8223b4d20..21d0890e436c5 100644 --- a/docs/reference/connector/docs/connectors-sharepoint-online.asciidoc +++ b/docs/reference/connector/docs/connectors-sharepoint-online.asciidoc @@ -87,14 +87,16 @@ Select an expiration date. (At this expiration date, you will need to generate a + ``` Graph API -- Sites.Read.All +- Sites.Selected - Files.Read.All - Group.Read.All - User.Read.All Sharepoint -- Sites.Read.All +- Sites.Selected ``` +NOTE: If the `Comma-separated list of sites` configuration is set to `*` or if a user enables the toggle button `Enumerate all sites`, the connector requires `Sites.Read.All` permission. + * **Grant admin consent**, using the `Grant Admin Consent` link from the permissions screen. * Save the tenant name (i.e. Domain name) of Azure platform. @@ -138,7 +140,7 @@ Refer to https://learn.microsoft.com/en-us/sharepoint/dev/general-development/ho Here's a summary of why we use these Graph API permissions: -* *Sites.Read.All* is used to fetch the sites and their metadata +* *Sites.Selected* is used to fetch the sites and their metadata * *Files.Read.All* is used to fetch Site Drives and files in these drives * *Groups.Read.All* is used to fetch groups for document-level permissions * *User.Read.All* is used to fetch user information for document-level permissions @@ -546,14 +548,16 @@ Select an expiration date. (At this expiration date, you will need to generate a + ``` Graph API -- Sites.Read.All +- Sites.Selected - Files.Read.All - Group.Read.All - User.Read.All Sharepoint -- Sites.Read.All +- Sites.Selected ``` +NOTE: If the `Comma-separated list of sites` configuration is set to `*` or if a user enables the toggle button `Enumerate all sites`, the connector requires `Sites.Read.All` permission. + * **Grant admin consent**, using the `Grant Admin Consent` link from the permissions screen. * Save the tenant name (i.e. Domain name) of Azure platform. @@ -597,7 +601,7 @@ Refer to https://learn.microsoft.com/en-us/sharepoint/dev/general-development/ho Here's a summary of why we use these Graph API permissions: -* *Sites.Read.All* is used to fetch the sites and their metadata +* *Sites.Selected* is used to fetch the sites and their metadata * *Files.Read.All* is used to fetch Site Drives and files in these drives * *Groups.Read.All* is used to fetch groups for document-level permissions * *User.Read.All* is used to fetch user information for document-level permissions From 7039a1dc8c886e23fda47a4b38cbab72746ac8cf Mon Sep 17 00:00:00 2001 From: Ying Mao Date: Tue, 12 Nov 2024 10:26:13 -0500 Subject: [PATCH 15/16] Adds support for `input_type` field to Vertex inference service (#116431) * Adding input type to google vertex ai service * Update docs/changelog/116431.yaml * PR feedback - backwards compatibility * Fix lint error --- docs/changelog/116431.yaml | 5 + .../org/elasticsearch/TransportVersions.java | 1 + .../GoogleVertexAiActionCreator.java | 6 +- .../GoogleVertexAiActionVisitor.java | 3 +- .../GoogleVertexAiEmbeddingsRequest.java | 2 +- ...GoogleVertexAiEmbeddingsRequestEntity.java | 37 +++- .../googlevertexai/GoogleVertexAiModel.java | 18 +- .../googlevertexai/GoogleVertexAiService.java | 4 +- .../GoogleVertexAiEmbeddingsModel.java | 51 +++++- ...VertexAiEmbeddingsRequestTaskSettings.java | 27 ++- .../GoogleVertexAiEmbeddingsTaskSettings.java | 99 ++++++++-- .../rerank/GoogleVertexAiRerankModel.java | 9 +- ...eVertexAiEmbeddingsRequestEntityTests.java | 96 ++++++++-- .../GoogleVertexAiEmbeddingsRequestTests.java | 36 +++- .../GoogleVertexAiServiceTests.java | 90 +++++++--- .../GoogleVertexAiEmbeddingsModelTests.java | 104 ++++++++++- ...xAiEmbeddingsRequestTaskSettingsTests.java | 45 ++++- ...leVertexAiEmbeddingsTaskSettingsTests.java | 170 ++++++++++++++++-- 18 files changed, 697 insertions(+), 106 deletions(-) create mode 100644 docs/changelog/116431.yaml diff --git a/docs/changelog/116431.yaml b/docs/changelog/116431.yaml new file mode 100644 index 0000000000000..50c6baf1d01c7 --- /dev/null +++ b/docs/changelog/116431.yaml @@ -0,0 +1,5 @@ +pr: 116431 +summary: Adds support for `input_type` field to Vertex inference service +area: Machine Learning +type: enhancement +issues: [] diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java index 5f3b466f9f7bd..6e62845383a14 100644 --- a/server/src/main/java/org/elasticsearch/TransportVersions.java +++ b/server/src/main/java/org/elasticsearch/TransportVersions.java @@ -193,6 +193,7 @@ static TransportVersion def(int id) { public static final TransportVersion ROLE_MONITOR_STATS = def(8_787_00_0); public static final TransportVersion DATA_STREAM_INDEX_VERSION_DEPRECATION_CHECK = def(8_788_00_0); public static final TransportVersion ADD_COMPATIBILITY_VERSIONS_TO_NODE_INFO = def(8_789_00_0); + public static final TransportVersion VERTEX_AI_INPUT_TYPE_ADDED = def(8_790_00_0); /* * STOP! READ THIS FIRST! No, really, diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionCreator.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionCreator.java index 27b3ae95f1aa4..99f535f81485c 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionCreator.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionCreator.java @@ -7,6 +7,7 @@ package org.elasticsearch.xpack.inference.external.action.googlevertexai; +import org.elasticsearch.inference.InputType; import org.elasticsearch.xpack.inference.external.action.ExecutableAction; import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction; import org.elasticsearch.xpack.inference.external.http.sender.GoogleVertexAiEmbeddingsRequestManager; @@ -33,9 +34,10 @@ public GoogleVertexAiActionCreator(Sender sender, ServiceComponents serviceCompo } @Override - public ExecutableAction create(GoogleVertexAiEmbeddingsModel model, Map taskSettings) { + public ExecutableAction create(GoogleVertexAiEmbeddingsModel model, Map taskSettings, InputType inputType) { + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, taskSettings, inputType); var requestManager = new GoogleVertexAiEmbeddingsRequestManager( - model, + overriddenModel, serviceComponents.truncator(), serviceComponents.threadPool() ); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionVisitor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionVisitor.java index def8f09ce06be..2b5cd5854c8ab 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionVisitor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionVisitor.java @@ -7,6 +7,7 @@ package org.elasticsearch.xpack.inference.external.action.googlevertexai; +import org.elasticsearch.inference.InputType; import org.elasticsearch.xpack.inference.external.action.ExecutableAction; import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsModel; import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel; @@ -15,7 +16,7 @@ public interface GoogleVertexAiActionVisitor { - ExecutableAction create(GoogleVertexAiEmbeddingsModel model, Map taskSettings); + ExecutableAction create(GoogleVertexAiEmbeddingsModel model, Map taskSettings, InputType inputType); ExecutableAction create(GoogleVertexAiRerankModel model, Map taskSettings); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequest.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequest.java index c0e36baf2e98f..75320bc762c8b 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequest.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequest.java @@ -40,7 +40,7 @@ public HttpRequest createHttpRequest() { HttpPost httpPost = new HttpPost(model.uri()); ByteArrayEntity byteEntity = new ByteArrayEntity( - Strings.toString(new GoogleVertexAiEmbeddingsRequestEntity(truncationResult.input(), model.getTaskSettings().autoTruncate())) + Strings.toString(new GoogleVertexAiEmbeddingsRequestEntity(truncationResult.input(), model.getTaskSettings())) .getBytes(StandardCharsets.UTF_8) ); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntity.java index 2fae999599ba2..fc33df0d63acd 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntity.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntity.java @@ -7,23 +7,35 @@ package org.elasticsearch.xpack.inference.external.request.googlevertexai; -import org.elasticsearch.core.Nullable; +import org.elasticsearch.inference.InputType; import org.elasticsearch.xcontent.ToXContentObject; import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings; import java.io.IOException; import java.util.List; import java.util.Objects; -public record GoogleVertexAiEmbeddingsRequestEntity(List inputs, @Nullable Boolean autoTruncation) implements ToXContentObject { +import static org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsTaskSettings.invalidInputTypeMessage; + +public record GoogleVertexAiEmbeddingsRequestEntity(List inputs, GoogleVertexAiEmbeddingsTaskSettings taskSettings) + implements + ToXContentObject { private static final String INSTANCES_FIELD = "instances"; private static final String CONTENT_FIELD = "content"; private static final String PARAMETERS_FIELD = "parameters"; private static final String AUTO_TRUNCATE_FIELD = "autoTruncate"; + private static final String TASK_TYPE_FIELD = "task_type"; + + private static final String CLASSIFICATION_TASK_TYPE = "CLASSIFICATION"; + private static final String CLUSTERING_TASK_TYPE = "CLUSTERING"; + private static final String RETRIEVAL_DOCUMENT_TASK_TYPE = "RETRIEVAL_DOCUMENT"; + private static final String RETRIEVAL_QUERY_TASK_TYPE = "RETRIEVAL_QUERY"; public GoogleVertexAiEmbeddingsRequestEntity { Objects.requireNonNull(inputs); + Objects.requireNonNull(taskSettings); } @Override @@ -35,16 +47,20 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.startObject(); { builder.field(CONTENT_FIELD, input); + + if (taskSettings.getInputType() != null) { + builder.field(TASK_TYPE_FIELD, convertToString(taskSettings.getInputType())); + } } builder.endObject(); } builder.endArray(); - if (autoTruncation != null) { + if (taskSettings.autoTruncate() != null) { builder.startObject(PARAMETERS_FIELD); { - builder.field(AUTO_TRUNCATE_FIELD, autoTruncation); + builder.field(AUTO_TRUNCATE_FIELD, taskSettings.autoTruncate()); } builder.endObject(); } @@ -52,4 +68,17 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws return builder; } + + static String convertToString(InputType inputType) { + return switch (inputType) { + case INGEST -> RETRIEVAL_DOCUMENT_TASK_TYPE; + case SEARCH -> RETRIEVAL_QUERY_TASK_TYPE; + case CLASSIFICATION -> CLASSIFICATION_TASK_TYPE; + case CLUSTERING -> CLUSTERING_TASK_TYPE; + default -> { + assert false : invalidInputTypeMessage(inputType); + yield null; + } + }; + } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiModel.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiModel.java index 17e6ec2152e7e..caa244f8af4f2 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiModel.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiModel.java @@ -7,13 +7,16 @@ package org.elasticsearch.xpack.inference.services.googlevertexai; +import org.elasticsearch.inference.InputType; import org.elasticsearch.inference.Model; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.ModelSecrets; import org.elasticsearch.inference.ServiceSettings; +import org.elasticsearch.inference.TaskSettings; import org.elasticsearch.xpack.inference.external.action.ExecutableAction; import org.elasticsearch.xpack.inference.external.action.googlevertexai.GoogleVertexAiActionVisitor; +import java.net.URI; import java.util.Map; import java.util.Objects; @@ -21,6 +24,8 @@ public abstract class GoogleVertexAiModel extends Model { private final GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings; + protected URI uri; + public GoogleVertexAiModel( ModelConfigurations configurations, ModelSecrets secrets, @@ -34,13 +39,24 @@ public GoogleVertexAiModel( public GoogleVertexAiModel(GoogleVertexAiModel model, ServiceSettings serviceSettings) { super(model, serviceSettings); + uri = model.uri(); + rateLimitServiceSettings = model.rateLimitServiceSettings(); + } + + public GoogleVertexAiModel(GoogleVertexAiModel model, TaskSettings taskSettings) { + super(model, taskSettings); + + uri = model.uri(); rateLimitServiceSettings = model.rateLimitServiceSettings(); } - public abstract ExecutableAction accept(GoogleVertexAiActionVisitor creator, Map taskSettings); + public abstract ExecutableAction accept(GoogleVertexAiActionVisitor creator, Map taskSettings, InputType inputType); public GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings() { return rateLimitServiceSettings; } + public URI uri() { + return uri; + } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java index 0b4da10e7130f..a05b1a937d376 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java @@ -210,7 +210,7 @@ protected void doInfer( var actionCreator = new GoogleVertexAiActionCreator(getSender(), getServiceComponents()); - var action = googleVertexAiModel.accept(actionCreator, taskSettings); + var action = googleVertexAiModel.accept(actionCreator, taskSettings, inputType); action.execute(inputs, timeout, listener); } @@ -235,7 +235,7 @@ protected void doChunkedInfer( ).batchRequestsWithListeners(listener); for (var request : batchedRequests) { - var action = googleVertexAiModel.accept(actionCreator, taskSettings); + var action = googleVertexAiModel.accept(actionCreator, taskSettings, inputType); action.execute(new DocumentsOnlyInput(request.batch().inputs()), timeout, request.listener()); } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModel.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModel.java index 1df8ee937497a..a5acbb80b76ec 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModel.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModel.java @@ -11,12 +11,14 @@ import org.elasticsearch.common.util.LazyInitializable; import org.elasticsearch.core.Nullable; import org.elasticsearch.inference.ChunkingSettings; +import org.elasticsearch.inference.InputType; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.ModelSecrets; import org.elasticsearch.inference.SettingsConfiguration; import org.elasticsearch.inference.TaskType; import org.elasticsearch.inference.configuration.SettingsConfigurationDisplayType; import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType; +import org.elasticsearch.inference.configuration.SettingsConfigurationSelectOption; import org.elasticsearch.xpack.inference.external.action.ExecutableAction; import org.elasticsearch.xpack.inference.external.action.googlevertexai.GoogleVertexAiActionVisitor; import org.elasticsearch.xpack.inference.external.request.googlevertexai.GoogleVertexAiUtils; @@ -29,13 +31,25 @@ import java.util.Collections; import java.util.HashMap; import java.util.Map; +import java.util.stream.Stream; import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings.AUTO_TRUNCATE; +import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE; public class GoogleVertexAiEmbeddingsModel extends GoogleVertexAiModel { - private URI uri; + public static GoogleVertexAiEmbeddingsModel of( + GoogleVertexAiEmbeddingsModel model, + Map taskSettings, + InputType inputType + ) { + var requestTaskSettings = GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap(taskSettings); + return new GoogleVertexAiEmbeddingsModel( + model, + GoogleVertexAiEmbeddingsTaskSettings.of(model.getTaskSettings(), requestTaskSettings, inputType) + ); + } public GoogleVertexAiEmbeddingsModel( String inferenceEntityId, @@ -62,6 +76,10 @@ public GoogleVertexAiEmbeddingsModel(GoogleVertexAiEmbeddingsModel model, Google super(model, serviceSettings); } + public GoogleVertexAiEmbeddingsModel(GoogleVertexAiEmbeddingsModel model, GoogleVertexAiEmbeddingsTaskSettings taskSettings) { + super(model, taskSettings); + } + // Should only be used directly for testing GoogleVertexAiEmbeddingsModel( String inferenceEntityId, @@ -126,13 +144,9 @@ public GoogleVertexAiEmbeddingsRateLimitServiceSettings rateLimitServiceSettings return (GoogleVertexAiEmbeddingsRateLimitServiceSettings) super.rateLimitServiceSettings(); } - public URI uri() { - return uri; - } - @Override - public ExecutableAction accept(GoogleVertexAiActionVisitor visitor, Map taskSettings) { - return visitor.create(this, taskSettings); + public ExecutableAction accept(GoogleVertexAiActionVisitor visitor, Map taskSettings, InputType inputType) { + return visitor.create(this, taskSettings, inputType); } public static URI buildUri(String location, String projectId, String modelId) throws URISyntaxException { @@ -161,11 +175,32 @@ public static Map get() { new LazyInitializable<>(() -> { var configurationMap = new HashMap(); + configurationMap.put( + INPUT_TYPE, + new SettingsConfiguration.Builder().setDisplay(SettingsConfigurationDisplayType.DROPDOWN) + .setLabel("Input Type") + .setOrder(1) + .setRequired(false) + .setSensitive(false) + .setTooltip("Specifies the type of input passed to the model.") + .setType(SettingsConfigurationFieldType.STRING) + .setOptions( + Stream.of( + InputType.CLASSIFICATION.toString(), + InputType.CLUSTERING.toString(), + InputType.INGEST.toString(), + InputType.SEARCH.toString() + ).map(v -> new SettingsConfigurationSelectOption.Builder().setLabelAndValue(v).build()).toList() + ) + .setValue("") + .build() + ); + configurationMap.put( AUTO_TRUNCATE, new SettingsConfiguration.Builder().setDisplay(SettingsConfigurationDisplayType.TOGGLE) .setLabel("Auto Truncate") - .setOrder(1) + .setOrder(2) .setRequired(false) .setSensitive(false) .setTooltip("Specifies if the API truncates inputs longer than the maximum token length automatically.") diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettings.java index 14a67a64377e2..e39c423582151 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettings.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettings.java @@ -9,29 +9,46 @@ import org.elasticsearch.common.ValidationException; import org.elasticsearch.core.Nullable; +import org.elasticsearch.inference.InputType; +import org.elasticsearch.inference.ModelConfigurations; import java.util.Map; import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalBoolean; +import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum; +import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE; +import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings.VALID_REQUEST_VALUES; -public record GoogleVertexAiEmbeddingsRequestTaskSettings(@Nullable Boolean autoTruncate) { +public record GoogleVertexAiEmbeddingsRequestTaskSettings(@Nullable Boolean autoTruncate, @Nullable InputType inputType) { - public static final GoogleVertexAiEmbeddingsRequestTaskSettings EMPTY_SETTINGS = new GoogleVertexAiEmbeddingsRequestTaskSettings(null); + public static final GoogleVertexAiEmbeddingsRequestTaskSettings EMPTY_SETTINGS = new GoogleVertexAiEmbeddingsRequestTaskSettings( + null, + null + ); public static GoogleVertexAiEmbeddingsRequestTaskSettings fromMap(Map map) { - if (map.isEmpty()) { - return GoogleVertexAiEmbeddingsRequestTaskSettings.EMPTY_SETTINGS; + if (map == null || map.isEmpty()) { + return EMPTY_SETTINGS; } ValidationException validationException = new ValidationException(); + InputType inputType = extractOptionalEnum( + map, + INPUT_TYPE, + ModelConfigurations.TASK_SETTINGS, + InputType::fromString, + VALID_REQUEST_VALUES, + validationException + ); + Boolean autoTruncate = extractOptionalBoolean(map, GoogleVertexAiEmbeddingsTaskSettings.AUTO_TRUNCATE, validationException); if (validationException.validationErrors().isEmpty() == false) { throw validationException; } - return new GoogleVertexAiEmbeddingsRequestTaskSettings(autoTruncate); + return new GoogleVertexAiEmbeddingsRequestTaskSettings(autoTruncate, inputType); } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettings.java index dcdbbda33575f..9b759a4661bce 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettings.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettings.java @@ -9,19 +9,24 @@ import org.elasticsearch.TransportVersion; import org.elasticsearch.TransportVersions; +import org.elasticsearch.common.Strings; import org.elasticsearch.common.ValidationException; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.core.Nullable; +import org.elasticsearch.inference.InputType; +import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.TaskSettings; import org.elasticsearch.xcontent.XContentBuilder; import java.io.IOException; +import java.util.EnumSet; import java.util.HashMap; import java.util.Map; import java.util.Objects; import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalBoolean; +import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum; public class GoogleVertexAiEmbeddingsTaskSettings implements TaskSettings { @@ -29,48 +34,108 @@ public class GoogleVertexAiEmbeddingsTaskSettings implements TaskSettings { public static final String AUTO_TRUNCATE = "auto_truncate"; - public static final GoogleVertexAiEmbeddingsTaskSettings EMPTY_SETTINGS = new GoogleVertexAiEmbeddingsTaskSettings( - Boolean.valueOf(null) + public static final String INPUT_TYPE = "input_type"; + + static final EnumSet VALID_REQUEST_VALUES = EnumSet.of( + InputType.INGEST, + InputType.SEARCH, + InputType.CLASSIFICATION, + InputType.CLUSTERING ); + public static final GoogleVertexAiEmbeddingsTaskSettings EMPTY_SETTINGS = new GoogleVertexAiEmbeddingsTaskSettings(null, null); + public static GoogleVertexAiEmbeddingsTaskSettings fromMap(Map map) { + if (map == null || map.isEmpty()) { + return EMPTY_SETTINGS; + } + ValidationException validationException = new ValidationException(); + InputType inputType = extractOptionalEnum( + map, + INPUT_TYPE, + ModelConfigurations.TASK_SETTINGS, + InputType::fromString, + VALID_REQUEST_VALUES, + validationException + ); + Boolean autoTruncate = extractOptionalBoolean(map, AUTO_TRUNCATE, validationException); if (validationException.validationErrors().isEmpty() == false) { throw validationException; } - return new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate); + return new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, inputType); } public static GoogleVertexAiEmbeddingsTaskSettings of( GoogleVertexAiEmbeddingsTaskSettings originalSettings, - GoogleVertexAiEmbeddingsRequestTaskSettings requestSettings + GoogleVertexAiEmbeddingsRequestTaskSettings requestSettings, + InputType requestInputType ) { + var inputTypeToUse = getValidInputType(originalSettings, requestSettings, requestInputType); var autoTruncate = requestSettings.autoTruncate() == null ? originalSettings.autoTruncate : requestSettings.autoTruncate(); - return new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate); + return new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, inputTypeToUse); + } + + private static InputType getValidInputType( + GoogleVertexAiEmbeddingsTaskSettings originalSettings, + GoogleVertexAiEmbeddingsRequestTaskSettings requestTaskSettings, + InputType requestInputType + ) { + InputType inputTypeToUse = originalSettings.inputType; + + if (VALID_REQUEST_VALUES.contains(requestInputType)) { + inputTypeToUse = requestInputType; + } else if (requestTaskSettings.inputType() != null) { + inputTypeToUse = requestTaskSettings.inputType(); + } + + return inputTypeToUse; } + private final InputType inputType; private final Boolean autoTruncate; - public GoogleVertexAiEmbeddingsTaskSettings(@Nullable Boolean autoTruncate) { + public GoogleVertexAiEmbeddingsTaskSettings(@Nullable Boolean autoTruncate, @Nullable InputType inputType) { + validateInputType(inputType); + this.inputType = inputType; this.autoTruncate = autoTruncate; } public GoogleVertexAiEmbeddingsTaskSettings(StreamInput in) throws IOException { this.autoTruncate = in.readOptionalBoolean(); + + var inputType = (in.getTransportVersion().onOrAfter(TransportVersions.VERTEX_AI_INPUT_TYPE_ADDED)) + ? in.readOptionalEnum(InputType.class) + : null; + + validateInputType(inputType); + this.inputType = inputType; + } + + private static void validateInputType(InputType inputType) { + if (inputType == null) { + return; + } + + assert VALID_REQUEST_VALUES.contains(inputType) : invalidInputTypeMessage(inputType); } @Override public boolean isEmpty() { - return autoTruncate == null; + return inputType == null && autoTruncate == null; } public Boolean autoTruncate() { return autoTruncate; } + public InputType getInputType() { + return inputType; + } + @Override public String getWriteableName() { return NAME; @@ -84,11 +149,19 @@ public TransportVersion getMinimalSupportedVersion() { @Override public void writeTo(StreamOutput out) throws IOException { out.writeOptionalBoolean(this.autoTruncate); + + if (out.getTransportVersion().onOrAfter(TransportVersions.VERTEX_AI_INPUT_TYPE_ADDED)) { + out.writeOptionalEnum(this.inputType); + } } @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); + if (inputType != null) { + builder.field(INPUT_TYPE, inputType); + } + if (autoTruncate != null) { builder.field(AUTO_TRUNCATE, autoTruncate); } @@ -101,19 +174,23 @@ public boolean equals(Object object) { if (this == object) return true; if (object == null || getClass() != object.getClass()) return false; GoogleVertexAiEmbeddingsTaskSettings that = (GoogleVertexAiEmbeddingsTaskSettings) object; - return Objects.equals(autoTruncate, that.autoTruncate); + return Objects.equals(inputType, that.inputType) && Objects.equals(autoTruncate, that.autoTruncate); } @Override public int hashCode() { - return Objects.hash(autoTruncate); + return Objects.hash(autoTruncate, inputType); + } + + public static String invalidInputTypeMessage(InputType inputType) { + return Strings.format("received invalid input type value [%s]", inputType.toString()); } @Override public TaskSettings updatedTaskSettings(Map newSettings) { - GoogleVertexAiEmbeddingsRequestTaskSettings requestSettings = GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap( + GoogleVertexAiEmbeddingsRequestTaskSettings updatedSettings = GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap( new HashMap<>(newSettings) ); - return of(this, requestSettings); + return of(this, updatedSettings, updatedSettings.inputType() != null ? updatedSettings.inputType() : this.inputType); } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/rerank/GoogleVertexAiRerankModel.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/rerank/GoogleVertexAiRerankModel.java index 3f9c4f7a66560..e73d8d2e2613a 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/rerank/GoogleVertexAiRerankModel.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/rerank/GoogleVertexAiRerankModel.java @@ -10,6 +10,7 @@ import org.apache.http.client.utils.URIBuilder; import org.elasticsearch.common.util.LazyInitializable; import org.elasticsearch.core.Nullable; +import org.elasticsearch.inference.InputType; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.ModelSecrets; import org.elasticsearch.inference.SettingsConfiguration; @@ -34,8 +35,6 @@ public class GoogleVertexAiRerankModel extends GoogleVertexAiModel { - private URI uri; - public GoogleVertexAiRerankModel( String inferenceEntityId, TaskType taskType, @@ -122,12 +121,8 @@ public GoogleDiscoveryEngineRateLimitServiceSettings rateLimitServiceSettings() return (GoogleDiscoveryEngineRateLimitServiceSettings) super.rateLimitServiceSettings(); } - public URI uri() { - return uri; - } - @Override - public ExecutableAction accept(GoogleVertexAiActionVisitor visitor, Map taskSettings) { + public ExecutableAction accept(GoogleVertexAiActionVisitor visitor, Map taskSettings, InputType inputType) { return visitor.create(this, taskSettings); } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntityTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntityTests.java index f4912e0862e60..18ae7425aaaf2 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntityTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntityTests.java @@ -8,10 +8,12 @@ package org.elasticsearch.xpack.inference.external.request.googlevertexai; import org.elasticsearch.common.Strings; +import org.elasticsearch.inference.InputType; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentFactory; import org.elasticsearch.xcontent.XContentType; +import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings; import java.io.IOException; import java.util.List; @@ -20,8 +22,11 @@ public class GoogleVertexAiEmbeddingsRequestEntityTests extends ESTestCase { - public void testToXContent_SingleEmbeddingRequest_WritesAutoTruncationIfDefined() throws IOException { - var entity = new GoogleVertexAiEmbeddingsRequestEntity(List.of("abc"), true); + public void testToXContent_SingleEmbeddingRequest_WritesAllFields() throws IOException { + var entity = new GoogleVertexAiEmbeddingsRequestEntity( + List.of("abc"), + new GoogleVertexAiEmbeddingsTaskSettings(true, InputType.SEARCH) + ); XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); entity.toXContent(builder, null); @@ -31,7 +36,8 @@ public void testToXContent_SingleEmbeddingRequest_WritesAutoTruncationIfDefined( { "instances": [ { - "content": "abc" + "content": "abc", + "task_type": "RETRIEVAL_QUERY" } ], "parameters": { @@ -42,7 +48,10 @@ public void testToXContent_SingleEmbeddingRequest_WritesAutoTruncationIfDefined( } public void testToXContent_SingleEmbeddingRequest_DoesNotWriteAutoTruncationIfNotDefined() throws IOException { - var entity = new GoogleVertexAiEmbeddingsRequestEntity(List.of("abc"), null); + var entity = new GoogleVertexAiEmbeddingsRequestEntity( + List.of("abc"), + new GoogleVertexAiEmbeddingsTaskSettings(null, InputType.INGEST) + ); XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); entity.toXContent(builder, null); @@ -52,15 +61,16 @@ public void testToXContent_SingleEmbeddingRequest_DoesNotWriteAutoTruncationIfNo { "instances": [ { - "content": "abc" + "content": "abc", + "task_type": "RETRIEVAL_DOCUMENT" } ] } """)); } - public void testToXContent_MultipleEmbeddingsRequest_WritesAutoTruncationIfDefined() throws IOException { - var entity = new GoogleVertexAiEmbeddingsRequestEntity(List.of("abc", "def"), true); + public void testToXContent_SingleEmbeddingRequest_DoesNotWriteInputTypeIfNotDefined() throws IOException { + var entity = new GoogleVertexAiEmbeddingsRequestEntity(List.of("abc"), new GoogleVertexAiEmbeddingsTaskSettings(false, null)); XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); entity.toXContent(builder, null); @@ -71,9 +81,35 @@ public void testToXContent_MultipleEmbeddingsRequest_WritesAutoTruncationIfDefin "instances": [ { "content": "abc" + } + ], + "parameters": { + "autoTruncate": false + } + } + """)); + } + + public void testToXContent_MultipleEmbeddingsRequest_WritesAllFields() throws IOException { + var entity = new GoogleVertexAiEmbeddingsRequestEntity( + List.of("abc", "def"), + new GoogleVertexAiEmbeddingsTaskSettings(true, InputType.CLUSTERING) + ); + + XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); + entity.toXContent(builder, null); + String xContentResult = Strings.toString(builder); + + assertThat(xContentResult, equalToIgnoringWhitespaceInJsonString(""" + { + "instances": [ + { + "content": "abc", + "task_type": "CLUSTERING" }, { - "content": "def" + "content": "def", + "task_type": "CLUSTERING" } ], "parameters": { @@ -83,8 +119,8 @@ public void testToXContent_MultipleEmbeddingsRequest_WritesAutoTruncationIfDefin """)); } - public void testToXContent_MultipleEmbeddingsRequest_DoesNotWriteAutoTruncationIfNotDefined() throws IOException { - var entity = new GoogleVertexAiEmbeddingsRequestEntity(List.of("abc", "def"), null); + public void testToXContent_MultipleEmbeddingsRequest_DoesNotWriteInputTypeIfNotDefined() throws IOException { + var entity = new GoogleVertexAiEmbeddingsRequestEntity(List.of("abc", "def"), new GoogleVertexAiEmbeddingsTaskSettings(true, null)); XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); entity.toXContent(builder, null); @@ -99,8 +135,48 @@ public void testToXContent_MultipleEmbeddingsRequest_DoesNotWriteAutoTruncationI { "content": "def" } + ], + "parameters": { + "autoTruncate": true + } + } + """)); + } + + public void testToXContent_MultipleEmbeddingsRequest_DoesNotWriteAutoTruncationIfNotDefined() throws IOException { + var entity = new GoogleVertexAiEmbeddingsRequestEntity( + List.of("abc", "def"), + new GoogleVertexAiEmbeddingsTaskSettings(null, InputType.CLASSIFICATION) + ); + + XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); + entity.toXContent(builder, null); + String xContentResult = Strings.toString(builder); + + assertThat(xContentResult, equalToIgnoringWhitespaceInJsonString(""" + { + "instances": [ + { + "content": "abc", + "task_type": "CLASSIFICATION" + }, + { + "content": "def", + "task_type": "CLASSIFICATION" + } ] } """)); } + + public void testToXContent_ThrowsIfInputIsNull() { + expectThrows( + NullPointerException.class, + () -> new GoogleVertexAiEmbeddingsRequestEntity(null, new GoogleVertexAiEmbeddingsTaskSettings(null, InputType.CLASSIFICATION)) + ); + } + + public void testToXContent_ThrowsIfTaskSettingsIsNull() { + expectThrows(NullPointerException.class, () -> new GoogleVertexAiEmbeddingsRequestEntity(List.of("abc", "def"), null)); + } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestTests.java index b28fd8d3a0cf9..a26d3496bed6b 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestTests.java @@ -10,6 +10,7 @@ import org.apache.http.HttpHeaders; import org.apache.http.client.methods.HttpPost; import org.elasticsearch.core.Nullable; +import org.elasticsearch.inference.InputType; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.common.Truncator; @@ -31,11 +32,11 @@ public class GoogleVertexAiEmbeddingsRequestTests extends ESTestCase { private static final String AUTH_HEADER_VALUE = "foo"; - public void testCreateRequest_WithoutDimensionsSet_And_WithoutAutoTruncateSet() throws IOException { + public void testCreateRequest_WithoutDimensionsSet_And_WithoutAutoTruncateSet_And_WithoutInputTypeSet() throws IOException { var model = "model"; var input = "input"; - var request = createRequest(model, input, null); + var request = createRequest(model, input, null, null); var httpRequest = request.createHttpRequest(); assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class)); @@ -54,7 +55,7 @@ public void testCreateRequest_WithAutoTruncateSet() throws IOException { var input = "input"; var autoTruncate = true; - var request = createRequest(model, input, autoTruncate); + var request = createRequest(model, input, autoTruncate, null); var httpRequest = request.createHttpRequest(); assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class)); @@ -68,11 +69,29 @@ public void testCreateRequest_WithAutoTruncateSet() throws IOException { assertThat(requestMap, is(Map.of("instances", List.of(Map.of("content", "input")), "parameters", Map.of("autoTruncate", true)))); } + public void testCreateRequest_WithInputTypeSet() throws IOException { + var model = "model"; + var input = "input"; + + var request = createRequest(model, input, null, InputType.SEARCH); + var httpRequest = request.createHttpRequest(); + + assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class)); + var httpPost = (HttpPost) httpRequest.httpRequestBase(); + + assertThat(httpPost.getLastHeader(HttpHeaders.CONTENT_TYPE).getValue(), is(XContentType.JSON.mediaType())); + assertThat(httpPost.getLastHeader(HttpHeaders.AUTHORIZATION).getValue(), is(AUTH_HEADER_VALUE)); + + var requestMap = entityAsMap(httpPost.getEntity().getContent()); + assertThat(requestMap, aMapWithSize(1)); + assertThat(requestMap, is(Map.of("instances", List.of(Map.of("content", "input", "task_type", "RETRIEVAL_QUERY"))))); + } + public void testTruncate_ReducesInputTextSizeByHalf() throws IOException { var model = "model"; var input = "abcd"; - var request = createRequest(model, input, null); + var request = createRequest(model, input, null, null); var truncatedRequest = request.truncate(); var httpRequest = truncatedRequest.createHttpRequest(); @@ -87,8 +106,13 @@ public void testTruncate_ReducesInputTextSizeByHalf() throws IOException { assertThat(requestMap, is(Map.of("instances", List.of(Map.of("content", "ab"))))); } - private static GoogleVertexAiEmbeddingsRequest createRequest(String modelId, String input, @Nullable Boolean autoTruncate) { - var embeddingsModel = GoogleVertexAiEmbeddingsModelTests.createModel(modelId, autoTruncate); + private static GoogleVertexAiEmbeddingsRequest createRequest( + String modelId, + String input, + @Nullable Boolean autoTruncate, + @Nullable InputType inputType + ) { + var embeddingsModel = GoogleVertexAiEmbeddingsModelTests.createModel(modelId, autoTruncate, inputType); return new GoogleVertexAiEmbeddingsWithoutAuthRequest( TruncatorTests.createTruncator(), diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiServiceTests.java index 6f28301078853..906a825e49561 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiServiceTests.java @@ -13,8 +13,10 @@ import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.core.Nullable; import org.elasticsearch.inference.ChunkingSettings; import org.elasticsearch.inference.InferenceServiceConfiguration; +import org.elasticsearch.inference.InputType; import org.elasticsearch.inference.Model; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.TaskType; @@ -109,7 +111,7 @@ public void testParseRequestConfig_CreatesGoogleVertexAiEmbeddingsModel() throws projectId ) ), - new HashMap<>(Map.of()), + getTaskSettingsMap(true, InputType.INGEST), getSecretSettingsMap(serviceAccountJson) ), modelListener @@ -154,7 +156,7 @@ public void testParseRequestConfig_CreatesAGoogleVertexAiEmbeddingsModelWhenChun projectId ) ), - new HashMap<>(Map.of()), + getTaskSettingsMap(true, InputType.INGEST), createRandomChunkingSettingsMap(), getSecretSettingsMap(serviceAccountJson) ), @@ -200,7 +202,7 @@ public void testParseRequestConfig_CreatesAGoogleVertexAiEmbeddingsModelWhenChun projectId ) ), - new HashMap<>(Map.of()), + getTaskSettingsMap(false, InputType.SEARCH), getSecretSettingsMap(serviceAccountJson) ), modelListener @@ -281,7 +283,7 @@ public void testParseRequestConfig_ThrowsWhenAnExtraKeyExistsInConfig() throws I "project" ) ), - getTaskSettingsMap(true), + getTaskSettingsMap(true, InputType.SEARCH), getSecretSettingsMap("{}") ); config.put("extra_key", "value"); @@ -308,7 +310,7 @@ public void testParseRequestConfig_ThrowsWhenAnExtraKeyExistsInServiceSettingsMa ); serviceSettings.put("extra_key", "value"); - var config = getRequestConfigMap(serviceSettings, getTaskSettingsMap(true), getSecretSettingsMap("{}")); + var config = getRequestConfigMap(serviceSettings, getTaskSettingsMap(true, InputType.CLUSTERING), getSecretSettingsMap("{}")); var failureListener = getModelListenerForException( ElasticsearchStatusException.class, @@ -362,7 +364,7 @@ public void testParseRequestConfig_ThrowsWhenAnExtraKeyExistsInSecretSettingsMap "project" ) ), - getTaskSettingsMap(true), + getTaskSettingsMap(true, null), secretSettings ); @@ -399,7 +401,7 @@ public void testParsePersistedConfigWithSecrets_CreatesGoogleVertexAiEmbeddingsM true ) ), - getTaskSettingsMap(autoTruncate), + getTaskSettingsMap(autoTruncate, InputType.SEARCH), getSecretSettingsMap(serviceAccountJson) ); @@ -417,7 +419,7 @@ public void testParsePersistedConfigWithSecrets_CreatesGoogleVertexAiEmbeddingsM assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, InputType.SEARCH))); assertThat(embeddingsModel.getSecretSettings().serviceAccountJson().toString(), is(serviceAccountJson)); } } @@ -447,7 +449,7 @@ public void testParsePersistedConfigWithSecrets_CreatesAGoogleVertexAiEmbeddings true ) ), - getTaskSettingsMap(autoTruncate), + getTaskSettingsMap(autoTruncate, null), createRandomChunkingSettingsMap(), getSecretSettingsMap(serviceAccountJson) ); @@ -466,7 +468,7 @@ public void testParsePersistedConfigWithSecrets_CreatesAGoogleVertexAiEmbeddings assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, null))); assertThat(embeddingsModel.getConfigurations().getChunkingSettings(), instanceOf(ChunkingSettings.class)); assertThat(embeddingsModel.getSecretSettings().serviceAccountJson().toString(), is(serviceAccountJson)); } @@ -497,7 +499,7 @@ public void testParsePersistedConfigWithSecrets_CreatesAnEmbeddingsModelWhenChun true ) ), - getTaskSettingsMap(autoTruncate), + getTaskSettingsMap(autoTruncate, null), getSecretSettingsMap(serviceAccountJson) ); @@ -515,7 +517,7 @@ public void testParsePersistedConfigWithSecrets_CreatesAnEmbeddingsModelWhenChun assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, null))); assertThat(embeddingsModel.getConfigurations().getChunkingSettings(), instanceOf(ChunkingSettings.class)); assertThat(embeddingsModel.getSecretSettings().serviceAccountJson().toString(), is(serviceAccountJson)); } @@ -573,7 +575,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists true ) ), - getTaskSettingsMap(autoTruncate), + getTaskSettingsMap(autoTruncate, InputType.INGEST), getSecretSettingsMap(serviceAccountJson) ); persistedConfig.config().put("extra_key", "value"); @@ -592,7 +594,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, InputType.INGEST))); assertThat(embeddingsModel.getSecretSettings().serviceAccountJson().toString(), is(serviceAccountJson)); } } @@ -625,7 +627,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists true ) ), - getTaskSettingsMap(autoTruncate), + getTaskSettingsMap(autoTruncate, null), secretSettingsMap ); @@ -643,7 +645,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, null))); assertThat(embeddingsModel.getSecretSettings().serviceAccountJson().toString(), is(serviceAccountJson)); } } @@ -676,7 +678,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists var persistedConfig = getPersistedConfigMap( serviceSettingsMap, - getTaskSettingsMap(autoTruncate), + getTaskSettingsMap(autoTruncate, InputType.CLUSTERING), getSecretSettingsMap(serviceAccountJson) ); @@ -694,7 +696,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, InputType.CLUSTERING))); assertThat(embeddingsModel.getSecretSettings().serviceAccountJson().toString(), is(serviceAccountJson)); } } @@ -711,7 +713,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists """; try (var service = createGoogleVertexAiService()) { - var taskSettings = getTaskSettingsMap(autoTruncate); + var taskSettings = getTaskSettingsMap(autoTruncate, InputType.SEARCH); taskSettings.put("extra_key", "value"); var persistedConfig = getPersistedConfigMap( @@ -745,7 +747,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, InputType.SEARCH))); assertThat(embeddingsModel.getSecretSettings().serviceAccountJson().toString(), is(serviceAccountJson)); } } @@ -770,7 +772,7 @@ public void testParsePersistedConfig_CreatesAGoogleVertexAiEmbeddingsModelWhenCh true ) ), - getTaskSettingsMap(autoTruncate), + getTaskSettingsMap(autoTruncate, null), createRandomChunkingSettingsMap() ); @@ -783,7 +785,7 @@ public void testParsePersistedConfig_CreatesAGoogleVertexAiEmbeddingsModelWhenCh assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, null))); assertThat(embeddingsModel.getConfigurations().getChunkingSettings(), instanceOf(ChunkingSettings.class)); } } @@ -808,7 +810,7 @@ public void testParsePersistedConfig_CreatesAnEmbeddingsModelWhenChunkingSetting true ) ), - getTaskSettingsMap(autoTruncate) + getTaskSettingsMap(autoTruncate, null) ); var model = service.parsePersistedConfig("id", TaskType.TEXT_EMBEDDING, persistedConfig.config()); @@ -820,7 +822,7 @@ public void testParsePersistedConfig_CreatesAnEmbeddingsModelWhenChunkingSetting assertThat(embeddingsModel.getServiceSettings().location(), is(location)); assertThat(embeddingsModel.getServiceSettings().projectId(), is(projectId)); assertThat(embeddingsModel.getServiceSettings().dimensionsSetByUser(), is(Boolean.TRUE)); - assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(embeddingsModel.getTaskSettings(), is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, null))); assertThat(embeddingsModel.getConfigurations().getChunkingSettings(), instanceOf(ChunkingSettings.class)); } } @@ -838,12 +840,44 @@ public void testGetConfiguration() throws Exception { { "task_type": "text_embedding", "configuration": { + "input_type": { + "default_value": null, + "depends_on": [], + "display": "dropdown", + "label": "Input Type", + "options": [ + { + "label": "classification", + "value": "classification" + }, + { + "label": "clustering", + "value": "clustering" + }, + { + "label": "ingest", + "value": "ingest" + }, + { + "label": "search", + "value": "search" + } + ], + "order": 1, + "required": false, + "sensitive": false, + "tooltip": "Specifies the type of input passed to the model.", + "type": "str", + "ui_restrictions": [], + "validations": [], + "value": "" + }, "auto_truncate": { "default_value": null, "depends_on": [], "display": "toggle", "label": "Auto Truncate", - "order": 1, + "order": 2, "required": false, "sensitive": false, "tooltip": "Specifies if the API truncates inputs longer than the maximum token length automatically.", @@ -1005,11 +1039,15 @@ private static ActionListener getModelListenerForException(Class excep }); } - private static Map getTaskSettingsMap(Boolean autoTruncate) { + private static Map getTaskSettingsMap(Boolean autoTruncate, @Nullable InputType inputType) { var taskSettings = new HashMap(); taskSettings.put(GoogleVertexAiEmbeddingsTaskSettings.AUTO_TRUNCATE, autoTruncate); + if (inputType != null) { + taskSettings.put(GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE, inputType.toString()); + } + return taskSettings; } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModelTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModelTests.java index 68d03d350d06e..7836c5c15cfb1 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModelTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsModelTests.java @@ -10,14 +10,18 @@ import org.elasticsearch.common.Strings; import org.elasticsearch.common.settings.SecureString; import org.elasticsearch.core.Nullable; +import org.elasticsearch.inference.InputType; import org.elasticsearch.inference.SimilarityMeasure; import org.elasticsearch.inference.TaskType; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiSecretSettings; +import org.hamcrest.MatcherAssert; import java.net.URI; import java.net.URISyntaxException; +import java.util.Map; +import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettingsTests.getTaskSettingsMap; import static org.hamcrest.Matchers.is; public class GoogleVertexAiEmbeddingsModelTests extends ESTestCase { @@ -45,6 +49,75 @@ public void testBuildUri() throws URISyntaxException { ); } + public void testOverrideWith_DoesNotOverrideAndModelRemainsEqual_WhenSettingsAreEmpty_AndInputTypeIsInvalid() { + var model = createModel("model", Boolean.FALSE, InputType.SEARCH); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, Map.of(), InputType.UNSPECIFIED); + + MatcherAssert.assertThat(overriddenModel, is(model)); + } + + public void testOverrideWith_DoesNotOverrideAndModelRemainsEqual_WhenSettingsAreNull_AndInputTypeIsInvalid() { + var model = createModel("model", Boolean.FALSE, InputType.SEARCH); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, null, InputType.UNSPECIFIED); + + MatcherAssert.assertThat(overriddenModel, is(model)); + } + + public void testOverrideWith_SetsInputTypeToOverride_WhenFieldIsNullInModelTaskSettings_AndNullInRequestTaskSettings() { + var model = createModel("model", Boolean.FALSE, null); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, getTaskSettingsMap(null, null), InputType.SEARCH); + + var expectedModel = createModel("model", Boolean.FALSE, InputType.SEARCH); + MatcherAssert.assertThat(overriddenModel, is(expectedModel)); + } + + public void testOverrideWith_SetsInputType_FromRequest_IfValid_OverridingStoredTaskSettings() { + var model = createModel("model", Boolean.FALSE, InputType.INGEST); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, getTaskSettingsMap(null, null), InputType.SEARCH); + + var expectedModel = createModel("model", Boolean.FALSE, InputType.SEARCH); + MatcherAssert.assertThat(overriddenModel, is(expectedModel)); + } + + public void testOverrideWith_SetsInputType_FromRequest_IfValid_OverridingRequestTaskSettings() { + var model = createModel("model", Boolean.FALSE, null); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, getTaskSettingsMap(null, InputType.CLUSTERING), InputType.SEARCH); + + var expectedModel = createModel("model", Boolean.FALSE, InputType.SEARCH); + MatcherAssert.assertThat(overriddenModel, is(expectedModel)); + } + + public void testOverrideWith_OverridesInputType_WithRequestTaskSettingsSearch_WhenRequestInputTypeIsInvalid() { + var model = createModel("model", Boolean.FALSE, InputType.INGEST); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, getTaskSettingsMap(null, InputType.SEARCH), InputType.UNSPECIFIED); + + var expectedModel = createModel("model", Boolean.FALSE, InputType.SEARCH); + MatcherAssert.assertThat(overriddenModel, is(expectedModel)); + } + + public void testOverrideWith_DoesNotSetInputType_FromRequest_IfInputTypeIsInvalid() { + var model = createModel("model", Boolean.FALSE, null); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, getTaskSettingsMap(null, null), InputType.UNSPECIFIED); + + var expectedModel = createModel("model", Boolean.FALSE, null); + MatcherAssert.assertThat(overriddenModel, is(expectedModel)); + } + + public void testOverrideWith_DoesNotSetInputType_WhenRequestTaskSettingsIsNull_AndRequestInputTypeIsInvalid() { + var model = createModel("model", Boolean.FALSE, InputType.INGEST); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, getTaskSettingsMap(null, null), InputType.UNSPECIFIED); + + var expectedModel = createModel("model", Boolean.FALSE, InputType.INGEST); + MatcherAssert.assertThat(overriddenModel, is(expectedModel)); + } + + public void testOverrideWith_DoesNotOverrideModelUri() { + var model = createModel("model", Boolean.FALSE, InputType.SEARCH); + var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, Map.of(), null); + + MatcherAssert.assertThat(overriddenModel.uri(), is(model.uri())); + } + public static GoogleVertexAiEmbeddingsModel createModel( String location, String projectId, @@ -58,12 +131,37 @@ public static GoogleVertexAiEmbeddingsModel createModel( "service", uri, new GoogleVertexAiEmbeddingsServiceSettings(location, projectId, modelId, false, null, null, null, null), - new GoogleVertexAiEmbeddingsTaskSettings(Boolean.FALSE), + new GoogleVertexAiEmbeddingsTaskSettings(Boolean.FALSE, null), new GoogleVertexAiSecretSettings(new SecureString(serviceAccountJson.toCharArray())) ); } - public static GoogleVertexAiEmbeddingsModel createModel(String modelId, @Nullable Boolean autoTruncate) { + public static GoogleVertexAiEmbeddingsModel createModel(String modelId, @Nullable Boolean autoTruncate, @Nullable InputType inputType) { + return new GoogleVertexAiEmbeddingsModel( + "id", + TaskType.TEXT_EMBEDDING, + "service", + new GoogleVertexAiEmbeddingsServiceSettings( + "location", + "projectId", + modelId, + false, + null, + null, + SimilarityMeasure.DOT_PRODUCT, + null + ), + new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, inputType), + null, + new GoogleVertexAiSecretSettings(new SecureString("testString".toCharArray())) + ); + } + + public static GoogleVertexAiEmbeddingsModel createRandomizedModel( + String modelId, + @Nullable Boolean autoTruncate, + @Nullable InputType inputType + ) { return new GoogleVertexAiEmbeddingsModel( "id", TaskType.TEXT_EMBEDDING, @@ -78,7 +176,7 @@ public static GoogleVertexAiEmbeddingsModel createModel(String modelId, @Nullabl SimilarityMeasure.DOT_PRODUCT, null ), - new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate), + new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, inputType), null, new GoogleVertexAiSecretSettings(new SecureString(randomAlphaOfLength(8).toCharArray())) ); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettingsTests.java index 1e9a2f435cb08..a49e0f2e3f57d 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettingsTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsRequestTaskSettingsTests.java @@ -7,6 +7,8 @@ package org.elasticsearch.xpack.inference.services.googlevertexai.embeddings; +import org.elasticsearch.common.ValidationException; +import org.elasticsearch.inference.InputType; import org.elasticsearch.test.ESTestCase; import java.util.HashMap; @@ -21,9 +23,14 @@ public void testFromMap_ReturnsEmptySettings_IfMapEmpty() { assertThat(requestTaskSettings, is(GoogleVertexAiEmbeddingsRequestTaskSettings.EMPTY_SETTINGS)); } + public void testFromMap_ReturnsEmptySettings_IfMapNull() { + var requestTaskSettings = GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap(null); + assertThat(requestTaskSettings, is(GoogleVertexAiEmbeddingsRequestTaskSettings.EMPTY_SETTINGS)); + } + public void testFromMap_DoesNotThrowValidationException_IfAutoTruncateIsMissing() { var requestTaskSettings = GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap(new HashMap<>(Map.of("unrelated", true))); - assertThat(requestTaskSettings, is(new GoogleVertexAiEmbeddingsRequestTaskSettings(null))); + assertThat(requestTaskSettings, is(new GoogleVertexAiEmbeddingsRequestTaskSettings(null, null))); } public void testFromMap_ExtractsAutoTruncate() { @@ -31,6 +38,40 @@ public void testFromMap_ExtractsAutoTruncate() { var requestTaskSettings = GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap( new HashMap<>(Map.of(GoogleVertexAiEmbeddingsTaskSettings.AUTO_TRUNCATE, autoTruncate)) ); - assertThat(requestTaskSettings, is(new GoogleVertexAiEmbeddingsRequestTaskSettings(autoTruncate))); + assertThat(requestTaskSettings, is(new GoogleVertexAiEmbeddingsRequestTaskSettings(autoTruncate, null))); + } + + public void testFromMap_ThrowsValidationException_IfAutoTruncateIsInvalidValue() { + expectThrows( + ValidationException.class, + () -> GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap( + new HashMap<>(Map.of(GoogleVertexAiEmbeddingsTaskSettings.AUTO_TRUNCATE, "invalid")) + ) + ); + } + + public void testFromMap_ExtractsInputType() { + var requestTaskSettings = GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap( + new HashMap<>(Map.of(GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE, InputType.INGEST.toString())) + ); + assertThat(requestTaskSettings, is(new GoogleVertexAiEmbeddingsRequestTaskSettings(null, InputType.INGEST))); + } + + public void testFromMap_ThrowsValidationException_IfInputTypeIsInvalidValue() { + expectThrows( + ValidationException.class, + () -> GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap( + new HashMap<>(Map.of(GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE, "abc")) + ) + ); + } + + public void testFromMap_ThrowsValidationException_IfInputTypeIsUnspecified() { + expectThrows( + ValidationException.class, + () -> GoogleVertexAiEmbeddingsRequestTaskSettings.fromMap( + new HashMap<>(Map.of(GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE, InputType.UNSPECIFIED.toString())) + ) + ); } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettingsTests.java index 5b87bbc3c42c8..0a390b114702c 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettingsTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/embeddings/GoogleVertexAiEmbeddingsTaskSettingsTests.java @@ -8,21 +8,30 @@ package org.elasticsearch.xpack.inference.services.googlevertexai.embeddings; import org.elasticsearch.TransportVersion; +import org.elasticsearch.TransportVersions; import org.elasticsearch.common.Strings; import org.elasticsearch.common.ValidationException; import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.core.Nullable; +import org.elasticsearch.inference.InputType; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentFactory; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase; +import org.hamcrest.MatcherAssert; import java.io.IOException; +import java.util.Arrays; import java.util.Collections; +import java.util.EnumSet; import java.util.HashMap; +import java.util.Locale; import java.util.Map; +import static org.elasticsearch.xpack.inference.InputTypeTests.randomWithoutUnspecified; import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings.AUTO_TRUNCATE; +import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE; +import static org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings.VALID_REQUEST_VALUES; import static org.hamcrest.Matchers.is; public class GoogleVertexAiEmbeddingsTaskSettingsTests extends AbstractBWCWireSerializationTestCase { @@ -39,6 +48,9 @@ public void testUpdatedTaskSettings() { if (newSettings.autoTruncate() != null) { newSettingsMap.put(GoogleVertexAiEmbeddingsTaskSettings.AUTO_TRUNCATE, newSettings.autoTruncate()); } + if (newSettings.getInputType() != null) { + newSettingsMap.put(GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE, newSettings.getInputType().toString()); + } GoogleVertexAiEmbeddingsTaskSettings updatedSettings = (GoogleVertexAiEmbeddingsTaskSettings) initialSettings.updatedTaskSettings( Collections.unmodifiableMap(newSettingsMap) ); @@ -47,56 +59,144 @@ public void testUpdatedTaskSettings() { } else { assertEquals(newSettings.autoTruncate(), updatedSettings.autoTruncate()); } + if (newSettings.getInputType() == null) { + assertEquals(initialSettings.getInputType(), updatedSettings.getInputType()); + } else { + assertEquals(newSettings.getInputType(), updatedSettings.getInputType()); + } + } + + public void testFromMap_CreatesEmptySettings_WhenAllFieldsAreNull() { + MatcherAssert.assertThat( + GoogleVertexAiEmbeddingsTaskSettings.fromMap(new HashMap<>()), + is(new GoogleVertexAiEmbeddingsTaskSettings(null, null)) + ); + assertNull(GoogleVertexAiEmbeddingsTaskSettings.fromMap(new HashMap<>()).autoTruncate()); + assertNull(GoogleVertexAiEmbeddingsTaskSettings.fromMap(new HashMap<>()).getInputType()); + } + + public void testFromMap_CreatesEmptySettings_WhenMapIsNull() { + MatcherAssert.assertThat( + GoogleVertexAiEmbeddingsTaskSettings.fromMap(null), + is(new GoogleVertexAiEmbeddingsTaskSettings(null, null)) + ); + assertNull(GoogleVertexAiEmbeddingsTaskSettings.fromMap(null).autoTruncate()); + assertNull(GoogleVertexAiEmbeddingsTaskSettings.fromMap(null).getInputType()); } public void testFromMap_AutoTruncateIsSet() { var autoTruncate = true; - var taskSettingsMap = getTaskSettingsMap(autoTruncate); + var taskSettingsMap = getTaskSettingsMap(autoTruncate, null); var taskSettings = GoogleVertexAiEmbeddingsTaskSettings.fromMap(taskSettingsMap); - assertThat(taskSettings, is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate))); + assertThat(taskSettings, is(new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, null))); } public void testFromMap_ThrowsValidationException_IfAutoTruncateIsInvalidValue() { - var taskSettings = getTaskSettingsMap("invalid"); + var taskSettings = getTaskSettingsMap("invalid", null); expectThrows(ValidationException.class, () -> GoogleVertexAiEmbeddingsTaskSettings.fromMap(taskSettings)); } public void testFromMap_AutoTruncateIsNull() { - var taskSettingsMap = getTaskSettingsMap(null); + var taskSettingsMap = getTaskSettingsMap(null, null); var taskSettings = GoogleVertexAiEmbeddingsTaskSettings.fromMap(taskSettingsMap); // needed, because of constructors being ambiguous otherwise Boolean nullBoolean = null; - assertThat(taskSettings, is(new GoogleVertexAiEmbeddingsTaskSettings(nullBoolean))); + assertThat(taskSettings, is(new GoogleVertexAiEmbeddingsTaskSettings(nullBoolean, null))); } - public void testFromMap_DoesNotThrow_WithEmptyMap() { - assertNull(GoogleVertexAiEmbeddingsTaskSettings.fromMap(new HashMap<>()).autoTruncate()); + public void testFromMap_ReturnsFailure_WhenInputTypeIsInvalid() { + var exception = expectThrows( + ValidationException.class, + () -> GoogleVertexAiEmbeddingsTaskSettings.fromMap( + new HashMap<>(Map.of(GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE, "abc")) + ) + ); + + assertThat( + exception.getMessage(), + is( + Strings.format( + "Validation Failed: 1: [task_settings] Invalid value [abc] received. [input_type] must be one of [%s];", + getValidValuesSortedAndCombined(VALID_REQUEST_VALUES) + ) + ) + ); + } + + public void testFromMap_ReturnsFailure_WhenInputTypeIsUnspecified() { + var exception = expectThrows( + ValidationException.class, + () -> GoogleVertexAiEmbeddingsTaskSettings.fromMap( + new HashMap<>(Map.of(GoogleVertexAiEmbeddingsTaskSettings.INPUT_TYPE, InputType.UNSPECIFIED.toString())) + ) + ); + + assertThat( + exception.getMessage(), + is( + Strings.format( + "Validation Failed: 1: [task_settings] Invalid value [unspecified] received. [input_type] must be one of [%s];", + getValidValuesSortedAndCombined(VALID_REQUEST_VALUES) + ) + ) + ); } public void testOf_UseRequestSettings() { var originalAutoTruncate = true; - var originalSettings = new GoogleVertexAiEmbeddingsTaskSettings(originalAutoTruncate); + var originalSettings = new GoogleVertexAiEmbeddingsTaskSettings(originalAutoTruncate, null); var requestAutoTruncate = originalAutoTruncate == false; - var requestTaskSettings = new GoogleVertexAiEmbeddingsRequestTaskSettings(requestAutoTruncate); + var requestTaskSettings = new GoogleVertexAiEmbeddingsRequestTaskSettings(requestAutoTruncate, null); - assertThat(GoogleVertexAiEmbeddingsTaskSettings.of(originalSettings, requestTaskSettings).autoTruncate(), is(requestAutoTruncate)); + assertThat( + GoogleVertexAiEmbeddingsTaskSettings.of(originalSettings, requestTaskSettings, null).autoTruncate(), + is(requestAutoTruncate) + ); + } + + public void testOf_UseRequestSettings_AndRequestInputType() { + var originalAutoTruncate = true; + var originalSettings = new GoogleVertexAiEmbeddingsTaskSettings(originalAutoTruncate, InputType.SEARCH); + + var requestAutoTruncate = originalAutoTruncate == false; + var requestTaskSettings = new GoogleVertexAiEmbeddingsRequestTaskSettings(requestAutoTruncate, null); + + assertThat( + GoogleVertexAiEmbeddingsTaskSettings.of(originalSettings, requestTaskSettings, InputType.INGEST).getInputType(), + is(InputType.INGEST) + ); } public void testOf_UseOriginalSettings() { var originalAutoTruncate = true; - var originalSettings = new GoogleVertexAiEmbeddingsTaskSettings(originalAutoTruncate); + var originalSettings = new GoogleVertexAiEmbeddingsTaskSettings(originalAutoTruncate, null); - var requestTaskSettings = new GoogleVertexAiEmbeddingsRequestTaskSettings(null); + var requestTaskSettings = new GoogleVertexAiEmbeddingsRequestTaskSettings(null, null); - assertThat(GoogleVertexAiEmbeddingsTaskSettings.of(originalSettings, requestTaskSettings).autoTruncate(), is(originalAutoTruncate)); + assertThat( + GoogleVertexAiEmbeddingsTaskSettings.of(originalSettings, requestTaskSettings, null).autoTruncate(), + is(originalAutoTruncate) + ); + } + + public void testOf_UseOriginalSettings_WithInputType() { + var originalAutoTruncate = true; + var originalSettings = new GoogleVertexAiEmbeddingsTaskSettings(originalAutoTruncate, InputType.INGEST); + + var requestTaskSettings = new GoogleVertexAiEmbeddingsRequestTaskSettings(null, null); + + assertThat( + GoogleVertexAiEmbeddingsTaskSettings.of(originalSettings, requestTaskSettings, null).autoTruncate(), + is(originalAutoTruncate) + ); } public void testToXContent_WritesAutoTruncateIfNotNull() throws IOException { - var settings = GoogleVertexAiEmbeddingsTaskSettings.fromMap(getTaskSettingsMap(true)); + var settings = GoogleVertexAiEmbeddingsTaskSettings.fromMap(getTaskSettingsMap(true, null)); XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); settings.toXContent(builder, null); @@ -107,7 +207,7 @@ public void testToXContent_WritesAutoTruncateIfNotNull() throws IOException { } public void testToXContent_DoesNotWriteAutoTruncateIfNull() throws IOException { - var settings = GoogleVertexAiEmbeddingsTaskSettings.fromMap(getTaskSettingsMap(null)); + var settings = GoogleVertexAiEmbeddingsTaskSettings.fromMap(getTaskSettingsMap(null, null)); XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); settings.toXContent(builder, null); @@ -117,6 +217,25 @@ public void testToXContent_DoesNotWriteAutoTruncateIfNull() throws IOException { {}""")); } + public void testToXContent_WritesInputTypeIfNotNull() throws IOException { + var settings = GoogleVertexAiEmbeddingsTaskSettings.fromMap(getTaskSettingsMap(true, InputType.INGEST)); + + XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON); + settings.toXContent(builder, null); + String xContentResult = Strings.toString(builder); + + assertThat(xContentResult, is(""" + {"input_type":"ingest","auto_truncate":true}""")); + } + + public void testToXContent_ThrowsAssertionFailure_WhenInputTypeIsUnspecified() { + var thrownException = expectThrows( + AssertionError.class, + () -> new GoogleVertexAiEmbeddingsTaskSettings(false, InputType.UNSPECIFIED) + ); + assertThat(thrownException.getMessage(), is("received invalid input type value [unspecified]")); + } + @Override protected Writeable.Reader instanceReader() { return GoogleVertexAiEmbeddingsTaskSettings::new; @@ -137,20 +256,37 @@ protected GoogleVertexAiEmbeddingsTaskSettings mutateInstanceForVersion( GoogleVertexAiEmbeddingsTaskSettings instance, TransportVersion version ) { + if (version.before(TransportVersions.VERTEX_AI_INPUT_TYPE_ADDED)) { + // default to null input type if node is on a version before input type was introduced + return new GoogleVertexAiEmbeddingsTaskSettings(instance.autoTruncate(), null); + } return instance; } private static GoogleVertexAiEmbeddingsTaskSettings createRandom() { - return new GoogleVertexAiEmbeddingsTaskSettings(randomFrom(new Boolean[] { null, randomBoolean() })); + var inputType = randomBoolean() ? randomWithoutUnspecified() : null; + var autoTruncate = randomFrom(new Boolean[] { null, randomBoolean() }); + return new GoogleVertexAiEmbeddingsTaskSettings(autoTruncate, inputType); + } + + private static > String getValidValuesSortedAndCombined(EnumSet validValues) { + var validValuesAsStrings = validValues.stream().map(value -> value.toString().toLowerCase(Locale.ROOT)).toArray(String[]::new); + Arrays.sort(validValuesAsStrings); + + return String.join(", ", validValuesAsStrings); } - private static Map getTaskSettingsMap(@Nullable Object autoTruncate) { + public static Map getTaskSettingsMap(@Nullable Object autoTruncate, @Nullable InputType inputType) { var map = new HashMap(); if (autoTruncate != null) { map.put(AUTO_TRUNCATE, autoTruncate); } + if (inputType != null) { + map.put(INPUT_TYPE, inputType.toString()); + } + return map; } } From 778ab8fee362e5a17195f24d23dbbf6ee88557c4 Mon Sep 17 00:00:00 2001 From: Salvatore Campagna <93581129+salvatore-campagna@users.noreply.github.com> Date: Tue, 12 Nov 2024 16:39:41 +0100 Subject: [PATCH 16/16] Re-structure document ID generation favoring `_id` inverted index compression (#104683) This implementation restructures auto-generated document IDs to maximize compression within Lucene's terms dictionary. The key insight is placing stable or slowly-changing components at the start of the ID - the most significant bytes of the timestamp change very gradually (the first byte shifts only every 35 years, the second every 50 days). This careful ordering means that large sequences of IDs generated close in time will share common prefixes, allowing Lucene's Finite State Transducer (FST) to store terms more compactly. To maintain uniqueness while preserving these compression benefits, the ID combines three elements: a timestamp that ensures time-based ordering, the coordinator's MAC address for cluster-wide uniqueness, and a sequence number for handling high-throughput scenarios. The timestamp handling is particularly robust, using atomic operations to prevent backwards movement even if the system clock shifts. For high-volume indices generating millions of documents, this optimization can lead to substantial storage savings while maintaining strict guarantees about ID uniqueness and ordering. --- docs/changelog/104683.yaml | 5 + .../action/bulk/BulkIntegrationIT.java | 6 +- .../action/index/IndexRequest.java | 29 +++- .../cluster/routing/IndexRouting.java | 12 +- .../TimeBasedKOrderedUUIDGenerator.java | 73 ++++++++++ .../common/TimeBasedUUIDGenerator.java | 4 +- .../java/org/elasticsearch/common/UUIDs.java | 10 ++ .../elasticsearch/index/IndexVersions.java | 1 + .../action/index/IndexRequestTests.java | 6 + .../org/elasticsearch/common/UUIDTests.java | 134 ++++++++++++------ 10 files changed, 225 insertions(+), 55 deletions(-) create mode 100644 docs/changelog/104683.yaml create mode 100644 server/src/main/java/org/elasticsearch/common/TimeBasedKOrderedUUIDGenerator.java diff --git a/docs/changelog/104683.yaml b/docs/changelog/104683.yaml new file mode 100644 index 0000000000000..d4f40b59cfd91 --- /dev/null +++ b/docs/changelog/104683.yaml @@ -0,0 +1,5 @@ +pr: 104683 +summary: "Feature: re-structure document ID generation favoring _id inverted index compression" +area: Logs +type: enhancement +issues: [] diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkIntegrationIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkIntegrationIT.java index 34170d7c0f747..e45555b1dec19 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkIntegrationIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkIntegrationIT.java @@ -99,7 +99,11 @@ public void testBulkWithWriteIndexAndRouting() { // allowing the auto-generated timestamp to externally be set would allow making the index inconsistent with duplicate docs public void testExternallySetAutoGeneratedTimestamp() { IndexRequest indexRequest = new IndexRequest("index1").source(Collections.singletonMap("foo", "baz")); - indexRequest.autoGenerateId(); + if (randomBoolean()) { + indexRequest.autoGenerateId(); + } else { + indexRequest.autoGenerateTimeBasedId(); + } if (randomBoolean()) { indexRequest.id("test"); } diff --git a/server/src/main/java/org/elasticsearch/action/index/IndexRequest.java b/server/src/main/java/org/elasticsearch/action/index/IndexRequest.java index d0785a60dd0f5..c0811e7424b0d 100644 --- a/server/src/main/java/org/elasticsearch/action/index/IndexRequest.java +++ b/server/src/main/java/org/elasticsearch/action/index/IndexRequest.java @@ -51,6 +51,7 @@ import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.function.Supplier; import static org.elasticsearch.action.ValidateActions.addValidationError; import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_PRIMARY_TERM; @@ -76,6 +77,9 @@ public class IndexRequest extends ReplicatedWriteRequest implement private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(IndexRequest.class); private static final TransportVersion PIPELINES_HAVE_RUN_FIELD_ADDED = TransportVersions.V_8_10_X; + private static final Supplier ID_GENERATOR = UUIDs::base64UUID; + private static final Supplier K_SORTED_TIME_BASED_ID_GENERATOR = UUIDs::base64TimeBasedKOrderedUUID; + /** * Max length of the source document to include into string() * @@ -692,10 +696,18 @@ public void process(IndexRouting indexRouting) { * request compatible with the append-only optimization. */ public void autoGenerateId() { - assert id == null; - assert autoGeneratedTimestamp == UNSET_AUTO_GENERATED_TIMESTAMP : "timestamp has already been generated!"; - assert ifSeqNo == UNASSIGNED_SEQ_NO; - assert ifPrimaryTerm == UNASSIGNED_PRIMARY_TERM; + assertBeforeGeneratingId(); + autoGenerateTimestamp(); + id(ID_GENERATOR.get()); + } + + public void autoGenerateTimeBasedId() { + assertBeforeGeneratingId(); + autoGenerateTimestamp(); + id(K_SORTED_TIME_BASED_ID_GENERATOR.get()); + } + + private void autoGenerateTimestamp() { /* * Set the auto generated timestamp so the append only optimization * can quickly test if this request *must* be unique without reaching @@ -704,8 +716,13 @@ public void autoGenerateId() { * never work before 1970, but that's ok. It's after 1970. */ autoGeneratedTimestamp = Math.max(0, System.currentTimeMillis()); - String uid = UUIDs.base64UUID(); - id(uid); + } + + private void assertBeforeGeneratingId() { + assert id == null; + assert autoGeneratedTimestamp == UNSET_AUTO_GENERATED_TIMESTAMP : "timestamp has already been generated!"; + assert ifSeqNo == UNASSIGNED_SEQ_NO; + assert ifPrimaryTerm == UNASSIGNED_PRIMARY_TERM; } /** diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/IndexRouting.java b/server/src/main/java/org/elasticsearch/cluster/routing/IndexRouting.java index 3fb3c182f89cd..1c89d3bf259b5 100644 --- a/server/src/main/java/org/elasticsearch/cluster/routing/IndexRouting.java +++ b/server/src/main/java/org/elasticsearch/cluster/routing/IndexRouting.java @@ -24,6 +24,8 @@ import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.core.Nullable; import org.elasticsearch.features.NodeFeature; +import org.elasticsearch.index.IndexMode; +import org.elasticsearch.index.IndexVersion; import org.elasticsearch.index.IndexVersions; import org.elasticsearch.index.mapper.TimeSeriesRoutingHashFieldMapper; import org.elasticsearch.transport.Transports; @@ -147,11 +149,15 @@ public void checkIndexSplitAllowed() {} private abstract static class IdAndRoutingOnly extends IndexRouting { private final boolean routingRequired; + private final IndexVersion creationVersion; + private final IndexMode indexMode; IdAndRoutingOnly(IndexMetadata metadata) { super(metadata); + this.creationVersion = metadata.getCreationVersion(); MappingMetadata mapping = metadata.mapping(); this.routingRequired = mapping == null ? false : mapping.routingRequired(); + this.indexMode = metadata.getIndexMode(); } protected abstract int shardId(String id, @Nullable String routing); @@ -161,7 +167,11 @@ public void process(IndexRequest indexRequest) { // generate id if not already provided final String id = indexRequest.id(); if (id == null) { - indexRequest.autoGenerateId(); + if (creationVersion.onOrAfter(IndexVersions.TIME_BASED_K_ORDERED_DOC_ID) && indexMode == IndexMode.LOGSDB) { + indexRequest.autoGenerateTimeBasedId(); + } else { + indexRequest.autoGenerateId(); + } } else if (id.isEmpty()) { throw new IllegalArgumentException("if _id is specified it must not be empty"); } diff --git a/server/src/main/java/org/elasticsearch/common/TimeBasedKOrderedUUIDGenerator.java b/server/src/main/java/org/elasticsearch/common/TimeBasedKOrderedUUIDGenerator.java new file mode 100644 index 0000000000000..9c97cb8fe7e85 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/common/TimeBasedKOrderedUUIDGenerator.java @@ -0,0 +1,73 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.common; + +import java.nio.ByteBuffer; +import java.util.Base64; + +/** + * Generates a base64-encoded, k-ordered UUID string optimized for compression and efficient indexing. + *

+ * This method produces a time-based UUID where slowly changing components like the timestamp appear first, + * improving prefix-sharing and compression during indexing. It ensures uniqueness across nodes by incorporating + * a timestamp, a MAC address, and a sequence ID. + *

+ * Timestamp: Represents the current time in milliseconds, ensuring ordering and uniqueness. + *
+ * MAC Address: Ensures uniqueness across different coordinators. + *
+ * Sequence ID: Differentiates UUIDs generated within the same millisecond, ensuring uniqueness even at high throughput. + *

+ * The result is a compact base64-encoded string, optimized for efficient compression of the _id field in an inverted index. + */ +public class TimeBasedKOrderedUUIDGenerator extends TimeBasedUUIDGenerator { + private static final Base64.Encoder BASE_64_NO_PADDING = Base64.getEncoder().withoutPadding(); + + @Override + public String getBase64UUID() { + final int sequenceId = this.sequenceNumber.incrementAndGet() & 0x00FF_FFFF; + + // Calculate timestamp to ensure ordering and avoid backward movement in case of time shifts. + // Uses AtomicLong to guarantee that timestamp increases even if the system clock moves backward. + // If the sequenceId overflows (reaches 0 within the same millisecond), the timestamp is incremented + // to ensure strict ordering. + long timestamp = this.lastTimestamp.accumulateAndGet( + currentTimeMillis(), + sequenceId == 0 ? (lastTimestamp, currentTimeMillis) -> Math.max(lastTimestamp, currentTimeMillis) + 1 : Math::max + ); + + final byte[] uuidBytes = new byte[15]; + final ByteBuffer buffer = ByteBuffer.wrap(uuidBytes); + + buffer.put((byte) (timestamp >>> 40)); // changes every 35 years + buffer.put((byte) (timestamp >>> 32)); // changes every ~50 days + buffer.put((byte) (timestamp >>> 24)); // changes every ~4.5h + buffer.put((byte) (timestamp >>> 16)); // changes every ~65 secs + + // MAC address of the coordinator might change if there are many coordinators in the cluster + // and the indexing api does not necessarily target the same coordinator. + byte[] macAddress = macAddress(); + assert macAddress.length == 6; + buffer.put(macAddress, 0, macAddress.length); + + buffer.put((byte) (sequenceId >>> 16)); + + // From hereinafter everything is almost like random and does not compress well + // due to unlikely prefix-sharing + buffer.put((byte) (timestamp >>> 8)); + buffer.put((byte) (sequenceId >>> 8)); + buffer.put((byte) timestamp); + buffer.put((byte) sequenceId); + + assert buffer.position() == uuidBytes.length; + + return BASE_64_NO_PADDING.encodeToString(uuidBytes); + } +} diff --git a/server/src/main/java/org/elasticsearch/common/TimeBasedUUIDGenerator.java b/server/src/main/java/org/elasticsearch/common/TimeBasedUUIDGenerator.java index 73528ed0d3866..2ed979ae66ffa 100644 --- a/server/src/main/java/org/elasticsearch/common/TimeBasedUUIDGenerator.java +++ b/server/src/main/java/org/elasticsearch/common/TimeBasedUUIDGenerator.java @@ -24,10 +24,10 @@ class TimeBasedUUIDGenerator implements UUIDGenerator { // We only use bottom 3 bytes for the sequence number. Paranoia: init with random int so that if JVM/OS/machine goes down, clock slips // backwards, and JVM comes back up, we are less likely to be on the same sequenceNumber at the same time: - private final AtomicInteger sequenceNumber = new AtomicInteger(SecureRandomHolder.INSTANCE.nextInt()); + protected final AtomicInteger sequenceNumber = new AtomicInteger(SecureRandomHolder.INSTANCE.nextInt()); // Used to ensure clock moves forward: - private final AtomicLong lastTimestamp = new AtomicLong(0); + protected final AtomicLong lastTimestamp = new AtomicLong(0); private static final byte[] SECURE_MUNGED_ADDRESS = MacAddressProvider.getSecureMungedAddress(); diff --git a/server/src/main/java/org/elasticsearch/common/UUIDs.java b/server/src/main/java/org/elasticsearch/common/UUIDs.java index 61ee4bd5d64ab..0f73b8172c10f 100644 --- a/server/src/main/java/org/elasticsearch/common/UUIDs.java +++ b/server/src/main/java/org/elasticsearch/common/UUIDs.java @@ -16,6 +16,8 @@ public class UUIDs { private static final RandomBasedUUIDGenerator RANDOM_UUID_GENERATOR = new RandomBasedUUIDGenerator(); + + private static final UUIDGenerator TIME_BASED_K_ORDERED_GENERATOR = new TimeBasedKOrderedUUIDGenerator(); private static final UUIDGenerator TIME_UUID_GENERATOR = new TimeBasedUUIDGenerator(); /** @@ -33,6 +35,14 @@ public static String base64UUID() { return TIME_UUID_GENERATOR.getBase64UUID(); } + public static String base64TimeBasedKOrderedUUID() { + return TIME_BASED_K_ORDERED_GENERATOR.getBase64UUID(); + } + + public static String base64TimeBasedUUID() { + return TIME_UUID_GENERATOR.getBase64UUID(); + } + /** * The length of a UUID string generated by {@link #randomBase64UUID} and {@link #randomBase64UUIDSecureString}. */ diff --git a/server/src/main/java/org/elasticsearch/index/IndexVersions.java b/server/src/main/java/org/elasticsearch/index/IndexVersions.java index 440613263d441..9264b9e1c3a20 100644 --- a/server/src/main/java/org/elasticsearch/index/IndexVersions.java +++ b/server/src/main/java/org/elasticsearch/index/IndexVersions.java @@ -132,6 +132,7 @@ private static Version parseUnchecked(String version) { public static final IndexVersion LOGSDB_DEFAULT_IGNORE_DYNAMIC_BEYOND_LIMIT_BACKPORT = def(8_519_00_0, Version.LUCENE_9_12_0); public static final IndexVersion UPGRADE_TO_LUCENE_10_0_0 = def(9_000_00_0, Version.LUCENE_10_0_0); public static final IndexVersion LOGSDB_DEFAULT_IGNORE_DYNAMIC_BEYOND_LIMIT = def(9_001_00_0, Version.LUCENE_10_0_0); + public static final IndexVersion TIME_BASED_K_ORDERED_DOC_ID = def(9_002_00_0, Version.LUCENE_10_0_0); /* * STOP! READ THIS FIRST! No, really, * ____ _____ ___ ____ _ ____ _____ _ ____ _____ _ _ ___ ____ _____ ___ ____ ____ _____ _ diff --git a/server/src/test/java/org/elasticsearch/action/index/IndexRequestTests.java b/server/src/test/java/org/elasticsearch/action/index/IndexRequestTests.java index 32297e0c09b8f..9d74c2069ec10 100644 --- a/server/src/test/java/org/elasticsearch/action/index/IndexRequestTests.java +++ b/server/src/test/java/org/elasticsearch/action/index/IndexRequestTests.java @@ -128,6 +128,12 @@ public void testAutoGenerateId() { assertTrue("expected > 0 but got: " + request.getAutoGeneratedTimestamp(), request.getAutoGeneratedTimestamp() > 0); } + public void testAutoGenerateTimeBasedId() { + IndexRequest request = new IndexRequest("index"); + request.autoGenerateTimeBasedId(); + assertTrue("expected > 0 but got: " + request.getAutoGeneratedTimestamp(), request.getAutoGeneratedTimestamp() > 0); + } + public void testIndexResponse() { ShardId shardId = new ShardId(randomAlphaOfLengthBetween(3, 10), randomAlphaOfLengthBetween(3, 10), randomIntBetween(0, 1000)); String id = randomAlphaOfLengthBetween(3, 10); diff --git a/server/src/test/java/org/elasticsearch/common/UUIDTests.java b/server/src/test/java/org/elasticsearch/common/UUIDTests.java index 2e7dbb00aa2de..9fbeaf1c6c081 100644 --- a/server/src/test/java/org/elasticsearch/common/UUIDTests.java +++ b/server/src/test/java/org/elasticsearch/common/UUIDTests.java @@ -35,6 +35,7 @@ public class UUIDTests extends ESTestCase { static UUIDGenerator timeUUIDGen = new TimeBasedUUIDGenerator(); static UUIDGenerator randomUUIDGen = new RandomBasedUUIDGenerator(); + static UUIDGenerator kOrderedUUIDGen = new TimeBasedKOrderedUUIDGenerator(); public void testRandomUUID() { verifyUUIDSet(100000, randomUUIDGen); @@ -44,14 +45,49 @@ public void testTimeUUID() { verifyUUIDSet(100000, timeUUIDGen); } - public void testThreadedTimeUUID() { - testUUIDThreaded(timeUUIDGen); + public void testKOrderedUUID() { + verifyUUIDSet(100000, kOrderedUUIDGen); } public void testThreadedRandomUUID() { testUUIDThreaded(randomUUIDGen); } + public void testThreadedTimeUUID() { + testUUIDThreaded(timeUUIDGen); + } + + public void testThreadedKOrderedUUID() { + testUUIDThreaded(kOrderedUUIDGen); + } + + public void testCompression() throws Exception { + Logger logger = LogManager.getLogger(UUIDTests.class); + + assertThat(testCompression(timeUUIDGen, 100000, 10000, 3, logger), Matchers.lessThan(14d)); + assertThat(testCompression(timeUUIDGen, 100000, 1000, 3, logger), Matchers.lessThan(15d)); + assertThat(testCompression(timeUUIDGen, 100000, 100, 3, logger), Matchers.lessThan(21d)); + + assertThat(testCompression(kOrderedUUIDGen, 100000, 10000, 3, logger), Matchers.lessThan(13d)); + assertThat(testCompression(kOrderedUUIDGen, 100000, 1000, 3, logger), Matchers.lessThan(14d)); + assertThat(testCompression(kOrderedUUIDGen, 100000, 100, 3, logger), Matchers.lessThan(19d)); + } + + public void testComparativeCompression() throws Exception { + Logger logger = LogManager.getLogger(UUIDTests.class); + + int numDocs = 100000; + int docsPerSecond = 1000; + int nodes = 3; + + double randomCompression = testCompression(randomUUIDGen, numDocs, docsPerSecond, nodes, logger); + double baseCompression = testCompression(timeUUIDGen, numDocs, docsPerSecond, nodes, logger); + double kOrderedCompression = testCompression(kOrderedUUIDGen, numDocs, docsPerSecond, nodes, logger); + + assertThat(kOrderedCompression, Matchers.lessThanOrEqualTo(baseCompression)); + assertThat(kOrderedCompression, Matchers.lessThanOrEqualTo(randomCompression)); + } + Set verifyUUIDSet(int count, UUIDGenerator uuidSource) { HashSet uuidSet = new HashSet<>(); for (int i = 0; i < count; ++i) { @@ -109,49 +145,62 @@ public void testUUIDThreaded(UUIDGenerator uuidSource) { assertEquals(count * uuids, globalSet.size()); } - public void testCompression() throws Exception { - Logger logger = LogManager.getLogger(UUIDTests.class); - // Low number so that the test runs quickly, but the results are more interesting with larger numbers - // of indexed documents - assertThat(testCompression(100000, 10000, 3, logger), Matchers.lessThan(14d)); // ~12 in practice - assertThat(testCompression(100000, 1000, 3, logger), Matchers.lessThan(15d)); // ~13 in practice - assertThat(testCompression(100000, 100, 3, logger), Matchers.lessThan(21d)); // ~20 in practice - } - - private static double testCompression(int numDocs, int numDocsPerSecond, int numNodes, Logger logger) throws Exception { - final double intervalBetweenDocs = 1000. / numDocsPerSecond; // milliseconds + private static double testCompression(final UUIDGenerator generator, int numDocs, int numDocsPerSecond, int numNodes, Logger logger) + throws Exception { + final double intervalBetweenDocs = 1000. / numDocsPerSecond; final byte[][] macAddresses = new byte[numNodes][]; Random r = random(); for (int i = 0; i < macAddresses.length; ++i) { macAddresses[i] = new byte[6]; random().nextBytes(macAddresses[i]); } - UUIDGenerator generator = new TimeBasedUUIDGenerator() { - double currentTimeMillis = TestUtil.nextLong(random(), 0L, 10000000000L); - @Override - protected long currentTimeMillis() { - currentTimeMillis += intervalBetweenDocs * 2 * r.nextDouble(); - return (long) currentTimeMillis; + UUIDGenerator uuidSource = generator; + if (generator instanceof TimeBasedUUIDGenerator) { + if (generator instanceof TimeBasedKOrderedUUIDGenerator) { + uuidSource = new TimeBasedKOrderedUUIDGenerator() { + double currentTimeMillis = TestUtil.nextLong(random(), 0L, 10000000000L); + + @Override + protected long currentTimeMillis() { + currentTimeMillis += intervalBetweenDocs * 2 * r.nextDouble(); + return (long) currentTimeMillis; + } + + @Override + protected byte[] macAddress() { + return RandomPicks.randomFrom(r, macAddresses); + } + }; + } else { + uuidSource = new TimeBasedUUIDGenerator() { + double currentTimeMillis = TestUtil.nextLong(random(), 0L, 10000000000L); + + @Override + protected long currentTimeMillis() { + currentTimeMillis += intervalBetweenDocs * 2 * r.nextDouble(); + return (long) currentTimeMillis; + } + + @Override + protected byte[] macAddress() { + return RandomPicks.randomFrom(r, macAddresses); + } + }; } + } - @Override - protected byte[] macAddress() { - return RandomPicks.randomFrom(r, macAddresses); - } - }; - // Avoid randomization which will slow down things without improving - // the quality of this test Directory dir = newFSDirectory(createTempDir()); IndexWriterConfig config = new IndexWriterConfig().setCodec(Codec.forName(Lucene.LATEST_CODEC)) - .setMergeScheduler(new SerialMergeScheduler()); // for reproducibility + .setMergeScheduler(new SerialMergeScheduler()); + IndexWriter w = new IndexWriter(dir, config); Document doc = new Document(); StringField id = new StringField("_id", "", Store.NO); doc.add(id); long start = System.nanoTime(); for (int i = 0; i < numDocs; ++i) { - id.setStringValue(generator.getBase64UUID()); + id.setStringValue(uuidSource.getBase64UUID()); w.addDocument(doc); } w.forceMerge(1); @@ -164,30 +213,25 @@ protected byte[] macAddress() { dir.close(); double bytesPerDoc = (double) size / numDocs; logger.info( - numDocs - + " docs indexed at " - + numDocsPerSecond - + " docs/s required " - + ByteSizeValue.ofBytes(size) - + " bytes of disk space, or " - + bytesPerDoc - + " bytes per document. Took: " - + new TimeValue(time) - + "." + "{} - {} docs indexed at {} docs/s required {} bytes of disk space, or {} bytes per document. Took: {}.", + uuidSource.getClass().getSimpleName(), + numDocs, + numDocsPerSecond, + ByteSizeValue.ofBytes(size), + bytesPerDoc, + new TimeValue(time) ); return bytesPerDoc; } public void testStringLength() { assertEquals(UUIDs.RANDOM_BASED_UUID_STRING_LENGTH, getUnpaddedBase64StringLength(RandomBasedUUIDGenerator.SIZE_IN_BYTES)); - assertEquals(UUIDs.RANDOM_BASED_UUID_STRING_LENGTH, UUIDs.randomBase64UUID().length()); - assertEquals(UUIDs.RANDOM_BASED_UUID_STRING_LENGTH, UUIDs.randomBase64UUID(random()).length()); - try (var secureString = UUIDs.randomBase64UUIDSecureString()) { - assertEquals(UUIDs.RANDOM_BASED_UUID_STRING_LENGTH, secureString.toString().length()); - } - assertEquals(UUIDs.TIME_BASED_UUID_STRING_LENGTH, getUnpaddedBase64StringLength(TimeBasedUUIDGenerator.SIZE_IN_BYTES)); - assertEquals(UUIDs.TIME_BASED_UUID_STRING_LENGTH, UUIDs.base64UUID().length()); + assertEquals(UUIDs.TIME_BASED_UUID_STRING_LENGTH, getUnpaddedBase64StringLength(TimeBasedKOrderedUUIDGenerator.SIZE_IN_BYTES)); + + assertEquals(UUIDs.RANDOM_BASED_UUID_STRING_LENGTH, randomUUIDGen.getBase64UUID().length()); + assertEquals(UUIDs.TIME_BASED_UUID_STRING_LENGTH, timeUUIDGen.getBase64UUID().length()); + assertEquals(UUIDs.TIME_BASED_UUID_STRING_LENGTH, kOrderedUUIDGen.getBase64UUID().length()); } private static int getUnpaddedBase64StringLength(int sizeInBytes) {