[Feat] Add Raycluster metrics cleanup

Author: dc-ying

ray-operator/controllers/ray/metrics/ray_cluster_metrics.go

@@ -2,6 +2,8 @@ package metrics
 
 import (
 	"context"
+	"sync"
+	"time"
 	"strconv"
 
 	"github.com/go-logr/logr"
@@ -19,16 +21,28 @@ type RayClusterMetricsObserver interface {
 	ObserveRayClusterProvisionedDuration(name, namespace string, duration float64)
 }
 
+// RayClusterMetricCleanupItem represents an item in the RayCluster metric cleanup queue
+ type RayClusterMetricCleanupItem struct {
+	Name      string
+	Namespace string
+	DeleteAt  time.Time
+}
+
 // RayClusterMetricsManager implements the prometheus.Collector and RayClusterMetricsObserver interface to collect ray cluster metrics.
 type RayClusterMetricsManager struct {
 	rayClusterProvisionedDurationSeconds *prometheus.GaugeVec
 	rayClusterInfo                       *prometheus.Desc
 	rayClusterConditionProvisioned       *prometheus.Desc
 	client                               client.Client
 	log                                  logr.Logger
+
+	// Cleanup queue and related fields specific to RayCluster metrics
+	cleanupQueue []RayClusterMetricCleanupItem
+	queueMutex   sync.Mutex
+	metricTTL    time.Duration
 }
 
-// NewRayClusterMetricsManager creates a new RayClusterManager instance.
+// NewRayClusterMetricsManager creates a new RayClusterMetricsManager instance.
 func NewRayClusterMetricsManager(ctx context.Context, client client.Client) *RayClusterMetricsManager {
 	manager := &RayClusterMetricsManager{
 		rayClusterProvisionedDurationSeconds: prometheus.NewGaugeVec(
@@ -56,9 +70,14 @@ func NewRayClusterMetricsManager(ctx context.Context, client client.Client) *Ray
 			[]string{"name", "namespace", "condition"},
 			nil,
 		),
-		client: client,
-		log:    ctrl.LoggerFrom(ctx),
+		client:       client,
+		log:          ctrl.LoggerFrom(ctx),
+		cleanupQueue: make([]RayClusterMetricCleanupItem, 0),
+		metricTTL:    5 * time.Minute, // Keep metrics for 5 minutes
 	}
+	
+	// Start the cleanup goroutine
+	go manager.startRayClusterCleanupLoop(ctx)
 	return manager
 }
 
@@ -85,10 +104,65 @@ func (r *RayClusterMetricsManager) Collect(ch chan<- prometheus.Metric) {
 	}
 }
 
+// ObserveRayClusterProvisionedDuration records the provisioned duration of a RayCluster
 func (r *RayClusterMetricsManager) ObserveRayClusterProvisionedDuration(name, namespace string, duration float64) {
 	r.rayClusterProvisionedDurationSeconds.WithLabelValues(name, namespace).Set(duration)
 }
 
+// ScheduleRayClusterMetricForCleanup schedules a RayCluster metric for cleanup after the TTL
+func (r *RayClusterMetricsManager) ScheduleRayClusterMetricForCleanup(name, namespace string) {
+	r.queueMutex.Lock()
+	defer r.queueMutex.Unlock()
+
+	// Add to cleanup queue
+	item := RayClusterMetricCleanupItem{
+		Name:      name,
+		Namespace: namespace,
+		DeleteAt:  time.Now().Add(r.metricTTL),
+	}
+	r.cleanupQueue = append(r.cleanupQueue, item)
+	r.log.Info("Scheduled RayCluster metric for cleanup", "name", name, "namespace", namespace, "deleteAt", item.DeleteAt)
+}
+
+// startRayClusterCleanupLoop starts a loop to clean up expired RayCluster metrics
+func (r *RayClusterMetricsManager) startRayClusterCleanupLoop(ctx context.Context) {
+	// Check for expired metrics every minute
+	ticker := time.NewTicker(1 * time.Minute)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			r.cleanupExpiredRayClusterMetrics()
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+// cleanupExpiredRayClusterMetrics removes RayCluster metrics that have expired
+func (r *RayClusterMetricsManager) cleanupExpiredRayClusterMetrics() {
+	r.queueMutex.Lock()
+	defer r.queueMutex.Unlock()
+
+	now := time.Now()
+	remainingItems := make([]RayClusterMetricCleanupItem, 0)
+
+	for _, item := range r.cleanupQueue {
+		if now.After(item.DeleteAt) {
+			// Remove expired metric
+			r.rayClusterProvisionedDurationSeconds.DeleteLabelValues(item.Name, item.Namespace)
+			r.log.Info("Cleaned up expired RayCluster metric", "name", item.Name, "namespace", item.Namespace)
+		} else {
+			// Keep non-expired items
+			remainingItems = append(remainingItems, item)
+		}
+	}
+
+	// Update queue
+	r.cleanupQueue = remainingItems
+}
+
 func (r *RayClusterMetricsManager) collectRayClusterInfo(cluster *rayv1.RayCluster, ch chan<- prometheus.Metric) {
 	ownerKind := "None"
 	if v, ok := cluster.Labels[utils.RayOriginatedFromCRDLabelKey]; ok {

ray-operator/controllers/ray/metrics/ray_cluster_metrics_test.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -96,6 +97,106 @@ func TestRayClusterInfo(t *testing.T) {
 	}
 }
 
+// TestScheduleRayClusterMetricForCleanup tests scheduling a metric for cleanup
+func TestScheduleRayClusterMetricForCleanup(t *testing.T) {
+	ctx := context.Background()
+	k8sScheme := runtime.NewScheme()
+	require.NoError(t, rayv1.AddToScheme(k8sScheme))
+
+	client := fake.NewClientBuilder().WithScheme(k8sScheme).Build()
+	manager := NewRayClusterMetricsManager(ctx, client)
+
+	// Schedule a metric for cleanup
+	manager.ScheduleRayClusterMetricForCleanup("test-cluster", "test-namespace")
+
+	// Verify the cleanup queue has one item
+	assert.Len(t, manager.cleanupQueue, 1)
+	assert.Equal(t, "test-cluster", manager.cleanupQueue[0].Name)
+	assert.Equal(t, "test-namespace", manager.cleanupQueue[0].Namespace)
+	assert.WithinDuration(t, time.Now().Add(5*time.Minute), manager.cleanupQueue[0].DeleteAt, 1*time.Second)
+}
+
+// TestCleanupExpiredRayClusterMetrics tests cleaning up expired metrics
+func TestCleanupExpiredRayClusterMetrics(t *testing.T) {
+	ctx := context.Background()
+	k8sScheme := runtime.NewScheme()
+	require.NoError(t, rayv1.AddToScheme(k8sScheme))
+
+	client := fake.NewClientBuilder().WithScheme(k8sScheme).Build()
+	manager := NewRayClusterMetricsManager(ctx, client)
+
+	// Set up a metric
+	manager.ObserveRayClusterProvisionedDuration("expired-cluster", "test-namespace", 123.45)
+
+	// Add an expired item to the cleanup queue
+	manager.queueMutex.Lock()
+	manager.cleanupQueue = append(manager.cleanupQueue, RayClusterMetricCleanupItem{
+		Name:      "expired-cluster",
+		Namespace: "test-namespace",
+		DeleteAt:  time.Now().Add(-1 * time.Minute), // Expired 1 minute ago
+	})
+	manager.queueMutex.Unlock()
+
+	// Clean up expired metrics
+	manager.cleanupExpiredRayClusterMetrics()
+
+	// Verify the metric was deleted by checking the registry
+	reg := prometheus.NewRegistry()
+	reg.MustRegister(manager)
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "/metrics", nil)
+	require.NoError(t, err)
+	recorder := httptest.NewRecorder()
+	handler := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
+	handler.ServeHTTP(recorder, req)
+	assert.Equal(t, http.StatusOK, recorder.Code)
+	body := recorder.Body.String()
+	assert.NotContains(t, body, `kuberay_cluster_provisioned_duration_seconds{name="expired-cluster",namespace="test-namespace"}`)
+
+	// Verify the cleanup queue is empty
+	assert.Empty(t, manager.cleanupQueue)
+}
+
+// TestCleanupLoop tests the background cleanup loop
+func TestCleanupLoop(t *testing.T) {
+	// Create a context that we can cancel
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	k8sScheme := runtime.NewScheme()
+	require.NoError(t, rayv1.AddToScheme(k8sScheme))
+
+	client := fake.NewClientBuilder().WithScheme(k8sScheme).Build()
+	manager := NewRayClusterMetricsManager(ctx, client)
+
+	// Set up a metric
+	manager.ObserveRayClusterProvisionedDuration("test-cluster", "test-namespace", 123.45)
+
+	// Add an item to the cleanup queue with a very short TTL
+	manager.queueMutex.Lock()
+	manager.metricTTL = 1 * time.Second // Set TTL to 1 second for testing
+	manager.cleanupQueue = append(manager.cleanupQueue, RayClusterMetricCleanupItem{
+		Name:      "test-cluster",
+		Namespace: "test-namespace",
+		DeleteAt:  time.Now().Add(manager.metricTTL),
+	})
+	manager.queueMutex.Unlock()
+
+	// Wait for the cleanup loop to run and process the item
+	time.Sleep(2 * time.Second)
+
+	// Verify the metric was deleted by checking the registry
+	reg := prometheus.NewRegistry()
+	reg.MustRegister(manager)
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "/metrics", nil)
+	require.NoError(t, err)
+	recorder := httptest.NewRecorder()
+	handler := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
+	handler.ServeHTTP(recorder, req)
+	assert.Equal(t, http.StatusOK, recorder.Code)
+	body := recorder.Body.String()
+	assert.NotContains(t, body, `kuberay_cluster_provisioned_duration_seconds{name="test-cluster",namespace="test-namespace"}`)
+}
+
 func TestRayClusterConditionProvisioned(t *testing.T) {
 	tests := []struct {
 		name            string

ray-operator/controllers/ray/raycluster_controller.go

@@ -195,6 +195,13 @@ func (r *RayClusterReconciler) rayClusterReconcile(ctx context.Context, instance
 	// Please do NOT modify `originalRayClusterInstance` in the following code.
 	originalRayClusterInstance := instance.DeepCopy()
 
+	if instance.DeletionTimestamp != nil && !instance.DeletionTimestamp.IsZero() {
+		logger.Info("RayCluster is being deleted, scheduling metric for cleanup")
+		if r.options.RayClusterMetricsManager != nil {
+			r.options.RayClusterMetricsManager.ScheduleRayClusterMetricForCleanup(instance.Name, instance.Namespace)
+		}
+	}
+
 	// The `enableGCSFTRedisCleanup` is a feature flag introduced in KubeRay v1.0.0. It determines whether
 	// the Redis cleanup job should be activated. Users can disable the feature by setting the environment
 	// variable `ENABLE_GCS_FT_REDIS_CLEANUP` to `false`, and undertake the Redis storage namespace cleanup