Skip to content

Commit dbd86b2

Browse files
EtiennePerotgvisor-bot
authored andcommitted
testcluster: Support overwriting runtime selection on a pod.
Previously, calling `RuntimeType.ApplyPodSpec` on a pod spec for which this function was already called may add duplicate tolerations, or incorrectly not unset the `RuntimeClassName` option in the `RuntimeTypeUnsandboxed` case. Now this is avoided. PiperOrigin-RevId: 775899751
1 parent f8854c4 commit dbd86b2

File tree

1 file changed

+16
-5
lines changed

1 file changed

+16
-5
lines changed

test/kubernetes/testcluster/objects.go

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -433,25 +433,36 @@ func SetNodePlacementPolicyCompact(nodepool *cspb.NodePool, tpuTopology string)
433433
return nil
434434
}
435435

436+
// addToleration adds a toleration to the pod if it doesn't already exist.
437+
func addToleration(podSpec *v13.PodSpec, toleration v13.Toleration) {
438+
for _, t := range podSpec.Tolerations {
439+
if t == toleration {
440+
return
441+
}
442+
}
443+
podSpec.Tolerations = append(podSpec.Tolerations, toleration)
444+
}
445+
436446
// ApplyPodSpec modifies a PodSpec to use this runtime.
437447
func (t RuntimeType) ApplyPodSpec(podSpec *v13.PodSpec) {
438448
switch t {
439449
case RuntimeTypeGVisor:
440450
podSpec.RuntimeClassName = proto.String(gvisorRuntimeClass)
441451
podSpec.NodeSelector[NodepoolRuntimeKey] = string(RuntimeTypeGVisor)
442-
podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
452+
addToleration(podSpec, v13.Toleration{
443453
Key: "nvidia.com/gpu",
444454
Operator: v13.TolerationOpExists,
445455
})
446456
case RuntimeTypeUnsandboxed:
457+
podSpec.RuntimeClassName = nil
447458
podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
448459
Key: "nvidia.com/gpu",
449460
Operator: v13.TolerationOpExists,
450461
})
451462
// Allow the pod to schedule on gVisor nodes as well.
452463
// This enables the use of `--test-nodepool-runtime=runc` to run
453464
// unsandboxed benchmarks on gVisor test clusters.
454-
podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
465+
addToleration(podSpec, v13.Toleration{
455466
Effect: v13.TaintEffectNoSchedule,
456467
Key: gvisorNodepoolKey,
457468
Operator: v13.TolerationOpEqual,
@@ -460,16 +471,16 @@ func (t RuntimeType) ApplyPodSpec(podSpec *v13.PodSpec) {
460471
case RuntimeTypeGVisorTPU:
461472
podSpec.RuntimeClassName = proto.String(gvisorRuntimeClass)
462473
podSpec.NodeSelector[NodepoolRuntimeKey] = string(RuntimeTypeGVisorTPU)
463-
podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
474+
addToleration(podSpec, v13.Toleration{
464475
Key: "google.com/tpu",
465476
Operator: v13.TolerationOpExists,
466477
})
467478
case RuntimeTypeUnsandboxedTPU:
468-
podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
479+
addToleration(podSpec, v13.Toleration{
469480
Key: "google.com/tpu",
470481
Operator: v13.TolerationOpExists,
471482
})
472-
podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
483+
addToleration(podSpec, v13.Toleration{
473484
Effect: v13.TaintEffectNoSchedule,
474485
Key: gvisorNodepoolKey,
475486
Operator: v13.TolerationOpEqual,

0 commit comments

Comments
 (0)