diff --git a/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/CreateFunctionRequest.java b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/CreateFunctionRequest.java index 6b1d6b22b0..6ef3a3cde6 100644 --- a/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/CreateFunctionRequest.java +++ b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/CreateFunctionRequest.java @@ -341,8 +341,8 @@ static void validateModelFields(List models, throw new BadRequestException(MESG_MISSING_LLM_MODEL_URIS); } if (isLlmFunction) { - LlmConfigValidator.validateRoutingMethod( - model.getName(), model.getLlmConfig().getRoutingMethod()); + model.getLlmConfig().setRoutingMethod(LlmRoutingMethodValidator.validate( + model.getName(), model.getLlmConfig().getRoutingMethod())); LlmConfigValidator.validateTokenRateLimit( model.getName(), model.getLlmConfig().getTokenRateLimit()); } diff --git a/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidator.java b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidator.java index a77fb9a9ae..4dd0d7811a 100644 --- a/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidator.java +++ b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidator.java @@ -6,60 +6,28 @@ import com.nvidia.boot.exceptions.BadRequestException; import jakarta.annotation.Nullable; -import java.util.Locale; -import java.util.Set; import java.util.regex.Pattern; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; /** - * Rejects invalid {@code llmConfig} routingMethod/tokenRateLimit at create/update, so callers - * get a 400 up front instead of a late failure at invocation. + * Rejects an invalid {@code llmConfig} tokenRateLimit at create/update, so callers get a 400 up + * front instead of a late failure at invocation. */ @Slf4j public final class LlmConfigValidator { private LlmConfigValidator() {} - // Stargate LoadBalancerAlgorithm values; keep in sync. Blank = router default. - private static final Set VALID_ROUTING_METHODS = Set.of( - "power-of-two", - "wait-and-widen", - "round-robin", - "random", - "pulsar", - "pulsar-wait-and-widen", - // Deprecated Stargate aliases retained for existing deployments. - "groq-multiregion", - "pulsar-multiregion"); - // Comma-separated '-' entries, no unit repeated. private static final Pattern TOKEN_RATE_LIMIT_PATTERN = Pattern.compile( "^(?!.*-([SMHDW]).*-\\1)[1-9]\\d*-[SMHDW](,\\s*[1-9]\\d*-[SMHDW])*$"); - private static final String MESG_INVALID_ROUTING_METHOD = - "Invalid request: 'llmConfig.routingMethod' for model '%s' is invalid; supported " - + "values are [power-of-two, wait-and-widen, round-robin, random, pulsar, " - + "pulsar-wait-and-widen, groq-multiregion, pulsar-multiregion]"; private static final String MESG_INVALID_TOKEN_RATE_LIMIT = "Invalid request: 'llmConfig.tokenRateLimit' for model '%s' is invalid; expected " + "comma-separated '-' entries with unit in [S, M, H, D, W] " + "(for example '100000-S' or '10-M,5-S')"; - /** Rejects a routingMethod that is not one of the supported router algorithms. */ - public static void validateRoutingMethod(String modelName, @Nullable String routingMethod) { - if (StringUtils.isBlank(routingMethod)) { - return; - } - // Match the router: lowercase, '_' -> '-'. - var normalized = routingMethod.trim().toLowerCase(Locale.ROOT).replace('_', '-'); - if (!VALID_ROUTING_METHODS.contains(normalized)) { - var mesg = MESG_INVALID_ROUTING_METHOD.formatted(modelName); - log.error(mesg); - throw new BadRequestException(mesg); - } - } - /** Rejects a tokenRateLimit that is not '-' fragments. */ public static void validateTokenRateLimit(String modelName, @Nullable String tokenRateLimit) { if (StringUtils.isBlank(tokenRateLimit)) { @@ -67,7 +35,7 @@ public static void validateTokenRateLimit(String modelName, @Nullable String tok } if (!TOKEN_RATE_LIMIT_PATTERN.matcher(tokenRateLimit).matches()) { var mesg = MESG_INVALID_TOKEN_RATE_LIMIT.formatted(modelName); - log.error(mesg); + log.warn(mesg); throw new BadRequestException(mesg); } } diff --git a/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/LlmRoutingMethodValidator.java b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/LlmRoutingMethodValidator.java new file mode 100644 index 0000000000..c07f273cb3 --- /dev/null +++ b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/rest/function/management/dto/LlmRoutingMethodValidator.java @@ -0,0 +1,135 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.nvcf.rest.function.management.dto; + +import com.nvidia.boot.exceptions.BadRequestException; +import jakarta.annotation.Nullable; +import java.nio.charset.StandardCharsets; +import java.util.HashSet; +import java.util.Set; +import java.util.regex.Pattern; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; + +/** + * Validates the syntax of {@code llmConfig.routingMethod} at create/update: a method name + * optionally followed by {@code ;key=value} parameters. Methods and parameters are not + * interpreted here; the router owns their semantics and the value is stored as received. + */ +@Slf4j +public final class LlmRoutingMethodValidator { + + private static final int MAX_EXPRESSION_BYTES = 1024; + private static final int MAX_PARAMETERS = 32; + private static final Pattern METHOD_PATTERN = Pattern.compile("[A-Za-z][A-Za-z0-9_-]*"); + private static final Pattern PARAMETER_PATTERN = + Pattern.compile(" *([a-z][a-z0-9_]*)=(\\S(?:.*\\S)?)"); + private static final Pattern INTEGER_PATTERN = Pattern.compile("-?[0-9]{1,15}"); + private static final Pattern DECIMAL_PATTERN = Pattern.compile("-?[0-9]{1,12}\\.[0-9]{1,3}"); + private static final Pattern TOKEN_PATTERN = + Pattern.compile("[A-Za-z*][A-Za-z0-9!#$%&'*+.^_`|~:/-]*"); + private static final Pattern STRING_PATTERN = Pattern.compile( + "\"(?:[\\x20\\x21\\x23-\\x2b\\x2d-\\x3a\\x3c-\\x5b\\x5d-\\x7e]|\\\\[\"\\\\])*\""); + // \p{Cc} covers the C1 controls such as U+0085 (NEL), which \p{Cntrl} does not. + private static final Pattern CONTROL_CHARACTERS = + Pattern.compile("[\\p{Cc}\\p{Zl}\\p{Zp}]"); + + private static final String MESG_INVALID_ROUTING_METHOD = + "Invalid request: 'llmConfig.routingMethod' for model '%s' is invalid: %s"; + private static final String MESG_EXPRESSION_TOO_LONG = + "expression exceeds %d bytes".formatted(MAX_EXPRESSION_BYTES); + private static final String MESG_COMMAS_NOT_ALLOWED = "commas are not allowed"; + private static final String MESG_INVALID_METHOD_NAME = + "method name must match [A-Za-z][A-Za-z0-9_-]*"; + private static final String MESG_TOO_MANY_PARAMETERS = + "at most %d parameters are allowed".formatted(MAX_PARAMETERS); + private static final String MESG_INVALID_PARAMETER = + "parameter '%s' must be key=value with key matching [a-z][a-z0-9_]*"; + private static final String MESG_INVALID_VALUE = + "value for '%s' must be an integer, decimal, token, or quoted string"; + private static final String MESG_DUPLICATE_PARAMETER = "duplicate parameter '%s'"; + + private LlmRoutingMethodValidator() {} + + /** + * Rejects a routingMethod whose syntax the router could not parse and returns the value to + * store: the input without outer spaces, so validated and stored bytes are identical. Any + * other outer character, including tabs and line breaks, fails the grammar. + */ + @Nullable + public static String validate(String modelName, @Nullable String routingMethod) { + if (routingMethod == null) { + return null; + } + var value = StringUtils.strip(routingMethod, " "); + if (value.isEmpty()) { + return value; + } + if (value.getBytes(StandardCharsets.UTF_8).length > MAX_EXPRESSION_BYTES) { + reject(modelName, MESG_EXPRESSION_TOO_LONG); + } + if (value.contains(",")) { + reject(modelName, MESG_COMMAS_NOT_ALLOWED); + } + var segments = value.split(";", -1); + if (!METHOD_PATTERN.matcher(segments[0]).matches()) { + reject(modelName, MESG_INVALID_METHOD_NAME); + } + if (segments.length - 1 > MAX_PARAMETERS) { + reject(modelName, MESG_TOO_MANY_PARAMETERS); + } + var keys = new HashSet(); + for (var index = 1; index < segments.length; index++) { + validateParameter(modelName, segments[index], keys); + } + return value; + } + + private static void validateParameter(String modelName, String segment, Set keys) { + var parameter = PARAMETER_PATTERN.matcher(segment); + if (!parameter.matches()) { + reject(modelName, MESG_INVALID_PARAMETER.formatted(withoutControlCharacters(segment))); + } + var key = parameter.group(1); + if (!isValidBareValue(parameter.group(2))) { + reject(modelName, MESG_INVALID_VALUE.formatted(key)); + } + if (!keys.add(key)) { + reject(modelName, MESG_DUPLICATE_PARAMETER.formatted(key)); + } + } + + // The segment is raw request text echoed in the log and the 400 body; a line break in it + // could forge a log line. + private static String withoutControlCharacters(String segment) { + return CONTROL_CHARACTERS.matcher(segment).replaceAll("?"); + } + + private static boolean isValidBareValue(String value) { + return INTEGER_PATTERN.matcher(value).matches() + || DECIMAL_PATTERN.matcher(value).matches() + || TOKEN_PATTERN.matcher(value).matches() + || STRING_PATTERN.matcher(value).matches(); + } + + // A malformed client value is not an operator problem, so it is logged below error level. + private static void reject(String modelName, String rule) { + var mesg = MESG_INVALID_ROUTING_METHOD.formatted(withoutControlCharacters(modelName), rule); + log.warn(mesg); + throw new BadRequestException(mesg); + } +} diff --git a/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/service/function/FunctionLlmService.java b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/service/function/FunctionLlmService.java index bea23531fc..c5a3da7fe8 100644 --- a/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/service/function/FunctionLlmService.java +++ b/src/control-plane-services/cloud-functions/nvcf-core/src/main/java/com/nvidia/nvcf/service/function/FunctionLlmService.java @@ -26,6 +26,7 @@ import com.nvidia.nvcf.rest.function.management.dto.FunctionTypeEnum; import com.nvidia.nvcf.rest.function.management.dto.LlmConfigValidator; import com.nvidia.nvcf.rest.function.management.dto.LlmInvocationConfigDto; +import com.nvidia.nvcf.rest.function.management.dto.LlmRoutingMethodValidator; import com.nvidia.nvcf.rest.function.management.dto.UpdateFunctionRequest; import jakarta.annotation.Nullable; import java.util.Comparator; @@ -264,7 +265,8 @@ private Map propagateModelUpdatesToSiblings( UpdateFunctionRequest.ModelUpdateDto::modelName, u -> FunctionModelDto.LlmConfigDto.builder() .tokenRateLimit(u.llmConfig().tokenRateLimit()) - .routingMethod(u.llmConfig().routingMethod()) + .routingMethod(LlmRoutingMethodValidator.validate( + u.modelName(), u.llmConfig().routingMethod())) .build())); if (overrides.isEmpty()) { return Map.of(); @@ -403,9 +405,8 @@ private void updateModels( llmConfig.setTokenRateLimit(llmConfigUpdate.tokenRateLimit()); } if (llmConfigUpdate.routingMethod() != null) { - LlmConfigValidator.validateRoutingMethod( - modelUpdate.modelName(), llmConfigUpdate.routingMethod()); - llmConfig.setRoutingMethod(llmConfigUpdate.routingMethod()); + llmConfig.setRoutingMethod(LlmRoutingMethodValidator.validate( + modelUpdate.modelName(), llmConfigUpdate.routingMethod())); } updated = true; break; diff --git a/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/FunctionsWithLlmModelsTest.java b/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/FunctionsWithLlmModelsTest.java index af00e91836..1ccaa7bd7b 100644 --- a/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/FunctionsWithLlmModelsTest.java +++ b/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/FunctionsWithLlmModelsTest.java @@ -595,6 +595,55 @@ private FunctionDto createAdditionalLlmFunctionVersion( return response.getBody().function(); } + @Test + void shouldStoreRoutingExpressionsWithoutOuterSpacesOnCreateAndUpdate() { + var functionName = TEST_FUNCTION_NAME + "-" + Instant.now().toEpochMilli(); + var storedRoutingMethod = "Pulsar_Wait_And_Widen; seed=stable-a;n=2"; + var function = createInitialLlmFunction( + functionName, "1-M", " " + storedRoutingMethod + " "); + + assertThat(function.models().getFirst().getLlmConfig().getRoutingMethod()) + .isEqualTo(storedRoutingMethod); + assertLlmConfigPersisted(function.versionId(), "1-M", storedRoutingMethod); + + // A different value forces the sibling write, which copies the converter's trimmed DTO. + var secondRoutingMethod = "wait-and-widen;n=3"; + var secondVersion = createAdditionalLlmFunctionVersion( + function.id(), functionName, "1-M", " " + secondRoutingMethod + " "); + assertThat(secondVersion.models().getFirst().getLlmConfig().getRoutingMethod()) + .isEqualTo(secondRoutingMethod); + assertLlmConfigPersisted(function.versionId(), "1-M", secondRoutingMethod); + assertLlmConfigPersisted(secondVersion.versionId(), "1-M", secondRoutingMethod); + + // Unknown method and parameter: well formed, so it persists; the router owns semantics. + var updatedRoutingMethod = "fastest;widen=2"; + var updateToken = MOCK_OAUTH2_TOKEN_SERVER.getJwt( + TEST_CLIENT_SUBJECT, List.of(SCOPE_UPDATE_FUNCTION), 100); + var updateRequest = UpdateFunctionRequest.builder() + .modelUpdates(List.of(UpdateFunctionRequest.ModelUpdateDto.builder() + .modelName(TEST_LLM_MODEL_NAME) + .llmConfig(UpdateFunctionRequest.LlmConfigUpdateDto.builder() + .routingMethod(" " + updatedRoutingMethod + " ") + .build()) + .build())) + .build(); + var updateEntity = RequestEntity.put(URI.create("/v2/nvcf/functions/" + function.id() + + "/versions/" + function.versionId())) + .contentType(MediaType.APPLICATION_JSON) + .header("Authorization", "Bearer " + updateToken) + .body(updateRequest); + + var response = testRestTemplate.exchange(updateEntity, FunctionResponse.class); + + assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK); + assertThat(response.getBody()).isNotNull(); + var updatedModel = response.getBody().function().models().getFirst(); + assertThat(updatedModel.getLlmConfig().getRoutingMethod()) + .isEqualTo(updatedRoutingMethod); + assertLlmConfigPersisted(function.versionId(), "1-M", updatedRoutingMethod); + assertLlmConfigPersisted(secondVersion.versionId(), "1-M", updatedRoutingMethod); + } + @Test void shouldRejectCreateWithInvalidRoutingMethod() { var createToken = MOCK_OAUTH2_TOKEN_SERVER.getJwt(TEST_CLIENT_SUBJECT, @@ -605,7 +654,7 @@ void shouldRejectCreateWithInvalidRoutingMethod() { .inferenceUrl(TEST_INFERENCE_URL) .inferencePort(TEST_INFERENCE_PORT) .functionType(FunctionTypeEnum.LLM) - .models(List.of(llmModel("1-M", "not-a-method"))) + .models(List.of(llmModel("1-M", "pulsar,seed=x"))) .build(); var createEntity = RequestEntity.post(URI.create("/v2/nvcf/functions")) .contentType(MediaType.APPLICATION_JSON) @@ -615,7 +664,36 @@ void shouldRejectCreateWithInvalidRoutingMethod() { var response = testRestTemplate.exchange(createEntity, String.class); assertThat(response.getStatusCode()).isEqualTo(HttpStatus.BAD_REQUEST); - assertThat(response.getBody()).contains("llmConfig.routingMethod"); + assertThat(response.getBody()).contains( + "llmConfig.routingMethod", TEST_LLM_MODEL_NAME, "commas are not allowed"); + } + + @Test + void shouldRejectVersionCreateWithInvalidRoutingMethod() { + var functionName = TEST_FUNCTION_NAME + "-" + Instant.now().toEpochMilli(); + var function = createInitialLlmFunction(functionName, "1-M", "round-robin"); + + var createToken = MOCK_OAUTH2_TOKEN_SERVER.getJwt(TEST_CLIENT_SUBJECT, + List.of(SCOPE_REGISTER_FUNCTION), 100); + var createRequest = CreateFunctionRequest.builder() + .name(functionName) + .containerImage(TEST_NGC_CONTAINER_IMAGE) + .inferenceUrl(TEST_INFERENCE_URL) + .inferencePort(TEST_INFERENCE_PORT) + .functionType(FunctionTypeEnum.LLM) + .models(List.of(llmModel("1-M", "pulsar;seed="))) + .build(); + var createEntity = RequestEntity.post(URI.create( + "/v2/nvcf/functions/" + function.id() + "/versions")) + .contentType(MediaType.APPLICATION_JSON) + .header("Authorization", "Bearer " + createToken) + .body(createRequest); + + var response = testRestTemplate.exchange(createEntity, String.class); + + assertThat(response.getStatusCode()).isEqualTo(HttpStatus.BAD_REQUEST); + assertThat(response.getBody()).contains( + "llmConfig.routingMethod", TEST_LLM_MODEL_NAME, "must be key=value"); } @Test @@ -652,7 +730,7 @@ void shouldRejectUpdateWithInvalidRoutingMethod() { .modelUpdates(List.of(UpdateFunctionRequest.ModelUpdateDto.builder() .modelName(TEST_LLM_MODEL_NAME) .llmConfig(UpdateFunctionRequest.LlmConfigUpdateDto.builder() - .routingMethod("not-a-method") + .routingMethod("pulsar;n=?1") .build()) .build())) .build(); @@ -665,7 +743,8 @@ void shouldRejectUpdateWithInvalidRoutingMethod() { var response = testRestTemplate.exchange(updateEntity, String.class); assertThat(response.getStatusCode()).isEqualTo(HttpStatus.BAD_REQUEST); - assertThat(response.getBody()).contains("llmConfig.routingMethod"); + assertThat(response.getBody()).contains( + "llmConfig.routingMethod", TEST_LLM_MODEL_NAME, "value for 'n'"); } @Test diff --git a/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidatorTest.java b/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidatorTest.java index 62afbb9a91..83f7ff5d1f 100644 --- a/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidatorTest.java +++ b/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/dto/LlmConfigValidatorTest.java @@ -8,7 +8,6 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import com.nvidia.boot.exceptions.BadRequestException; -import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.NullAndEmptySource; import org.junit.jupiter.params.provider.ValueSource; @@ -17,36 +16,6 @@ class LlmConfigValidatorTest { private static final String MODEL = "meta/llama-3.1-8b-instruct"; - @ParameterizedTest - @ValueSource(strings = { - "power-of-two", "wait-and-widen", "round-robin", "random", "pulsar", - "pulsar-wait-and-widen", "groq-multiregion", "pulsar-multiregion", - // Router normalizes case and '_' to '-', so these are accepted too. - "Power-Of-Two", "power_of_two", "wait_and_widen", "pulsar_wait_and_widen", - "groq_multiregion", "pulsar_multiregion", " pulsar " - }) - void validRoutingMethodsAccepted(String routingMethod) { - assertThatCode(() -> LlmConfigValidator.validateRoutingMethod(MODEL, routingMethod)) - .doesNotThrowAnyException(); - } - - @ParameterizedTest - @NullAndEmptySource - @ValueSource(strings = {" "}) - void blankRoutingMethodAccepted(String routingMethod) { - assertThatCode(() -> LlmConfigValidator.validateRoutingMethod(MODEL, routingMethod)) - .doesNotThrowAnyException(); - } - - @ParameterizedTest - @ValueSource(strings = {"weighted", "sticky", "not-a-method", "round robin", "power-of-3"}) - void invalidRoutingMethodsRejected(String routingMethod) { - assertThatThrownBy(() -> LlmConfigValidator.validateRoutingMethod(MODEL, routingMethod)) - .isInstanceOf(BadRequestException.class) - .hasMessageContaining("routingMethod") - .hasMessageContaining(MODEL); - } - @ParameterizedTest @ValueSource(strings = {"100000-S", "10-M", "5-H", "1-D", "2-W", "10-M,5-S", "10-M, 5-S"}) void validTokenRateLimitsAccepted(String tokenRateLimit) { diff --git a/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/dto/LlmRoutingMethodValidatorTest.java b/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/dto/LlmRoutingMethodValidatorTest.java new file mode 100644 index 0000000000..e5ff225f75 --- /dev/null +++ b/src/control-plane-services/cloud-functions/nvcf-core/src/test/java/com/nvidia/nvcf/rest/function/management/dto/LlmRoutingMethodValidatorTest.java @@ -0,0 +1,186 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.nvcf.rest.function.management.dto; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import com.nvidia.boot.exceptions.BadRequestException; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.NullAndEmptySource; +import org.junit.jupiter.params.provider.ValueSource; + +class LlmRoutingMethodValidatorTest { + + private static final String MODEL = "meta/llama-3.1-8b-instruct"; + + @ParameterizedTest + @ValueSource(strings = { + "power-of-two", "wait-and-widen", "round-robin", "random", "pulsar", + "pulsar-wait-and-widen", "groq-multiregion", "pulsar-multiregion", + "Power-Of-Two", "power_of_two", "wait_and_widen", "pulsar_wait_and_widen", + "groq_multiregion", "pulsar_multiregion", " pulsar ", + "pulsar;seed=stable-a", + "pulsar; seed=stable-a; consider_kv_free_tokens=true", + "pulsar-wait-and-widen;seed=stable-a;n=2;max_queue_time_floor_ms=100;" + + "max_queue_time_ceil_ms=500", + "wait-and-widen;next_bucket_unlock_factor=\"0.0625\"", + "power-of-n;sample_count=4;comparator=queue-time", + "wait-and-widen;max_input_work_seconds=1.5", + "wait-and-widen;n=-1", "fastest;widen=2", + "pulsar; seed=x", " pulsar;seed=x ", + "pulsar;n=999999999999999", "pulsar;n=-999999999999999", + "pulsar;n=999999999999.999", "pulsar;n=-999999999999.999", + "pulsar;seed=*", "pulsar;seed=a!#$%&'*+.^_`|~:/-", + "pulsar;seed=\"\"", "pulsar;seed=\"a b\"", "pulsar;seed=\"a\\\"b\\\\c\"" + }) + @MethodSource("routingMethodsAtLimits") + void validRoutingMethodsAccepted(String routingMethod) { + assertThatCode(() -> LlmRoutingMethodValidator.validate(MODEL, routingMethod)) + .doesNotThrowAnyException(); + } + + @ParameterizedTest + @NullAndEmptySource + @ValueSource(strings = {" "}) + void blankRoutingMethodAccepted(String routingMethod) { + assertThatCode(() -> LlmRoutingMethodValidator.validate(MODEL, routingMethod)) + .doesNotThrowAnyException(); + } + + @ParameterizedTest(name = "{index}: {0} violates {1}") + @MethodSource("invalidRoutingMethods") + void invalidRoutingMethodsRejected(String routingMethod, String rule) { + assertThatThrownBy(() -> LlmRoutingMethodValidator.validate(MODEL, routingMethod)) + .isInstanceOf(BadRequestException.class) + .hasMessageContaining("routingMethod") + .hasMessageContaining(MODEL) + .hasMessageContaining(rule); + } + + @Test + void controlCharactersInRejectedSegmentReplaced() { + var routingMethod = "pulsar;seed=a\nb\rc"; + + assertThatThrownBy(() -> LlmRoutingMethodValidator.validate(MODEL, routingMethod)) + .isInstanceOf(BadRequestException.class) + .hasMessageContaining("parameter 'seed=a?b?c' must be key=value") + .hasMessageNotContaining("\n") + .hasMessageNotContaining("\r"); + } + + @Test + void nextLineControlCharacterInRejectedSegmentReplaced() { + var nextLine = String.valueOf((char) 0x85); + var routingMethod = "pulsar;seed=a" + nextLine + "b"; + + assertThatThrownBy(() -> LlmRoutingMethodValidator.validate(MODEL, routingMethod)) + .isInstanceOf(BadRequestException.class) + .hasMessageContaining("parameter 'seed=a?b' must be key=value") + .hasMessageNotContaining(nextLine); + } + + @Test + void controlCharactersInModelNameReplaced() { + assertThatThrownBy(() -> LlmRoutingMethodValidator.validate("m\nx", "pulsar,seed=x")) + .isInstanceOf(BadRequestException.class) + .hasMessageContaining("for model 'm?x'") + .hasMessageNotContaining("\n"); + } + + @ParameterizedTest(name = "{index}: [{0}] is stored as [{1}]") + @MethodSource("storedValues") + void returnsValueWithoutOuterSpaces(String routingMethod, String stored) { + assertThat(LlmRoutingMethodValidator.validate(MODEL, routingMethod)).isEqualTo(stored); + } + + private static Stream storedValues() { + return Stream.of( + Arguments.of(null, null), + Arguments.of("", ""), + Arguments.of(" ", ""), + Arguments.of("pulsar", "pulsar"), + Arguments.of(" pulsar ", "pulsar"), + Arguments.of(" Power_Of_Two;seed=x ", "Power_Of_Two;seed=x"), + Arguments.of("pulsar; seed=x", "pulsar; seed=x")); + } + + private static Stream routingMethodsAtLimits() { + return Stream.of( + "a".repeat(1024), " " + "a".repeat(1024) + " ", "pulsar" + parameters(32)); + } + + private static Stream invalidRoutingMethods() { + return Stream.of( + Arguments.of("round robin", "method name must match"), + Arguments.of("power-of-3!", "method name must match"), + Arguments.of(";seed=x", "method name must match"), + Arguments.of("pulsar\n", "method name must match"), + Arguments.of("\tpulsar", "method name must match"), + Arguments.of("\n", "method name must match"), + Arguments.of("pulsar;seed=x\n", "must be key=value"), + Arguments.of("pulsar,seed=x", "commas are not allowed"), + Arguments.of("pulsar;seed=\"a,b\"", "commas are not allowed"), + Arguments.of("pulsar;seed=\"a;b\"", "value for 'seed'"), + Arguments.of("pulsar ;seed=x", "method name must match"), + Arguments.of("pulsar;seed = x", "must be key=value"), + Arguments.of("pulsar;seed= x", "must be key=value"), + Arguments.of("pulsar;seed=x ;n=1", "must be key=value"), + Arguments.of("pulsar;\tseed=x", "must be key=value"), + Arguments.of("pulsar;seed", "must be key=value"), + Arguments.of("pulsar;seed=", "must be key=value"), + Arguments.of("pulsar;=x", "must be key=value"), + Arguments.of("pulsar;Seed=x", "must be key=value"), + Arguments.of("pulsar;2n=1", "must be key=value"), + Arguments.of("pulsar;max-queued=1", "must be key=value"), + Arguments.of("pulsar;_n=1", "must be key=value"), + Arguments.of("pulsar;seed=1abc", "value for 'seed'"), + Arguments.of("pulsar;seed=-abc", "value for 'seed'"), + Arguments.of("pulsar;seed=a\tb", "value for 'seed'"), + Arguments.of("pulsar;seed=a" + (char) 0x7f + "b", "value for 'seed'"), + Arguments.of("pulsar;n=?1", "value for 'n'"), + Arguments.of("pulsar;n=?0", "value for 'n'"), + Arguments.of("pulsar;seed=:YQ==:", "value for 'seed'"), + Arguments.of("pulsar;n=1.2345", "value for 'n'"), + Arguments.of("pulsar;n=1000000000000000", "value for 'n'"), + Arguments.of("pulsar;n=1000000000000.1", "value for 'n'"), + Arguments.of("pulsar;n=1.", "value for 'n'"), + Arguments.of("pulsar;n=1;n=2", "duplicate parameter 'n'"), + Arguments.of("pulsar" + parameters(33), "at most 32 parameters"), + Arguments.of("a".repeat(1025), "expression exceeds 1024 bytes"), + Arguments.of("pulsar;seed=\"" + "\u00e9".repeat(506) + "\"", + "expression exceeds 1024 bytes"), + Arguments.of("pulsar;seed=x;", "must be key=value"), + Arguments.of("pulsar;seed=\"unterminated", "value for 'seed'"), + Arguments.of("pulsar;seed=\"a\\nb\"", "value for 'seed'"), + Arguments.of("pulsar;seed=\"\u00e9\"", "value for 'seed'"), + Arguments.of("pulsar;seed=\"a\tb\"", "value for 'seed'")); + } + + private static String parameters(int count) { + return IntStream.range(0, count) + .mapToObj(index -> ";p" + index + "=x") + .collect(Collectors.joining()); + } +}