initial commit for coordinator support

api/jobset/v1alpha2/jobset_types.go

@@ -44,6 +44,11 @@ const (
 	// JobSetControllerName is the reserved value for the managedBy field for the built-in
 	// JobSet controller.
 	JobSetControllerName = "jobset.sigs.k8s.io/jobset-controller"
+
+	// CoordinatorKey is used as an annotation and label on Jobs and Pods. If the JobSet spec
+	// defines the .spec.coordinator field, this annotation/label will be added to store a stable
+	// network endpoint where the coordinator pod can be reached.
+	CoordinatorKey = "jobset.sigs.k8s.io/coordinator"
 )
 
 type JobSetConditionType string
@@ -95,6 +100,14 @@ type JobSetSpec struct {
 	// Suspend suspends all running child Jobs when set to true.
 	Suspend *bool `json:"suspend,omitempty"`
 
+	// Coordinator can be used to assign a specific pod as the coordinator for
+	// the JobSet. If defined, an annotation will be added to all Jobs and pods with
+	// coordinator pod, which contains the stable network endpoint where the
+	// coordinator pod can be reached.
+	// jobset.sigs.k8s.io/coordinator=<pod hostname>.<headless service>
+	// +optional
+	Coordinator *Coordinator `json:"coordinator,omitempty"`
+
 	// ManagedBy is used to indicate the controller or entity that manages a JobSet.
 	// The built-in JobSet controller reconciles JobSets which don't have this
 	// field at all or the field value is the reserved string
@@ -323,6 +336,22 @@ type StartupPolicy struct {
 	StartupPolicyOrder StartupPolicyOptions `json:"startupPolicyOrder"`
 }
 
+// Coordinator defines which pod can be marked as the coordinator for the JobSet workload.
+type Coordinator struct {
+	// ReplicatedJob is the name of the ReplicatedJob which contains
+	// the coordinator pod.
+	ReplicatedJob string `json:"replicatedJob"`
+
+	// JobIndex is the index of Job which contains the coordinator pod
+	// (i.e., for a ReplicatedJob with N replicas, there are Job indexes 0 to N-1).
+	// Defaults to 0 if unset.
+	JobIndex int `json:"jobIndex,omitempty"`
+
+	// PodIndex is the Job completion index of the coordinator pod.
+	// Defaults to 0 if unset.
+	PodIndex int `json:"podIndex,omitempty"`
+}
+
 func init() {
 	SchemeBuilder.Register(&JobSet{}, &JobSetList{})
 }

config/components/crd/bases/jobset.x-k8s.io_jobsets.yaml

@@ -59,6 +59,33 @@ spec:
           spec:
             description: JobSetSpec defines the desired state of JobSet
             properties:
+              coordinator:
+                description: |-
+                  Coordinator can be used to assign a specific pod as the coordinator for
+                  the JobSet. If defined, an annotation will be added to all Jobs and pods with
+                  coordinator pod, which contains the stable network endpoint where the
+                  coordinator pod can be reached.
+                  jobset.sigs.k8s.io/coordinator=<pod hostname>.<headless service>
+                properties:
+                  jobIndex:
+                    description: |-
+                      JobIndex is the index of Job which contains the coordinator pod
+                      (i.e., for a ReplicatedJob with N replicas, there are Job indexes 0 to N-1).
+                      Defaults to 0 if unset.
+                    type: integer
+                  podIndex:
+                    description: |-
+                      PodIndex is the Job completion index of the coordinator pod.
+                      Defaults to 0 if unset.
+                    type: integer
+                  replicatedJob:
+                    description: |-
+                      ReplicatedJob is the name of the ReplicatedJob which contains
+                      the coordinator pod.
+                    type: string
+                required:
+                - replicatedJob
+                type: object
               failurePolicy:
                 description: |-
                   FailurePolicy, if set, configures when to declare the JobSet as

hack/python-sdk/swagger.json

@@ -7,6 +7,30 @@
   },
   "paths": {},
   "definitions": {
+    "jobset.v1alpha2.Coordinator": {
+      "description": "Coordinator defines which pod can be marked as the coordinator for the JobSet workload.",
+      "type": "object",
+      "required": [
+        "replicatedJob"
+      ],
+      "properties": {
+        "jobIndex": {
+          "description": "JobIndex is the index of Job which contains the coordinator pod (i.e., for a ReplicatedJob with N replicas, there are Job indexes 0 to N-1). Defaults to 0 if unset.",
+          "type": "integer",
+          "format": "int32"
+        },
+        "podIndex": {
+          "description": "PodIndex is the Job completion index of the coordinator pod. Defaults to 0 if unset.",
+          "type": "integer",
+          "format": "int32"
+        },
+        "replicatedJob": {
+          "description": "ReplicatedJob is the name of the ReplicatedJob which contains the coordinator pod.",
+          "type": "string",
+          "default": ""
+        }
+      }
+    },
     "jobset.v1alpha2.FailurePolicy": {
       "type": "object",
       "properties": {
@@ -120,6 +144,10 @@
       "description": "JobSetSpec defines the desired state of JobSet",
       "type": "object",
       "properties": {
+        "coordinator": {
+          "description": "Coordinator can be used to assign a specific pod as the coordinator for the JobSet. If defined, an annotation will be added to all Jobs and pods with coordinator pod, which contains the stable network endpoint where the coordinator pod can be reached. jobset.sigs.k8s.io/coordinator=\u003cpod hostname\u003e.\u003cheadless service\u003e",
+          "$ref": "#/definitions/jobset.v1alpha2.Coordinator"
+        },
         "failurePolicy": {
           "description": "FailurePolicy, if set, configures when to declare the JobSet as failed. The JobSet is always declared failed if any job in the set finished with status failed.",
           "$ref": "#/definitions/jobset.v1alpha2.FailurePolicy"

pkg/controllers/jobset_controller.go

@@ -746,6 +746,12 @@ func labelAndAnnotateObject(obj metav1.Object, js *jobset.JobSet, rjob *jobset.R
 	annotations[jobset.JobIndexKey] = strconv.Itoa(jobIdx)
 	annotations[jobset.JobKey] = jobHashKey(js.Namespace, jobName)
 
+	// Apply coordinator annotation/label if a coordinator is defined in the JobSet spec.
+	if js.Spec.Coordinator != nil {
+		labels[jobset.CoordinatorKey] = coordinatorEndpoint(js)
+		annotations[jobset.CoordinatorKey] = coordinatorEndpoint(js)
+	}
+
 	// Check for JobSet level exclusive placement.
 	if topologyDomain, exists := js.Annotations[jobset.ExclusiveKey]; exists {
 		annotations[jobset.ExclusiveKey] = topologyDomain
@@ -1024,3 +1030,9 @@ func exclusiveConditions(cond1, cond2 metav1.Condition) bool {
 		cond2.Type == string(jobset.JobSetStartupPolicyInProgress)
 	return inProgressAndCompleted || completedAndInProgress
 }
+
+// coordinatorEndpoint returns the stable network endpoint where the coordinator pod can be reached.
+// This function assumes the caller has validated that jobset.Spec.Coordinator != nil.
+func coordinatorEndpoint(js *jobset.JobSet) string {
+	return fmt.Sprintf("%s-%s-%d-%d.%s", js.Name, js.Spec.Coordinator.ReplicatedJob, js.Spec.Coordinator.JobIndex, js.Spec.Coordinator.PodIndex, GetSubdomain(js))
+}