More container fields for SuggestionConfig (#2000)

* More container fields for SuggestionConfig * Inline corev1.Container into SuggestionConfig * Set default value for suggestion container name * Append suggestion volume and port only if not present * Deep-Copy base suggestion container * Check for suggestion container port number as well * Prohibit suggestion port to be set in suggestion config
kubeflow · Jan 25, 2023 · ff6441b · ff6441b
1 parent 55e6e34
commit ff6441b
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 43 deletions.
diff --git a/pkg/controller.v1beta1/suggestion/composer/composer.go b/pkg/controller.v1beta1/suggestion/composer/composer.go
@@ -76,6 +76,11 @@ func (g *General) DesiredDeployment(s *suggestionsv1beta1.Suggestion) (*appsv1.D
 	if err != nil {
 		return nil, err
 	}
+	if containsContainerPortWithName(suggestionConfigData.Ports, consts.DefaultSuggestionPortName) ||
+		containsContainerPort(suggestionConfigData.Ports, consts.DefaultSuggestionPort) {
+		return nil, fmt.Errorf("invalid suggestion config: a port with name %q or number %d must not be specified",
+			consts.DefaultSuggestionPortName, consts.DefaultSuggestionPort)
+	}
 
 	// If early stopping is used, get the config data.
 	earlyStoppingConfigData := katibconfig.EarlyStoppingConfig{}
@@ -181,21 +186,27 @@ func (g *General) desiredContainers(s *suggestionsv1beta1.Suggestion,
 	suggestionConfigData katibconfig.SuggestionConfig,
 	earlyStoppingConfigData katibconfig.EarlyStoppingConfig) []corev1.Container {
 
-	containers := []corev1.Container{}
-	suggestionContainer := corev1.Container{
-		Name:            consts.ContainerSuggestion,
-		Image:           suggestionConfigData.Image,
-		ImagePullPolicy: suggestionConfigData.ImagePullPolicy,
-		Ports: []corev1.ContainerPort{
-			{
-				Name:          consts.DefaultSuggestionPortName,
-				ContainerPort: consts.DefaultSuggestionPort,
-			},
-		},
-		Resources: suggestionConfigData.Resource,
+	var (
+		containers          []corev1.Container
+		suggestionContainer corev1.Container
+	)
+
+	suggestionConfigData.Container.DeepCopyInto(&suggestionContainer)
+
+	// Assign default values for suggestionContainer fields that are not set via
+	// the suggestion config.
+
+	if suggestionContainer.Name == "" {
+		suggestionContainer.Name = consts.ContainerSuggestion
 	}
 
-	if viper.GetBool(consts.ConfigEnableGRPCProbeInSuggestion) {
+	suggestionPort := corev1.ContainerPort{
+		Name:          consts.DefaultSuggestionPortName,
+		ContainerPort: consts.DefaultSuggestionPort,
+	}
+	suggestionContainer.Ports = append(suggestionContainer.Ports, suggestionPort)
+
+	if viper.GetBool(consts.ConfigEnableGRPCProbeInSuggestion) && suggestionContainer.ReadinessProbe == nil {
 		suggestionContainer.ReadinessProbe = &corev1.Probe{
 			ProbeHandler: corev1.ProbeHandler{
 				Exec: &corev1.ExecAction{
@@ -209,6 +220,8 @@ func (g *General) desiredContainers(s *suggestionsv1beta1.Suggestion,
 			InitialDelaySeconds: defaultInitialDelaySeconds,
 			PeriodSeconds:       defaultPeriodForReady,
 		}
+	}
+	if viper.GetBool(consts.ConfigEnableGRPCProbeInSuggestion) && suggestionContainer.LivenessProbe == nil {
 		suggestionContainer.LivenessProbe = &corev1.Probe{
 			ProbeHandler: corev1.ProbeHandler{
 				Exec: &corev1.ExecAction{
@@ -226,15 +239,14 @@ func (g *General) desiredContainers(s *suggestionsv1beta1.Suggestion,
 		}
 	}
 
-	// Attach volume mounts to the suggestion container if ResumePolicy = FromVolume
-	if s.Spec.ResumePolicy == experimentsv1beta1.FromVolume {
-		suggestionContainer.VolumeMounts = []corev1.VolumeMount{
-			{
-				Name:      consts.ContainerSuggestionVolumeName,
-				MountPath: suggestionConfigData.VolumeMountPath,
-			},
+	if s.Spec.ResumePolicy == experimentsv1beta1.FromVolume && !containsVolumeMountWithName(suggestionContainer.VolumeMounts, consts.ContainerSuggestionVolumeName) {
+		suggestionVolume := corev1.VolumeMount{
+			Name:      consts.ContainerSuggestionVolumeName,
+			MountPath: suggestionConfigData.VolumeMountPath,
 		}
+		suggestionContainer.VolumeMounts = append(suggestionContainer.VolumeMounts, suggestionVolume)
 	}
+
 	containers = append(containers, suggestionContainer)
 
 	if s.Spec.EarlyStopping != nil && s.Spec.EarlyStopping.AlgorithmName != "" {
@@ -256,6 +268,36 @@ func (g *General) desiredContainers(s *suggestionsv1beta1.Suggestion,
 	return containers
 }
 
+func containsVolumeMountWithName(volumeMounts []corev1.VolumeMount, name string) bool {
+	for i := range volumeMounts {
+		if volumeMounts[i].Name == name {
+			return true
+		}
+	}
+
+	return false
+}
+
+func containsContainerPortWithName(ports []corev1.ContainerPort, name string) bool {
+	for i := range ports {
+		if ports[i].Name == name {
+			return true
+		}
+	}
+
+	return false
+}
+
+func containsContainerPort(ports []corev1.ContainerPort, port int32) bool {
+	for i := range ports {
+		if ports[i].ContainerPort == port {
+			return true
+		}
+	}
+
+	return false
+}
+
 // DesiredVolume returns desired PVC and PV for Suggestion.
 // If PV doesn't exist in Katib config return nil for PV.
 func (g *General) DesiredVolume(s *suggestionsv1beta1.Suggestion) (*corev1.PersistentVolumeClaim, *corev1.PersistentVolume, error) {

diff --git a/pkg/controller.v1beta1/suggestion/composer/composer_test.go b/pkg/controller.v1beta1/suggestion/composer/composer_test.go
@@ -646,18 +646,20 @@ func newFakeSuggestionConfig() katibconfig.SuggestionConfig {
 	diskQ, _ := resource.ParseQuantity(disk)
 
 	return katibconfig.SuggestionConfig{
-		Image:           image,
-		ImagePullPolicy: imagePullPolicy,
-		Resource: corev1.ResourceRequirements{
-			Limits: corev1.ResourceList{
-				corev1.ResourceCPU:              cpuQ,
-				corev1.ResourceMemory:           memoryQ,
-				corev1.ResourceEphemeralStorage: diskQ,
-			},
-			Requests: corev1.ResourceList{
-				corev1.ResourceCPU:              cpuQ,
-				corev1.ResourceMemory:           memoryQ,
-				corev1.ResourceEphemeralStorage: diskQ,
+		Container: corev1.Container{
+			Image:           image,
+			ImagePullPolicy: imagePullPolicy,
+			Resources: corev1.ResourceRequirements{
+				Limits: corev1.ResourceList{
+					corev1.ResourceCPU:              cpuQ,
+					corev1.ResourceMemory:           memoryQ,
+					corev1.ResourceEphemeralStorage: diskQ,
+				},
+				Requests: corev1.ResourceList{
+					corev1.ResourceCPU:              cpuQ,
+					corev1.ResourceMemory:           memoryQ,
+					corev1.ResourceEphemeralStorage: diskQ,
+				},
 			},
 		},
 		ServiceAccountName: serviceAccount,

diff --git a/pkg/controller.v1beta1/suggestion/suggestion_controller_test.go b/pkg/controller.v1beta1/suggestion/suggestion_controller_test.go
@@ -435,7 +435,9 @@ func newKatibConfigMapInstance() *corev1.ConfigMap {
 	// Create suggestion config
 	suggestionConfig := map[string]katibconfig.SuggestionConfig{
 		"random": {
-			Image: suggestionImage,
+			Container: corev1.Container{
+				Image: suggestionImage,
+			},
 		},
 	}
 	bSuggestionConfig, _ := json.Marshal(suggestionConfig)

diff --git a/pkg/util/v1beta1/katibconfig/config.go b/pkg/util/v1beta1/katibconfig/config.go
@@ -34,9 +34,7 @@ import (
 
 // SuggestionConfig is the JSON suggestion structure in Katib config.
 type SuggestionConfig struct {
-	Image                     string                           `json:"image"`
-	ImagePullPolicy           corev1.PullPolicy                `json:"imagePullPolicy,omitempty"`
-	Resource                  corev1.ResourceRequirements      `json:"resources,omitempty"`
+	corev1.Container          `json:",inline"`
 	ServiceAccountName        string                           `json:"serviceAccountName,omitempty"`
 	VolumeMountPath           string                           `json:"volumeMountPath,omitempty"`
 	PersistentVolumeClaimSpec corev1.PersistentVolumeClaimSpec `json:"persistentVolumeClaimSpec,omitempty"`
@@ -99,7 +97,7 @@ func GetSuggestionConfigData(algorithmName string, client client.Client) (Sugges
 	suggestionConfigData.ImagePullPolicy = setImagePullPolicy(suggestionConfigData.ImagePullPolicy)
 
 	// Set resource requirements for suggestion
-	suggestionConfigData.Resource = setResourceRequirements(suggestionConfigData.Resource)
+	suggestionConfigData.Resources = setResourceRequirements(suggestionConfigData.Resources)
 
 	// Set default suggestion container volume mount path
 	if suggestionConfigData.VolumeMountPath == "" {

diff --git a/pkg/util/v1beta1/katibconfig/config_test.go b/pkg/util/v1beta1/katibconfig/config_test.go
@@ -54,13 +54,13 @@ func TestGetSuggestionConfigData(t *testing.T) {
 			katibConfig: func() *katibConfig {
 				kc := &katibConfig{suggestion: map[string]*SuggestionConfig{testAlgorithmName: newFakeSuggestionConfig()}}
 				kc.suggestion[testAlgorithmName].ImagePullPolicy = corev1.PullAlways
-				kc.suggestion[testAlgorithmName].Resource = *newFakeCustomResourceRequirements()
+				kc.suggestion[testAlgorithmName].Resources = *newFakeCustomResourceRequirements()
 				return kc
 			}(),
 			expected: func() *SuggestionConfig {
 				c := newFakeSuggestionConfig()
 				c.ImagePullPolicy = corev1.PullAlways
-				c.Resource = *newFakeCustomResourceRequirements()
+				c.Resources = *newFakeCustomResourceRequirements()
 				return c
 			}(),
 			inputAlgorithmName: testAlgorithmName,
@@ -107,7 +107,7 @@ func TestGetSuggestionConfigData(t *testing.T) {
 			testDescription: "GetSuggestionConfigData sets resource.requests and resource.limits for the suggestion service",
 			katibConfig: func() *katibConfig {
 				kc := &katibConfig{suggestion: map[string]*SuggestionConfig{testAlgorithmName: newFakeSuggestionConfig()}}
-				kc.suggestion[testAlgorithmName].Resource = corev1.ResourceRequirements{}
+				kc.suggestion[testAlgorithmName].Resources = corev1.ResourceRequirements{}
 				return kc
 			}(),
 			expected:           newFakeSuggestionConfig(),
@@ -402,9 +402,11 @@ func newFakeSuggestionConfig() *SuggestionConfig {
 	defaultVolumeStorage, _ := resource.ParseQuantity(consts.DefaultSuggestionVolumeStorage)
 
 	return &SuggestionConfig{
-		Image:           "suggestion-image",
-		ImagePullPolicy: consts.DefaultImagePullPolicy,
-		Resource:        *setFakeResourceRequirements(),
+		Container: corev1.Container{
+			Image:           "suggestion-image",
+			ImagePullPolicy: consts.DefaultImagePullPolicy,
+			Resources:       *setFakeResourceRequirements(),
+		},
 		VolumeMountPath: consts.DefaultContainerSuggestionVolumeMountPath,
 		PersistentVolumeClaimSpec: corev1.PersistentVolumeClaimSpec{
 			AccessModes: []corev1.PersistentVolumeAccessMode{