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.

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

Move exclusive placement annotation to ReplicatedJob template #389

Merged
merged 6 commits into from
Jan 28, 2024
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
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
4 changes: 2 additions & 2 deletions examples/simple/exclusive-placement.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@ apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: exclusive-placement
annotations:
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool # 1:1 job replica to node pool assignment
spec:
failurePolicy:
maxRestarts: 3
replicatedJobs:
- name: workers
replicas: 3 # set to number of node pools
template:
annotations:
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool # 1:1 job replica to node pool assignment
spec:
parallelism: 3
completions: 3
Expand Down
12 changes: 8 additions & 4 deletions pkg/controllers/jobset_controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -612,11 +612,11 @@ func constructJob(js *jobset.JobSet, rjob *jobset.ReplicatedJob, jobIdx int) (*b
}

// If this job should be exclusive per topology, configure the scheduling constraints accordingly.
if _, ok := js.Annotations[jobset.ExclusiveKey]; ok {
if _, ok := rjob.Template.Annotations[jobset.ExclusiveKey]; ok {
danielvegamyhre marked this conversation as resolved.
Show resolved Hide resolved
// If user has set the nodeSelectorStrategy annotation flag, add the job name label as a
// nodeSelector, and add a toleration for the no schedule taint.
// The node label and node taint must be added separately by a user/script.
if _, exists := js.Annotations[jobset.NodeSelectorStrategyKey]; exists {
if _, exists := rjob.Template.Annotations[jobset.NodeSelectorStrategyKey]; exists {
addNamespacedJobNodeSelector(job)
addTaintToleration(job)
}
Expand Down Expand Up @@ -675,8 +675,12 @@ func labelAndAnnotateObject(obj metav1.Object, js *jobset.JobSet, rjob *jobset.R

// Only set exclusive key label/annotation on jobs/pods if the parent JobSet
// is using exclusive placement.
if _, exists := js.Annotations[jobset.ExclusiveKey]; exists {
annotations[jobset.ExclusiveKey] = js.Annotations[jobset.ExclusiveKey]
if topologyDomain, exists := rjob.Template.Annotations[jobset.ExclusiveKey]; exists {
danielvegamyhre marked this conversation as resolved.
Show resolved Hide resolved
annotations[jobset.ExclusiveKey] = topologyDomain
// Add nodeSelectorStrategy annotation if it is being used.
if value, ok := rjob.Template.Annotations[jobset.NodeSelectorStrategyKey]; ok {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This labelAndAnnotate function is invoked for both the job template and the pod template. The job template inherits gets this already because the job is created out of rjob.Template anyways, and we don't need to set this annotation on the pod.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

anyways, we can clean this up when we add the canonical placement policy API.

annotations[jobset.NodeSelectorStrategyKey] = value
}
ahg-g marked this conversation as resolved.
Show resolved Hide resolved
}

obj.SetLabels(labels)
Expand Down
52 changes: 30 additions & 22 deletions pkg/controllers/jobset_controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -304,9 +304,10 @@ func TestConstructJobsFromTemplate(t *testing.T) {
{
name: "exclusive affinities",
js: testutils.MakeJobSet(jobSetName, ns).
SetAnnotations(map[string]string{jobset.ExclusiveKey: topologyDomain}).
ReplicatedJob(testutils.MakeReplicatedJob(replicatedJobName).
Job(testutils.MakeJobTemplate(jobName, ns).Obj()).
Job(testutils.MakeJobTemplate(jobName, ns).
SetAnnotations(map[string]string{jobset.ExclusiveKey: topologyDomain}).
Obj()).
Replicas(1).
Obj()).
Obj(),
Expand Down Expand Up @@ -400,28 +401,30 @@ func TestConstructJobsFromTemplate(t *testing.T) {
{
name: "node selector exclusive placement strategy enabled",
js: testutils.MakeJobSet(jobSetName, ns).
SetAnnotations(map[string]string{
jobset.ExclusiveKey: "topology",
jobset.NodeSelectorStrategyKey: "true",
}).
EnableDNSHostnames(true).
NetworkSubdomain(jobSetName).
ReplicatedJob(testutils.MakeReplicatedJob(replicatedJobName).
Job(testutils.MakeJobTemplate(jobName, ns).Obj()).
Job(testutils.MakeJobTemplate(jobName, ns).
SetAnnotations(map[string]string{
jobset.ExclusiveKey: topologyDomain,
jobset.NodeSelectorStrategyKey: "true",
}).
Obj()).
Subdomain(jobSetName).
Replicas(1).
Obj()).
Obj(),
ownedJobs: &childJobs{},
want: []*batchv1.Job{
makeJob(&makeJobArgs{
jobSetName: jobSetName,
replicatedJobName: replicatedJobName,
jobName: "test-jobset-replicated-job-0",
ns: ns,
replicas: 1,
jobIdx: 0,
topology: "topology"}).
jobSetName: jobSetName,
replicatedJobName: replicatedJobName,
jobName: "test-jobset-replicated-job-0",
ns: ns,
replicas: 1,
jobIdx: 0,
topology: topologyDomain,
nodeSelectorStrategy: true}).
Suspend(false).
Subdomain(jobSetName).
NodeSelector(map[string]string{
Expand Down Expand Up @@ -899,14 +902,15 @@ func TestCalculateReplicatedJobStatuses(t *testing.T) {
}

type makeJobArgs struct {
jobSetName string
replicatedJobName string
jobName string
ns string
replicas int
jobIdx int
restarts int
topology string
jobSetName string
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can do this as a followup, but we are using the wrapper pattern, we shouldn't need a makeJobArgs struct, we should be able to do something like

makeJob(name, namespace)
  .ReplicatedJobName(replicatedJobName)
  .Replicas(n)

replicatedJobName string
jobName string
ns string
replicas int
jobIdx int
restarts int
topology string
nodeSelectorStrategy bool
}

func makeJob(args *makeJobArgs) *testutils.JobWrapper {
Expand All @@ -929,6 +933,10 @@ func makeJob(args *makeJobArgs) *testutils.JobWrapper {
// Only set exclusive key if we are using exclusive placement per topology.
if args.topology != "" {
annotations[jobset.ExclusiveKey] = args.topology
// Exclusive placement topology domain must be set in order to use the node selector strategy.
if args.nodeSelectorStrategy {
annotations[jobset.NodeSelectorStrategyKey] = "true"
}
}
jobWrapper := testutils.MakeJob(args.jobName, args.ns).
JobLabels(labels).
Expand Down
6 changes: 6 additions & 0 deletions pkg/util/testing/wrappers.go
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,12 @@ func (j *JobTemplateWrapper) PodSpec(podSpec corev1.PodSpec) *JobTemplateWrapper
return j
}

// SetAnnotations sets the annotations on the Job template.
func (j *JobTemplateWrapper) SetAnnotations(annotations map[string]string) *JobTemplateWrapper {
j.Annotations = annotations
return j
}

// Obj returns the inner batchv1.JobTemplateSpec
func (j *JobTemplateWrapper) Obj() batchv1.JobTemplateSpec {
return j.JobTemplateSpec
Expand Down