Merge branch 'master' into add_metadata_write_logic_in_history_server

Signed-off-by: Jie-Kai Chang <jiekaichang@apache.org>

Author: 400Ping

historyserver/cmd/collector/main.go

@@ -1,11 +1,11 @@
 package main
 
 import (
-	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
 	"path"
 	"strconv"
 	"time"
@@ -108,18 +108,36 @@ func main() {
 		panic("Failed to create writer for runtime class name: " + runtimeClassName + " for role: " + role + ".")
 	}
 
+	var wg sync.WaitGroup
+
+	sigChan := make(chan os.Signal, 1)
+	stop := make(chan struct{}, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+
+	wg.Add(1)
 	// Create and initialize EventServer
-	eventServer := eventserver.NewEventServer(writer, rayRootDir, sessionDir, rayNodeId, rayClusterName, rayClusterId, sessionName)
-	eventServer.InitServer(eventsPort)
-
-	collector := runtime.NewCollector(&globalConfig, writer)
-	_ = collector.Start(context.TODO().Done())
-
-	eventStop := eventServer.WaitForStop()
-	logStop := collector.WaitForStop()
-	<-eventStop
-	logrus.Info("Event server shutdown")
-	<-logStop
-	logrus.Info("Log server shutdown")
-	logrus.Info("All servers shutdown")
+	go func() {
+		defer wg.Done()
+		eventServer := eventserver.NewEventServer(writer, rayRootDir, sessionDir, rayNodeId, rayClusterName, rayClusterId, sessionName)
+		eventServer.InitServer(stop, eventsPort)
+		logrus.Info("Event server shutdown")
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		collector := runtime.NewCollector(&globalConfig, writer)
+		collector.Run(stop)
+		logrus.Info("Log server shutdown")
+	}()
+
+	<-sigChan
+	logrus.Info("Received shutdown signal, initiating graceful shutdown...")
+
+	// Stop both the event server and the collector
+	close(stop)
+
+	// Wait for both goroutines to complete
+	wg.Wait()
+	logrus.Info("Graceful shutdown complete")
 }

historyserver/cmd/historyserver/main.go

@@ -63,24 +63,23 @@ func main() {
 	// WaitGroup to track goroutine completion
 	var wg sync.WaitGroup
 
+	sigChan := make(chan os.Signal, 1)
+	stop := make(chan struct{}, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+
 	// Start EventHandler in background goroutine
-	eventStop := make(chan struct{}, 1)
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
 		logrus.Info("Starting EventHandler in background...")
-		if err := eventHandler.Run(eventStop, 2); err != nil {
+		if err := eventHandler.Run(stop, 2); err != nil {
 			logrus.Errorf("EventHandler stopped with error: %v", err)
 		}
 		logrus.Info("EventHandler shutdown complete")
 	}()
 
 	handler := historyserver.NewServerHandler(&globalConfig, dashboardDir, reader, cliMgr, eventHandler)
 
-	sigChan := make(chan os.Signal, 1)
-	stop := make(chan struct{}, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
@@ -92,8 +91,7 @@ func main() {
 	logrus.Info("Received shutdown signal, initiating graceful shutdown...")
 
 	// Stop both the server and the event handler
-	stop <- struct{}{}
-	eventStop <- struct{}{}
+	close(stop)
 
 	// Wait for both goroutines to complete
 	wg.Wait()

historyserver/pkg/collector/eventserver/eventserver.go

@@ -7,11 +7,9 @@ import (
 	"io"
 	"net/http"
 	"os"
-	"os/signal"
 	"path"
 	"strings"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/emicklei/go-restful/v3"
@@ -83,7 +81,7 @@ func NewEventServer(writer storage.StorageWriter, rootDir, sessionDir, nodeID, c
 	return server
 }
 
-func (es *EventServer) InitServer(port int) {
+func (es *EventServer) InitServer(stop <-chan struct{}, port int) {
 	ws := new(restful.WebService)
 	ws.Path("/v1")
 	ws.Consumes(restful.MIME_JSON)
@@ -101,16 +99,10 @@ func (es *EventServer) InitServer(port int) {
 		es.periodicFlush()
 	}()
 
-	// Handle SIGTERM signal
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
-
-	go func() {
-		<-sigChan
-		logrus.Info("Received SIGTERM, flushing events to storage")
-		es.flushEvents()
-		close(es.stopped)
-	}()
+	<-stop
+	logrus.Info("Received stop signal, flushing events to storage")
+	es.flushEvents()
+	close(es.stopped)
 }
 
 // watchNodeIDFile watches /tmp/ray/raylet_node_id for content changes
@@ -297,10 +289,6 @@ func (es *EventServer) periodicFlush() {
 	}
 }
 
-func (es *EventServer) WaitForStop() <-chan struct{} {
-	return es.stopped
-}
-
 func (es *EventServer) flushEvents() {
 	es.mutex.Lock()
 	if len(es.events) == 0 {

historyserver/pkg/collector/logcollector/runtime/interface.go

@@ -1,6 +1,5 @@
 package runtime
 
 type RayLogCollector interface {
-	Start(stop <-chan struct{}) error
-	WaitForStop() <-chan struct{}
+	Run(stop <-chan struct{}) error
 }

historyserver/pkg/collector/logcollector/runtime/logcollector/collector.go

@@ -6,12 +6,10 @@ import (
 	"io/fs"
 	"net/http"
 	"os"
-	"os/signal"
 	"path"
 	"path/filepath"
 	"strings"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/fsnotify/fsnotify"
@@ -47,11 +45,6 @@ type RayLogHandler struct {
 	JobsUrlInfo                  *types.UrlInfo
 }
 
-func (r *RayLogHandler) Start(stop <-chan struct{}) error {
-	go r.Run(stop)
-	return nil
-}
-
 func (r *RayLogHandler) Run(stop <-chan struct{}) error {
 	// watchPath := r.LogDir
 	r.prevLogsDir = filepath.Join("/tmp", "ray", "prev-logs")
@@ -66,10 +59,6 @@ func (r *RayLogHandler) Run(stop <-chan struct{}) error {
 	}
 	defer watcher.Close()
 
-	// Setup signal handling for SIGTERM
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGTERM)
-
 	// WatchPrevLogsLoops performs an initial scan of the prev-logs directory on startup
 	// to process leftover log files in prev-logs/{sessionID}/{nodeID}/logs/ directories.
 	// After scanning, it watches for new directories and files. This ensures incomplete
@@ -83,22 +72,12 @@ func (r *RayLogHandler) Run(stop <-chan struct{}) error {
 		go r.PersistMetaLoop(stop)
 	}
 
-	select {
-	case <-sigChan:
-		logrus.Info("Received SIGTERM, processing all logs...")
-		r.processSessionLatestLogs()
-		close(r.ShutdownChan)
-	case <-stop:
-		logrus.Info("Received stop signal, processing all logs...")
-		r.processSessionLatestLogs()
-		close(r.ShutdownChan)
-	}
-	logrus.Warnf("Receive stop single, so stop ray collector ")
-	return nil
-}
+	<-stop
+	logrus.Info("Received stop signal, processing all logs...")
+	r.processSessionLatestLogs()
+	close(r.ShutdownChan)
 
-func (r *RayLogHandler) WaitForStop() <-chan struct{} {
-	return r.ShutdownChan
+	return nil
 }
 
 // processSessionLatestLogs processes logs in /tmp/ray/session_latest/logs directory

historyserver/pkg/historyserver/server.go

@@ -39,7 +39,7 @@ func NewServerHandler(c *types.RayHistoryServerConfig, dashboardDir string, read
 	}
 }
 
-func (s *ServerHandler) Run(stop chan struct{}) error {
+func (s *ServerHandler) Run(stop <-chan struct{}) error {
 	s.RegisterRouter()
 	port := ":8080"
 	server := &http.Server{

ray-operator/controllers/ray/batchscheduler/kai-scheduler/kai_scheduler.go

@@ -36,9 +36,18 @@ func GetPluginName() string { return "kai-scheduler" }
 func (k *KaiScheduler) Name() string { return GetPluginName() }
 
 func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, object metav1.Object) error {
-	_, ok := object.(*rayv1.RayCluster)
-	if !ok {
-		return fmt.Errorf("currently only RayCluster is supported, got %T", object)
+	// In K8sJobMode and InteractiveMode, RayJob first creates a RayCluster,
+	// and then creates the submitter pod after the RayCluster is ready.
+	// KAI-Scheduler does not handle this two-phase creation pattern.
+	// Other schedulers like Yunikorn and Volcano support this by pre-creating PodGroups.
+	// For more details, see https://github.com/ray-project/kuberay/pull/4418#pullrequestreview-3751609041
+	if rayJob, ok := object.(*rayv1.RayJob); ok {
+		switch rayJob.Spec.SubmissionMode {
+		case rayv1.K8sJobMode:
+			return fmt.Errorf("KAI-Scheduler does not support RayJob with K8sJobMode: the submitter pod is created after RayCluster is ready, preventing proper gang scheduling")
+		case rayv1.InteractiveMode:
+			return fmt.Errorf("KAI-Scheduler does not support RayJob with InteractiveMode: the user-provided submitter runs after RayCluster is ready, preventing proper gang scheduling")
+		}
 	}
 	return nil
 }

ray-operator/controllers/ray/batchscheduler/kai-scheduler/kai_scheduler_test.go

@@ -21,6 +21,19 @@ func createTestRayCluster(labels map[string]string) *rayv1.RayCluster {
 	}
 }
 
+func createTestRayJob(labels map[string]string) *rayv1.RayJob {
+	return &rayv1.RayJob{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-job",
+			Namespace: "default",
+			Labels:    labels,
+		},
+		Spec: rayv1.RayJobSpec{
+			RayClusterSpec: &rayv1.RayClusterSpec{},
+		},
+	}
+}
+
 func createTestPod() *corev1.Pod {
 	return &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
@@ -139,3 +152,39 @@ func TestAddMetadataToChildResource_PreservesExistingPodLabels(t *testing.T) {
 	a.Equal("existing-value", pod.Labels["existing-label"])
 	a.Equal("ray", pod.Labels["app"])
 }
+
+func TestAddMetadataToChildResource_WithRayJob(t *testing.T) {
+	a := assert.New(t)
+	scheduler := &KaiScheduler{}
+	ctx := context.Background()
+
+	rayJob := createTestRayJob(map[string]string{
+		QueueLabelName: "test-queue",
+	})
+	pod := createTestPod()
+
+	scheduler.AddMetadataToChildResource(ctx, rayJob, pod, "test-group")
+
+	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
+
+	a.NotNil(pod.Labels)
+	a.Equal("test-queue", pod.Labels[QueueLabelName])
+}
+
+func TestAddMetadataToChildResource_WithRayJob_WithoutQueueLabel(t *testing.T) {
+	a := assert.New(t)
+	scheduler := &KaiScheduler{}
+	ctx := context.Background()
+
+	rayJob := createTestRayJob(map[string]string{})
+	pod := createTestPod()
+
+	scheduler.AddMetadataToChildResource(ctx, rayJob, pod, "test-group")
+
+	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
+
+	if pod.Labels != nil {
+		_, exists := pod.Labels[QueueLabelName]
+		a.False(exists)
+	}
+}