Adds shouldUseTriton methods that decides whether to use Triton for ONNX inference. Used by both Triton sidecar and TritonOnnxRuntime, making connection between them more explicit. (#35784)

Author: glebashnik

config-model/src/main/java/com/yahoo/vespa/model/container/xml/ContainerModelBuilder.java

@@ -248,15 +248,18 @@ private void addClusterContent(ApplicationContainerCluster cluster, Element spec
 
         addParameterStoreValidationHandler(cluster, deployState);
     }
-
-    private List<SidecarSpec> getSidecars(ApplicationContainerCluster cluster, DeployState deployState, NodesSpecification nodesSpecification) {
-        var sidecars = new ArrayList<SidecarSpec>();
-
+    
+    private boolean shouldUseTriton(ApplicationContainerCluster cluster, DeployState deployState) {
         var isPublicCloud = deployState.zone().system().isPublicCloudLike();
         var hasOnnxModels =  !cluster.onnxModelCostCalculator().models().isEmpty();
-        var useTritonFlagValue = deployState.featureFlags().useTriton();
+        var useTritonFeatureFlagValue = deployState.featureFlags().useTriton();
+        return useTritonFeatureFlagValue && isPublicCloud && hasOnnxModels;
+    }
 
-        if (useTritonFlagValue && isPublicCloud && hasOnnxModels) {
+    private List<SidecarSpec> getSidecars(ApplicationContainerCluster cluster, DeployState deployState, NodesSpecification nodesSpecification) {
+        var sidecars = new ArrayList<SidecarSpec>();
+        
+        if (shouldUseTriton(cluster, deployState)) {
             var hasGpu = !nodesSpecification.minResources().nodeResources().gpuResources().isZero();
 
             // Hardcoded values for changes to be reviewed and tested
@@ -974,7 +977,7 @@ protected void addModelEvaluationRuntime(DeployState deployState, ApplicationCon
         cluster.addPlatformBundle(ContainerModelEvaluation.MODEL_INTEGRATION_BUNDLE_FILE);
         cluster.addPlatformBundle(ContainerModelEvaluation.ONNXRUNTIME_BUNDLE_FILE);
         /* The ONNX runtime is always available for injection to any component */
-        if (deployState.featureFlags().useTriton()) {
+        if (shouldUseTriton(cluster, deployState)) {
             cluster.addSimpleComponent(
                     ContainerModelEvaluation.TRITON_ONNX_RUNTIME_CLASS, null, ContainerModelEvaluation.INTEGRATION_BUNDLE_NAME);
         } else {

config-model/src/test/java/com/yahoo/vespa/model/application/validation/change/RestartOnDeployForTritonOnnxRuntimeValidatorTest.java

@@ -3,14 +3,21 @@
 
 import com.yahoo.config.model.api.ConfigChangeAction;
 import com.yahoo.config.model.api.ConfigChangeRestartAction.ConfigChange;
+import com.yahoo.config.model.api.OnnxModelCost;
+import com.yahoo.config.model.api.OnnxModelOptions;
 import com.yahoo.config.model.deploy.DeployState;
 import com.yahoo.config.model.deploy.TestProperties;
+import com.yahoo.config.provision.Environment;
+import com.yahoo.config.provision.RegionName;
+import com.yahoo.config.provision.SystemName;
+import com.yahoo.config.provision.Zone;
 import com.yahoo.vespa.model.VespaModel;
 import com.yahoo.vespa.model.application.validation.ValidationTester;
 import com.yahoo.vespa.model.test.utils.VespaModelCreatorWithMockPkg;
 import org.junit.jupiter.api.Test;
 
 import java.util.List;
+import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -21,39 +28,39 @@
 public class RestartOnDeployForTritonOnnxRuntimeValidatorTest {
 
     private static final String SERVICES_XML_WITH_ONE_CLUSTER = """
-            <services version='1.0'>
-              <container id='cluster1' version='1.0'>
-              </container>
-            </services>
-            """;
+        <services version='1.0'>
+          <container id='cluster1' version='1.0'>
+          </container>
+        </services>
+        """;
 
     // Clusters must have different ports
     private static final String SERVICES_XML_TWO_CLUSTERS = """
-            <services version='1.0'>
-              <container id='cluster1' version='1.0'>
-                <http>
-                  <server id='server1' port='8080'/>
-                </http>
-              </container>
-              <container id='other-cluster' version='1.0'>
-                <http>
-                  <server id='server2' port='8081'/>
-                </http>
-              </container>
-            </services>
-            """;
+        <services version='1.0'>
+          <container id='cluster1' version='1.0'>
+            <http>
+              <server id='server1' port='8080'/>
+            </http>
+          </container>
+          <container id='other-cluster' version='1.0'>
+            <http>
+              <server id='server2' port='8081'/>
+            </http>
+          </container>
+        </services>
+        """;
 
     private static final String SERVICES_XML_OTHER_CLUSTER = """
-            <services version='1.0'>
-              <container id='other-cluster' version='1.0'>
-              </container>
-            </services>
-            """;
+        <services version='1.0'>
+          <container id='other-cluster' version='1.0'>
+          </container>
+        </services>
+        """;
 
     @Test
     void restart_when_triton_runtime_enabled() {
-        var previous = createModel(SERVICES_XML_WITH_ONE_CLUSTER,false);
-        var next = createModel(SERVICES_XML_WITH_ONE_CLUSTER,true);
+        var previous = createModel(SERVICES_XML_WITH_ONE_CLUSTER, false);
+        var next = createModel(SERVICES_XML_WITH_ONE_CLUSTER, true);
         var result = validateModel(previous, next);
 
         assertEquals(1, result.size());
@@ -66,7 +73,7 @@ void restart_when_triton_runtime_enabled() {
 
     @Test
     void restart_when_triton_runtime_disabled() {
-        var previous = createModel(SERVICES_XML_WITH_ONE_CLUSTER,true);
+        var previous = createModel(SERVICES_XML_WITH_ONE_CLUSTER, true);
         var next = createModel(SERVICES_XML_WITH_ONE_CLUSTER, false);
         var result = validateModel(previous, next);
 
@@ -89,7 +96,7 @@ void no_restart_when_triton_runtime_remains_enabled() {
 
     @Test
     void no_restart_when_triton_runtime_remains_disabled() {
-        var previous = createModel(SERVICES_XML_WITH_ONE_CLUSTER,false);
+        var previous = createModel(SERVICES_XML_WITH_ONE_CLUSTER, false);
         var next = createModel(SERVICES_XML_WITH_ONE_CLUSTER, false);
         var result = validateModel(previous, next);
 
@@ -115,19 +122,32 @@ void no_restart_when_cluster_with_triton_runtime_is_removed() {
     }
 
     private static List<ConfigChangeAction> validateModel(VespaModel current, VespaModel next) {
-        return ValidationTester.validateChanges(new RestartOnDeployForTritonOnnxRuntimeValidator(),
-                                                next,
-                                                deployStateBuilder(false).previousModel(current).build());
+        return ValidationTester.validateChanges(
+                new RestartOnDeployForTritonOnnxRuntimeValidator(),
+                next,
+                deployStateBuilder(false).previousModel(current).build());
     }
-    
+
     private static VespaModel createModel(String servicesXml, boolean useTriton) {
         var builder = deployStateBuilder(useTriton);
         return new VespaModelCreatorWithMockPkg(null, servicesXml).create(builder);
     }
 
     private static DeployState.Builder deployStateBuilder(boolean useTriton) {
-        return new DeployState.Builder()
-                .properties(new TestProperties().setUseTriton(useTriton));
+        var deployStateBuilder = new DeployState.Builder().properties(new TestProperties().setUseTriton(useTriton));
+
+        // Need a model and cloud to enable Triton.
+        if (useTriton) {
+            var mockModelCost = new OnnxModelCost.DisabledOnnxModelCost() {
+                @Override
+                public Map<String, ModelInfo> models() {
+                    return Map.of("modernbert", new ModelInfo("modernbert", 1, 1, OnnxModelOptions.empty()));
+                }
+            };
+            deployStateBuilder.onnxModelCost(mockModelCost);
+            deployStateBuilder.zone(new Zone(SystemName.PublicCd, Environment.dev, RegionName.defaultName()));
+        }
+
+        return deployStateBuilder;
     }
-
 }