Enable to set dynamic watch from controller flags

Author: andreyvelich

cmd/katib-controller/v1beta1/main.go

@@ -30,6 +30,7 @@ import (
 	apis "github.com/kubeflow/katib/pkg/apis/controller"
 	controller "github.com/kubeflow/katib/pkg/controller.v1beta1"
 	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
+	trialutil "github.com/kubeflow/katib/pkg/controller.v1beta1/trial/util"
 	webhook "github.com/kubeflow/katib/pkg/webhook/v1beta1"
 )
 
@@ -44,6 +45,7 @@ func main() {
 	var injectSecurityContext bool
 	var serviceName string
 	var enableGRPCProbeInSuggestion bool
+	var trialResources trialutil.GvkListFlag
 
 	flag.StringVar(&experimentSuggestionName, "experiment-suggestion-name",
 		"default", "The implementation of suggestion interface in experiment controller (default)")
@@ -53,6 +55,7 @@ func main() {
 	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.Var(&trialResources, "trial-resources", "The list of resources that can be used as trial template, in the form: Kind.version.group (e.g. TFJob.v1.kubeflow.org)")
 
 	flag.Parse()
 
@@ -61,6 +64,7 @@ func main() {
 	viper.Set(consts.ConfigCertLocalFS, certLocalFS)
 	viper.Set(consts.ConfigInjectSecurityContext, injectSecurityContext)
 	viper.Set(consts.ConfigEnableGRPCProbeInSuggestion, enableGRPCProbeInSuggestion)
+	viper.Set(consts.ConfigTrialResources, trialResources)
 
 	log.Info("Config:",
 		consts.ConfigExperimentSuggestionName,
@@ -75,6 +79,8 @@ func main() {
 		viper.GetBool(consts.ConfigInjectSecurityContext),
 		consts.ConfigEnableGRPCProbeInSuggestion,
 		viper.GetBool(consts.ConfigEnableGRPCProbeInSuggestion),
+		"trial-resources",
+		viper.Get(consts.ConfigTrialResources),
 	)
 
 	// Get a config to talk to the apiserver

pkg/controller.v1beta1/consts/const.go

@@ -22,6 +22,9 @@ const (
 	// ConfigEnableGRPCProbeInSuggestion is the config name which indicates
 	// if we should set GRPC probe in suggestion deployments.
 	ConfigEnableGRPCProbeInSuggestion = "enable-grpc-probe-in-suggestion"
+	// ConfigTrialResources is the config name which indicates
+	// resources list which can be used as trial template
+	ConfigTrialResources = "trial-resources"
 
 	// LabelExperimentName is the label of experiment name.
 	LabelExperimentName = "experiment"

pkg/controller.v1beta1/trial/trial_controller.go

@@ -42,10 +42,12 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/source"
 
 	trialsv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/trials/v1beta1"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
 	"github.com/kubeflow/katib/pkg/controller.v1beta1/trial/managerclient"
 	trialutil "github.com/kubeflow/katib/pkg/controller.v1beta1/trial/util"
 	"github.com/kubeflow/katib/pkg/controller.v1beta1/util"
 	jobv1beta1 "github.com/kubeflow/katib/pkg/job/v1beta1"
+	"github.com/spf13/viper"
 )
 
 const (
@@ -124,6 +126,35 @@ func add(mgr manager.Manager, r reconcile.Reconciler) error {
 			log.Info("Job watch added successfully", "CRD Kind", gvk.Kind)
 		}
 	}
+
+	trialResources := viper.Get(consts.ConfigTrialResources)
+	if trialResources != nil {
+		// Cast interface to gvk slice object
+		gvkList := trialResources.(trialutil.GvkListFlag)
+
+		// Watch for changes in custom resources
+		for _, gvk := range gvkList {
+			unstructuredJob := &unstructured.Unstructured{}
+			unstructuredJob.SetGroupVersionKind(gvk)
+			err = c.Watch(
+				&source.Kind{Type: unstructuredJob},
+				&handler.EnqueueRequestForOwner{
+					IsController: true,
+					OwnerType:    &trialsv1beta1.Trial{},
+				})
+			if err != nil {
+				if meta.IsNoMatchError(err) {
+					log.Info("Job watch error. CRD might be missing. Please install CRD and restart katib-controller",
+						"CRD Group", gvk.Group, "CRD Version", gvk.Version, "CRD Kind", gvk.Kind)
+					continue
+				}
+				return err
+			}
+			log.Info("Job watch added successfully",
+				"CRD Group", gvk.Group, "CRD Version", gvk.Version, "CRD Kind", gvk.Kind)
+		}
+	}
+
 	log.Info("Trial controller created")
 	return nil
 }

pkg/controller.v1beta1/trial/trial_controller_test.go

@@ -8,6 +8,7 @@ import (
 	tfv1 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1"
 	"github.com/onsi/gomega"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/spf13/viper"
 	"golang.org/x/net/context"
 	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/api/core/v1"
@@ -22,6 +23,7 @@ import (
 	commonv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
 	trialsv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/trials/v1beta1"
 	api_pb "github.com/kubeflow/katib/pkg/apis/manager/v1beta1"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
 	trialutil "github.com/kubeflow/katib/pkg/controller.v1beta1/trial/util"
 	util "github.com/kubeflow/katib/pkg/controller.v1beta1/util"
 	managerclientmock "github.com/kubeflow/katib/pkg/mock/v1beta1/trial/managerclient"
@@ -50,6 +52,22 @@ func TestAdd(t *testing.T) {
 	mgr, err := manager.New(cfg, manager.Options{})
 	g.Expect(err).NotTo(gomega.HaveOccurred())
 
+	// Set fake trial resources
+	trialResources := trialutil.GvkListFlag{
+		{
+			Group:   "kubeflow.org",
+			Version: "v1",
+			Kind:    "TFJob",
+		},
+		{
+			Group:   "kubeflow.org",
+			Version: "v1",
+			Kind:    "MPIJob",
+		},
+	}
+
+	viper.Set(consts.ConfigTrialResources, trialResources)
+
 	// Test - Try to add Trial controller to the manager
 	g.Expect(Add(mgr)).NotTo(gomega.HaveOccurred())
 }

pkg/controller.v1beta1/trial/util/flag_util.go

@@ -0,0 +1,30 @@
+package util
+
+import (
+	"fmt"
+	"strings"
+
+	"k8s.io/apimachinery/pkg/runtime/schema"
+)
+
+// GvkListFlag is the custom flag to parse GroupVersionKind list for trial resources.
+type GvkListFlag []schema.GroupVersionKind
+
+// Set is the method to convert gvk to string value
+func (flag *GvkListFlag) String() string {
+	gvkStrings := []string{}
+	for _, x := range []schema.GroupVersionKind(*flag) {
+		gvkStrings = append(gvkStrings, x.String())
+	}
+	return strings.Join(gvkStrings, ",")
+}
+
+// Set is the method to set gvk from string flag value
+func (flag *GvkListFlag) Set(value string) error {
+	gvk, _ := schema.ParseKindArg(value)
+	if gvk == nil {
+		return fmt.Errorf("Invalid GroupVersionKind: %v", value)
+	}
+	*flag = append(*flag, *gvk)
+	return nil
+}

pkg/controller.v1beta1/trial/util/prometheus_metrics.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package trial
+package util
 
 import (
 	"context"