Skip to content
This repository was archived by the owner on Sep 19, 2022. It is now read-only.

Add master role label for PyTorchJob #116

Merged
merged 2 commits into from
Dec 18, 2018
Merged

Add master role label for PyTorchJob #116

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9a1a483
Add master role label
johnugeorge Dec 18, 2018
577dea8
Adding comments
johnugeorge Dec 18, 2018
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
1 change: 1 addition & 0 deletions pkg/controller.v1beta1/pytorch/controller.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ const (
replicaIndexLabel = "pytorch-replica-index"
labelGroupName = "group_name"
labelPyTorchJobName = "pytorch_job_name"
labelPyTorchJobRole = "pytorch_job_role"
)

var (
Expand Down
14 changes: 12 additions & 2 deletions pkg/controller.v1beta1/pytorch/pod.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,17 +58,24 @@ func (pc *PyTorchController) reconcilePods(
}
replicas := int(*spec.Replicas)
restart := false
masterRole := false

initializePyTorchReplicaStatuses(job, rtype)

podSlices := getPodSlices(pods, replicas, logger)
for index, podSlice := range podSlices {
masterRole = false
if len(podSlice) > 1 {
logger.Warningf("We have too many pods for %s %d", rt, index)
// TODO(gaocegege): Kill some pods.
} else if len(podSlice) == 0 {
logger.Infof("Need to create new pod: %s-%d", rt, index)
err = pc.createNewPod(job, rtype, strconv.Itoa(index), spec)

//Pytorch Job will have exactly one Master pod available
if rtype == v1beta1.PyTorchReplicaTypeMaster {
masterRole = true
}
err = pc.createNewPod(job, rtype, strconv.Itoa(index), spec, masterRole)
if err != nil {
return err
}
Expand Down Expand Up @@ -125,7 +132,7 @@ func getPodSlices(pods []*v1.Pod, replicas int, logger *log.Entry) [][]*v1.Pod {
}

// createNewPod creates a new pod for the given index and type.
func (pc *PyTorchController) createNewPod(job *v1beta1.PyTorchJob, rtype v1beta1.PyTorchReplicaType, index string, spec *common.ReplicaSpec) error {
func (pc *PyTorchController) createNewPod(job *v1beta1.PyTorchJob, rtype v1beta1.PyTorchReplicaType, index string, spec *common.ReplicaSpec, masterRole bool) error {
rt := strings.ToLower(string(rtype))
jobKey, err := KeyFunc(job)
if err != nil {
Expand All @@ -146,6 +153,9 @@ func (pc *PyTorchController) createNewPod(job *v1beta1.PyTorchJob, rtype v1beta1
labels[replicaTypeLabel] = rt
labels[replicaIndexLabel] = index

if masterRole {
labels[labelPyTorchJobRole] = "master"
}
podTemplate := spec.Template.DeepCopy()
totalReplicas := getTotalReplicas(job)
// Set name for the template.
Expand Down