[Feat] Add RayJob metrics cleanup

Author: dc-ying

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

@@ -156,8 +156,8 @@ func TestCleanupExpiredRayClusterMetrics(t *testing.T) {
 	assert.Empty(t, manager.cleanupQueue)
 }
 
-// TestCleanupLoop tests the background cleanup loop
-func TestCleanupLoop(t *testing.T) {
+// TestRayClusterCleanupLoop tests the background cleanup loop
+func TestRayClusterCleanupLoop(t *testing.T) {
 	// Create a context that we can cancel
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

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

@@ -3,6 +3,8 @@ package metrics
 import (
 	"context"
 	"strconv"
+	"sync"
+	"time"
 
 	"github.com/go-logr/logr"
 	"github.com/prometheus/client_golang/prometheus"
@@ -17,18 +19,30 @@ type RayJobMetricsObserver interface {
 	ObserveRayJobExecutionDuration(name, namespace string, jobDeploymentStatus rayv1.JobDeploymentStatus, retryCount int, duration float64)
 }
 
+// RayJobMetricCleanupItem represents an item in the RayJob metric cleanup queue
+ type RayJobMetricCleanupItem struct {
+	Name      string
+	Namespace string
+	DeleteAt  time.Time
+}
+
 // RayJobMetricsManager implements the prometheus.Collector and RayJobMetricsObserver interface to collect ray job metrics.
 type RayJobMetricsManager struct {
 	rayJobExecutionDurationSeconds *prometheus.GaugeVec
 	rayJobInfo                     *prometheus.Desc
 	rayJobDeploymentStatus         *prometheus.Desc
 	client                         client.Client
 	log                            logr.Logger
+
+	// Cleanup queue and related fields specific to RayJob metrics
+	cleanupQueue []RayJobMetricCleanupItem
+	queueMutex   sync.Mutex
+	metricTTL    time.Duration
 }
 
 // NewRayJobMetricsManager creates a new RayJobMetricsManager instance.
 func NewRayJobMetricsManager(ctx context.Context, client client.Client) *RayJobMetricsManager {
-	collector := &RayJobMetricsManager{
+	manager := &RayJobMetricsManager{
 		rayJobExecutionDurationSeconds: prometheus.NewGaugeVec(
 			prometheus.GaugeOpts{
 				Name: "kuberay_job_execution_duration_seconds",
@@ -50,10 +64,15 @@ func NewRayJobMetricsManager(ctx context.Context, client client.Client) *RayJobM
 			[]string{"name", "namespace", "deployment_status"},
 			nil,
 		),
-		client: client,
-		log:    ctrl.LoggerFrom(ctx),
+		client:       client,
+		log:          ctrl.LoggerFrom(ctx),
+		cleanupQueue: make([]RayJobMetricCleanupItem, 0),
+		metricTTL:    5 * time.Minute, // Keep metrics for 5 minutes
 	}
-	return collector
+	
+	// Start the cleanup goroutine
+	go manager.startRayJobCleanupLoop(ctx)
+	return manager
 }
 
 // Describe implements prometheus.Collector interface Describe method.
@@ -84,6 +103,63 @@ func (r *RayJobMetricsManager) ObserveRayJobExecutionDuration(name, namespace st
 	r.rayJobExecutionDurationSeconds.WithLabelValues(name, namespace, string(jobDeploymentStatus), strconv.Itoa(retryCount)).Set(duration)
 }
 
+// ScheduleRayJobMetricForCleanup schedules a RayJob metric for cleanup after the TTL
+func (r *RayJobMetricsManager) ScheduleRayJobMetricForCleanup(name, namespace string) {
+	r.queueMutex.Lock()
+	defer r.queueMutex.Unlock()
+
+	// Add to cleanup queue
+	item := RayJobMetricCleanupItem{
+		Name:      name,
+		Namespace: namespace,
+		DeleteAt:  time.Now().Add(r.metricTTL),
+	}
+	r.cleanupQueue = append(r.cleanupQueue, item)
+	r.log.Info("Scheduled RayJob metric for cleanup", "name", name, "namespace", namespace, "deleteAt", item.DeleteAt)
+}
+
+// startRayJobCleanupLoop starts a loop to clean up expired RayJob metrics
+func (r *RayJobMetricsManager) startRayJobCleanupLoop(ctx context.Context) {
+	// Check for expired metrics every minute
+	ticker := time.NewTicker(1 * time.Minute)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			r.cleanupExpiredRayJobMetrics()
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+// cleanupExpiredRayJobMetrics removes RayJob metrics that have expired
+func (r *RayJobMetricsManager) cleanupExpiredRayJobMetrics() {
+	r.queueMutex.Lock()
+	defer r.queueMutex.Unlock()
+
+	now := time.Now()
+	remainingItems := make([]RayJobMetricCleanupItem, 0)
+
+	for _, item := range r.cleanupQueue {
+		if now.After(item.DeleteAt) {
+			// Remove all metrics associated with this job
+			r.rayJobExecutionDurationSeconds.DeletePartialMatch(prometheus.Labels{
+				"name":      item.Name,
+				"namespace": item.Namespace,
+			})
+			r.log.Info("Cleaned up expired RayJob metrics", "name", item.Name, "namespace", item.Namespace)
+		} else {
+			// Keep non-expired items
+			remainingItems = append(remainingItems, item)
+		}
+	}
+
+	// Update queue
+	r.cleanupQueue = remainingItems
+}
+
 func (r *RayJobMetricsManager) collectRayJobInfo(rayJob *rayv1.RayJob, ch chan<- prometheus.Metric) {
 	ch <- prometheus.MustNewConstMetric(
 		r.rayJobInfo,

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

@@ -4,7 +4,9 @@ import (
 	"context"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -171,3 +173,191 @@ func TestMetricRayJobDeploymentStatus(t *testing.T) {
 		})
 	}
 }
+
+// TestRayJobMetricCleanupItem tests the structure of RayJobMetricCleanupItem
+func TestRayJobMetricCleanupItem(t *testing.T) {
+	// Create a test cleanup item
+	item := RayJobMetricCleanupItem{
+		Name:      "test-job",
+		Namespace: "default",
+		DeleteAt:  time.Now().Add(5 * time.Minute),
+	}
+
+	// Verify the fields are correctly set
+	assert.Equal(t, "test-job", item.Name)
+	assert.Equal(t, "default", item.Namespace)
+	assert.WithinDuration(t, time.Now().Add(5*time.Minute), item.DeleteAt, time.Second)
+}
+
+// TestScheduleRayJobMetricForCleanup tests adding items to the cleanup queue
+func TestScheduleRayJobMetricForCleanup(t *testing.T) {
+	// Create a fake client
+	k8sScheme := runtime.NewScheme()
+	require.NoError(t, rayv1.AddToScheme(k8sScheme))
+	client := fake.NewClientBuilder().WithScheme(k8sScheme).Build()
+	// Create a metrics manager
+	manager := NewRayJobMetricsManager(context.Background(), client)
+
+	// Schedule a cleanup for a job
+	manager.ScheduleRayJobMetricForCleanup("test-job", "default")
+
+	// Verify the item was added to the queue
+	manager.queueMutex.Lock()
+	defer manager.queueMutex.Unlock()
+
+	assert.Len(t, manager.cleanupQueue, 1)
+	assert.Equal(t, "test-job", manager.cleanupQueue[0].Name)
+	assert.Equal(t, "default", manager.cleanupQueue[0].Namespace)
+	assert.WithinDuration(t, time.Now().Add(5*time.Minute), manager.cleanupQueue[0].DeleteAt, time.Second)
+}
+
+// TestCleanupExpiredRayJobMetrics tests the cleanup of expired metrics
+func TestCleanupExpiredRayJobMetrics(t *testing.T) {
+	// Create a registry, fake client and metrics manager
+	registry := prometheus.NewRegistry()
+	k8sScheme := runtime.NewScheme()
+	require.NoError(t, rayv1.AddToScheme(k8sScheme))
+	client := fake.NewClientBuilder().WithScheme(k8sScheme).Build()
+	manager := NewRayJobMetricsManager(context.Background(), client)
+
+	// Register the manager with the registry
+	registry.MustRegister(manager)
+
+	// Record a metric for a job
+	manager.ObserveRayJobExecutionDuration("test-job", "default", rayv1.JobDeploymentStatusComplete, 0, 10.5)
+
+	// Verify the metric exists
+	metrics, err := registry.Gather()
+	require.NoError(t, err)
+	assert.Len(t, metrics, 1)
+	assert.Equal(t, "kuberay_job_execution_duration_seconds", metrics[0].GetName())
+
+	// Add the job to cleanup queue with a past delete time
+	manager.queueMutex.Lock()
+	manager.cleanupQueue = append(manager.cleanupQueue, RayJobMetricCleanupItem{
+		Name:      "test-job",
+		Namespace: "default",
+		DeleteAt:  time.Now().Add(-1 * time.Minute), // Expired
+	})
+	manager.queueMutex.Unlock()
+
+	// Run cleanup
+	manager.cleanupExpiredRayJobMetrics()
+
+	// Verify the metric was deleted
+	metrics, err = registry.Gather()
+	require.NoError(t, err)
+	// The metric should still exist but have no samples
+	assert.Len(t, metrics, 1)
+	assert.Equal(t, 0, len(metrics[0].GetMetric()))
+}
+
+// TestRayJobCleanupLoop tests the background cleanup loop
+func TestRayJobCleanupLoop(t *testing.T) {
+	// Create a fake client and metrics manager
+	k8sScheme := runtime.NewScheme()
+	require.NoError(t, rayv1.AddToScheme(k8sScheme))
+	client := fake.NewClientBuilder().WithScheme(k8sScheme).Build()
+	manager := NewRayJobMetricsManager(context.Background(), client)
+
+	// Start the cleanup loop
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		manager.startRayJobCleanupLoop(ctx)
+	}()
+
+	// Schedule a cleanup for a job
+	manager.ScheduleRayJobMetricForCleanup("test-job", "default")
+
+	// Wait for the TTL to expire and cleanup to run
+	time.Sleep(2 * time.Second)
+
+	// Verify the cleanup queue is empty
+	manager.queueMutex.Lock()
+	defer manager.queueMutex.Unlock()
+	assert.Len(t, manager.cleanupQueue, 0)
+
+	// Stop the cleanup loop
+	cancel()
+	wg.Wait()
+}
+
+// TestRayJobConditionProvisioned tests metrics when job is provisioned and cleaned up
+func TestRayJobConditionProvisioned(t *testing.T) {
+	// Create a registry, fake client and metrics manager
+	registry := prometheus.NewRegistry()
+	k8sScheme := runtime.NewScheme()
+	require.NoError(t, rayv1.AddToScheme(k8sScheme))
+	client := fake.NewClientBuilder().WithScheme(k8sScheme).Build()
+	manager := NewRayJobMetricsManager(context.Background(), client)
+	registry.MustRegister(manager)
+
+	// Simulate a job becoming provisioned
+	job := &rayv1.RayJob{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-job",
+			Namespace: "default",
+		},
+		Status: rayv1.RayJobStatus{
+			JobDeploymentStatus: rayv1.JobDeploymentStatusRunning,
+			JobStatus:           rayv1.JobStatusRunning,
+			StartTime:           &metav1.Time{Time: time.Now().Add(-10 * time.Second)},
+		},
+	}
+
+	// Simulate job completion and emit metrics
+	oldStatus := job.Status
+	job.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusComplete
+	job.Status.JobStatus = rayv1.JobStatusSucceeded
+
+	// Emit metrics and schedule cleanup
+	emitRayJobExecutionDuration(manager, job.Name, job.Namespace, oldStatus, job.Status)
+
+	// Verify the metric was recorded
+	metrics, err := registry.Gather()
+	require.NoError(t, err)
+	assert.Len(t, metrics, 1)
+	assert.Equal(t, "kuberay_job_execution_duration_seconds", metrics[0].GetName())
+
+	// Fast-forward time and run cleanup
+	manager.queueMutex.Lock()
+	for i := range manager.cleanupQueue {
+		manager.cleanupQueue[i].DeleteAt = time.Now().Add(-1 * time.Minute) // Force expire
+	}
+	manager.queueMutex.Unlock()
+
+	manager.cleanupExpiredRayJobMetrics()
+
+	// Verify the metric was cleaned up
+	metrics, err = registry.Gather()
+	require.NoError(t, err)
+	assert.Len(t, metrics[0].GetMetric(), 0)
+}
+
+// Helper function to match the one in controller
+func emitRayJobExecutionDuration(rayJobMetricsObserver RayJobMetricsObserver, rayJobName, rayJobNamespace string, originalRayJobStatus, rayJobStatus rayv1.RayJobStatus) {
+	if rayJobStatus.StartTime == nil {
+		// Set a default start time if not provided
+		now := time.Now()
+		rayJobStatus.StartTime = &metav1.Time{Time: now.Add(-10 * time.Second)}
+	}
+	if !rayv1.IsJobDeploymentTerminal(originalRayJobStatus.JobDeploymentStatus) && (rayv1.IsJobDeploymentTerminal(rayJobStatus.JobDeploymentStatus) ||
+		rayJobStatus.JobDeploymentStatus == rayv1.JobDeploymentStatusRetrying) {
+
+		retryCount := 0
+		if originalRayJobStatus.Failed != nil {
+			retryCount += int(*originalRayJobStatus.Failed)
+		}
+
+		rayJobMetricsObserver.ObserveRayJobExecutionDuration(
+			rayJobName,
+			rayJobNamespace,
+			rayJobStatus.JobDeploymentStatus,
+			retryCount,
+			time.Since(rayJobStatus.StartTime.Time).Seconds(),
+		)
+	}
+}

ray-operator/controllers/ray/rayjob_controller.go

@@ -465,6 +465,10 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 		return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 	}
 	emitRayJobMetrics(r.options.RayJobMetricsManager, rayJobInstance.Name, rayJobInstance.Namespace, originalRayJobInstance.Status, rayJobInstance.Status)
+	// Schedule metrics cleanup when job reaches terminal status
+	if !rayv1.IsJobDeploymentTerminal(originalRayJobInstance.Status.JobDeploymentStatus) && rayv1.IsJobDeploymentTerminal(rayJobInstance.Status.JobDeploymentStatus) {
+		r.options.RayJobMetricsManager.ScheduleRayJobMetricForCleanup(rayJobInstance.Name, rayJobInstance.Namespace)
+	}
 	return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, nil
 }