main.go

package main

import (
	"flag"
	"os"

	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"

	mlv1alpha1 "github.com/ataime/k8s-lab/ml-platform/operator/api/v1alpha1"
	"github.com/ataime/k8s-lab/ml-platform/operator/controller"
)

var scheme = runtime.NewScheme()

func init() {
	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
	utilruntime.Must(mlv1alpha1.AddToScheme(scheme))
}

func main() {
	var inferenceImage string
	flag.StringVar(&inferenceImage, "inference-image", "107.148.176.193:30180/ml-platform/inference:v1", "推理服务镜像")
	flag.Parse()

	ctrl.SetLogger(zap.New(zap.UseDevMode(true)))
	logger := ctrl.Log.WithName("ml-operator")

	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:                 scheme,
		HealthProbeBindAddress: ":8081",
		LeaderElection:         false, // 单副本，不需要选举
	})
	if err != nil {
		logger.Error(err, "创建 Manager 失败")
		os.Exit(1)
	}

	// 注册控制器
	reconciler := &controller.MLModelReconciler{
		Client:         mgr.GetClient(),
		Scheme:         mgr.GetScheme(),
		InferenceImage: inferenceImage,
	}
	if err := reconciler.SetupWithManager(mgr); err != nil {
		logger.Error(err, "注册控制器失败")
		os.Exit(1)
	}

	// 健康检查端点
	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
		logger.Error(err, "添加健康检查失败")
		os.Exit(1)
	}
	if err := mgr.AddReadyzCheck("readyz", healthz.Ping); err != nil {
		logger.Error(err, "添加就绪检查失败")
		os.Exit(1)
	}

	logger.Info("启动 ML Operator", "inference-image", inferenceImage)
	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		logger.Error(err, "Manager 运行失败")
		os.Exit(1)
	}
}