Enable imagePullPolicy from katib-config (#1013)

kubeflow · Jan 15, 2020 · 96df045 · 96df045
1 parent 489d52b
commit 96df045
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 6 deletions.
diff --git a/manifests/v1alpha3/katib-controller/katib-config.yaml b/manifests/v1alpha3/katib-controller/katib-config.yaml
@@ -40,6 +40,7 @@ data:
       },
       "nasrl": {
         "image": "gcr.io/kubeflow-images-public/katib/v1alpha3/suggestion-nasrl",
+        "imagePullPolicy": "Always",
         "resources": {
           "limits": {
             "memory": "200Mi"

diff --git a/pkg/controller.v1alpha3/consts/const.go b/pkg/controller.v1alpha3/consts/const.go
@@ -66,6 +66,10 @@ const (
 	LabelSuggestionDiskRequestTag = "diskRequest"
 	// DefaultDiskRequest is the default value for disk request.
 	DefaultDiskRequest = "500Mi"
+	// LabelSuggestionImagePullPolicy is the name of suggestion image pull policy in configmap.
+	LabelSuggestionImagePullPolicy = "imagePullPolicy"
+	// DefaultImagePullPolicy is the default value for image pull policy.
+	DefaultImagePullPolicy = "IfNotPresent"
 	// LabelMetricsCollectorSidecar is the name of metrics collector config in configmap.
 	LabelMetricsCollectorSidecar = "metrics-collector-sidecar"
 	// LabelMetricsCollectorSidecarImage is the name of metrics collector image config in configmap.
@@ -82,6 +86,8 @@ const (
 	LabelMetricsCollectorDiskLimitTag = "diskLimit"
 	// LabelMetricsCollectorDiskRequestTag is the name of metrics collector Disk Request config in configmap.
 	LabelMetricsCollectorDiskRequestTag = "diskRequest"
+	// LabelMetricsCollectorImagePullPolicy is the name of metrics collector image pull policy in configmap.
+	LabelMetricsCollectorImagePullPolicy = "imagePullPolicy"
 
 	// ReconcileErrorReason is the reason when there is a reconcile error.
 	ReconcileErrorReason = "ReconcileError"

diff --git a/pkg/controller.v1alpha3/suggestion/composer/composer.go b/pkg/controller.v1alpha3/suggestion/composer/composer.go
@@ -118,6 +118,7 @@ func (g *General) desiredContainer(s *suggestionsv1alpha3.Suggestion) (*corev1.C
 	}
 	// Get Suggestion data from config
 	suggestionContainerImage := suggestionConfigData[consts.LabelSuggestionImageTag]
+	suggestionImagePullPolicy := suggestionConfigData[consts.LabelSuggestionImagePullPolicy]
 	suggestionCPULimit := suggestionConfigData[consts.LabelSuggestionCPULimitTag]
 	suggestionCPURequest := suggestionConfigData[consts.LabelSuggestionCPURequestTag]
 	suggestionMemLimit := suggestionConfigData[consts.LabelSuggestionMemLimitTag]
@@ -128,7 +129,7 @@ func (g *General) desiredContainer(s *suggestionsv1alpha3.Suggestion) (*corev1.C
 		Name: consts.ContainerSuggestion,
 	}
 	c.Image = suggestionContainerImage
-	c.ImagePullPolicy = corev1.PullIfNotPresent
+	c.ImagePullPolicy = corev1.PullPolicy(suggestionImagePullPolicy)
 	c.Ports = []corev1.ContainerPort{
 		{
 			Name:          consts.DefaultSuggestionPortName,

diff --git a/pkg/util/v1alpha3/katibconfig/config.go b/pkg/util/v1alpha3/katibconfig/config.go
@@ -15,13 +15,15 @@ import (
 )
 
 type suggestionConfigJSON struct {
-	Image    string                      `json:"image"`
-	Resource corev1.ResourceRequirements `json:"resources"`
+	Image           string                      `json:"image"`
+	ImagePullPolicy corev1.PullPolicy           `json:"imagePullPolicy"`
+	Resource        corev1.ResourceRequirements `json:"resources"`
 }
 
 type metricsCollectorConfigJSON struct {
-	Image    string                      `json:"image"`
-	Resource corev1.ResourceRequirements `json:"resources"`
+	Image           string                      `json:"image"`
+	ImagePullPolicy corev1.PullPolicy           `json:"imagePullPolicy"`
+	Resource        corev1.ResourceRequirements `json:"resources"`
 }
 
 // GetSuggestionConfigData gets the config data for the given algorithm name.
@@ -50,6 +52,14 @@ func GetSuggestionConfigData(algorithmName string, client client.Client) (map[st
 				return map[string]string{}, errors.New("Required value for " + consts.LabelSuggestionImageTag + " configuration of algorithm name " + algorithmName)
 			}
 
+			// Get Image Pull Policy
+			imagePullPolicy := suggestionConfig.ImagePullPolicy
+			if imagePullPolicy == corev1.PullAlways || imagePullPolicy == corev1.PullIfNotPresent || imagePullPolicy == corev1.PullNever {
+				suggestionConfigData[consts.LabelSuggestionImagePullPolicy] = string(imagePullPolicy)
+			} else {
+				suggestionConfigData[consts.LabelSuggestionImagePullPolicy] = consts.DefaultImagePullPolicy
+			}
+
 			// Set default values for CPU, Memory and Disk
 			suggestionConfigData[consts.LabelSuggestionCPURequestTag] = consts.DefaultCPURequest
 			suggestionConfigData[consts.LabelSuggestionMemRequestTag] = consts.DefaultMemRequest
@@ -123,6 +133,14 @@ func GetMetricsCollectorConfigData(cKind common.CollectorKind, client client.Cli
 				return metricsCollectorConfigData, errors.New("Required value for " + consts.LabelMetricsCollectorSidecarImage + "configuration of metricsCollector kind " + kind)
 			}
 
+			// Get Image Pull Policy
+			imagePullPolicy := metricsCollectorConfig.ImagePullPolicy
+			if imagePullPolicy == corev1.PullAlways || imagePullPolicy == corev1.PullIfNotPresent || imagePullPolicy == corev1.PullNever {
+				metricsCollectorConfigData[consts.LabelMetricsCollectorImagePullPolicy] = string(imagePullPolicy)
+			} else {
+				metricsCollectorConfigData[consts.LabelMetricsCollectorImagePullPolicy] = consts.DefaultImagePullPolicy
+			}
+
 			// Set default values for CPU, Memory and Disk
 			metricsCollectorConfigData[consts.LabelMetricsCollectorCPURequestTag] = consts.DefaultCPURequest
 			metricsCollectorConfigData[consts.LabelMetricsCollectorMemRequestTag] = consts.DefaultMemRequest

diff --git a/pkg/webhook/v1alpha3/pod/inject_webhook.go b/pkg/webhook/v1alpha3/pod/inject_webhook.go
@@ -186,6 +186,7 @@ func (s *sidecarInjector) getMetricsCollectorContainer(trial *trialsv1alpha3.Tri
 
 	// Get metricsCollector data from config
 	metricsCollectorContainerImage := metricsCollectorConfigData[consts.LabelMetricsCollectorSidecarImage]
+	metricsCollectorImagePullPolicy := metricsCollectorConfigData[consts.LabelMetricsCollectorImagePullPolicy]
 	metricsCollectorCPULimit := metricsCollectorConfigData[consts.LabelMetricsCollectorCPULimitTag]
 	metricsCollectorCPURequest := metricsCollectorConfigData[consts.LabelMetricsCollectorCPURequestTag]
 	metricsCollectorMemLimit := metricsCollectorConfigData[consts.LabelMetricsCollectorMemLimitTag]
@@ -222,7 +223,7 @@ func (s *sidecarInjector) getMetricsCollectorContainer(trial *trialsv1alpha3.Tri
 		Name:            sidecarContainerName,
 		Image:           metricsCollectorContainerImage,
 		Args:            args,
-		ImagePullPolicy: v1.PullIfNotPresent,
+		ImagePullPolicy: v1.PullPolicy(metricsCollectorImagePullPolicy),
 		Resources: v1.ResourceRequirements{
 			Limits: v1.ResourceList{
 				v1.ResourceCPU:              cpuLimitQuantity,