chore(controller): make host path for pip cache optional

star-whale · Sep 27, 2022 · e7403a5 · e7403a5
1 parent bc5c50c
commit e7403a5
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 9 deletions.
diff --git a/server/controller/src/main/java/ai/starwhale/mlops/schedule/k8s/K8sJobTemplate.java b/server/controller/src/main/java/ai/starwhale/mlops/schedule/k8s/K8sJobTemplate.java
@@ -17,6 +17,7 @@
 package ai.starwhale.mlops.schedule.k8s;
 
 import io.kubernetes.client.openapi.models.V1Container;
+import io.kubernetes.client.openapi.models.V1EmptyDirVolumeSource;
 import io.kubernetes.client.openapi.models.V1Job;
 import io.kubernetes.client.openapi.models.V1JobSpec;
 import io.kubernetes.client.openapi.models.V1PodSpec;
@@ -46,20 +47,23 @@ public class K8sJobTemplate {
 
     public static final String pipCacheVolumeName = "pip-cache";
 
-    @Value("${sw.infra.k8s.host-path-for-cache}")
-    private String pipCacheHostPath;
+    private final String pipCacheHostPath;
 
     final String template;
     final V1Job v1Job;
 
-    public K8sJobTemplate(@Value("${sw.infra.k8s.job-template-path}") String templatePath)
+    public K8sJobTemplate(
+            @Value("${sw.infra.k8s.job-template-path}") String templatePath,
+            @Value("${sw.infra.k8s.host-path-for-cache}") String pipCacheHostPath
+    )
             throws IOException {
         if (!StringUtils.hasText(templatePath)) {
             this.template = getJobDefaultTemplate();
         } else {
             this.template = Files.readString(Paths.get(templatePath));
         }
         v1Job = Yaml.loadAs(template, V1Job.class);
+        this.pipCacheHostPath = pipCacheHostPath;
     }
 
     public List<V1Container> getInitContainerTemplates() {
@@ -110,10 +114,19 @@ public V1Job renderJob(String jobName,
 
         });
 
-        // replace host path
+        // patch pip cache volume
         List<V1Volume> volumes = job.getSpec().getTemplate().getSpec().getVolumes();
-        volumes.stream().filter(v -> v.getName().equals(pipCacheVolumeName))
-                .findFirst().ifPresent(volume -> volume.getHostPath().path(pipCacheHostPath));
+        var volume = volumes.stream().filter(v -> v.getName().equals(pipCacheVolumeName))
+                .findFirst().orElse(null);
+        if (volume != null) {
+            if (pipCacheHostPath.isEmpty()) {
+                // make volume emptyDir
+                volume.setHostPath(null);
+                volume.emptyDir(new V1EmptyDirVolumeSource());
+            } else {
+                volume.getHostPath().path(pipCacheHostPath);
+            }
+        }
 
         return job;
     }

diff --git a/server/controller/src/main/resources/application.yaml b/server/controller/src/main/resources/application.yaml
@@ -30,7 +30,7 @@ sw:
   infra:
     k8s:
       name-space: ${SW_K8S_NAME_SPACE:default}
-      host-path-for-cache: ${SW_K8S_HOST_PATH_FOR_CACHE:/mnt/data}
+      host-path-for-cache: ${SW_K8S_HOST_PATH_FOR_CACHE:}
       job-template-path: ${SW_K8S_JOB_TEMPLATE_PATH:}
   storage:
     type: ${SW_STORAGE_TYPE:minio}

diff --git a/server/controller/src/test/java/ai/starwhale/mlops/schedule/k8s/K8sJobTemplateTest.java b/server/controller/src/test/java/ai/starwhale/mlops/schedule/k8s/K8sJobTemplateTest.java
@@ -32,7 +32,7 @@
 
 public class K8sJobTemplateTest {
 
-    K8sJobTemplate k8sJobTemplate = new K8sJobTemplate("");
+    K8sJobTemplate k8sJobTemplate = new K8sJobTemplate("", "/path");
 
     public K8sJobTemplateTest() throws IOException {
     }
@@ -97,4 +97,20 @@ private Map<String, ContainerOverwriteSpec> buildContainerSpecMap() {
         return Map.of("worker", containerOverwriteSpecWorker, "data-provider", containerOverwriteSpecDp);
     }
 
+    @Test
+    public void testPipCache() throws IOException {
+        Map<String, ContainerOverwriteSpec> containerSpecMap = buildContainerSpecMap();
+        var job = k8sJobTemplate.renderJob("foo", containerSpecMap, Map.of());
+        var volume = job.getSpec().getTemplate().getSpec().getVolumes().stream()
+                .filter(v -> v.getName().equals(K8sJobTemplate.pipCacheVolumeName)).findFirst().orElse(null);
+        Assertions.assertEquals(volume.getHostPath().getPath(), "/path");
+
+        // empty host path
+        var template = new K8sJobTemplate("", "");
+        job = template.renderJob("foo", containerSpecMap, Map.of());
+        volume = job.getSpec().getTemplate().getSpec().getVolumes().stream()
+                .filter(v -> v.getName().equals(K8sJobTemplate.pipCacheVolumeName)).findFirst().orElse(null);
+        Assertions.assertNull(volume.getHostPath());
+        Assertions.assertNotNull(volume.getEmptyDir());
+    }
 }
diff --git a/server/controller/src/test/java/ai/starwhale/mlops/schedule/k8s/K8sTaskSchedulerTest.java b/server/controller/src/test/java/ai/starwhale/mlops/schedule/k8s/K8sTaskSchedulerTest.java
@@ -123,7 +123,7 @@ private Task mockTask() {
     public static class K8sJobTemplateMock extends K8sJobTemplate {
 
         public K8sJobTemplateMock(String templatePath) throws IOException {
-            super("");
+            super("", "/path");
         }
 
         @Override