Add replica_role label to metrics and replica-role field to logs

Author: IrvingMg

cmd/kueue/main.go

@@ -79,6 +79,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/scheduler/preemption/fairsharing"
 	"sigs.k8s.io/kueue/pkg/util/cert"
 	"sigs.k8s.io/kueue/pkg/util/kubeversion"
+	"sigs.k8s.io/kueue/pkg/util/roletracker"
 	"sigs.k8s.io/kueue/pkg/util/useragent"
 	"sigs.k8s.io/kueue/pkg/util/waitforpodsready"
 	"sigs.k8s.io/kueue/pkg/version"
@@ -233,6 +234,15 @@ func main() {
 		os.Exit(1)
 	}
 
+	var roleTracker *roletracker.RoleTracker
+	if cfg.LeaderElection != nil && ptr.Deref(cfg.LeaderElection.LeaderElect, false) {
+		roleTracker = roletracker.NewRoleTracker(mgr.Elected())
+		go roleTracker.Start(ctx, setupLog)
+		setupLog.Info("RoleTracker: leader election enabled")
+	} else {
+		setupLog.Info("RoleTracker: running in standalone mode")
+	}
+
 	certsReady := make(chan struct{})
 
 	if cfg.InternalCertManagement != nil && *cfg.InternalCertManagement.Enable {
@@ -243,8 +253,8 @@ func main() {
 	} else {
 		close(certsReady)
 	}
-	cacheOptions := []schdcache.Option{schdcache.WithPodsReadyTracking(blockForPodsReady(&cfg))}
-	queueOptions := []qcache.Option{qcache.WithPodsReadyRequeuingTimestamp(podsReadyRequeuingTimestamp(&cfg))}
+	cacheOptions := []schdcache.Option{schdcache.WithPodsReadyTracking(blockForPodsReady(&cfg)), schdcache.WithRoleTracker(roleTracker)}
+	queueOptions := []qcache.Option{qcache.WithPodsReadyRequeuingTimestamp(podsReadyRequeuingTimestamp(&cfg)), qcache.WithRoleTracker(roleTracker)}
 	if cfg.Resources != nil && len(cfg.Resources.ExcludeResourcePrefixes) > 0 {
 		cacheOptions = append(cacheOptions, schdcache.WithExcludedResourcePrefixes(cfg.Resources.ExcludeResourcePrefixes))
 		queueOptions = append(queueOptions, qcache.WithExcludedResourcePrefixes(cfg.Resources.ExcludeResourcePrefixes))
@@ -287,12 +297,12 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err := setupControllers(ctx, mgr, cCache, queues, &cfg, serverVersionFetcher); err != nil {
+	if err := setupControllers(ctx, mgr, cCache, queues, &cfg, serverVersionFetcher, roleTracker); err != nil {
 		setupLog.Error(err, "Unable to setup controllers")
 		os.Exit(1)
 	}
 
-	if failedWebhook, err := webhooks.Setup(mgr); err != nil {
+	if failedWebhook, err := webhooks.Setup(mgr, webhooks.WithRoleTracker(roleTracker)); err != nil {
 		setupLog.Error(err, "Unable to create webhook", "webhook", failedWebhook)
 		os.Exit(1)
 	}
@@ -309,7 +319,7 @@ func main() {
 		}()
 	}
 
-	if err := setupScheduler(mgr, cCache, queues, &cfg); err != nil {
+	if err := setupScheduler(mgr, cCache, queues, &cfg, roleTracker); err != nil {
 		setupLog.Error(err, "Could not setup scheduler")
 		os.Exit(1)
 	}
@@ -352,8 +362,8 @@ func setupIndexes(ctx context.Context, mgr ctrl.Manager, cfg *configapi.Configur
 	return jobframework.SetupIndexes(ctx, mgr.GetFieldIndexer(), opts...)
 }
 
-func setupControllers(ctx context.Context, mgr ctrl.Manager, cCache *schdcache.Cache, queues *qcache.Manager, cfg *configapi.Configuration, serverVersionFetcher *kubeversion.ServerVersionFetcher) error {
-	if failedCtrl, err := core.SetupControllers(mgr, queues, cCache, cfg); err != nil {
+func setupControllers(ctx context.Context, mgr ctrl.Manager, cCache *schdcache.Cache, queues *qcache.Manager, cfg *configapi.Configuration, serverVersionFetcher *kubeversion.ServerVersionFetcher, roleTracker *roletracker.RoleTracker) error {
+	if failedCtrl, err := core.SetupControllers(mgr, queues, cCache, cfg, core.WithSetupRoleTracker(roleTracker)); err != nil {
 		return fmt.Errorf("unable to create controller %s: %w", failedCtrl, err)
 	}
 	if features.Enabled(features.FailureRecoveryPolicy) {
@@ -413,7 +423,7 @@ func setupControllers(ctx context.Context, mgr ctrl.Manager, cCache *schdcache.C
 	}
 
 	if features.Enabled(features.TopologyAwareScheduling) {
-		if failedCtrl, err := tas.SetupControllers(mgr, queues, cCache, cfg); err != nil {
+		if failedCtrl, err := tas.SetupControllers(mgr, queues, cCache, cfg, tas.WithRoleTracker(roleTracker)); err != nil {
 			return fmt.Errorf("could not setup TAS controller %s: %w", failedCtrl, err)
 		}
 	}
@@ -429,6 +439,7 @@ func setupControllers(ctx context.Context, mgr ctrl.Manager, cCache *schdcache.C
 		jobframework.WithCache(cCache),
 		jobframework.WithQueues(queues),
 		jobframework.WithObjectRetentionPolicies(cfg.ObjectRetentionPolicies),
+		jobframework.WithRoleTracker(roleTracker),
 	}
 	nsSelector, err := metav1.LabelSelectorAsSelector(cfg.ManagedJobsNamespaceSelector)
 	if err != nil {
@@ -472,7 +483,7 @@ func setupProbeEndpoints(mgr ctrl.Manager, certsReady <-chan struct{}) error {
 	return nil
 }
 
-func setupScheduler(mgr ctrl.Manager, cCache *schdcache.Cache, queues *qcache.Manager, cfg *configapi.Configuration) error {
+func setupScheduler(mgr ctrl.Manager, cCache *schdcache.Cache, queues *qcache.Manager, cfg *configapi.Configuration, roleTracker *roletracker.RoleTracker) error {
 	sched := scheduler.New(
 		queues,
 		cCache,
@@ -481,6 +492,7 @@ func setupScheduler(mgr ctrl.Manager, cCache *schdcache.Cache, queues *qcache.Ma
 		scheduler.WithPodsReadyRequeuingTimestamp(podsReadyRequeuingTimestamp(cfg)),
 		scheduler.WithFairSharing(cfg.FairSharing),
 		scheduler.WithAdmissionFairSharing(cfg.AdmissionFairSharing),
+		scheduler.WithRoleTracker(roleTracker),
 	)
 	if err := mgr.Add(sched); err != nil {
 		return fmt.Errorf("unable to add scheduler to manager: %w", err)

pkg/cache/queue/manager.go

@@ -40,6 +40,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/metrics"
 	afs "sigs.k8s.io/kueue/pkg/util/admissionfairsharing"
 	"sigs.k8s.io/kueue/pkg/util/queue"
+	"sigs.k8s.io/kueue/pkg/util/roletracker"
 	"sigs.k8s.io/kueue/pkg/workload"
 )
 
@@ -52,6 +53,13 @@ var (
 // Option configures the manager.
 type Option func(*Manager)
 
+// WithRoleTracker sets the role tracker for logs and metrics. Can be nil if leader election is disabled.
+func WithRoleTracker(tracker *roletracker.RoleTracker) Option {
+	return func(m *Manager) {
+		m.roleTracker = tracker
+	}
+}
+
 // WithClock allows to specify a custom clock
 func WithClock(c clock.WithDelayedExecution) Option {
 	return func(m *Manager) {
@@ -105,6 +113,7 @@ type Manager struct {
 
 	clock         clock.WithDelayedExecution
 	client        client.Client
+	roleTracker   *roletracker.RoleTracker
 	statusChecker StatusChecker
 	localQueues   map[queue.LocalQueueReference]*LocalQueue
 
@@ -717,7 +726,7 @@ func (m *Manager) reportLQPendingWorkloads(lq *LocalQueue) {
 	metrics.ReportLocalQueuePendingWorkloads(metrics.LocalQueueReference{
 		Name:      lqName,
 		Namespace: namespace,
-	}, active, inadmissible)
+	}, active, inadmissible, m.roleTracker)
 }
 
 func (m *Manager) reportPendingWorkloads(cqName kueue.ClusterQueueReference, cq *ClusterQueue) {
@@ -726,7 +735,7 @@ func (m *Manager) reportPendingWorkloads(cqName kueue.ClusterQueueReference, cq
 		inadmissible += active
 		active = 0
 	}
-	metrics.ReportPendingWorkloads(cqName, active, inadmissible)
+	metrics.ReportPendingWorkloads(cqName, active, inadmissible, m.roleTracker)
 }
 
 func (m *Manager) GetClusterQueueNames() []kueue.ClusterQueueReference {

pkg/cache/scheduler/cache.go

@@ -41,6 +41,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/metrics"
 	"sigs.k8s.io/kueue/pkg/resources"
 	"sigs.k8s.io/kueue/pkg/util/queue"
+	"sigs.k8s.io/kueue/pkg/util/roletracker"
 	"sigs.k8s.io/kueue/pkg/workload"
 )
 
@@ -61,6 +62,13 @@ const (
 // Option configures the reconciler.
 type Option func(*Cache)
 
+// WithRoleTracker sets the role tracker for logs and metrics. Can be nil if leader election is disabled.
+func WithRoleTracker(tracker *roletracker.RoleTracker) Option {
+	return func(c *Cache) {
+		c.roleTracker = tracker
+	}
+}
+
 // WithPodsReadyTracking indicates the cache controller tracks the PodsReady
 // condition for admitted workloads, and allows to block admission of new
 // workloads until all admitted workloads are in the PodsReady condition.
@@ -101,6 +109,7 @@ type Cache struct {
 	podsReadyCond sync.Cond
 
 	client               client.Client
+	roleTracker          *roletracker.RoleTracker
 	assumedWorkloads     map[workload.Reference]kueue.ClusterQueueReference
 	resourceFlavors      map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor
 	podsReadyTracking    bool
@@ -142,6 +151,7 @@ func (c *Cache) newClusterQueue(log logr.Logger, cq *kueue.ClusterQueue) (*clust
 		resourceNode:        NewResourceNode(),
 		tasCache:            &c.tasCache,
 		AdmissionScope:      cq.Spec.AdmissionScope,
+		roleTracker:         c.roleTracker,
 
 		workloadsNotAccountedForTAS: sets.New[workload.Reference](),
 	}
@@ -359,7 +369,7 @@ func (c *Cache) TerminateClusterQueue(name kueue.ClusterQueueReference) {
 	defer c.Unlock()
 	if cq := c.hm.ClusterQueue(name); cq != nil {
 		cq.Status = terminating
-		metrics.ReportClusterQueueStatus(cq.Name, cq.Status)
+		metrics.ReportClusterQueueStatus(cq.Name, cq.Status, c.roleTracker)
 	}
 }
 

pkg/cache/scheduler/clusterqueue.go

@@ -39,6 +39,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	"sigs.k8s.io/kueue/pkg/util/api"
 	"sigs.k8s.io/kueue/pkg/util/queue"
+	"sigs.k8s.io/kueue/pkg/util/roletracker"
 	stringsutils "sigs.k8s.io/kueue/pkg/util/strings"
 	"sigs.k8s.io/kueue/pkg/workload"
 )
@@ -88,6 +89,7 @@ type clusterQueue struct {
 
 	workloadsNotAccountedForTAS sets.Set[workload.Reference]
 	AdmissionScope              *kueue.AdmissionScope
+	roleTracker                 *roletracker.RoleTracker
 }
 
 func (c *clusterQueue) GetName() kueue.ClusterQueueReference {
@@ -228,7 +230,7 @@ func (c *clusterQueue) updateQueueStatus(log logr.Logger) {
 	if status != c.Status {
 		log.V(3).Info("Updating status in cache", "clusterQueue", c.Name, "newStatus", status, "oldStatus", c.Status)
 		c.Status = status
-		metrics.ReportClusterQueueStatus(c.Name, c.Status)
+		metrics.ReportClusterQueueStatus(c.Name, c.Status, c.roleTracker)
 	}
 }
 
@@ -464,14 +466,16 @@ func (c *clusterQueue) deleteWorkload(log logr.Logger, w *kueue.Workload) {
 }
 
 func (c *clusterQueue) reportActiveWorkloads() {
-	metrics.AdmittedActiveWorkloads.WithLabelValues(string(c.Name)).Set(float64(c.admittedWorkloadsCount))
-	metrics.ReservingActiveWorkloads.WithLabelValues(string(c.Name)).Set(float64(len(c.Workloads)))
+	role := roletracker.GetMetricsRole(c.roleTracker)
+	metrics.AdmittedActiveWorkloads.WithLabelValues(string(c.Name), role).Set(float64(c.admittedWorkloadsCount))
+	metrics.ReservingActiveWorkloads.WithLabelValues(string(c.Name), role).Set(float64(len(c.Workloads)))
 }
 
-func (q *LocalQueue) reportActiveWorkloads() {
+func (q *LocalQueue) reportActiveWorkloads(roleTracker *roletracker.RoleTracker) {
 	namespace, name := queue.MustParseLocalQueueReference(q.key)
-	metrics.LocalQueueAdmittedActiveWorkloads.WithLabelValues(string(name), namespace).Set(float64(q.admittedWorkloads))
-	metrics.LocalQueueReservingActiveWorkloads.WithLabelValues(string(name), namespace).Set(float64(q.reservingWorkloads))
+	role := roletracker.GetMetricsRole(roleTracker)
+	metrics.LocalQueueAdmittedActiveWorkloads.WithLabelValues(string(name), namespace, role).Set(float64(q.admittedWorkloads))
+	metrics.LocalQueueReservingActiveWorkloads.WithLabelValues(string(name), namespace, role).Set(float64(q.reservingWorkloads))
 }
 
 // updateWorkloadUsage updates the usage of the ClusterQueue for the workload
@@ -501,7 +505,7 @@ func (c *clusterQueue) updateWorkloadUsage(log logr.Logger, wi *workload.Info, o
 			lq.admittedWorkloads += op.asSignedOne()
 		}
 		if features.Enabled(features.LocalQueueMetrics) {
-			lq.reportActiveWorkloads()
+			lq.reportActiveWorkloads(c.roleTracker)
 		}
 	}
 }
@@ -571,7 +575,7 @@ func (c *clusterQueue) addLocalQueue(q *kueue.LocalQueue) error {
 	}
 	c.localQueues[qKey] = qImpl
 	if features.Enabled(features.LocalQueueMetrics) {
-		qImpl.reportActiveWorkloads()
+		qImpl.reportActiveWorkloads(c.roleTracker)
 	}
 	return nil
 }

pkg/controller/core/admissioncheck_controller.go

@@ -40,6 +40,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
 	qcache "sigs.k8s.io/kueue/pkg/cache/queue"
 	schdcache "sigs.k8s.io/kueue/pkg/cache/scheduler"
+	"sigs.k8s.io/kueue/pkg/util/roletracker"
 	"sigs.k8s.io/kueue/pkg/util/slices"
 )
 
@@ -49,12 +50,13 @@ type AdmissionCheckUpdateWatcher interface {
 
 // AdmissionCheckReconciler reconciles a AdmissionCheck object
 type AdmissionCheckReconciler struct {
-	log        logr.Logger
-	qManager   *qcache.Manager
-	client     client.Client
-	cache      *schdcache.Cache
-	cqUpdateCh chan event.GenericEvent
-	watchers   []AdmissionCheckUpdateWatcher
+	log         logr.Logger
+	qManager    *qcache.Manager
+	client      client.Client
+	cache       *schdcache.Cache
+	cqUpdateCh  chan event.GenericEvent
+	watchers    []AdmissionCheckUpdateWatcher
+	roleTracker *roletracker.RoleTracker
 }
 
 var _ reconcile.Reconciler = (*AdmissionCheckReconciler)(nil)
@@ -64,13 +66,15 @@ func NewAdmissionCheckReconciler(
 	client client.Client,
 	qMgr *qcache.Manager,
 	cache *schdcache.Cache,
+	roleTracker *roletracker.RoleTracker,
 ) *AdmissionCheckReconciler {
 	return &AdmissionCheckReconciler{
-		log:        ctrl.Log.WithName("admissioncheck-reconciler"),
-		qManager:   qMgr,
-		client:     client,
-		cache:      cache,
-		cqUpdateCh: make(chan event.GenericEvent, updateChBuffer),
+		log:         ctrl.Log.WithName("admissioncheck-reconciler"),
+		qManager:    qMgr,
+		client:      client,
+		cache:       cache,
+		cqUpdateCh:  make(chan event.GenericEvent, updateChBuffer),
+		roleTracker: roleTracker,
 	}
 }
 
@@ -239,6 +243,7 @@ func (r *AdmissionCheckReconciler) SetupWithManager(mgr ctrl.Manager, cfg *confi
 		WithOptions(controller.Options{
 			NeedLeaderElection:      ptr.To(false),
 			MaxConcurrentReconciles: mgr.GetControllerOptions().GroupKindConcurrency[kueue.GroupVersion.WithKind("AdmissionCheck").GroupKind().String()],
+			LogConstructor:          roletracker.NewLogConstructor(r.roleTracker, "admissioncheck-reconciler"),
 		}).
 		WatchesRawSource(source.Channel(r.cqUpdateCh, &h)).
 		Complete(WithLeadingManager(mgr, r, &kueue.AdmissionCheck{}, cfg))