feat: Add a flag to control the logic about sc (#994)

* feat: Add a flag to control the logic about sc Signed-off-by: Ce Gao <[email protected]> * feat: Add a nil pointer check Signed-off-by: Ce Gao <[email protected]>
kubeflow · Jan 2, 2020 · 685c0a3 · 685c0a3
1 parent 9728cc4
commit 685c0a3
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 14 deletions.
diff --git a/cmd/katib-controller/v1alpha3/main.go b/cmd/katib-controller/v1alpha3/main.go
@@ -41,24 +41,33 @@ func main() {
 	var metricsAddr string
 	var webhookPort int
 	var certLocalFS bool
+	var injectSecurityContext bool
 
 	flag.StringVar(&experimentSuggestionName, "experiment-suggestion-name",
 		"default", "The implementation of suggestion interface in experiment controller (default|fake)")
 	flag.StringVar(&metricsAddr, "metrics-addr", ":8080", "The address the metric endpoint binds to.")
 	flag.IntVar(&webhookPort, "webhook-port", 8443, "The port number to be used for admission webhook server.")
 	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.Parse()
 
 	// Set the config in viper.
 	viper.Set(consts.ConfigExperimentSuggestionName, experimentSuggestionName)
 	viper.Set(consts.ConfigCertLocalFS, certLocalFS)
+	viper.Set(consts.ConfigInjectSecurityContext, injectSecurityContext)
 
 	log.Info("Config:",
 		consts.ConfigExperimentSuggestionName,
 		viper.GetString(consts.ConfigExperimentSuggestionName),
 		consts.ConfigCertLocalFS,
 		viper.GetBool(consts.ConfigCertLocalFS),
+		"webhook-port",
+		webhookPort,
+		"metrics-addr",
+		metricsAddr,
+		consts.ConfigInjectSecurityContext,
+		viper.GetBool(consts.ConfigInjectSecurityContext),
 	)
 
 	// Get a config to talk to the apiserver

diff --git a/pkg/controller.v1alpha3/consts/const.go b/pkg/controller.v1alpha3/consts/const.go
@@ -9,6 +9,10 @@ const (
 	// ConfigCertLocalFS is the config name which indicates if we
 	// should store the cert in file system.
 	ConfigCertLocalFS = "cert-local-filesystem"
+	// ConfigInjectSecurityContext is the config name which indicates
+	// if we should inject the security context into the metrics collector
+	// sidecar.
+	ConfigInjectSecurityContext = "inject-security-context"
 
 	// LabelExperimentName is the label of experiment name.
 	LabelExperimentName = "experiment"

diff --git a/pkg/webhook/v1alpha3/pod/inject_webhook.go b/pkg/webhook/v1alpha3/pod/inject_webhook.go
@@ -23,20 +23,21 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/spf13/viper"
+	v1 "k8s.io/api/core/v1"
+	k8serrors "k8s.io/apimachinery/pkg/api/errors"
+	apitypes "k8s.io/apimachinery/pkg/types"
+	"k8s.io/utils/pointer"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/runtime/inject"
 	logf "sigs.k8s.io/controller-runtime/pkg/runtime/log"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission/types"
 
-	v1 "k8s.io/api/core/v1"
-	k8serrors "k8s.io/apimachinery/pkg/api/errors"
-	apitypes "k8s.io/apimachinery/pkg/types"
-	"k8s.io/utils/pointer"
-
 	common "github.com/kubeflow/katib/pkg/apis/controller/common/v1alpha3"
 	trialsv1alpha3 "github.com/kubeflow/katib/pkg/apis/controller/trials/v1alpha3"
 	katibmanagerv1alpha3 "github.com/kubeflow/katib/pkg/common/v1alpha3"
+	"github.com/kubeflow/katib/pkg/controller.v1alpha3/consts"
 	jobv1alpha3 "github.com/kubeflow/katib/pkg/job/v1alpha3"
 	mccommon "github.com/kubeflow/katib/pkg/metricscollector/v1alpha3/common"
 	"github.com/kubeflow/katib/pkg/util/v1alpha3/katibconfig"
@@ -46,9 +47,12 @@ var log = logf.Log.WithName("injector-webhook")
 
 // sidecarInjector that inject metrics collect sidecar into master pod
 type sidecarInjector struct {
-	client         client.Client
-	decoder        types.Decoder
-	managerService string
+	client  client.Client
+	decoder types.Decoder
+
+	// injectSecurityContext indicates if we should inject the security
+	// context into the metrics collector sidecar.
+	injectSecurityContext bool
 }
 
 var _ admission.Handler = &sidecarInjector{}
@@ -96,10 +100,11 @@ func (s *sidecarInjector) InjectDecoder(d types.Decoder) error {
 	return nil
 }
 
-func NewSidecarInjector(c client.Client, ms string) *sidecarInjector {
+// NewSidecarInjector returns a new sidecar injector.
+func NewSidecarInjector(c client.Client) *sidecarInjector {
 	return &sidecarInjector{
-		client:         c,
-		managerService: ms,
+		injectSecurityContext: viper.GetBool(consts.ConfigInjectSecurityContext),
+		client:                c,
 	}
 }
 
@@ -177,14 +182,22 @@ func (s *sidecarInjector) getMetricsCollectorContainer(trial *trialsv1alpha3.Tri
 	}
 	args := getMetricsCollectorArgs(trial.Name, metricName, mc)
 	sidecarContainerName := getSidecarContainerName(trial.Spec.MetricsCollector.Collector.Kind)
-	securityContext := originalPod.Spec.Containers[0].SecurityContext.DeepCopy()
+
 	injectContainer := v1.Container{
 		Name:            sidecarContainerName,
 		Image:           image,
 		Args:            args,
 		ImagePullPolicy: v1.PullIfNotPresent,
-		SecurityContext: securityContext,
 	}
+
+	// Inject the security context when the flag is enabled.
+	if s.injectSecurityContext {
+		if len(originalPod.Spec.Containers) != 0 &&
+			originalPod.Spec.Containers[0].SecurityContext != nil {
+			injectContainer.SecurityContext = originalPod.Spec.Containers[0].SecurityContext.DeepCopy()
+		}
+	}
+
 	return &injectContainer, nil
 }
 

diff --git a/pkg/webhook/v1alpha3/webhook.go b/pkg/webhook/v1alpha3/webhook.go
@@ -112,7 +112,7 @@ func register(manager manager.Manager, server *webhook.Server) error {
 		Operations(admissionregistrationv1beta1.Create).
 		WithManager(manager).
 		ForType(&v1.Pod{}).
-		Handlers(pod.NewSidecarInjector(manager.GetClient(), manager.GetConfig().Host)).
+		Handlers(pod.NewSidecarInjector(manager.GetClient())).
 		Build()
 	if err != nil {
 		return err