feat: Support disable GRPC probe (#1020)

Signed-off-by: Ce Gao <gaoce@caicloud.io>
kubeflow · Jan 20, 2020 · 4d46158 · 4d46158
1 parent db45574
commit 4d46158
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 23 deletions.
diff --git a/cmd/katib-controller/v1alpha3/main.go b/cmd/katib-controller/v1alpha3/main.go
@@ -43,6 +43,7 @@ func main() {
 	var certLocalFS bool
 	var injectSecurityContext bool
 	var serviceName string
+	var enableGRPCProbeInSuggestion bool
 
 	flag.StringVar(&experimentSuggestionName, "experiment-suggestion-name",
 		"default", "The implementation of suggestion interface in experiment controller (default|fake)")
@@ -51,13 +52,15 @@ func main() {
 	flag.BoolVar(&certLocalFS, "cert-localfs", false, "Store the webhook cert in local file system")
 	flag.BoolVar(&injectSecurityContext, "webhook-inject-securitycontext", false, "Inject the securityContext of container[0] in the sidecar")
 	flag.StringVar(&serviceName, "webhook-service-name", "katib-controller", "The service name which will be used in webhook")
+	flag.BoolVar(&enableGRPCProbeInSuggestion, "enable-grpc-probe-in-suggestion", true, "enable grpc probe in suggestions")
 
 	flag.Parse()
 
 	// Set the config in viper.
 	viper.Set(consts.ConfigExperimentSuggestionName, experimentSuggestionName)
 	viper.Set(consts.ConfigCertLocalFS, certLocalFS)
 	viper.Set(consts.ConfigInjectSecurityContext, injectSecurityContext)
+	viper.Set(consts.ConfigEnableGRPCProbeInSuggestion, enableGRPCProbeInSuggestion)
 
 	log.Info("Config:",
 		consts.ConfigExperimentSuggestionName,
@@ -70,6 +73,8 @@ func main() {
 		metricsAddr,
 		consts.ConfigInjectSecurityContext,
 		viper.GetBool(consts.ConfigInjectSecurityContext),
+		consts.ConfigEnableGRPCProbeInSuggestion,
+		viper.GetBool(consts.ConfigEnableGRPCProbeInSuggestion),
 	)
 
 	// Get a config to talk to the apiserver

diff --git a/pkg/controller.v1alpha3/consts/const.go b/pkg/controller.v1alpha3/consts/const.go
@@ -13,6 +13,9 @@ const (
 	// if we should inject the security context into the metrics collector
 	// sidecar.
 	ConfigInjectSecurityContext = "inject-security-context"
+	// ConfigEnableGRPCProbeInSuggestion is the config name which indicates
+	// if we should set GRPC probe in suggestion deployments.
+	ConfigEnableGRPCProbeInSuggestion = "enable-grpc-probe-in-suggestion"
 
 	// LabelExperimentName is the label of experiment name.
 	LabelExperimentName = "experiment"

diff --git a/pkg/controller.v1alpha3/suggestion/composer/composer.go b/pkg/controller.v1alpha3/suggestion/composer/composer.go
@@ -3,6 +3,7 @@ package composer
 import (
 	"fmt"
 
+	"github.com/spf13/viper"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -175,33 +176,35 @@ func (g *General) desiredContainer(s *suggestionsv1alpha3.Suggestion) (*corev1.C
 		},
 	}
 
-	c.ReadinessProbe = &corev1.Probe{
-		Handler: corev1.Handler{
-			Exec: &corev1.ExecAction{
-				Command: []string{
-					defaultGRPCHealthCheckProbe,
-					fmt.Sprintf("-addr=:%d", consts.DefaultSuggestionPort),
-					fmt.Sprintf("-service=%s", consts.DefaultGRPCService),
+	if viper.GetBool(consts.ConfigEnableGRPCProbeInSuggestion) {
+		c.ReadinessProbe = &corev1.Probe{
+			Handler: corev1.Handler{
+				Exec: &corev1.ExecAction{
+					Command: []string{
+						defaultGRPCHealthCheckProbe,
+						fmt.Sprintf("-addr=:%d", consts.DefaultSuggestionPort),
+						fmt.Sprintf("-service=%s", consts.DefaultGRPCService),
+					},
 				},
 			},
-		},
-		InitialDelaySeconds: defaultInitialDelaySeconds,
-		PeriodSeconds:       defaultPeriodForReady,
-	}
-	c.LivenessProbe = &corev1.Probe{
-		Handler: corev1.Handler{
-			Exec: &corev1.ExecAction{
-				Command: []string{
-					defaultGRPCHealthCheckProbe,
-					fmt.Sprintf("-addr=:%d", consts.DefaultSuggestionPort),
-					fmt.Sprintf("-service=%s", consts.DefaultGRPCService),
+			InitialDelaySeconds: defaultInitialDelaySeconds,
+			PeriodSeconds:       defaultPeriodForReady,
+		}
+		c.LivenessProbe = &corev1.Probe{
+			Handler: corev1.Handler{
+				Exec: &corev1.ExecAction{
+					Command: []string{
+						defaultGRPCHealthCheckProbe,
+						fmt.Sprintf("-addr=:%d", consts.DefaultSuggestionPort),
+						fmt.Sprintf("-service=%s", consts.DefaultGRPCService),
+					},
 				},
 			},
-		},
-		// Ref https://srcco.de/posts/kubernetes-liveness-probes-are-dangerous.html
-		InitialDelaySeconds: defaultInitialDelaySeconds,
-		PeriodSeconds:       defaultPeriodForLive,
-		FailureThreshold:    defaultFailureThreshold,
+			// Ref https://srcco.de/posts/kubernetes-liveness-probes-are-dangerous.html
+			InitialDelaySeconds: defaultInitialDelaySeconds,
+			PeriodSeconds:       defaultPeriodForLive,
+			FailureThreshold:    defaultFailureThreshold,
+		}
 	}
 	return c, nil
 }