github.com/kubeflow/training-operator@v1.7.0/pkg/controller.v1/pytorch/master.go (about)

     1  package pytorch
     2  
     3  import (
     4  	"strconv"
     5  	"sync"
     6  
     7  	corev1 "k8s.io/api/core/v1"
     8  
     9  	kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
    10  )
    11  
    12  var (
    13  	masterGenerator EnvVarGenerator
    14  	onceMaster      sync.Once
    15  	EnvMasterPort   = "MASTER_PORT"
    16  	EnvMasterAddr   = "MASTER_ADDR"
    17  
    18  	PETMasterPort = "PET_MASTER_PORT"
    19  	PETMasterAddr = "PET_MASTER_ADDR"
    20  )
    21  
    22  // MasterEnvVarGenerator is the environment variable generator for Master related arguments.
    23  type MasterEnvVarGenerator struct {
    24  }
    25  
    26  func GetMasterEnvVarGenerator() EnvVarGenerator {
    27  	onceMaster.Do(func() {
    28  		masterGenerator = &MasterEnvVarGenerator{}
    29  	})
    30  	return masterGenerator
    31  }
    32  
    33  func (e MasterEnvVarGenerator) Generate(
    34  	job *kubeflowv1.PyTorchJob) ([]corev1.EnvVar, error) {
    35  	envVars := []corev1.EnvVar{}
    36  	if job.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeMaster] != nil {
    37  		masterPort, err := getPortFromPyTorchJob(job, kubeflowv1.PyTorchJobReplicaTypeMaster)
    38  		if err != nil {
    39  			return nil, err
    40  		}
    41  
    42  		masterAddr := replicaName(job.Name, kubeflowv1.PyTorchJobReplicaTypeMaster, 0)
    43  
    44  		envVars = append(envVars, corev1.EnvVar{
    45  			Name:  EnvMasterPort,
    46  			Value: strconv.Itoa(int(masterPort)),
    47  		})
    48  		envVars = append(envVars, corev1.EnvVar{
    49  			Name:  PETMasterPort,
    50  			Value: strconv.Itoa(int(masterPort)),
    51  		})
    52  		envVars = append(envVars, corev1.EnvVar{
    53  			Name:  EnvMasterAddr,
    54  			Value: masterAddr,
    55  		})
    56  		envVars = append(envVars, corev1.EnvVar{
    57  			Name:  PETMasterAddr,
    58  			Value: masterAddr,
    59  		})
    60  	}
    61  	return envVars, nil
    62  }