[Feat] [history server] Add actor task endpoint

historyserver/cmd/historyserver/main.go

@@ -21,14 +21,20 @@ func main() {
 	kubeconfigs := ""
 	runtimeClassConfigPath := "/var/collector-config/data"
 	dashboardDir := ""
+	useKubernetesProxy := false
 	flag.StringVar(&runtimeClassName, "runtime-class-name", "", "")
 	flag.StringVar(&rayRootDir, "ray-root-dir", "", "")
 	flag.StringVar(&kubeconfigs, "kubeconfigs", "", "")
 	flag.StringVar(&dashboardDir, "dashboard-dir", "/dashboard", "")
 	flag.StringVar(&runtimeClassConfigPath, "runtime-class-config-path", "", "") //"/var/collector-config/data"
+	flag.BoolVar(&useKubernetesProxy, "use-kubernetes-proxy", false, "")
 	flag.Parse()
 
-	cliMgr := historyserver.NewClientManager(kubeconfigs)
+	cliMgr, err := historyserver.NewClientManager(kubeconfigs, useKubernetesProxy)
+	if err != nil {
+		logrus.Errorf("Failed to create client manager: %v", err)
+		os.Exit(1)
+	}
 
 	jsonData := make(map[string]interface{})
 	if runtimeClassConfigPath != "" {
@@ -75,7 +81,11 @@ func main() {
 		logrus.Info("EventHandler shutdown complete")
 	}()
 
-	handler := historyserver.NewServerHandler(&globalConfig, dashboardDir, reader, cliMgr, eventHandler)
+	handler, err := historyserver.NewServerHandler(&globalConfig, dashboardDir, reader, cliMgr, eventHandler, useKubernetesProxy)
+	if err != nil {
+		logrus.Errorf("Failed to create server handler: %v", err)
+		os.Exit(1)
+	}
 
 	sigChan := make(chan os.Signal, 1)
 	stop := make(chan struct{}, 1)

historyserver/pkg/eventserver/eventserver.go

@@ -198,6 +198,7 @@ func (h *EventHandler) storeEvent(eventMap map[string]any) error {
 	}
 	eventType := types.EventType(eventTypeStr)
 
+	// clusterNameVal is actually the cluster session key.
 	clusterNameVal, ok := eventMap["clusterName"]
 	if !ok {
 		return fmt.Errorf("event missing 'clusterName' field")
@@ -210,126 +211,9 @@ func (h *EventHandler) storeEvent(eventMap map[string]any) error {
 	logrus.Infof("current eventType: %v", eventType)
 	switch eventType {
 	case types.TASK_DEFINITION_EVENT:
-		taskDef, ok := eventMap["taskDefinitionEvent"]
-		if !ok {
-			return fmt.Errorf("event does not have 'taskDefinitionEvent'")
-		}
-		jsonTaskDefinition, err := json.Marshal(taskDef)
-		if err != nil {
-			return err
-		}
-
-		var currTask types.Task
-		if err := json.Unmarshal(jsonTaskDefinition, &currTask); err != nil {
-			return err
-		}
-
-		taskMap := h.ClusterTaskMap.GetOrCreateTaskMap(currentClusterName)
-		taskMap.CreateOrMergeAttempt(currTask.TaskID, currTask.AttemptNumber, func(t *types.Task) {
-			// Merge definition fields (preserve existing Events if any)
-			existingEvents := t.Events
-			*t = currTask
-			if len(existingEvents) > 0 {
-				t.Events = existingEvents
-				t.State = existingEvents[len(existingEvents)-1].State
-			}
-		})
-
+		return h.handleTaskDefinitionEvent(eventMap, currentClusterName, false)
 	case types.TASK_LIFECYCLE_EVENT:
-		lifecycleEvent, ok := eventMap["taskLifecycleEvent"].(map[string]any)
-		if !ok {
-			return fmt.Errorf("invalid taskLifecycleEvent format")
-		}
-
-		taskId, _ := lifecycleEvent["taskId"].(string)
-		taskAttempt, _ := lifecycleEvent["taskAttempt"].(float64)
-		transitions, _ := lifecycleEvent["stateTransitions"].([]any)
-
-		nodeId, _ := lifecycleEvent["nodeId"].(string)
-		workerId, _ := lifecycleEvent["workerId"].(string)
-
-		if len(transitions) == 0 || taskId == "" {
-			return nil
-		}
-
-		// Parse state transitions
-		var stateEvents []types.StateEvent
-		for _, transition := range transitions {
-			tr, ok := transition.(map[string]any)
-			if !ok {
-				continue
-			}
-			state, _ := tr["state"].(string)
-			timestampStr, _ := tr["timestamp"].(string)
-
-			var timestamp time.Time
-			if timestampStr != "" {
-				timestamp, _ = time.Parse(time.RFC3339Nano, timestampStr)
-			}
-
-			stateEvents = append(stateEvents, types.StateEvent{
-				State:     types.TaskStatus(state),
-				Timestamp: timestamp,
-			})
-		}
-
-		if len(stateEvents) == 0 {
-			return nil
-		}
-
-		taskMap := h.ClusterTaskMap.GetOrCreateTaskMap(currentClusterName)
-		taskMap.CreateOrMergeAttempt(taskId, int(taskAttempt), func(t *types.Task) {
-			// --- DEDUPLICATION using (State + Timestamp) as unique key ---
-			// Build a set of existing event keys to detect duplicates
-			type eventKey struct {
-				State     string
-				Timestamp int64
-			}
-			existingKeys := make(map[eventKey]bool)
-			for _, e := range t.Events {
-				existingKeys[eventKey{string(e.State), e.Timestamp.UnixNano()}] = true
-			}
-
-			// Only append events that haven't been seen before
-			for _, e := range stateEvents {
-				key := eventKey{string(e.State), e.Timestamp.UnixNano()}
-				if !existingKeys[key] {
-					t.Events = append(t.Events, e)
-					existingKeys[key] = true
-				}
-			}
-
-			// Sort events by timestamp to ensure correct order
-			sort.Slice(t.Events, func(i, j int) bool {
-				return t.Events[i].Timestamp.Before(t.Events[j].Timestamp)
-			})
-
-			if len(t.Events) == 0 {
-				return
-			}
-
-			t.State = t.Events[len(t.Events)-1].State
-
-			if nodeId != "" {
-				t.NodeID = nodeId
-			}
-			if workerId != "" {
-				t.WorkerID = workerId
-			}
-			if t.StartTime.IsZero() {
-				for _, e := range t.Events {
-					if e.State == types.RUNNING {
-						t.StartTime = e.Timestamp
-						break
-					}
-				}
-			}
-			lastEvent := t.Events[len(t.Events)-1]
-			if lastEvent.State == types.FINISHED || lastEvent.State == types.FAILED {
-				t.EndTime = lastEvent.Timestamp
-			}
-		})
-
+		return h.handleTaskLifecycleEvent(eventMap, currentClusterName)
 	case types.ACTOR_DEFINITION_EVENT:
 		actorDef, ok := eventMap["actorDefinitionEvent"]
 		if !ok {
@@ -525,12 +409,8 @@ func (h *EventHandler) storeEvent(eventMap map[string]any) error {
 			}
 			a.NumRestarts = restartCount
 		})
-
 	case types.ACTOR_TASK_DEFINITION_EVENT:
-		// TODO: Handle actor task definition event
-		// This is related to GET /api/v0/tasks (type=ACTOR_TASK)
-		logrus.Debugf("ACTOR_TASK_DEFINITION_EVENT received, not yet implemented")
-
+		return h.handleTaskDefinitionEvent(eventMap, currentClusterName, true)
 	case types.DRIVER_JOB_DEFINITION_EVENT:
 		// NOTE: When event comes in, JobID will be in base64, processing will convert it to Hex
 		jobDef, ok := eventMap["driverJobDefinitionEvent"]
@@ -743,29 +623,31 @@ func (h *EventHandler) getAllNodeEventFiles(clusterInfo utils.ClusterInfo) []str
 	return nodeEventFiles
 }
 
-// GetTasks returns a thread-safe deep copy of all tasks (including all attempts) for a given cluster.
-// Each task attempt is returned as a separate element in the slice.
+// GetTasks returns a slice of thread-safe deep copies of all task attempts for a given cluster session.
 // Deep copy ensures the returned data is safe to use after the lock is released.
-func (h *EventHandler) GetTasks(clusterName string) []types.Task {
+func (h *EventHandler) GetTasks(clusterSessionKey string) []types.Task {
 	h.ClusterTaskMap.RLock()
 	defer h.ClusterTaskMap.RUnlock()
 
-	taskMap, ok := h.ClusterTaskMap.ClusterTaskMap[clusterName]
+	taskMap, ok := h.ClusterTaskMap.ClusterTaskMap[clusterSessionKey]
 	if !ok {
+		// TODO(jwj): Add error handling.
+		logrus.Errorf("Task map not found for cluster session: %s", clusterSessionKey)
 		return []types.Task{}
 	}
 
 	taskMap.Lock()
 	defer taskMap.Unlock()
 
-	// Flatten all attempts into a single slice with deep copy
-	var tasks []types.Task
-	for _, attempts := range taskMap.TaskMap {
-		for _, task := range attempts {
-			tasks = append(tasks, task.DeepCopy())
+	// Flatten all task attempts into a single slice with deep copy.
+	allTasks := make([]types.Task, 0)
+	for _, taskAttempts := range taskMap.TaskMap {
+		for _, taskAttempt := range taskAttempts {
+			allTasks = append(allTasks, taskAttempt.DeepCopy())
 		}
 	}
-	return tasks
+
+	return allTasks
 }
 
 // GetTaskByID returns all attempts for a specific task ID in a given cluster.
@@ -915,3 +797,136 @@ func (h *EventHandler) GetJobByJobID(clusterName, jobID string) (types.Job, bool
 	}
 	return job.DeepCopy(), true
 }
+
+// handleTaskDefinitionEvent processes TASK_DEFINITION_EVENT or ACTOR_TASK_DEFINITION_EVENT and preserves the task attempt ordering.
+func (h *EventHandler) handleTaskDefinitionEvent(eventMap map[string]any, clusterSessionKey string, isActorTask bool) error {
+	var taskDefField string
+	if isActorTask {
+		taskDefField = "actorTaskDefinitionEvent"
+	} else {
+		taskDefField = "taskDefinitionEvent"
+	}
+
+	taskDef, ok := eventMap[taskDefField]
+	if !ok {
+		return fmt.Errorf("event does not have '%s' field", taskDefField)
+	}
+
+	jsonTaskDefinition, err := json.Marshal(taskDef)
+	if err != nil {
+		return fmt.Errorf("failed to marshal %s event: %w", taskDefField, err)
+	}
+
+	var currTask types.Task
+	if err := json.Unmarshal(jsonTaskDefinition, &currTask); err != nil {
+		return fmt.Errorf("failed to unmarshal %s event: %w", taskDefField, err)
+	}
+
+	// Manually set the task type for an actor task.
+	if isActorTask {
+		currTask.TaskType = types.ACTOR_TASK
+	}
+
+	taskMap := h.ClusterTaskMap.GetOrCreateTaskMap(clusterSessionKey)
+	taskMap.CreateOrMergeAttempt(currTask.TaskID, currTask.TaskAttempt, func(task *types.Task) {
+		// Preserve existing state transitions.
+		existingStateTransitions := task.StateTransitions
+
+		*task = currTask
+
+		if len(existingStateTransitions) > 0 {
+			task.StateTransitions = existingStateTransitions
+			task.State = task.GetLastState()
+		}
+	})
+
+	return nil
+}
+
+// handleTaskLifecycleEvent processes TASK_LIFECYCLE_EVENT and merges state transitions for a given task attempt.
+func (h *EventHandler) handleTaskLifecycleEvent(eventMap map[string]any, clusterSessionKey string) error {
+	taskLifecycle, ok := eventMap["taskLifecycleEvent"]
+	if !ok {
+		return fmt.Errorf("event does not have 'taskLifecycleEvent' field")
+	}
+
+	jsonTaskLifecycle, err := json.Marshal(taskLifecycle)
+	if err != nil {
+		return fmt.Errorf("failed to marshal task lifecycle event: %w", err)
+	}
+
+	var currTask types.Task
+	if err := json.Unmarshal(jsonTaskLifecycle, &currTask); err != nil {
+		return fmt.Errorf("failed to unmarshal task lifecycle event: %w", err)
+	}
+	if currTask.TaskID == "" {
+		return fmt.Errorf("task ID is empty")
+	}
+
+	// TODO(jwj): Clarify if there must be at least one state transition. Can one task have more than one state transition?
+	if len(currTask.StateTransitions) == 0 {
+		return fmt.Errorf("TASK_LIFECYCLE_EVENT must have at least one state transition")
+	}
+
+	taskMap := h.ClusterTaskMap.GetOrCreateTaskMap(clusterSessionKey)
+	taskMap.CreateOrMergeAttempt(currTask.TaskID, currTask.TaskAttempt, func(task *types.Task) {
+		// --- DEDUPLICATION using (State + Timestamp) as unique key ---
+		// Build a set of existing event keys to detect duplicates
+		type eventKey struct {
+			State     string
+			Timestamp int64
+		}
+		existingKeys := make(map[eventKey]bool)
+		for _, e := range task.StateTransitions {
+			existingKeys[eventKey{string(e.State), e.Timestamp.UnixNano()}] = true
+		}
+
+		// Only append events that haven't been seen before
+		for _, e := range currTask.StateTransitions {
+			key := eventKey{string(e.State), e.Timestamp.UnixNano()}
+			if !existingKeys[key] {
+				task.StateTransitions = append(task.StateTransitions, e)
+				existingKeys[key] = true
+			}
+		}
+
+		// Sort events by timestamp to ensure correct order
+		sort.Slice(task.StateTransitions, func(i, j int) bool {
+			return task.StateTransitions[i].Timestamp.Before(task.StateTransitions[j].Timestamp)
+		})
+
+		if len(task.StateTransitions) == 0 {
+			return
+		}
+
+		// TODO(jwj): Before beta, the lifecycle-related fields are overwritten.
+		// In beta, the complete historical replay will be supported.
+		task.JobID = currTask.JobID
+		task.NodeID = currTask.NodeID
+		task.WorkerID = currTask.WorkerID
+		task.WorkerPID = currTask.WorkerPID
+		task.IsDebuggerPaused = currTask.IsDebuggerPaused
+		task.ActorReprName = currTask.ActorReprName
+		task.State = task.GetLastState()
+
+		// Derive creation time, start time and end time from state transitions.
+		// Ref: https://github.com/ray-project/ray/blob/d0b1d151d8ea964a711e451d0ae736f8bf95b629/python/ray/util/state/common.py#L1660-L1685
+		for _, tr := range task.StateTransitions {
+			switch tr.State {
+			case types.PENDING_ARGS_AVAIL:
+				if task.CreationTime.IsZero() {
+					task.CreationTime = tr.Timestamp
+				}
+			case types.RUNNING:
+				if task.StartTime.IsZero() {
+					task.StartTime = tr.Timestamp
+				}
+			case types.FINISHED, types.FAILED:
+				// Take the latest timestamp as the end time.
+				task.EndTime = tr.Timestamp
+			}
+		}
+	})
+
+	return nil
+}