Add config for number of GPUs available to ONNX runtime

1 files changed, 21 insertions, 6 deletions

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java b/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java
index f901bf3c826..9e21fd2d23a 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java
@@ -9,6 +9,7 @@ import com.yahoo.config.model.producer.TreeConfigProducer;
 import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.search.config.QrStartConfig;
+import com.yahoo.vespa.config.search.core.OnnxModelsConfig;
 import com.yahoo.vespa.model.container.component.SimpleComponent;
 import java.time.Duration;
 import java.util.Optional;
@@ -20,6 +21,7 @@ import java.util.Optional;
  */
 public final class ApplicationContainer extends Container implements
         QrStartConfig.Producer,
+        OnnxModelsConfig.Producer,
         ZookeeperServerConfig.Producer {
 
     private final boolean isHostedVespa;
@@ -42,12 +44,15 @@ public final class ApplicationContainer extends Container implements
 
     @Override
     public void getConfig(QrStartConfig.Builder builder) {
-        if (getHostResource() != null) {
-            NodeResources nodeResources = getHostResource().realResources();
-            if ( ! nodeResources.isUnspecified()) {
-                builder.jvm.availableProcessors(Math.max(2, (int)Math.ceil(nodeResources.vcpu())));
-            }
-        }
+        realResources().ifPresent(r -> builder.jvm.availableProcessors(Math.max(2, (int) Math.ceil(r.vcpu()))));
+    }
+
+    @Override
+    public void getConfig(OnnxModelsConfig.Builder builder) {
+        realResources().ifPresent(r -> {
+            int count = r.gpuResources().count();
+            if (count >= 0) builder.gpu.count(count);
+        });
     }
 
     @Override
@@ -84,4 +89,14 @@ public final class ApplicationContainer extends Container implements
 
     @Override public Optional<String> getPreShutdownCommand() { return Optional.of(prepareStopCommand(Duration.ofMinutes(6))); }
 
+    private Optional<NodeResources> realResources() {
+        if (getHostResource() != null) {
+            NodeResources nodeResources = getHostResource().realResources();
+            if ( ! nodeResources.isUnspecified()) {
+                return Optional.of(nodeResources);
+            }
+        }
+        return Optional.empty();
+    }
+
 }