Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Sep 27, 2022
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 19 additions & 6 deletions server/controller/src/main/java/ai/starwhale/mlops/schedule/k8s/K8sJobTemplate.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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;
}
Expand Down
2 changes: 1 addition & 1 deletion server/controller/src/main/resources/application.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
18 changes: 17 additions & 1 deletion server/controller/src/test/java/ai/starwhale/mlops/schedule/k8s/K8sJobTemplateTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@

public class K8sJobTemplateTest {

K8sJobTemplate k8sJobTemplate = new K8sJobTemplate("");
K8sJobTemplate k8sJobTemplate = new K8sJobTemplate("", "/path");

public K8sJobTemplateTest() throws IOException {
}
Expand Down Expand Up @@ -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());
}
}
2 changes: 1 addition & 1 deletion server/controller/src/test/java/ai/starwhale/mlops/schedule/k8s/K8sTaskSchedulerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ private Task mockTask() {
public static class K8sJobTemplateMock extends K8sJobTemplate {

public K8sJobTemplateMock(String templatePath) throws IOException {
super("");
super("", "/path");
}

@Override
Expand Down