main.go

/*
Copyright 2024 The Kubeflow Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"errors"
	"flag"
	"net/http"
	"os"
	"sync/atomic"

	zaplog "go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	apiruntime "k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	ctrl "sigs.k8s.io/controller-runtime"
	ctrlpkg "sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
	schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"
	volcanov1beta1 "volcano.sh/apis/pkg/apis/scheduling/v1beta1"

	configapi "github.com/kubeflow/trainer/v2/pkg/apis/config/v1alpha1"
	trainer "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
	"github.com/kubeflow/trainer/v2/pkg/config"
	"github.com/kubeflow/trainer/v2/pkg/controller"
	"github.com/kubeflow/trainer/v2/pkg/features"
	"github.com/kubeflow/trainer/v2/pkg/runtime"
	runtimecore "github.com/kubeflow/trainer/v2/pkg/runtime/core"
	"github.com/kubeflow/trainer/v2/pkg/statusserver"
	"github.com/kubeflow/trainer/v2/pkg/util/cert"
	"github.com/kubeflow/trainer/v2/pkg/webhooks"
)

const (
	webhookConfigurationName = "validator.trainer.kubeflow.org"
)

var (
	scheme   = apiruntime.NewScheme()
	setupLog = ctrl.Log.WithName("setup")
)

func init() {
	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
	utilruntime.Must(configapi.AddToScheme(scheme))
	utilruntime.Must(trainer.AddToScheme(scheme))
	utilruntime.Must(jobsetv1alpha2.AddToScheme(scheme))
	utilruntime.Must(schedulerpluginsv1alpha1.AddToScheme(scheme))
	utilruntime.Must(volcanov1beta1.AddToScheme(scheme))
}

func main() {
	var configFile string
	var enableHTTP2 bool
	var featureGates string

	flag.StringVar(&configFile, "config", "",
		"The controller will load its initial configuration from this file. "+
			"Omit this flag to use the default configuration values. "+
			"Command-line flags override configuration from this file.")
	// if the enable-http2 flag is false (the default), http/2 should be disabled
	// due to its vulnerabilities. More specifically, disabling http/2 will
	// prevent from being vulnerable to the HTTP/2 Stream Cancellation and
	// Rapid Reset CVEs. For more information see:
	// - https://github.com/advisories/GHSA-qppj-fm5r-hxr3
	// - https://github.com/advisories/GHSA-4374-p667-p6c8
	flag.BoolVar(&enableHTTP2, "enable-http2", false,
		"If set, HTTP/2 will be enabled for the metrics and webhook servers")
	flag.StringVar(&featureGates, "feature-gates", "",
		"A comma-separated list of key=value pairs that describe feature gates. "+
			"Command-line feature gates override those specified in the config file.")

	zapOpts := zap.Options{
		TimeEncoder: zapcore.RFC3339NanoTimeEncoder,
		ZapOpts:     []zaplog.Option{zaplog.AddCaller()},
	}
	zapOpts.BindFlags(flag.CommandLine)
	flag.Parse()

	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&zapOpts)))

	setupLog.Info("Loading configuration", "configFile", configFile)
	options, cfg, err := config.Load(scheme, configFile, enableHTTP2)
	if err != nil {
		setupLog.Error(err, "Unable to load configuration")
		os.Exit(1)
	}

	// Set feature gates from config file first
	if err := utilfeature.DefaultMutableFeatureGate.SetFromMap(cfg.FeatureGates); err != nil {
		setupLog.Error(err, "Unable to set feature gates from config file")
		os.Exit(1)
	}

	// Command-line feature gates override config file settings
	if featureGates != "" {
		if err := utilfeature.DefaultMutableFeatureGate.Set(featureGates); err != nil {
			setupLog.Error(err, "Unable to set feature gates from command line")
			os.Exit(1)
		}
	}

	setupLog.Info("Creating manager")
	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), options)
	if err != nil {
		setupLog.Error(err, "unable to start manager")
		os.Exit(1)
	}

	certsReady := make(chan struct{})
	if config.IsCertManagementEnabled(&cfg) {
		setupLog.Info("Setting up certificate management")
		if err = cert.ManageCerts(mgr, cert.Config{
			WebhookSecretName:        cfg.CertManagement.WebhookSecretName,
			WebhookServiceName:       cfg.CertManagement.WebhookServiceName,
			WebhookConfigurationName: webhookConfigurationName,
		}, certsReady); err != nil {
			setupLog.Error(err, "unable to set up cert rotation")
			os.Exit(1)
		}
	} else {
		setupLog.Info("Certificate management is disabled, certificates must be provided externally")
		close(certsReady)
	}

	ctx := ctrl.SetupSignalHandler()

	setupProbeEndpoints(mgr, certsReady)
	runtimes, err := runtimecore.New(ctx, mgr.GetClient(), mgr.GetFieldIndexer(), &cfg)
	if err != nil {
		setupLog.Error(err, "Could not initialize runtimes")
		os.Exit(1)
	}

	// Register status server before goroutine starts and before mgr.Start().
	// AddHealthzCheck/AddReadyzCheck return errors if called after the manager
	// has already started, which is what was happening when this was inside
	// the setupManagerComponents goroutine.
	var statusServerReady atomic.Bool
	if features.Enabled(features.TrainJobStatus) {
		if err := mgr.AddHealthzCheck("status-server", healthz.Ping); err != nil {
			setupLog.Error(err, "unable to set up status server health check")
			os.Exit(1)
		}
		if err := mgr.AddReadyzCheck("status-server", func(_ *http.Request) error {
			if !statusServerReady.Load() {
				return errors.New("status server is not ready")
			}
			return nil
		}); err != nil {
			setupLog.Error(err, "unable to set up status server ready check")
			os.Exit(1)
		}
	}

	// Set up controllers and other components using goroutines to start the manager quickly.
	go setupManagerComponents(mgr, runtimes, &cfg, certsReady, enableHTTP2, &statusServerReady)

	setupLog.Info("Starting manager")
	if err = mgr.Start(ctx); err != nil {
		setupLog.Error(err, "Could not run manager")
		os.Exit(1)
	}
}

func setupManagerComponents(mgr ctrl.Manager, runtimes map[string]runtime.Runtime, cfg *configapi.Configuration, certsReady <-chan struct{}, enableHTTP2 bool, statusServerReady *atomic.Bool) {
	setupLog.Info("Waiting for certificate generation to complete")
	<-certsReady
	setupLog.Info("Certs ready")

	if failedCtrlName, err := controller.SetupControllers(mgr, runtimes, ctrlpkg.Options{}); err != nil {
		setupLog.Error(err, "Could not create controller", "controller", failedCtrlName)
		os.Exit(1)
	}
	if failedWebhook, err := webhooks.Setup(mgr, runtimes); err != nil {
		setupLog.Error(err, "Could not create webhook", "webhook", failedWebhook)
		os.Exit(1)
	}
	if features.Enabled(features.TrainJobStatus) {
		if err := statusserver.SetupServer(mgr, cfg.StatusServer, enableHTTP2); err != nil {
			setupLog.Error(err, "Could not create runtime status server")
			os.Exit(1)
		}
		statusServerReady.Store(true)
	}
}

func setupProbeEndpoints(mgr ctrl.Manager, certsReady <-chan struct{}) {
	defer setupLog.Info("Probe endpoints are configured on healthz and readyz")

	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
		setupLog.Error(err, "unable to set up health check")
		os.Exit(1)
	}
	// Wait for the webhook server to be listening before advertising the
	// training-operator replica as ready. This allows users to wait with sending the first
	// requests, requiring webhooks, until the training-operator deployment is available, so
	// that the early requests are not rejected during the training-operator's startup.
	// We wrap the call to GetWebhookServer in a closure to delay calling
	// the function, otherwise a not fully-initialized webhook server (without
	// ready certs) fails the start of the manager.
	if err := mgr.AddReadyzCheck("readyz", func(req *http.Request) error {
		select {
		case <-certsReady:
			return mgr.GetWebhookServer().StartedChecker()(req)
		default:
			return errors.New("certificates are not ready")
		}
	}); err != nil {
		setupLog.Error(err, "unable to set up ready check")
		os.Exit(1)
	}
}