feat: support container parameters overrides (#17)

epam · Jan 13, 2025 · c7eeb4e · c7eeb4e
1 parent e358549
commit c7eeb4e
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 57 deletions.
diff --git a/src/main/java/com/epam/aidial/config/AppConfiguration.java b/src/main/java/com/epam/aidial/config/AppConfiguration.java
@@ -1,6 +1,7 @@
 package com.epam.aidial.config;
 
 import com.epam.aidial.kubernetes.knative.V1Service;
+import io.kubernetes.client.openapi.models.V1Container;
 import io.kubernetes.client.openapi.models.V1Job;
 import io.kubernetes.client.openapi.models.V1Secret;
 import io.kubernetes.client.util.Yaml;
@@ -27,6 +28,18 @@ public class AppConfiguration {
     private V1Service serviceConfig;
     private String serviceConfigString;
 
+    @Getter
+    private V1Container templateContainer;
+    private String templateContainerString;
+
+    @Getter
+    private V1Container builderContainer;
+    private String builderContainerString;
+
+    @Getter
+    private V1Container serviceContainer;
+    private String serviceContainerString;
+
     @Getter
     @Setter
     private Map<String, RuntimeConfiguration> runtimes;
@@ -46,6 +59,21 @@ public void setServiceConfig(V1Service serviceConfig) {
         this.serviceConfigString = Yaml.dump(serviceConfig);
     }
 
+    public void setTemplateContainer(V1Container container) {
+        this.templateContainer = container;
+        this.templateContainerString = Yaml.dump(container);
+    }
+
+    public void setBuilderContainer(V1Container container) {
+        this.builderContainer = container;
+        this.builderContainerString = Yaml.dump(container);
+    }
+
+    public void setServiceContainer(V1Container container) {
+        this.serviceContainer = container;
+        this.serviceContainerString = Yaml.dump(container);
+    }
+
     public V1Secret cloneSecretConfig() {
         return Yaml.loadAs(secretConfigString, V1Secret.class);
     }
@@ -58,6 +86,18 @@ public V1Service cloneServiceConfig() {
         return Yaml.loadAs(serviceConfigString, V1Service.class);
     }
 
+    public V1Container cloneTemplateContainer() {
+        return Yaml.loadAs(templateContainerString, V1Container.class);
+    }
+
+    public V1Container cloneBuilderContainer() {
+        return Yaml.loadAs(builderContainerString, V1Container.class);
+    }
+
+    public V1Container cloneServiceContainer() {
+        return Yaml.loadAs(serviceContainerString, V1Container.class);
+    }
+
     @Data
     public static class RuntimeConfiguration {
         private String image;

diff --git a/src/main/java/com/epam/aidial/service/ConfigService.java b/src/main/java/com/epam/aidial/service/ConfigService.java
@@ -64,15 +64,6 @@ public class ConfigService {
     private final RegistryService registryService;
     private final AppConfiguration appconfig;
 
-    @Value("${app.template-container.name}")
-    private final String pullerContainer;
-
-    @Value("${app.builder-container.name}")
-    private final String builderContainer;
-
-    @Value("${app.service-container.name}")
-    private final String serviceContainer;
-
     @Value("${app.docker-config-path}")
     private final String dockerConfigPath;
 
@@ -113,21 +104,21 @@ public V1Job buildJobConfig(String name, String sources, String runtime) {
         MappingChain<V1PodSpec> podSpec = config.get(JOB_SPEC_FIELD)
                 .get(JOB_TEMPLATE_FIELD)
                 .get(JOB_TEMPLATE_SPEC_FIELD);
-        MappingChain<V1Container> puller = podSpec.getList(POD_INIT_CONTAINERS_FIELD, CONTAINER_NAME)
-                .get(pullerContainer);
-        ListMapper<V1EnvVar> pullerEnvs = puller.getList(CONTAINER_ENV_FIELD, ENV_VAR_NAME);
+        MappingChain<V1Container> template = podSpec.getList(POD_INIT_CONTAINERS_FIELD, CONTAINER_NAME)
+                .getOrDefault(appconfig.getTemplateContainer().getName(), appconfig::cloneTemplateContainer);
+        ListMapper<V1EnvVar> pullerEnvs = template.getList(CONTAINER_ENV_FIELD, ENV_VAR_NAME);
         pullerEnvs.get("SOURCES")
                 .data()
                 .setValue(sources);
         pullerEnvs.get("PROFILE")
                 .data()
                 .setValue(runtimeConfig.getProfile());
         String secretName = dialAuthSecretName(name);
-        puller.get(CONTAINER_ENV_FROM_FIELD)
+        template.get(CONTAINER_ENV_FROM_FIELD)
                 .data()
                 .add(new V1EnvFromSource().secretRef(new V1SecretEnvSource().name(secretName)));
         MappingChain<V1Container> builder = podSpec.getList(POD_CONTAINERS_FIELD, CONTAINER_NAME)
-                .get(builderContainer);
+                .getOrDefault(appconfig.getBuilderContainer().getName(), appconfig::cloneBuilderContainer);
         builder.get(CONTAINER_ARGS_FIELD)
                 .data()
                 .addAll(List.of(
@@ -178,7 +169,7 @@ public V1Service appServiceConfig(
         MappingChain<V1Container> container = template
                 .get(SERVICE_TEMPLATE_SPEC_FIELD)
                 .getList(TEMPLATE_CONTAINERS_FIELD, CONTAINER_NAME)
-                .get(serviceContainer);
+                .getOrDefault(appconfig.getServiceContainer().getName(), appconfig::cloneServiceContainer);
 
         container.data()
                 .setImage(Objects.requireNonNullElse(image, registryService.fullImageName(name)));

diff --git a/src/main/java/com/epam/aidial/util/mapping/ListMapper.java b/src/main/java/com/epam/aidial/util/mapping/ListMapper.java
@@ -1,33 +1,37 @@
-package com.epam.aidial.util.mapping;
-
-import java.util.List;
-import java.util.Map;
-import java.util.function.BiConsumer;
-import java.util.function.Function;
-import java.util.function.Supplier;
-import java.util.stream.Collectors;
-
-public class ListMapper<T> {
-    private final List<T> list;
-    private final Supplier<T> itemFactory;
-    private final BiConsumer<T, String> nameSetter;
-    private final Map<String, T> map;
-
-    public ListMapper(List<T> list, NamedItemMapper<T> itemMapper) {
-        this.list = list;
-        this.itemFactory = itemMapper.factory();
-        this.nameSetter = itemMapper.setter();
-        this.map = list.stream()
-                .collect(Collectors.toMap(itemMapper.getter(), Function.identity()));
-    }
-
-    public MappingChain<T> get(String name) {
-        return new MappingChain<>(map.computeIfAbsent(name, n -> {
-            T newItem = itemFactory.get();
-            nameSetter.accept(newItem, n);
-            list.add(newItem);
-
-            return newItem;
-        }));
-    }
-}
+package com.epam.aidial.util.mapping;
+
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+public class ListMapper<T> {
+    private final List<T> list;
+    private final Supplier<T> itemFactory;
+    private final BiConsumer<T, String> nameSetter;
+    private final Map<String, T> map;
+
+    public ListMapper(List<T> list, NamedItemMapper<T> itemMapper) {
+        this.list = list;
+        this.itemFactory = itemMapper.factory();
+        this.nameSetter = itemMapper.setter();
+        this.map = list.stream()
+                .collect(Collectors.toMap(itemMapper.getter(), Function.identity()));
+    }
+
+    public MappingChain<T> get(String name) {
+        return getOrDefault(name, itemFactory);
+    }
+
+    public MappingChain<T> getOrDefault(String name, Supplier<T> itemFactory) {
+        return new MappingChain<>(map.computeIfAbsent(name, n -> {
+            T newItem = itemFactory.get();
+            nameSetter.accept(newItem, n);
+            list.add(newItem);
+
+            return newItem;
+        }));
+    }
+}
diff --git a/src/main/resources/application.yaml b/src/main/resources/application.yaml
@@ -20,7 +20,7 @@ app:
   heartbeat-period-sec: 30
   max-error-log-lines: 20
   max-error-log-chars: 1000
-  template-container: &template-container
+  template-container:
     name: template-container
     image: ${app.template-image}
     imagePullPolicy: Always
@@ -40,7 +40,7 @@ app:
       runAsUser: 1001
       runAsNonRoot: true
       allowPrivilegeEscalation: false
-  builder-container: &builder-container
+  builder-container:
     name: builder-container
     image: ${app.builder-image}
     args:
@@ -55,7 +55,7 @@ app:
         subPath: templates
         mountPath: /templates
         readOnly: true
-  service-container: &service-container
+  service-container:
     name: app-container
     imagePullPolicy: Always
     resources:
@@ -126,8 +126,6 @@ app:
           idleTimeoutSeconds: 300
           containerConcurrency: 50
           automountServiceAccountToken: false
-          containers:
-            - <<: *service-container
   job-config:
     apiVersion: batch/v1
     kind: Job
@@ -136,10 +134,6 @@ app:
       template:
         spec:
           automountServiceAccountToken: false
-          initContainers:
-            - <<: *template-container
-          containers:
-            - <<: *builder-container
           restartPolicy: Never
           volumes:
             - name: volume