Skip to content

Commit

Permalink
Merge pull request #806 from at88mph/remove-pull-secret-when-not-used
Browse files Browse the repository at this point in the history 
fix: use java api to merge image pull secret if present
  • Loading branch information
@shinybrar
shinybrar authored Feb 6, 2025
2 parents 3c2889a + 5dde660 commit 80ab29b
Show file tree
Hide file tree
Showing 6 changed files with 84 additions and 70 deletions.
1 change: 0 additions & 1 deletion skaha/src/main/java/org/opencadc/skaha/SkahaAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,6 @@ public SkahaAction() {

log.debug("skaha.hostname=" + K8SUtil.getSkahaHostName());
log.debug("skaha.sessions.hostname=" + K8SUtil.getSessionsHostName());
log.debug("skaha.hostname=" + K8SUtil.getSkahaHostName());
log.debug("skaha.homedir=" + homedir);
log.debug("SKAHA_TLD=" + skahaTld);
log.debug("skaha.scratchdir=" + scratchdir);
Expand Down
28 changes: 11 additions & 17 deletions skaha/src/main/java/org/opencadc/skaha/session/PostAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,6 @@ public class PostAction extends SessionAction {
public static final String SOFTWARE_LIMITS_RAM = "software.limits.ram";
public static final String HEADLESS_PRIORITY = "headless.priority";
public static final String HEADLESS_IMAGE_BUNDLE = "headless.image.bundle";
public static final String DEFAULT_SOFTWARE_IMAGESECRET_VALUE = "notused";

// k8s rejects label size > 63. Since k8s appends a maximum of six characters
// to a job name to form a pod name, we limit the job name length to 57 characters.
Expand Down Expand Up @@ -586,22 +585,11 @@ public void createSession(
final String headlessPriority = getHeadlessPriority();
final String headlessImageBundle = getHeadlessImageBundle(image, cmd, args, envs);

final String imageRegistrySecretName;
// In the absence of the existence of a public image, assume Private. The validateImage() step above will have
// caught a non-existent Image already.
if (getPublicImage(image) == null) {
final ImageRepositoryAuth userRegistryAuth = getRegistryAuth(getRegistryHost(image));
imageRegistrySecretName = createRegistryImageSecret(userRegistryAuth);
} else {
imageRegistrySecretName = PostAction.DEFAULT_SOFTWARE_IMAGESECRET_VALUE;
}

String jobName = K8SUtil.getJobName(sessionID, type, posixPrincipal.username);

SessionJobBuilder sessionJobBuilder = SessionJobBuilder.fromPath(Paths.get(jobLaunchPath))
.withGPUEnabled(this.gpuEnabled)
.withGPUCount(gpus)
.withImageSecret(imageRegistrySecretName)
.withParameter(PostAction.SKAHA_SESSIONID, this.sessionID)
.withParameter(PostAction.SKAHA_SESSIONNAME, name.toLowerCase())
.withParameter(PostAction.SKAHA_SESSIONEXPIRY, K8SUtil.getSessionExpiry())
Expand All @@ -619,15 +607,22 @@ public void createSession(
.withParameter(PostAction.SOFTWARE_REQUESTS_RAM, ram.toString() + "Gi")
.withParameter(PostAction.SOFTWARE_LIMITS_CORES, cores.toString())
.withParameter(PostAction.SOFTWARE_LIMITS_RAM, ram + "Gi")
.withParameter(PostAction.SKAHA_TLD, this.skahaTld);

sessionJobBuilder = sessionJobBuilder.withParameter(
PostAction.SKAHA_SUPPLEMENTALGROUPS, StringUtil.hasText(supplementalGroups) ? supplementalGroups : "");
.withParameter(PostAction.SKAHA_TLD, this.skahaTld)
.withParameter(
PostAction.SKAHA_SUPPLEMENTALGROUPS,
StringUtil.hasText(supplementalGroups) ? supplementalGroups : "");

if (type.equals(SessionAction.SESSION_TYPE_DESKTOP)) {
sessionJobBuilder = sessionJobBuilder.withParameter(PostAction.DESKTOP_SESSION_APP_TOKEN, generateToken());
}

// In the absence of the existence of a public image, assume Private. The validateImage() step above will have
// caught a non-existent Image already.
if (getPublicImage(image) == null) {
final ImageRepositoryAuth userRegistryAuth = getRegistryAuth(getRegistryHost(image));
sessionJobBuilder = sessionJobBuilder.withImageSecret(createRegistryImageSecret(userRegistryAuth));
}

String jobLaunchString = sessionJobBuilder.build();
String jsonLaunchFile = super.stageFile(jobLaunchString);

Expand Down Expand Up @@ -795,7 +790,6 @@ public void attachDesktopApp(
final String launchSoftwarePath = K8SUtil.getUserHome() + "/config/launch-desktop-app.yaml";
SessionJobBuilder sessionJobBuilder = SessionJobBuilder.fromPath(Paths.get(launchSoftwarePath))
.withGPUEnabled(this.gpuEnabled)
.withImageSecret(PostAction.DEFAULT_SOFTWARE_IMAGESECRET_VALUE)
.withParameter(PostAction.SKAHA_SESSIONID, this.sessionID)
.withParameter(PostAction.SKAHA_SESSIONEXPIRY, K8SUtil.getSessionExpiry())
.withParameter(PostAction.SKAHA_SESSIONTYPE, SessionAction.TYPE_DESKTOP_APP)
Expand Down
50 changes: 35 additions & 15 deletions skaha/src/main/java/org/opencadc/skaha/session/SessionJobBuilder.java
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
package org.opencadc.skaha.session;

import ca.nrc.cadc.util.StringUtil;
import io.kubernetes.client.openapi.models.V1Affinity;
import io.kubernetes.client.openapi.models.V1Job;
import io.kubernetes.client.openapi.models.V1JobSpec;
import io.kubernetes.client.openapi.models.V1LocalObjectReference;
import io.kubernetes.client.openapi.models.V1NodeAffinity;
import io.kubernetes.client.openapi.models.V1NodeSelector;
import io.kubernetes.client.openapi.models.V1NodeSelectorRequirement;
Expand All @@ -21,25 +23,23 @@
import java.util.Map;
import org.apache.log4j.Logger;

/**
* Class to interface with Kubernetes.
*/
/** Class to interface with Kubernetes. */
public class SessionJobBuilder {
private static final Logger log = Logger.getLogger(SessionJobBuilder.class);
private static final String SOFTWARE_LIMITS_GPUS = "software.limits.gpus";
private static final String SOFTWARE_IMAGESECRET = "software.imagesecret";
private final Map<String, String> parameters = new HashMap<>();
private Path jobFilePath;
private boolean gpuEnabled;
private Integer gpuCount;
private String imageRegistrySecretName;

private SessionJobBuilder() {}

/**
* Create a new builder from the provided path.
*
* @param jobFilePath The Path of the template file.
* @return SessionJobBuilder instance. Never null.
* @return SessionJobBuilder instance. Never null.
*/
static SessionJobBuilder fromPath(final Path jobFilePath) {
final SessionJobBuilder sessionJobBuilder = new SessionJobBuilder();
Expand Down Expand Up @@ -111,7 +111,7 @@ SessionJobBuilder withGPUCount(final int gpuCount) {
/**
* Build a single parameter into this builder's parameter map.
*
* @param key The key to find.
* @param key The key to find.
* @param value The value to replace with.
* @return This SessionJobBuilder, never null.
*/
Expand All @@ -122,11 +122,12 @@ SessionJobBuilder withParameter(final String key, final String value) {

/**
* Use the provided Kubernetes secret to authenticate with the Image Registry to pull the Image.
* @param imageRegistrySecretName String existing secret name.
* @return This SessionJobBuilder, never null.
*
* @param imageRegistrySecretName String existing secret name.
* @return This SessionJobBuilder, never null.
*/
SessionJobBuilder withImageSecret(final String imageRegistrySecretName) {
this.withParameter(SessionJobBuilder.SOFTWARE_IMAGESECRET, imageRegistrySecretName);
this.imageRegistrySecretName = imageRegistrySecretName;
return this;
}

Expand All @@ -143,13 +144,20 @@ String build() throws IOException {
jobFileString = SessionJobBuilder.setConfigValue(jobFileString, entry.getKey(), entry.getValue());
}

return mergeAffinity(jobFileString);
return buildJob(jobFileString);
}

private String buildJob(final String jobFileString) throws IOException {
final V1Job launchJob = (V1Job) Yaml.load(jobFileString);
mergeAffinity(launchJob);
mergeImagePullSecret(launchJob);

return Yaml.dump(launchJob);
}

private String mergeAffinity(final String jobFileString) throws IOException {
private void mergeAffinity(final V1Job launchJob) {
final V1Affinity gpuAffinity = getGPUSchedulingAffinity();
if (gpuAffinity != null) {
final V1Job launchJob = (V1Job) Yaml.load(jobFileString);
final V1JobSpec podTemplate = launchJob.getSpec();
if (podTemplate != null) {
// spec.template.spec
Expand Down Expand Up @@ -232,11 +240,23 @@ private String mergeAffinity(final String jobFileString) throws IOException {
}
}
}

return Yaml.dump(launchJob);
}
}

return jobFileString;
private void mergeImagePullSecret(final V1Job launchJob) {
final V1JobSpec podTemplate = launchJob.getSpec();
if (podTemplate != null && StringUtil.hasText(this.imageRegistrySecretName)) {
final V1PodSpec podTemplateSpec = podTemplate.getTemplate().getSpec();
if (podTemplateSpec != null) {
final List<V1LocalObjectReference> imagePullSecrets = podTemplateSpec.getImagePullSecrets();
if (imagePullSecrets == null) {
podTemplateSpec.setImagePullSecrets(
Collections.singletonList(new V1LocalObjectReference().name(this.imageRegistrySecretName)));
} else {
imagePullSecrets.add(new V1LocalObjectReference().name(this.imageRegistrySecretName));
}
}
}
}

private V1Affinity getGPUSchedulingAffinity() {
Expand Down
49 changes: 30 additions & 19 deletions skaha/src/test/java/org/opencadc/skaha/session/SessionJobBuilderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,12 +8,9 @@
import io.kubernetes.client.util.Yaml;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.*;
import org.apache.commons.lang3.RandomStringUtils;
import org.jetbrains.annotations.NotNull;
import org.junit.Assert;
import org.junit.Test;

Expand Down Expand Up @@ -46,6 +43,11 @@ public void testParsing() throws Exception {
Assert.assertFalse("Entry not replaced.", output.contains(entry.getKey()));
Assert.assertTrue("Value not injected into file.", output.contains(entry.getValue()));
}

V1Job job = (V1Job) Yaml.load(output);
V1PodSpec podSpec = Objects.requireNonNull(job.getSpec()).getTemplate().getSpec();
Assert.assertNotNull("PodSpec should not be null", podSpec);
Assert.assertNull("PodSpec should have image pull secrets", podSpec.getImagePullSecrets());
}

@Test
Expand Down Expand Up @@ -76,24 +78,15 @@ public void testWithAffinityMerging() throws Exception {
}

final V1Job job = (V1Job) Yaml.load(output);
final V1PodSpec podSpec = job.getSpec().getTemplate().getSpec();
final V1NodeAffinity nodeAffinity = podSpec.getAffinity().getNodeAffinity();

final List<V1NodeSelectorRequirement> testMatchExpressions = new ArrayList<>();
final List<V1NodeSelectorRequirement> matchExpressions = nodeAffinity
.getRequiredDuringSchedulingIgnoredDuringExecution()
.getNodeSelectorTerms()
.get(0)
.getMatchExpressions();

if (matchExpressions != null) {
testMatchExpressions.addAll(matchExpressions);
}
final V1PodSpec podSpec =
Objects.requireNonNull(job.getSpec()).getTemplate().getSpec();
Assert.assertNotNull("PodSpec should not be null", podSpec);
final List<V1NodeSelectorRequirement> testMatchExpressions = getV1NodeSelectorRequirements(podSpec);

Assert.assertEquals(
"Wrong pull secret.",
"my-secret",
podSpec.getImagePullSecrets().get(0).getName());
Objects.requireNonNull(podSpec.getImagePullSecrets()).get(0).getName());

final V1NodeSelectorRequirement gpuRequirement = new V1NodeSelectorRequirement();
gpuRequirement.setKey("nvidia.com/gpu.count");
Expand All @@ -109,4 +102,22 @@ public void testWithAffinityMerging() throws Exception {
"Missing provided (custom) required match expression.",
testMatchExpressions.contains(providedRequirement));
}

@NotNull private static List<V1NodeSelectorRequirement> getV1NodeSelectorRequirements(V1PodSpec podSpec) {
assert podSpec != null;
final V1NodeAffinity nodeAffinity =
Objects.requireNonNull(podSpec.getAffinity()).getNodeAffinity();

final List<V1NodeSelectorRequirement> testMatchExpressions = new ArrayList<>();
final List<V1NodeSelectorRequirement> matchExpressions = Objects.requireNonNull(
Objects.requireNonNull(nodeAffinity).getRequiredDuringSchedulingIgnoredDuringExecution())
.getNodeSelectorTerms()
.get(0)
.getMatchExpressions();

if (matchExpressions != null) {
testMatchExpressions.addAll(matchExpressions);
}
return testMatchExpressions;
}
}
2 changes: 0 additions & 2 deletions skaha/src/test/resources/test-base-values-affinity.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,6 @@ spec:
automountServiceAccountToken: false
enableServiceLinks: false
restartPolicy: OnFailure
imagePullSecrets:
- name: ${software.imagesecret}
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
Expand Down
24 changes: 8 additions & 16 deletions skaha/src/test/resources/test-base-values.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ metadata:
spec:
parallelism: 1
completions: 1
activeDeadlineSeconds: ${skaha.sessionexpiry}
activeDeadlineSeconds: 86400
ttlSecondsAfterFinished: 86400
template:
metadata:
Expand All @@ -21,20 +21,13 @@ spec:
canfar-net-sessionName: "${skaha.sessionname}"
canfar-net-sessionType: "${skaha.sessiontype}"
canfar-net-userid: "${skaha.userid}"
job-name: "${skaha.jobname}"
job-name: "test-job"
spec:
automountServiceAccountToken: false
enableServiceLinks: false
restartPolicy: OnFailure
{{ template "skaha.job.nodeAffinity" . }}
imagePullSecrets:
- name: "${software.imagesecret}"
securityContext:
{{ template "skaha.job.securityContext" . }}
priorityClassName: uber-user-preempt-medium
hostname: "${software.hostname}"
initContainers:
{{ template "skaha.job.initContainers" . }}
containers:
- name: "${skaha.jobname}"
env:
Expand Down Expand Up @@ -83,14 +76,13 @@ spec:
imagePullPolicy: IfNotPresent
resources:
requests:
memory: "${software.requests.ram}"
cpu: "${software.requests.cores}"
ephemeral-storage: "{{ .Values.deployment.skaha.sessions.minEphemeralStorage }}"
memory: "1Gi"
cpu: "1"
ephemeral-storage: "10Gi"
limits:
memory: "${software.limits.ram}"
cpu: "${software.limits.cores}"
${software.limits.gpus}
ephemeral-storage: "{{ .Values.deployment.skaha.sessions.maxEphemeralStorage }}"
memory: "2Gi"
cpu: "2"
ephemeral-storage: "100Gi"
ports:
- containerPort: 8888
protocol: TCP
Expand Down

0 comments on commit 80ab29b

Please sign in to comment.