Fix dimensions_set_by_user for SageMaker OpenAI by daguimu · Pull Request #150834 · elastic/elasticsearch

daguimu · 2026-06-05T02:18:26Z

Summary

OpenAiTextEmbeddingPayload.ApiServiceSettings lost the dimensionsSetByUser flag across a persist → re-parse cycle. toXContent only wrote dimensions, so when an endpoint whose dimensions were auto-discovered (updateModelWithEmbeddingDetails(...) constructs ApiServiceSettings(dimensions, false)) was persisted and re-read, fromMap recomputed the flag as dimensions != null, flipping it to true.

Closes #146881

What changed

ConfigurationParseContext is threaded through SageMakerSchemaPayload#apiServiceSettings(...) (and SageMakerSchema / SageMakerServiceSettings.fromMap), with SageMakerModelBuilder passing REQUEST from fromRequest and PERSISTENT from fromStorage.
OpenAiTextEmbeddingPayload.ApiServiceSettings persists dimensions_set_by_user in toXContent. In a request the field is not parsed — a user-supplied value is left in the map and rejected as an unknown setting — and the flag is derived from whether dimensions was provided. In a persisted config the stored value is read back, defaulting to false when absent (configs persisted before the field existed), as BaseJinaAIEmbeddingsServiceSettings does.
ElasticTextEmbeddingPayload.ApiServiceSettings gets the same parse-context handling, and its request parsing no longer ignores whether dimensions was actually set. Since this class has always persisted the flag, a persisted config missing it now adds a validation error, as GoogleVertexAiEmbeddingsServiceSettings does.
The flag is internal: both payloads override getFilteredXContentObject() to omit it from the GET response, and SageMakerServiceSettings renders the api settings' filtered view rather than this.

Tests

Request context: supplying dimensions_set_by_user in parseRequestConfig returns a 400 (SageMakerServiceTests); the flag is derived from dimensions presence and the key is left unconsumed (SageMaker{OpenAi,ElasticTextEmbedding}ServiceSettingsTests).
Persisted context: the stored flag is read back; a missing flag defaults to false for OpenAI and adds a validation error for Elastic; parsePersistedConfig accepts stored configs containing the flag (SageMakerServiceTests).
testFilteredXContentObjectOmitsDimensionsSetByUser verifies the GET view omits the flag while the persisted form keeps it, and randomApiServiceSettings() was widened so the inherited XContent round-trip also exercises the auto-discovered case.

ApiServiceSettings.toXContent only wrote `dimensions`, so persisting an endpoint whose dimensions were auto-discovered (dimensionsSetByUser=false) and re-parsing via fromMap recomputed the flag as `dimensions != null`, flipping it to true. Persist `dimensions_set_by_user` and read it back (falling back to the previous inference when absent, for user requests and already-persisted configs). The flag is internal, so it is excluded from the filtered GET response: SageMakerServiceSettings now renders the api settings' getFilteredXContentObject(), and the OpenAI settings override it to omit the flag while keeping it in the persisted toXContent.

elasticsearchmachine · 2026-06-05T17:11:54Z

Pinging @elastic/search-inference-team (Team:Search - Inference)

DonalEvans

In order to fully correct the behaviour here, the ConfigurationParseContext needs to passed into the fromMap() implementation on OpenAiTextEmbeddingPayload.ApiServiceSettings. This allows us to know whether we're dealing with a user request (where dimensions_set_by_user should NOT be parsed from the map, because it's not a setting that should be directly user-settable) or parsing persisted objects (where dimensions_set_by_user should be a required field). With the code as it currently is, a user could specify dimensions in a request but also include "dimensions_set_by_user": false, which is obviously not right. To cover this behaviour, we should have new tests of the ApiServiceSettings.fromMap() method to confirm that requests are parsed correctly.

We should also add some tests to SageMakerServiceTests to confirm that specifying the dimensions results in the correct dimensions value being used, and that specifying dimensions_set_by_user in a request context (i.e. via calling parseRequestConfig()) results in an error. An example of how the second test could be written is:

public void testParseRequestConfig_TextEmbedding_DimensionsSetByUser_ReturnsError() {
    var modelBuilder = new SageMakerModelBuilder(new SageMakerSchemas());
    var sageMakerService = new SageMakerService(modelBuilder, client, schemas, sageMakerThreadPool, Map::of, mockClusterServiceEmpty());
    TestPlainActionFuture<Model> listener = new TestPlainActionFuture<>();
    var serviceSettings = new HashMap<>(Map.of(ACCESS_KEY_FIELD, "api_key",
        SECRET_KEY_FIELD, "secret_key",
        "endpoint_name", "endpoint",
        "api", "openai",
        "region", "region",
        DIMENSIONS_SET_BY_USER, true
    ));
    sageMakerService.parseRequestConfig("modelId", TaskType.TEXT_EMBEDDING, new HashMap<>(Map.of("service_settings", serviceSettings)), listener);

    var exception = expectThrows(ElasticsearchStatusException.class, () -> listener.actionGet(TEST_REQUEST_TIMEOUT));

    assertThat(exception.status(), is(RestStatus.BAD_REQUEST));
    assertThat(exception.getMessage(), is("Configuration contains settings [{dimensions_set_by_user=true}] unknown to the [sagemaker] service"));
}

While we're at it, it would be good to add tests for parsePersistedConfig() as well, to ensure that we don't return an error when dimensions_set_by_user is present in the map we parse for text_embedding models.

Regarding ElasticTextEmbeddingPayload.ApiServiceSettings, let's also fix that class. It has a slightly different bug in its fromMap() implementation where if only considers the parsed value of dimensionsSetByUser but not whether dimensions has actually been set, but will also need to be refactored to pass the appropriate ConfigurationParseContext in to the method and should have similar tests to the ones added for OpenAiTextEmbeddingPayload.

DonalEvans · 2026-06-05T17:20:44Z

@@ -134,23 +134,31 @@ public void writeTo(StreamOutput out) throws IOException {

    @Override
    public ToXContentObject getFilteredXContentObject() {


For consistency with other integrations, I think it would be better to have SageMakerServiceSettings extend FilteredXContentObject and implement the toXContentFragmentOfExposedFields() method instead of getFilteredXContentObject(). This is the approach used by almost all other service settings classes to exclude fields from the XContent.

Keeping the current structure per your call in the conversation — thanks!

DonalEvans · 2026-06-05T17:25:44Z

            return builder;
        }

+        @Override
+        public ToXContentObject getFilteredXContentObject() {


The ApiServiceSettings class could also be changed to extend FilteredXContentObject and implement toXContentFragmentOfExposedFields() instead of getFilteredXContentObject().

Keeping the getFilteredXContentObject() override per the same call — FilteredXContentObject#getFilteredXContentObject() wraps in startObject()/endObject(), which would break these flat fragments.

DonalEvans · 2026-06-05T17:32:11Z

+    public void testDimensionsNotSetByUserSurvivesXContentRoundTrip() throws IOException {
+        // An endpoint whose dimensions were auto-discovered (not set by the user) persists its dimensions, but must
+        // round-trip with dimensionsSetByUser=false; otherwise the stored config is wrongly re-read as user-set.
+        var autoDiscovered = new OpenAiTextEmbeddingPayload.ApiServiceSettings(randomIntBetween(1, 100), false);
+        var roundTripped = fromMutableMap(new HashMap<>(toMap(autoDiscovered)));
+        assertThat(roundTripped.dimensionsSetByUser(), equalTo(false));
+        assertThat(roundTripped, equalTo(autoDiscovered));
+    }


The behaviour in this test is already covered by InferenceSettingsTestCase.testXContentRoundTrip() on the base class, so I think this test is redundant.

Removed in 5a0945c.

DonalEvans · 2026-06-05T17:45:44Z

-            return new ApiServiceSettings(dimensions, dimensions != null);
+            // dimensions_set_by_user is persisted with stored configurations; it is absent from user requests, where it is
+            // inferred from whether dimensions were supplied. Reading it back preserves the distinction across a persist/parse cycle.
+            return new ApiServiceSettings(dimensions, dimensionsSetByUser != null ? dimensionsSetByUser : dimensions != null);


Since the dimensionsSetByUser field on ApiServiceSettings can never be null, it should probably be boolean rather than Boolean in the constructor.

Kept as a record Boolean component since the accessor has to override ServiceSettings#dimensionsSetByUser(), which returns Boolean — covered by your call to keep the structure as-is.

Thread ConfigurationParseContext through the SageMaker schema so the text-embedding payloads parse dimensions_set_by_user only from persisted configs; in a request it is left unparsed and rejected as an unknown setting. Apply the same fix to ElasticTextEmbeddingPayload (which also no longer ignores whether dimensions was set) and hide the flag from its GET response. Add request/persisted tests at the service and settings level, and drop the round-trip test already covered by the base case.

daguimu · 2026-06-08T02:25:50Z

Thanks for the thorough review @DonalEvans! Pushed in 5a0945c.

Done:

Threaded ConfigurationParseContext through SageMakerSchemaPayload.apiServiceSettings(...) → SageMakerSchema → SageMakerServiceSettings.fromMap(...), with SageMakerModelBuilder passing REQUEST from fromRequest and PERSISTENT from fromStorage. In a request the payloads no longer parse dimensions_set_by_user, so a user-supplied value is left in the map and rejected as an unknown setting; in a persisted config it is read back (falling back to dimensions != null for configs written before the field existed).
Added SageMakerServiceTests cases for both the request error and the persisted no-error path, plus request/persisted unit tests on ApiServiceSettings.fromMap for both payloads.
Fixed ElasticTextEmbeddingPayload the same way — its fromMap no longer ignores whether dimensions was actually set, and it now also keeps dimensions_set_by_user out of the GET response (it was previously rendering it).
Removed the redundant round-trip test (covered by InferenceSettingsTestCase#testXContentRoundTrip), and widened the OpenAI randomApiServiceSettings() so that base test now also exercises the auto-discovered case.

Two points I'd like your call on — I ran into a constraint with FilteredXContentObject:

The ApiServiceSettings types render as flat fragments inside SageMakerServiceSettings (isFragment() == true), but FilteredXContentObject#getFilteredXContentObject() wraps its output in startObject()/endObject(). Extending it would emit a nested object and break the flat layout, so I kept a getFilteredXContentObject() override that returns the exposed fragment (there is precedent for a record doing this in DeepSeekServiceSettings). For SageMakerServiceSettings itself there are no internal-only fields of its own — it just delegates to the api settings' filtered view — so extending FilteredXContentObject wouldn't change its output.
On boolean vs Boolean: ServiceSettings#dimensionsSetByUser() returns Boolean, and these are records, so a boolean component would generate a boolean accessor that can't override the Boolean method. I kept them as records, matching the other SageMaker payloads.

If you'd prefer I convert the payloads to classes to adopt the FilteredXContentObject + boolean shape despite the fragment wrapping, I'm glad to — just let me know.

DonalEvans · 2026-06-10T16:41:23Z

If you'd prefer I convert the payloads to classes to adopt the FilteredXContentObject + boolean shape despite the fragment wrapping, I'm glad to — just let me know.

I think we're fine keeping the structure as it is; refactoring these classes too heavily is outside the scope of this PR and the Sagemaker service is already doing a lot of things differently from other services, so some inconsistency is probably unavoidable.

I do think that SageMakerStoredServiceSchema (the interface implemented by the two ApiServiceSettings records) probably shouldn't be implementing ServiceSettings, since the ApiServiceSettings are not objects, but rather fragments, which should be handled using ToXContentFragment, but that can be left for another PR.

DonalEvans · 2026-06-10T17:10:28Z

@elasticsearchmachine test this

DonalEvans · 2026-06-10T16:53:57Z

+                dimensionsSetByUser = dimensions != null;
+            } else {
+                var storedDimensionsSetByUser = extractOptionalBoolean(serviceSettings, DIMENSIONS_SET_BY_USER_FIELD, validationException);
+                dimensionsSetByUser = storedDimensionsSetByUser != null ? storedDimensionsSetByUser : dimensions != null;


dimensions_set_by_user should never be null in a persisted context, since there has never been a time when it didn't exist for this class. If it's not present, then we should return an error, similar to what we do in GoogleVertexAiEmbeddingsServiceSettings.fromMap().

Done in 61467d2 — in a persisted context a missing value now adds a validation error, mirroring GoogleVertexAiEmbeddingsServiceSettings.fromMap().

DonalEvans · 2026-06-10T17:06:27Z

+            // In a persisted config we read the stored value, falling back to that same inference for configs written before it existed.
+            boolean dimensionsSetByUser;
+            if (ConfigurationParseContext.isRequestContext(context)) {
+                dimensionsSetByUser = dimensions != null;
+            } else {
+                var storedDimensionsSetByUser = extractOptionalBoolean(
+                    serviceSettings,
+                    ServiceFields.DIMENSIONS_SET_BY_USER,
+                    validationException
+                );
+                dimensionsSetByUser = storedDimensionsSetByUser != null ? storedDimensionsSetByUser : dimensions != null;
+            }


In a persisted context, dimensions should always be non-null, because even if a user doesn't set a value when creating the endpoint, we infer it from the validation call we do as part of endpoint creation. As such, the dimensions != null check will always resolve to true. In other integrations, we default to false when no value for dimensions_set_by_user was written (see BaseJinaAIEmbeddingsServiceSettings.fromMap()) so it would probably be better to do that here too.

Done in 61467d2 — defaults to false when absent from a persisted config, mirroring BaseJinaAIEmbeddingsServiceSettings.fromMap().

…dimensions-set-by-user

In a persisted config the OpenAI settings now default a missing dimensions_set_by_user to false (as BaseJinaAIEmbeddingsServiceSettings does), since dimensions is always present after endpoint validation and falling back to dimensions != null always resolved to true. The Elastic settings treat the field as required and add a validation error when it is missing (as GoogleVertexAiEmbeddingsServiceSettings does), since this class has always persisted the field.

daguimu · 2026-06-11T01:47:56Z

Pushed 61467d2 addressing both inline comments:

OpenAI: a missing dimensions_set_by_user in a persisted config now defaults to false (mirroring BaseJinaAIEmbeddingsServiceSettings.fromMap()) — you're right that the dimensions != null fallback always resolved to true there.
Elastic: the field is treated as required in a persisted context, adding a validation error when missing (mirroring GoogleVertexAiEmbeddingsServiceSettings.fromMap()), since this class has always persisted it.

Both are covered by new persisted-context tests (testFromStorage_MissingDimensionsSetByUser_DefaultsToFalse / ..._AddsValidationError).

I also merged main into the branch — it had fallen ~700 commits behind, and the bwc-snapshots failures in the last run didn't look related (this PR doesn't touch wire serialization and the bwc suites don't exercise SageMaker). Could you trigger CI again when you get a chance? Thanks!

DonalEvans · 2026-06-11T01:50:44Z

@elasticsearchmachine test this

DonalEvans · 2026-06-11T17:50:57Z

The serverless tests all passed, but the status isn't being reflected in GitHub for some reason, so I'll merge it.

@timestamp

* upstream/main: (523 commits) Upgrade netty to 4.1.135.Final (elastic#151099) ES|QL: Add benchmark for match at runtime (elastic#151071) ES|QL: Unmute test for elastic#150965 (elastic#151123) ES|QL: make the worker queue dynamic (elastic#151103) ES|QL: Add `briefSummary` element to FunctionInfo for all functions, add one-line summaries to overview lists (elastic#150546) Mute org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT test {csv-spec:approximation.Approximate stats by with zero variance} elastic#151120 Mute org.elasticsearch.compute.operator.topn.TopNOperatorTests testSplitOnSize elastic#151118 Mute org.elasticsearch.backwards.MixedClusterClientYamlTestSuiteIT test {p0=index/100_field_name_length_limit/Test field name length limit synthetic source} elastic#151117 Mute org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.DivTests testEvaluate {TestCase=<long>, <dense_vector>} elastic#151116 Mute org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.DivTests testEvaluate {TestCase=<double>, <dense_vector>} elastic#151115 Mute org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.DivTests testEvaluate {TestCase=<integer>, <dense_vector>} elastic#151114 Mute org.elasticsearch.telemetry.apm.internal.export.otelsdk.OtelSdkExportMeterSupplierTests testSpanProcessorSelfMonitoringMetricsFlowIntoHealthProvider elastic#151113 Fix ES94BloomFilterDocValuesFormatTests test failu (elastic#151084) Fix RandomizedTimeSeriesIT test failure (elastic#151082) Fix dimensions_set_by_user for SageMaker OpenAI (elastic#150834) Unmute some tests (elastic#151081) ESQL: Fix TSTEP auto-@timestamp import (elastic#151044) Reuse docs ids array in LuceneOperator (elastic#150828) Mute org.elasticsearch.test.apmintegration.OTelMetricsBufferingIT testOutageBuffersToDiskAndDrainsOnRecovery elastic#150387 Setting for trace sampling and batch processor (elastic#150876) ...

ApiServiceSettings.toXContent only wrote `dimensions`, so persisting an endpoint whose dimensions were auto-discovered (dimensionsSetByUser=false) and re-parsing via fromMap recomputed the flag as `dimensions != null`, flipping it to true. Persist `dimensions_set_by_user` and read it back (falling back to the previous inference when absent, for user requests and already-persisted configs). The flag is internal, so it is excluded from the filtered GET response: SageMakerServiceSettings now renders the api settings' getFilteredXContentObject(), and the OpenAI settings override it to omit the flag while keeping it in the persisted toXContent. Thread ConfigurationParseContext through the SageMaker schema so the text-embedding payloads parse dimensions_set_by_user only from persisted configs; in a request it is left unparsed and rejected as an unknown setting. Apply the same fix to ElasticTextEmbeddingPayload (which also no longer ignores whether dimensions was set) and hide the flag from its GET response. Add request/persisted tests at the service and settings level, and drop the round-trip test already covered by the base case. In a persisted config the OpenAI settings now default a missing dimensions_set_by_user to false (as BaseJinaAIEmbeddingsServiceSettings does), since dimensions is always present after endpoint validation and falling back to dimensions != null always resolved to true. The Elastic settings treat the field as required and add a validation error when it is missing (as GoogleVertexAiEmbeddingsServiceSettings does), since this class has always persisted the field.

elasticsearchmachine added v9.5.0 needs:triage Requires assignment of a team area label external-contributor Pull request authored by a developer outside the Elasticsearch team labels Jun 5, 2026

Add changelog for elastic#150834

50320e0

DonalEvans added Feature:GenAI Features around GenAI :SearchOrg/Inference Label for the Search Inference team Team:Search - Inference >bug labels Jun 5, 2026

DonalEvans self-assigned this Jun 5, 2026

elasticsearchmachine removed the needs:triage Requires assignment of a team area label label Jun 5, 2026

DonalEvans requested changes Jun 5, 2026

View reviewed changes

DonalEvans reviewed Jun 10, 2026

View reviewed changes

daguimu added 2 commits June 11, 2026 09:38

Merge remote-tracking branch 'origin/main' into fix/146881-sagemaker-…

88a0173

…dimensions-set-by-user

DonalEvans approved these changes Jun 11, 2026

View reviewed changes

DonalEvans merged commit ab9dab1 into elastic:main Jun 11, 2026
37 of 38 checks passed

		@@ -134,23 +134,31 @@ public void writeTo(StreamOutput out) throws IOException {

		@Override
		public ToXContentObject getFilteredXContentObject() {

Conversation

daguimu commented Jun 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Tests

Uh oh!

elasticsearchmachine commented Jun 5, 2026

Uh oh!

DonalEvans left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

daguimu commented Jun 8, 2026

Uh oh!

DonalEvans commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DonalEvans commented Jun 10, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

daguimu commented Jun 11, 2026

Uh oh!

DonalEvans commented Jun 11, 2026

Uh oh!

DonalEvans commented Jun 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

daguimu commented Jun 5, 2026 •

edited

Loading

DonalEvans commented Jun 10, 2026 •

edited

Loading