diff --git a/historyserver/cmd/collector/main.go b/historyserver/cmd/collector/main.go index 7648a2c9014..4c5f2306ea5 100644 --- a/historyserver/cmd/collector/main.go +++ b/historyserver/cmd/collector/main.go @@ -160,6 +160,8 @@ func main() { } } + // RAY_COLLECTOR_ADDITIONAL_ENDPOINTS is optional: the head collector always + // polls its built-in endpoints, and anything listed here is polled on top. var additionalEndpoints []string if epStr := os.Getenv("RAY_COLLECTOR_ADDITIONAL_ENDPOINTS"); epStr != "" { for _, ep := range strings.Split(epStr, ",") { @@ -170,16 +172,15 @@ func main() { } } + // Fall back instead of exiting: crash-looping this sidecar would take the head pod + // out of its Service endpoints. endpointPollInterval := 30 * time.Second - if intervalStr := os.Getenv("RAY_COLLECTOR_POLL_INTERVAL"); intervalStr != "" { - parsed, parseErr := time.ParseDuration(intervalStr) - if parseErr != nil { - logrus.Fatalf("Failed to parse RAY_COLLECTOR_POLL_INTERVAL: %v", parseErr) - } - if parsed <= 0 { - logrus.Fatalf("RAY_COLLECTOR_POLL_INTERVAL must be positive, got: %s", intervalStr) + if v := os.Getenv("RAY_COLLECTOR_POLL_INTERVAL"); v != "" { + if parsed, err := time.ParseDuration(v); err == nil && parsed > 0 { + endpointPollInterval = parsed + } else { + logrus.Warnf("Invalid RAY_COLLECTOR_POLL_INTERVAL=%s, using default %s", v, endpointPollInterval) } - endpointPollInterval = parsed } jsonData := make(map[string]interface{}) @@ -229,6 +230,12 @@ func main() { sessionName := path.Base(activeSessionDir) + // The head collector shares a pod with the dashboard; override for non-default ports. + dashboardAddress := "http://localhost:8265" + if v := os.Getenv("RAY_DASHBOARD_ADDRESS"); v != "" { + dashboardAddress = v + } + globalConfig := types.RayCollectorConfig{ RootDir: rayRootDir, SessionDir: activeSessionDir, @@ -238,7 +245,7 @@ func main() { RayClusterNamespace: rayClusterNamespace, PushInterval: pushInterval, LogBatching: logBatching, - DashboardAddress: os.Getenv("RAY_DASHBOARD_ADDRESS"), + DashboardAddress: dashboardAddress, OwnerKind: ownerKind, OwnerName: ownerName, diff --git a/historyserver/config/ray-data.yaml b/historyserver/config/ray-data.yaml new file mode 100644 index 00000000000..eb79f423650 --- /dev/null +++ b/historyserver/config/ray-data.yaml @@ -0,0 +1,112 @@ +apiVersion: ray.io/v1 +kind: RayJob +metadata: + name: rayjob-ray-data +spec: + # Self-contained: brings up its own cluster; also exercises the collector's shutdown path. + shutdownAfterJobFinishes: true + # Long enough for one polling cycle after the job succeeds. + ttlSecondsAfterFinished: 30 + entrypoint: | + python -c " + import ray + ray.init() + + # materialize() is required: an unexecuted Dataset produces no stats. + ds = ray.data.range(100).map_batches(lambda batch: batch).materialize() + print(f'Dataset rows: {ds.count()}') + " + rayClusterSpec: + # Head-only: worker collectors would need the generated head Service FQDN in FQ_RAY_IP. + headGroupSpec: + rayStartParams: + dashboard-host: 0.0.0.0 + serviceType: ClusterIP + template: + spec: + containers: + - env: + - name: RAY_TMP_ROOT + value: &rayTmpRoot /tmp/ray + - name: RAY_enable_ray_event + value: "true" + - name: RAY_enable_core_worker_ray_event_to_aggregator + value: "true" + - name: RAY_DASHBOARD_AGGREGATOR_AGENT_EVENTS_EXPORT_ADDR + value: "http://localhost:8084/v1/events" + - name: RAY_DASHBOARD_AGGREGATOR_AGENT_EXPOSABLE_EVENT_TYPES + value: "ALL" + image: rayproject/ray:2.56.0 + imagePullPolicy: IfNotPresent + name: ray-head + securityContext: + allowPrivilegeEscalation: true + privileged: true + resources: + limits: + cpu: "5" + memory: 10G + requests: + cpu: "50m" + memory: 1G + volumeMounts: + - name: historyserver + mountPath: *rayTmpRoot + - name: collector + image: collector:v0.1.0 + imagePullPolicy: IfNotPresent + env: + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + # KubeRay generates the cluster name; read it back from the pod label. + - name: RAY_CLUSTER_NAME + valueFrom: + fieldRef: + fieldPath: metadata.labels['ray.io/cluster'] + - name: RAY_CLUSTER_NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + # ray.io/originated-from-* labels are on the RayCluster, not the pod; must match metadata.name. + - name: OWNER_KIND + value: "RayJob" + - name: OWNER_NAME + value: "rayjob-ray-data" + # Only used to look up this pod's Ray NodeID; the dashboard is in this pod. + - name: FQ_RAY_IP + value: "localhost" + - name: RAY_TMP_ROOT + value: *rayTmpRoot + # Shorter than the 30s default so a cycle fits inside ttlSecondsAfterFinished. + - name: RAY_COLLECTOR_POLL_INTERVAL + value: "5s" + - name: S3DISABLE_SSL + value: "true" + - name: AWS_ACCESS_KEY_ID + value: minioadmin + - name: AWS_SECRET_ACCESS_KEY + value: minioadmin + - name: AWS_SESSION_TOKEN + value: "" + - name: S3_BUCKET + value: "ray-historyserver" + - name: S3_ENDPOINT + value: "minio-service.minio-dev:9000" + - name: S3_REGION + value: "test" + - name: S3FORCE_PATH_STYLE + value: "true" + command: + - collector + - --role=Head + - --runtime-class-name=s3 + - --ray-root-dir=log + - --events-port=8084 + volumeMounts: + - name: historyserver + mountPath: *rayTmpRoot + volumes: + - name: historyserver + emptyDir: {} diff --git a/historyserver/config/raycluster-azureblob.yaml b/historyserver/config/raycluster-azureblob.yaml index 80150df58db..d292a2e2be0 100644 --- a/historyserver/config/raycluster-azureblob.yaml +++ b/historyserver/config/raycluster-azureblob.yaml @@ -61,6 +61,20 @@ spec: value: $(RAY_CLUSTER_NAME)-head-svc.$(RAY_CLUSTER_NAMESPACE).svc.cluster.local - name: RAY_TMP_ROOT value: *rayTmpRoot + # Optional; defaults to http://localhost:8265 (head only). + # - name: RAY_DASHBOARD_ADDRESS + # value: "http://localhost:9265" + # Optional; defaults to 30s. + # - name: RAY_COLLECTOR_POLL_INTERVAL + # value: "1m" + # Optional extras on top of the built-in endpoints (Serve, placement groups, Ray Data). + # Paths must match the History Server replay request URI, query string included. + # For example, preserve Ray Train V2 run details for a future History Server view. + # Requires the Ray Train workload/driver (not the collector) to set + # RAY_TRAIN_V2_ENABLED=1 and RAY_TRAIN_ENABLE_STATE_TRACKING=1; + # verify this DeveloperAPI URI when changing Ray versions. + # - name: RAY_COLLECTOR_ADDITIONAL_ENDPOINTS + # value: "/api/train/v2/runs/v1" - name: RAY_ROLE value: "Head" - name: STORAGE_BACKEND diff --git a/historyserver/config/raycluster-gcs.yaml b/historyserver/config/raycluster-gcs.yaml index 6cd8dc50bbf..7794c3a0702 100644 --- a/historyserver/config/raycluster-gcs.yaml +++ b/historyserver/config/raycluster-gcs.yaml @@ -66,27 +66,20 @@ spec: value: *rayTmpRoot - name: GCS_BUCKET value: "${GCS_BUCKET}" - # RAY_DASHBOARD_ADDRESS is used by the head collector to fetch endpoints' results - # (e.g., /api/v0/cluster_metadata) from the Ray Dashboard running in the same pod. - # Only the head collector uses this; worker collectors do not need it. - # If your Ray Dashboard uses a non-default port (not 8265), update this value accordingly. - - name: RAY_DASHBOARD_ADDRESS - value: "http://localhost:8265" - # RAY_COLLECTOR_ADDITIONAL_ENDPOINTS is a comma-separated list of Ray Dashboard - # API endpoint paths that the head collector will periodically poll and store. - # Use this for endpoints whose data cannot be obtained via Ray events. - # You can add more endpoints, e.g., "/api/v0/placement_groups,/api/serve/applications/" - # Note: Only static endpoints (no dynamic path parameters like {job_id}) are supported. - - name: RAY_COLLECTOR_ADDITIONAL_ENDPOINTS - value: "/api/v0/placement_groups?detail=1&limit=10000" - # Query params must match the Ray Dashboard frontend request exactly - # (see https://github.com/ray-project/ray/blob/cb9c80fee6a700efe61ea97987248ce82e3fa2e2/python/ray/dashboard/client/src/service/placementGroup.ts): - # detail=1 — include bundles and stats fields required by PlacementGroupTable - # limit=10000 — match the frontend's default limit to avoid truncation - # RAY_COLLECTOR_POLL_INTERVAL controls how often the collector polls the additional - # endpoints above. Accepts Go duration format (e.g., "30s", "1m", "5m"). - - name: RAY_COLLECTOR_POLL_INTERVAL - value: "30s" + # Optional; defaults to http://localhost:8265 (head only). + # - name: RAY_DASHBOARD_ADDRESS + # value: "http://localhost:9265" + # Optional; defaults to 30s. + # - name: RAY_COLLECTOR_POLL_INTERVAL + # value: "1m" + # Optional extras on top of the built-in endpoints (Serve, placement groups, Ray Data). + # Paths must match the History Server replay request URI, query string included. + # For example, preserve Ray Train V2 run details for a future History Server view. + # Requires the Ray Train workload/driver (not the collector) to set + # RAY_TRAIN_V2_ENABLED=1 and RAY_TRAIN_ENABLE_STATE_TRACKING=1; + # verify this DeveloperAPI URI when changing Ray versions. + # - name: RAY_COLLECTOR_ADDITIONAL_ENDPOINTS + # value: "/api/train/v2/runs/v1" - name: RAY_ROLE value: "Head" - name: STORAGE_BACKEND diff --git a/historyserver/config/raycluster.yaml b/historyserver/config/raycluster.yaml index fc7fea4d32c..1d55b0611c0 100644 --- a/historyserver/config/raycluster.yaml +++ b/historyserver/config/raycluster.yaml @@ -65,27 +65,20 @@ spec: value: $(RAY_CLUSTER_NAME)-head-svc.$(RAY_CLUSTER_NAMESPACE).svc.cluster.local - name: RAY_TMP_ROOT value: *rayTmpRoot - # RAY_DASHBOARD_ADDRESS is used by the head collector to fetch endpoints' results - # (e.g., /api/v0/cluster_metadata) from the Ray Dashboard running in the same pod. - # Only the head collector uses this; worker collectors do not need it. - # If your Ray Dashboard uses a non-default port (not 8265), update this value accordingly. - - name: RAY_DASHBOARD_ADDRESS - value: "http://localhost:8265" - # RAY_COLLECTOR_ADDITIONAL_ENDPOINTS is a comma-separated list of Ray Dashboard - # API endpoint paths that the head collector will periodically poll and store. - # Use this for endpoints whose data cannot be obtained via Ray events. - # You can add more endpoints, e.g., "/api/v0/placement_groups,/api/serve/applications/" - # Note: Only static endpoints (no dynamic path parameters like {job_id}) are supported. - - name: RAY_COLLECTOR_ADDITIONAL_ENDPOINTS - # Query params must match the Ray Dashboard frontend request exactly - # (see https://github.com/ray-project/ray/blob/cb9c80fee6a700efe61ea97987248ce82e3fa2e2/python/ray/dashboard/client/src/service/placementGroup.ts): - # detail=1 — include bundles and stats fields required by PlacementGroupTable - # limit=10000 — match the frontend's default limit to avoid truncation - value: "/api/v0/placement_groups?detail=1&limit=10000" - # RAY_COLLECTOR_POLL_INTERVAL controls how often the collector polls the additional - # endpoints above. Accepts Go duration format (e.g., "30s", "1m", "5m"). - - name: RAY_COLLECTOR_POLL_INTERVAL - value: "30s" + # Optional; defaults to http://localhost:8265 (head only). + # - name: RAY_DASHBOARD_ADDRESS + # value: "http://localhost:9265" + # Optional; defaults to 30s. + # - name: RAY_COLLECTOR_POLL_INTERVAL + # value: "1m" + # Optional extras on top of the built-in endpoints (Serve, placement groups, Ray Data). + # Paths must match the History Server replay request URI, query string included. + # For example, preserve Ray Train V2 run details for a future History Server view. + # Requires the Ray Train workload/driver (not the collector) to set + # RAY_TRAIN_V2_ENABLED=1 and RAY_TRAIN_ENABLE_STATE_TRACKING=1; + # verify this DeveloperAPI URI when changing Ray versions. + # - name: RAY_COLLECTOR_ADDITIONAL_ENDPOINTS + # value: "/api/train/v2/runs/v1" - name: S3DISABLE_SSL value: "true" - name: AWS_ACCESS_KEY_ID diff --git a/historyserver/config/rayservice.yaml b/historyserver/config/rayservice.yaml new file mode 100644 index 00000000000..79cb0a5fa7e --- /dev/null +++ b/historyserver/config/rayservice.yaml @@ -0,0 +1,126 @@ +apiVersion: ray.io/v1 +kind: RayService +metadata: + name: rayservice-historyserver +spec: + serveConfigV2: | + proxy_location: EveryNode + applications: + - name: history-e2e + route_prefix: / + import_path: microbenchmarks.no_ops:app_builder + args: + num_forwards: 0 + runtime_env: + working_dir: https://github.com/ray-project/serve_workloads/archive/a9f184f4d9ddb7f9a578502ae106470f87a702ef.zip + deployments: + - name: NoOp + num_replicas: 1 + ray_actor_options: + num_cpus: 0.5 + rayClusterConfig: + # Head-only: worker collectors would need the generated head Service FQDN in FQ_RAY_IP. + headGroupSpec: + rayStartParams: + dashboard-host: 0.0.0.0 + serviceType: ClusterIP + template: + spec: + containers: + - env: + - name: RAY_TMP_ROOT + value: &rayTmpRoot /tmp/ray + - name: RAY_enable_ray_event + value: "true" + - name: RAY_enable_core_worker_ray_event_to_aggregator + value: "true" + - name: RAY_DASHBOARD_AGGREGATOR_AGENT_EVENTS_EXPORT_ADDR + value: "http://localhost:8084/v1/events" + # in ray 2.52.0, we need to set RAY_DASHBOARD_AGGREGATOR_AGENT_EXPOSABLE_EVENT_TYPES + # in ray 2.53.0 (noy yet done). we need to set RAY_DASHBOARD_AGGREGATOR_AGENT_PUBLISHER_HTTP_ENDPOINT_EXPOSABLE_EVENT_TYPES + - name: RAY_DASHBOARD_AGGREGATOR_AGENT_EXPOSABLE_EVENT_TYPES + value: "TASK_DEFINITION_EVENT,TASK_LIFECYCLE_EVENT,ACTOR_TASK_DEFINITION_EVENT, + TASK_PROFILE_EVENT,DRIVER_JOB_DEFINITION_EVENT,DRIVER_JOB_LIFECYCLE_EVENT, + ACTOR_DEFINITION_EVENT,ACTOR_LIFECYCLE_EVENT,NODE_DEFINITION_EVENT,NODE_LIFECYCLE_EVENT" + image: rayproject/ray:2.52.0 + imagePullPolicy: IfNotPresent + name: ray-head + securityContext: + allowPrivilegeEscalation: true + privileged: true + ports: + - containerPort: 6379 + name: gcs-server + - containerPort: 8265 + name: dashboard + - containerPort: 10001 + name: client + - containerPort: 8000 + name: serve + resources: + limits: + cpu: "5" + memory: 10G + requests: + cpu: "50m" + memory: 1G + volumeMounts: + - name: historyserver + mountPath: *rayTmpRoot + - name: collector + image: collector:v0.1.0 + imagePullPolicy: IfNotPresent + env: + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + # KubeRay generates the cluster name; read it back from the pod label. + - name: RAY_CLUSTER_NAME + valueFrom: + fieldRef: + fieldPath: metadata.labels['ray.io/cluster'] + - name: RAY_CLUSTER_NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + # ray.io/originated-from-* labels are on the RayCluster, not the pod; must match metadata.name. + - name: OWNER_KIND + value: "RayService" + - name: OWNER_NAME + value: "rayservice-historyserver" + # Only used to look up this pod's Ray NodeID; the dashboard is in this pod. + - name: FQ_RAY_IP + value: "localhost" + - name: RAY_TMP_ROOT + value: *rayTmpRoot + - name: S3DISABLE_SSL + value: "true" + - name: AWS_ACCESS_KEY_ID + value: minioadmin + - name: AWS_SECRET_ACCESS_KEY + value: minioadmin + - name: AWS_SESSION_TOKEN + value: "" + - name: S3_BUCKET + value: "ray-historyserver" + - name: S3_ENDPOINT + value: "minio-service.minio-dev:9000" + - name: S3_REGION + value: "test" + - name: S3FORCE_PATH_STYLE + value: "true" + - name: RAY_COLLECTOR_POLL_INTERVAL + value: "30s" + command: + - collector + - --role=Head + - --runtime-class-name=s3 + - --ray-root-dir=log + - --events-port=8084 + volumeMounts: + - name: historyserver + mountPath: *rayTmpRoot + volumes: + - name: historyserver + emptyDir: {} diff --git a/historyserver/pkg/collector/logcollector/runtime/logcollector/collector.go b/historyserver/pkg/collector/logcollector/runtime/logcollector/collector.go index 8d3e96274af..1a1183ac6ff 100644 --- a/historyserver/pkg/collector/logcollector/runtime/logcollector/collector.go +++ b/historyserver/pkg/collector/logcollector/runtime/logcollector/collector.go @@ -82,22 +82,35 @@ func (r *RayLogHandler) Run(stop <-chan struct{}) error { // uploads from previous runs are resumed. go r.WatchPrevLogsLoops() go r.PollActiveSessionChanges() + var periodicPollResults <-chan periodicPollResult if r.IsHead { go r.WatchSessionLatestLoops() // Watch session_latest symlink changes go r.FetchAndStoreClusterMetadata() go r.FetchAndStoreTimezone() - go r.PollAdditionalEndpointsPeriodically() + // Driven by stop rather than ShutdownChan, which closes only after the final + // poll below: a tick in between would overwrite that final snapshot. + periodicPollResults = r.startPeriodicEndpointPolling(stop) } <-stop logrus.Info("Received stop signal, processing all logs...") - r.processSessionLatestLogs() - // Perform one final poll of additional endpoints before shutting down. - // This must happen before close(r.ShutdownChan) because pollSingleEndpoint - // uses ShutdownChan to cancel in-flight HTTP requests. + + // Endpoint data dies with the dashboard; log files stay on disk, so poll concurrently. + var wg sync.WaitGroup if r.IsHead { - r.processAdditionalEndpoints() + // Reuse the periodic poller's state only after it exits and transfers ownership. + periodicResult, joined := waitForPeriodicPollResult(periodicPollResults, periodicPollJoinTimeout) + if !joined { + logrus.Warn("Periodic endpoint poller still busy, starting the final poll anyway") + } + wg.Go(func() { + r.processAdditionalEndpoints(periodicResult) + }) } + r.processSessionLatestLogs() + wg.Wait() + + // Signal background goroutines after the synchronous shutdown flushes complete. close(r.ShutdownChan) return nil diff --git a/historyserver/pkg/collector/logcollector/runtime/logcollector/endpoint_fetch_once.go b/historyserver/pkg/collector/logcollector/runtime/logcollector/endpoint_fetch_once.go index 93d9c92c5f9..5737eb2fac3 100644 --- a/historyserver/pkg/collector/logcollector/runtime/logcollector/endpoint_fetch_once.go +++ b/historyserver/pkg/collector/logcollector/runtime/logcollector/endpoint_fetch_once.go @@ -12,6 +12,7 @@ import ( "github.com/sirupsen/logrus" + "github.com/ray-project/kuberay/historyserver/pkg/storage/clusterlogs" "github.com/ray-project/kuberay/historyserver/pkg/utils" ) @@ -125,7 +126,7 @@ func (r *RayLogHandler) fetchAndStoreEndpoint(cfg endpointFetchConfig) { // Successfully fetched — store it under the session path storageKey := utils.EndpointPathToStorageKey(cfg.endpoint) - objectKey := path.Join(r.ClusterDir, sessionName, utils.RAY_SESSIONDIR_FETCHED_ENDPOINTS_NAME, storageKey) + objectKey := path.Join(clusterlogs.FetchedEndpointsDir(r.ClusterDir, sessionName), storageKey) if err := r.Writer.WriteFile(objectKey, bytes.NewReader(body)); err != nil { logrus.Errorf("Failed to store %s at %s: %v", cfg.endpoint, objectKey, err) if !r.sleepOrShutdown(retryInterval) { diff --git a/historyserver/pkg/collector/logcollector/runtime/logcollector/poll.go b/historyserver/pkg/collector/logcollector/runtime/logcollector/poll.go index 18179444624..34a78d616b9 100644 --- a/historyserver/pkg/collector/logcollector/runtime/logcollector/poll.go +++ b/historyserver/pkg/collector/logcollector/runtime/logcollector/poll.go @@ -3,140 +3,384 @@ package logcollector import ( "bytes" "context" + "encoding/json" + "fmt" "io" "net/http" "path" "path/filepath" + "slices" + "strings" "time" "github.com/sirupsen/logrus" + "github.com/ray-project/kuberay/historyserver/pkg/storage/clusterlogs" "github.com/ray-project/kuberay/historyserver/pkg/utils" ) -// PollAdditionalEndpointsPeriodically periodically fetches user-configured additional -// endpoints from the Ray Dashboard and stores their responses in storage. Unlike -// FetchAndStoreClusterMetadata (which runs once), this runs continuously on a timer -// until shutdown. -// -// Each endpoint response is stored at: {ClusterDir}/{sessionName}/fetched_endpoints/{storageKey} -// where storageKey is derived from the endpoint path (e.g., "/api/v0/nodes/summary" -// becomes "restful__api__v0__nodes__summary"). Each poll cycle overwrites the -// previous response. -func (r *RayLogHandler) PollAdditionalEndpointsPeriodically() { - if len(r.AdditionalEndpoints) == 0 { - logrus.Info("No additional endpoints configured, skipping polling") - return +// Each string must match the frontend's request URI, query string included: the +// storage key is derived from it, so a mismatch makes the stored object unreachable. +const ( + serveApplicationsEndpoint = "/api/serve/applications/" + placementGroupsEndpoint = "/api/v0/placement_groups?detail=1&limit=10000" + // Only used to discover job IDs; its response is not stored. + jobsEndpoint = "/api/jobs/" + // Per job, where job_id is the hex core job ID (e.g. "01000000"), not the submission ID. + dataDatasetsEndpointPrefix = "/api/data/datasets/" +) + +// dataDatasetsEndpointPrefix needs a job ID, so pollDataDatasets handles it. +var staticPolledEndpoints = []string{ + serveApplicationsEndpoint, + placementGroupsEndpoint, +} + +// polledEndpoints merges the built-in and configured endpoints, deduplicated. +func (r *RayLogHandler) polledEndpoints() []string { + all := slices.Concat(staticPolledEndpoints, r.AdditionalEndpoints) + endpoints := make([]string, 0, len(all)) + seen := make(map[string]struct{}, len(all)) + for _, endpoint := range all { + if _, ok := seen[endpoint]; ok { + continue + } + seen[endpoint] = struct{}{} + endpoints = append(endpoints, endpoint) + } + return endpoints +} + +// Statuses a job never leaves, so its datasets only need to be stored once. +// Ref: https://github.com/ray-project/ray/blob/ray-2.54.1/python/ray/dashboard/modules/job/common.py#L38-L50 +var terminalJobStatuses = map[string]bool{ + "SUCCEEDED": true, + "FAILED": true, + "STOPPED": true, +} + +// pollOutcome distinguishes "nothing worth storing" from "could not store". +type pollOutcome int + +const ( + pollFailed pollOutcome = iota + pollStored + pollSkippedEmpty +) + +// Stats can appear slightly after a job reports terminal, so one empty response is not final. +const terminalEmptyPollsBeforeGivingUp = 2 + +// The final poll shares the pod's termination grace period, so it gives up rather than overrun it. +const shutdownPollBudget = 10 * time.Second + +// Caps the shutdown join on the periodic poller: a storage write in flight is not +// cancelable, and waiting it out could eat the grace period the final poll needs. +const periodicPollJoinTimeout = 5 * time.Second + +// datasetPollState is owned by one goroutine at a time; no lock is needed. +type datasetPollState struct { + // terminalStored contains terminal jobs whose dataset response was successfully written. + terminalStored map[string]struct{} + // emptyRuns counts consecutive empty responses observed after a job becomes terminal. + emptyRuns map[string]int +} + +func newDatasetPollState() *datasetPollState { + return &datasetPollState{ + terminalStored: make(map[string]struct{}), + emptyRuns: make(map[string]int), + } +} + +// periodicPollResult transfers ownership of state after the periodic poller exits. +type periodicPollResult struct { + sessionName string + state *datasetPollState +} + +// PollAdditionalEndpointsPeriodically fetches the built-in endpoints, plus anything from +// RAY_COLLECTOR_ADDITIONAL_ENDPOINTS, on a timer; each cycle overwrites the previous one. +// It stops when stop closes and cancels any blocked resolve or in-flight request at that point. +func (r *RayLogHandler) PollAdditionalEndpointsPeriodically(stop <-chan struct{}) { + r.pollAdditionalEndpointsPeriodically(stop) +} + +func (r *RayLogHandler) startPeriodicEndpointPolling(stop <-chan struct{}) <-chan periodicPollResult { + // The core sends exactly once. A one-element buffer lets that send finish even if + // the shutdown join times out and no receiver remains. + results := make(chan periodicPollResult, 1) + go func() { + results <- r.pollAdditionalEndpointsPeriodically(stop) + }() + return results +} + +func waitForPeriodicPollResult(results <-chan periodicPollResult, timeout time.Duration) (periodicPollResult, bool) { + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case result := <-results: + return result, true + case <-timer.C: + return periodicPollResult{}, false } +} + +func (r *RayLogHandler) pollAdditionalEndpointsPeriodically(stop <-chan struct{}) periodicPollResult { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + select { + case <-stop: + cancel() + case <-ctx.Done(): + } + }() - // Note: session name staleness is not a concern here because the log collector runs as a - // sidecar in the Ray head pod. If the Ray head process restarts (creating a new session), - // the entire pod — including this sidecar container — restarts, so resolveSessionName() - // always runs in a fresh container lifecycle with the current session. - sessionName, err := r.resolveSessionName() + sessionName, err := waitForSessionName(ctx) if err != nil { - logrus.Errorf("Failed to resolve session name for additional endpoints polling: %v", err) - return + logrus.Errorf("Failed to resolve session name for endpoint polling: %v", err) + return periodicPollResult{} } - logrus.Infof("Starting additional endpoints polling (interval=%v, endpoints=%v)", r.EndpointPollInterval, r.AdditionalEndpoints) + logrus.Infof("Starting endpoint polling (interval=%v, endpoints=%v)", r.EndpointPollInterval, r.polledEndpoints()) - // Perform an initial poll immediately on startup. - r.pollAllEndpoints(sessionName) + state := newDatasetPollState() + r.pollAllEndpoints(ctx, sessionName, state, false) ticker := time.NewTicker(r.EndpointPollInterval) defer ticker.Stop() for { select { - case <-r.ShutdownChan: - logrus.Info("Shutdown signaled, stopping additional endpoints polling") - return + case <-stop: + logrus.Info("Shutdown signaled, stopping endpoint polling") + return periodicPollResult{sessionName: sessionName, state: state} case <-ticker.C: - r.pollAllEndpoints(sessionName) + sessionName, state = r.pollCycle(ctx, sessionName, state) } } } -// processAdditionalEndpoints performs a final poll of all additional endpoints -// before shutdown. This mirrors processSessionLatestLogs as a shutdown cleanup step. -// -// Unlike PollAdditionalEndpointsPeriodically, this does NOT retry session name -// resolution because it runs during shutdown — if session_latest is gone (e.g., -// Ray head already exited), retrying would hang forever since ShutdownChan has -// not been closed yet. -func (r *RayLogHandler) processAdditionalEndpoints() { - if len(r.AdditionalEndpoints) == 0 { - return +// waitForSessionName resolves session_latest, retrying until ctx is canceled: +// at startup the symlink appears only once Ray has bootstrapped. +func waitForSessionName(ctx context.Context) (string, error) { + for { + name, err := currentSessionName() + if err == nil { + return name, nil + } + logrus.Warnf("session_latest symlink not ready: %v, retrying in %v", err, defaultRetryInterval) + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-time.After(defaultRetryInterval): + } } - logrus.Info("Processing additional endpoints before shutdown") +} - // Resolve session name directly without retry — this is a shutdown path. - sessionLatestDir := utils.GetRaySessionLatestPath() - sessionRealDir, err := filepath.EvalSymlinks(sessionLatestDir) - if err != nil { - logrus.Errorf("Failed to resolve session name for final additional endpoints poll: %v", err) - return +// pollCycle runs one polling pass, re-resolving the session first: the Ray head container +// can restart alone, starting a new session (and new job IDs) while this sidecar lives on. +func (r *RayLogHandler) pollCycle(ctx context.Context, sessionName string, state *datasetPollState) (string, *datasetPollState) { + switch current, err := currentSessionName(); { + case err != nil: + logrus.Warnf("Failed to re-resolve session name, polling into %s: %v", sessionName, err) + case current != sessionName: + logrus.Infof("Session changed from %s to %s, resetting dataset polling state", sessionName, current) + sessionName, state = current, newDatasetPollState() } - sessionName := filepath.Base(sessionRealDir) - r.pollAllEndpoints(sessionName) - logrus.Info("Finished processing additional endpoints") + r.pollAllEndpoints(ctx, sessionName, state, false) + return sessionName, state } -// pollAllEndpoints fetches all configured additional endpoints and stores their responses. -func (r *RayLogHandler) pollAllEndpoints(sessionName string) { - for _, endpoint := range r.AdditionalEndpoints { - r.pollSingleEndpoint(endpoint, sessionName) +// currentSessionName resolves session_latest without retrying, for callers that must not block. +func currentSessionName() (string, error) { + sessionRealDir, err := filepath.EvalSymlinks(utils.GetRaySessionLatestPath()) + if err != nil { + return "", err } + return filepath.Base(sessionRealDir), nil } -// pollSingleEndpoint fetches a single endpoint from the Ray Dashboard and writes -// the response to storage. -func (r *RayLogHandler) pollSingleEndpoint(endpoint, sessionName string) { - url := r.DashboardAddress + endpoint +// processAdditionalEndpoints performs one final poll before shutdown. +func (r *RayLogHandler) processAdditionalEndpoints(previous periodicPollResult) { + logrus.Info("Processing polled endpoints before shutdown") - ctx, cancel := context.WithTimeout(context.Background(), defaultRequestTimeout) - go func() { - select { - case <-r.ShutdownChan: - cancel() - case <-ctx.Done(): + sessionName, err := currentSessionName() + if err != nil { + logrus.Errorf("Failed to resolve session name for final endpoint poll: %v", err) + return + } + + state := newDatasetPollState() + if previous.state != nil && previous.sessionName == sessionName { + state = previous.state + } + + // The budget cancels HTTP fetches. StorageWriter.WriteFile does not accept a + // context, so a write already in progress may finish after the deadline. + ctx, cancel := context.WithTimeout(context.Background(), shutdownPollBudget) + defer cancel() + + r.pollAllEndpoints(ctx, sessionName, state, true) + logrus.Info("Finished processing polled endpoints") +} + +// pollAllEndpoints fetches the polled endpoints plus per-job dataset endpoints and stores +// their responses, stopping early once ctx is canceled. +// finalPoll marks the shutdown pass, which distrusts empty Serve responses (see isEmptyPayload). +func (r *RayLogHandler) pollAllEndpoints(ctx context.Context, sessionName string, state *datasetPollState, finalPoll bool) { + for _, endpoint := range r.polledEndpoints() { + if ctx.Err() != nil { + logrus.Warnf("Stopped polling before %s: %v", endpoint, ctx.Err()) + return } - }() + r.pollSingleEndpoint(ctx, endpoint, sessionName, finalPoll) + } + r.pollDataDatasets(ctx, sessionName, state, finalPoll) +} - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) +// pollDataDatasets stores one datasets object per job discovered via jobsEndpoint. +// Settled terminal jobs are skipped so the per-cycle cost does not grow with job count. +func (r *RayLogHandler) pollDataDatasets(ctx context.Context, sessionName string, state *datasetPollState, finalPoll bool) { + body, err := r.fetchEndpoint(ctx, jobsEndpoint) if err != nil { - cancel() - logrus.Errorf("Failed to create request for additional endpoint %s: %v", endpoint, err) + logrus.Warnf("Failed to fetch %s for dataset polling: %v", jobsEndpoint, err) return } - resp, err := r.HttpClient.Do(req) - if err != nil { - cancel() - logrus.Warnf("Failed to fetch additional endpoint %s: %v", endpoint, err) + var jobs []struct { + JobID string `json:"job_id"` + Status string `json:"status"` + } + if err := json.Unmarshal(body, &jobs); err != nil { + logrus.Warnf("Failed to parse %s response: %v", jobsEndpoint, err) return } - body, err := io.ReadAll(resp.Body) - resp.Body.Close() - cancel() + for i, job := range jobs { + if ctx.Err() != nil { + logrus.Warnf("Stopped dataset polling after %d/%d jobs: %v", i, len(jobs), ctx.Err()) + return + } + // job_id is empty for submission jobs whose driver has not started yet. + if job.JobID == "" { + continue + } + if _, ok := state.terminalStored[job.JobID]; ok { + continue + } + if !finalPoll && terminalJobStatuses[job.Status] && + state.emptyRuns[job.JobID] >= terminalEmptyPollsBeforeGivingUp { + continue + } + + outcome := r.pollSingleEndpoint(ctx, dataDatasetsEndpointPrefix+job.JobID, sessionName, finalPoll) + if !terminalJobStatuses[job.Status] { + continue + } + switch outcome { + case pollStored: + state.terminalStored[job.JobID] = struct{}{} + delete(state.emptyRuns, job.JobID) + case pollSkippedEmpty: + state.emptyRuns[job.JobID]++ + case pollFailed: + // Retried next cycle. + } + } +} + +func (r *RayLogHandler) pollSingleEndpoint(ctx context.Context, endpoint, sessionName string, finalPoll bool) pollOutcome { + body, err := r.fetchEndpoint(ctx, endpoint) if err != nil { - logrus.Warnf("Failed to read response body for additional endpoint %s: %v", endpoint, err) - return + logrus.Warnf("Failed to poll endpoint %s: %v", endpoint, err) + return pollFailed } - if resp.StatusCode != http.StatusOK { - logrus.Warnf("Additional endpoint %s returned status %d", endpoint, resp.StatusCode) - return + // Do not start a storage write after the parent context is canceled. + if ctx.Err() != nil { + return pollFailed + } + + if isEmptyPayload(endpoint, body, finalPoll) { + logrus.Debugf("Skipping %s: nothing to store", endpoint) + return pollSkippedEmpty } storageKey := utils.EndpointPathToStorageKey(endpoint) - objectKey := path.Join(r.ClusterDir, sessionName, utils.RAY_SESSIONDIR_FETCHED_ENDPOINTS_NAME, storageKey) + objectKey := path.Join(clusterlogs.FetchedEndpointsDir(r.ClusterDir, sessionName), storageKey) if err := r.Writer.WriteFile(objectKey, bytes.NewReader(body)); err != nil { - logrus.Errorf("Failed to store additional endpoint %s at %s: %v", endpoint, objectKey, err) - return + logrus.Errorf("Failed to store endpoint %s at %s: %v", endpoint, objectKey, err) + return pollFailed } - logrus.Infof("Successfully stored additional endpoint %s at %s (%d bytes)", endpoint, objectKey, len(body)) + logrus.Infof("Successfully stored endpoint %s at %s (%d bytes)", endpoint, objectKey, len(body)) + return pollStored +} + +// isEmptyPayload reports whether a response carries nothing worth storing. +// Datasets: empty can mean stats-actor eviction, so it never overwrites a snapshot. +// Serve: empty from a healthy cluster is the live truth and is stored, but on the final +// shutdown poll it usually means the Serve controller died before the dashboard. +func isEmptyPayload(endpoint string, body []byte, finalPoll bool) bool { + switch { + case strings.HasPrefix(endpoint, dataDatasetsEndpointPrefix): + return !hasDatasets(body) + case finalPoll && endpoint == serveApplicationsEndpoint: + return !hasServeApplications(body) + default: + return false + } +} + +// Unparsable bodies count as non-empty so unexpected shapes are stored, not dropped. +func hasServeApplications(body []byte) bool { + var resp struct { + Applications map[string]json.RawMessage `json:"applications"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return true + } + return len(resp.Applications) > 0 +} + +// Unparsable bodies count as non-empty so unexpected shapes are stored, not dropped. +func hasDatasets(body []byte) bool { + var resp struct { + Datasets []json.RawMessage `json:"datasets"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return true + } + return len(resp.Datasets) > 0 +} + +// fetchEndpoint GETs one dashboard endpoint with the parent and per-request deadlines. +func (r *RayLogHandler) fetchEndpoint(parent context.Context, endpoint string) ([]byte, error) { + url := r.DashboardAddress + endpoint + + ctx, cancel := context.WithTimeout(parent, defaultRequestTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + resp, err := r.HttpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("unexpected status %d", resp.StatusCode) + } + return body, nil } diff --git a/historyserver/pkg/collector/logcollector/runtime/logcollector/poll_test.go b/historyserver/pkg/collector/logcollector/runtime/logcollector/poll_test.go new file mode 100644 index 00000000000..92fb8ec3522 --- /dev/null +++ b/historyserver/pkg/collector/logcollector/runtime/logcollector/poll_test.go @@ -0,0 +1,832 @@ +package logcollector + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "testing" + "time" + + . "github.com/onsi/gomega" + + "github.com/ray-project/kuberay/historyserver/pkg/utils" +) + +const trainRunsEndpoint = "/api/train/v2/runs/v1" + +// fakeDashboard stands in for the Ray Dashboard, recording every requested path. +type fakeDashboard struct { + mu sync.Mutex + requests []string + // jobs is the /api/jobs/ response body. + jobs string + // datasets maps a job ID to its /api/data/datasets/{job_id} response body. + datasets map[string]string +} + +func (f *fakeDashboard) start(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + f.mu.Lock() + f.requests = append(f.requests, r.URL.RequestURI()) + f.mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + switch { + case r.URL.Path == jobsEndpoint: + _, _ = w.Write([]byte(f.jobs)) + case r.URL.Path == serveApplicationsEndpoint: + // Non-empty by default: an empty Serve response is deliberately not stored. + _, _ = w.Write([]byte(`{"applications": {"app": {"status": "RUNNING"}}}`)) + case strings.HasPrefix(r.URL.Path, dataDatasetsEndpointPrefix): + jobID := strings.TrimPrefix(r.URL.Path, dataDatasetsEndpointPrefix) + body, ok := f.datasets[jobID] + if !ok { + body = `{"datasets": []}` + } + _, _ = w.Write([]byte(body)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func (f *fakeDashboard) requestsFor(prefix string) []string { + f.mu.Lock() + defer f.mu.Unlock() + var out []string + for _, req := range f.requests { + if strings.HasPrefix(req, prefix) { + out = append(out, req) + } + } + return out +} + +func newPollTestHandler(t *testing.T, dashboardAddr string) (*RayLogHandler, *MockStorageWriter) { + t.Helper() + writer := NewMockStorageWriter() + return &RayLogHandler{ + Writer: writer, + HttpClient: &http.Client{}, + ShutdownChan: make(chan struct{}), + ClusterDir: "cluster-dir", + DashboardAddress: dashboardAddr, + IsHead: true, + }, writer +} + +func writtenKeys(writer *MockStorageWriter) []string { + writer.mu.Lock() + defer writer.mu.Unlock() + keys := make([]string, 0, len(writer.writtenFiles)) + for k := range writer.writtenFiles { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +type blockingStorageWriter struct { + entered chan struct{} + release chan struct{} + once sync.Once +} + +func (w *blockingStorageWriter) CreateDirectory(string) error { + return nil +} + +func (w *blockingStorageWriter) WriteFile(string, io.ReadSeeker) error { + w.once.Do(func() { close(w.entered) }) + <-w.release + return nil +} + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +type cancelOnEOFBody struct { + io.Reader + cancel context.CancelFunc +} + +func (b *cancelOnEOFBody) Read(p []byte) (int, error) { + n, err := b.Reader.Read(p) + if err == io.EOF { + b.cancel() + } + return n, err +} + +func (b *cancelOnEOFBody) Close() error { + return nil +} + +// TestPollDataDatasetsFansOutPerJob verifies per-job fan-out from /api/jobs/, skipping blank IDs. +func TestPollDataDatasetsFansOutPerJob(t *testing.T) { + g := NewWithT(t) + + dash := &fakeDashboard{ + jobs: `[ + {"job_id": "01000000", "status": "RUNNING"}, + {"job_id": null, "status": "PENDING"}, + {"job_id": "02000000", "status": "RUNNING"} + ]`, + datasets: map[string]string{ + "01000000": `{"datasets": [{"dataset": "ds_a", "job_id": "01000000"}]}`, + "02000000": `{"datasets": [{"dataset": "ds_b", "job_id": "02000000"}]}`, + }, + } + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + + handler.pollDataDatasets(context.Background(), "session_1", newDatasetPollState(), false) + + // The blank job_id is skipped. + g.Expect(dash.requestsFor(dataDatasetsEndpointPrefix)).To(ConsistOf( + "/api/data/datasets/01000000", + "/api/data/datasets/02000000", + )) + g.Expect(writtenKeys(writer)).To(Equal([]string{ + "cluster-dir/session_1/fetched_endpoints/restful__api__data__datasets__01000000", + "cluster-dir/session_1/fetched_endpoints/restful__api__data__datasets__02000000", + })) +} + +// TestPollDataDatasetsSkipsEmptyResponse verifies an empty datasets response is never written, +// so a stats-actor eviction cannot replace datasets captured earlier. +func TestPollDataDatasetsSkipsEmptyResponse(t *testing.T) { + g := NewWithT(t) + + dash := &fakeDashboard{ + jobs: `[{"job_id": "01000000", "status": "RUNNING"}]`, + // No entry, so the fake dashboard replies {"datasets": []}. + datasets: map[string]string{}, + } + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + + handler.pollDataDatasets(context.Background(), "session_1", newDatasetPollState(), false) + + g.Expect(dash.requestsFor(dataDatasetsEndpointPrefix)).To(HaveLen(1)) + g.Expect(writtenKeys(writer)).To(BeEmpty()) +} + +// TestPollDataDatasetsDoesNotCountRunningEmpties verifies the retry cap begins only after +// a job becomes terminal, so early empty responses cannot suppress its first terminal fetch. +func TestPollDataDatasetsDoesNotCountRunningEmpties(t *testing.T) { + g := NewWithT(t) + + var mu sync.Mutex + status := "RUNNING" + attempts := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + defer mu.Unlock() + if r.URL.Path == jobsEndpoint { + _, _ = w.Write([]byte(fmt.Sprintf(`[{"job_id": "01000000", "status": %q}]`, status))) + return + } + attempts++ + if status == "RUNNING" { + _, _ = w.Write([]byte(`{"datasets": []}`)) + return + } + _, _ = w.Write([]byte(`{"datasets": [{"dataset": "complete"}]}`)) + })) + t.Cleanup(srv.Close) + + handler, writer := newPollTestHandler(t, srv.URL) + state := newDatasetPollState() + handler.pollDataDatasets(context.Background(), "session_1", state, false) + handler.pollDataDatasets(context.Background(), "session_1", state, false) + g.Expect(state.emptyRuns).NotTo(HaveKey("01000000")) + + mu.Lock() + status = "SUCCEEDED" + mu.Unlock() + handler.pollDataDatasets(context.Background(), "session_1", state, false) + + mu.Lock() + defer mu.Unlock() + g.Expect(attempts).To(Equal(3)) + g.Expect(state.terminalStored).To(HaveKey("01000000")) + g.Expect(writtenKeys(writer)).To(HaveLen(1)) +} + +// TestPollDataDatasetsStopsPollingTerminalJobs verifies a terminal job is fetched once +// while a running job keeps being refreshed. +func TestPollDataDatasetsStopsPollingTerminalJobs(t *testing.T) { + g := NewWithT(t) + + dash := &fakeDashboard{ + jobs: `[ + {"job_id": "01000000", "status": "SUCCEEDED"}, + {"job_id": "02000000", "status": "RUNNING"} + ]`, + datasets: map[string]string{ + "01000000": `{"datasets": [{"dataset": "ds_done"}]}`, + "02000000": `{"datasets": [{"dataset": "ds_live"}]}`, + }, + } + srv := dash.start(t) + handler, _ := newPollTestHandler(t, srv.URL) + + state := newDatasetPollState() + handler.pollDataDatasets(context.Background(), "session_1", state, false) + handler.pollDataDatasets(context.Background(), "session_1", state, false) + handler.pollDataDatasets(context.Background(), "session_1", state, false) + + g.Expect(state.terminalStored).To(HaveKey("01000000")) + g.Expect(state.terminalStored).NotTo(HaveKey("02000000")) + // The terminal job is fetched only on the first cycle; the running one every cycle. + g.Expect(dash.requestsFor("/api/data/datasets/01000000")).To(HaveLen(1)) + g.Expect(dash.requestsFor("/api/data/datasets/02000000")).To(HaveLen(3)) +} + +// TestPollDataDatasetsRetriesFailedTerminalJob verifies a failed fetch is retried, so a +// transient dashboard error does not lose a terminal job's datasets. +func TestPollDataDatasetsRetriesFailedTerminalJob(t *testing.T) { + g := NewWithT(t) + + var mu sync.Mutex + attempts := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsEndpoint { + _, _ = w.Write([]byte(`[{"job_id": "01000000", "status": "SUCCEEDED"}]`)) + return + } + mu.Lock() + attempts++ + first := attempts == 1 + mu.Unlock() + if first { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + _, _ = w.Write([]byte(`{"datasets": [{"dataset": "ds_done"}]}`)) + })) + t.Cleanup(srv.Close) + + handler, writer := newPollTestHandler(t, srv.URL) + + state := newDatasetPollState() + handler.pollDataDatasets(context.Background(), "session_1", state, false) + g.Expect(state.terminalStored).To(BeEmpty()) + g.Expect(writtenKeys(writer)).To(BeEmpty()) + + handler.pollDataDatasets(context.Background(), "session_1", state, false) + g.Expect(state.terminalStored).To(HaveKey("01000000")) + g.Expect(writtenKeys(writer)).To(Equal([]string{ + "cluster-dir/session_1/fetched_endpoints/restful__api__data__datasets__01000000", + })) +} + +// TestPollDataDatasetsRetriesTerminalJobWithLateStats verifies an empty first response is +// polled again, because Ray Data registers stats slightly after the job reports SUCCEEDED. +func TestPollDataDatasetsRetriesTerminalJobWithLateStats(t *testing.T) { + g := NewWithT(t) + + var mu sync.Mutex + attempts := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsEndpoint { + _, _ = w.Write([]byte(`[{"job_id": "01000000", "status": "SUCCEEDED"}]`)) + return + } + mu.Lock() + attempts++ + first := attempts == 1 + mu.Unlock() + if first { + _, _ = w.Write([]byte(`{"datasets": []}`)) + return + } + _, _ = w.Write([]byte(`{"datasets": [{"dataset": "ds_late"}]}`)) + })) + t.Cleanup(srv.Close) + + handler, writer := newPollTestHandler(t, srv.URL) + + state := newDatasetPollState() + handler.pollDataDatasets(context.Background(), "session_1", state, false) + g.Expect(state.terminalStored).To(BeEmpty()) + g.Expect(writtenKeys(writer)).To(BeEmpty()) + + handler.pollDataDatasets(context.Background(), "session_1", state, false) + g.Expect(state.terminalStored).To(HaveKey("01000000")) + g.Expect(writtenKeys(writer)).To(Equal([]string{ + "cluster-dir/session_1/fetched_endpoints/restful__api__data__datasets__01000000", + })) +} + +// TestPollDataDatasetsGivesUpOnRepeatedlyEmptyTerminalJob verifies the retry is bounded. +func TestPollDataDatasetsGivesUpOnRepeatedlyEmptyTerminalJob(t *testing.T) { + g := NewWithT(t) + + dash := &fakeDashboard{ + jobs: `[{"job_id": "01000000", "status": "SUCCEEDED"}]`, + // No entry, so the fake dashboard always replies {"datasets": []}. + datasets: map[string]string{}, + } + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + + state := newDatasetPollState() + for i := 0; i < 4; i++ { + handler.pollDataDatasets(context.Background(), "session_1", state, false) + } + + g.Expect(state.terminalStored).NotTo(HaveKey("01000000")) + g.Expect(state.emptyRuns).To(HaveKeyWithValue("01000000", terminalEmptyPollsBeforeGivingUp)) + g.Expect(dash.requestsFor("/api/data/datasets/01000000")).To(HaveLen(terminalEmptyPollsBeforeGivingUp)) + g.Expect(writtenKeys(writer)).To(BeEmpty()) +} + +// TestFinalDatasetPollReusesPeriodicState verifies the final poll skips terminal jobs +// already stored, while retrying terminal jobs that were empty or failed periodically. +func TestFinalDatasetPollReusesPeriodicState(t *testing.T) { + g := NewWithT(t) + + var mu sync.Mutex + attempts := map[string]int{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsEndpoint { + _, _ = w.Write([]byte(`[ + {"job_id": "01000000", "status": "SUCCEEDED"}, + {"job_id": "02000000", "status": "SUCCEEDED"}, + {"job_id": "03000000", "status": "SUCCEEDED"} + ]`)) + return + } + jobID := strings.TrimPrefix(r.URL.Path, dataDatasetsEndpointPrefix) + mu.Lock() + attempts[jobID]++ + attempt := attempts[jobID] + mu.Unlock() + + switch jobID { + case "01000000": + _, _ = w.Write([]byte(`{"datasets": [{"dataset": "stored"}]}`)) + case "02000000": + if attempt <= terminalEmptyPollsBeforeGivingUp { + _, _ = w.Write([]byte(`{"datasets": []}`)) + return + } + _, _ = w.Write([]byte(`{"datasets": [{"dataset": "late"}]}`)) + case "03000000": + if attempt <= 2 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + _, _ = w.Write([]byte(`{"datasets": [{"dataset": "retried"}]}`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + t.Cleanup(srv.Close) + + handler, writer := newPollTestHandler(t, srv.URL) + state := newDatasetPollState() + handler.pollDataDatasets(context.Background(), "session_1", state, false) + handler.pollDataDatasets(context.Background(), "session_1", state, false) + + g.Expect(state.terminalStored).To(HaveKey("01000000")) + g.Expect(state.terminalStored).NotTo(HaveKey("02000000")) + g.Expect(state.terminalStored).NotTo(HaveKey("03000000")) + g.Expect(state.emptyRuns).To(HaveKeyWithValue("02000000", terminalEmptyPollsBeforeGivingUp)) + + handler.pollDataDatasets(context.Background(), "session_1", state, true) + + mu.Lock() + defer mu.Unlock() + g.Expect(attempts).To(Equal(map[string]int{ + "01000000": 1, + "02000000": 3, + "03000000": 3, + })) + g.Expect(state.terminalStored).To(HaveKey("01000000")) + g.Expect(state.terminalStored).To(HaveKey("02000000")) + g.Expect(state.terminalStored).To(HaveKey("03000000")) + g.Expect(state.emptyRuns).NotTo(HaveKey("02000000")) + g.Expect(writtenKeys(writer)).To(ConsistOf( + "cluster-dir/session_1/fetched_endpoints/restful__api__data__datasets__01000000", + "cluster-dir/session_1/fetched_endpoints/restful__api__data__datasets__02000000", + "cluster-dir/session_1/fetched_endpoints/restful__api__data__datasets__03000000", + )) +} + +// TestPollAllEndpointsStoresStaticEndpoints verifies the built-in endpoints are stored under +// the exact URIs the frontend requests. +func TestPollAllEndpointsStoresStaticEndpoints(t *testing.T) { + g := NewWithT(t) + + dash := &fakeDashboard{jobs: `[]`} + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + + handler.pollAllEndpoints(context.Background(), "session_1", newDatasetPollState(), false) + + g.Expect(dash.requestsFor("/api/serve/applications/")).To(HaveLen(1)) + g.Expect(dash.requestsFor("/api/v0/placement_groups?detail=1&limit=10000")).To(HaveLen(1)) + g.Expect(writtenKeys(writer)).To(Equal([]string{ + "cluster-dir/session_1/fetched_endpoints/restful__api__serve__applications", + "cluster-dir/session_1/fetched_endpoints/restful__api__v0__placement_groups?detail=1&limit=10000", + })) +} + +// TestPollCycleFollowsSessionChange verifies a Ray head restart moves polling to the new +// session and resets state, since job IDs restart with the session. +func TestPollCycleFollowsSessionChange(t *testing.T) { + g := NewWithT(t) + + tmpRoot := t.TempDir() + t.Setenv("RAY_TMP_ROOT", tmpRoot) + symlink := filepath.Join(tmpRoot, "session_latest") + pointAt := func(session string) { + g.Expect(os.MkdirAll(filepath.Join(tmpRoot, session), 0o755)).To(Succeed()) + _ = os.Remove(symlink) + g.Expect(os.Symlink(filepath.Join(tmpRoot, session), symlink)).To(Succeed()) + } + + dash := &fakeDashboard{ + jobs: `[{"job_id": "01000000", "status": "SUCCEEDED"}]`, + datasets: map[string]string{"01000000": `{"datasets": [{"dataset": "ds"}]}`}, + } + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + + pointAt("session_old") + session, state := handler.pollCycle(context.Background(), "session_old", newDatasetPollState()) + g.Expect(session).To(Equal("session_old")) + g.Expect(state.terminalStored).To(HaveKey("01000000")) + + pointAt("session_new") + session, state = handler.pollCycle(context.Background(), session, state) + g.Expect(session).To(Equal("session_new")) + g.Expect(state.terminalStored).To(HaveKey("01000000"), "the new session's job must be captured, not skipped") + + // The same job ID is stored once per session, not once overall. + g.Expect(dash.requestsFor("/api/data/datasets/01000000")).To(HaveLen(2)) + g.Expect(writtenKeys(writer)).To(ContainElements( + "cluster-dir/session_old/fetched_endpoints/restful__api__data__datasets__01000000", + "cluster-dir/session_new/fetched_endpoints/restful__api__data__datasets__01000000", + )) +} + +// TestFinalPollReusesStateOnlyForCurrentSession verifies a restarted Ray head cannot +// inherit terminal job IDs from the previous session. +func TestFinalPollReusesStateOnlyForCurrentSession(t *testing.T) { + for _, tc := range []struct { + name string + previousSession string + wantRequests int + }{ + {name: "same session", previousSession: "session_new", wantRequests: 0}, + {name: "changed session", previousSession: "session_old", wantRequests: 1}, + } { + t.Run(tc.name, func(t *testing.T) { + g := NewWithT(t) + tmpRoot := t.TempDir() + t.Setenv("RAY_TMP_ROOT", tmpRoot) + g.Expect(os.MkdirAll(filepath.Join(tmpRoot, "session_new"), 0o755)).To(Succeed()) + g.Expect(os.Symlink(filepath.Join(tmpRoot, "session_new"), filepath.Join(tmpRoot, "session_latest"))).To(Succeed()) + + dash := &fakeDashboard{ + jobs: `[{"job_id": "01000000", "status": "SUCCEEDED"}]`, + datasets: map[string]string{"01000000": `{"datasets": [{"dataset": "new"}]}`}, + } + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + state := newDatasetPollState() + state.terminalStored["01000000"] = struct{}{} + + handler.processAdditionalEndpoints(periodicPollResult{ + sessionName: tc.previousSession, + state: state, + }) + + g.Expect(dash.requestsFor("/api/data/datasets/01000000")).To(HaveLen(tc.wantRequests)) + dataKey := "cluster-dir/session_new/fetched_endpoints/restful__api__data__datasets__01000000" + if tc.wantRequests == 0 { + g.Expect(writtenKeys(writer)).NotTo(ContainElement(dataKey)) + } else { + g.Expect(writtenKeys(writer)).To(ContainElement(dataKey)) + } + }) + } +} + +// TestPeriodicPollingStopsOnShutdownSignal verifies the loop exits on the shutdown signal, +// not ShutdownChan, so no tick can overwrite the final shutdown snapshot. +func TestPeriodicPollingStopsOnShutdownSignal(t *testing.T) { + g := NewWithT(t) + + tmpRoot := t.TempDir() + t.Setenv("RAY_TMP_ROOT", tmpRoot) + g.Expect(os.MkdirAll(filepath.Join(tmpRoot, "session_1"), 0o755)).To(Succeed()) + g.Expect(os.Symlink(filepath.Join(tmpRoot, "session_1"), filepath.Join(tmpRoot, "session_latest"))).To(Succeed()) + + dash := &fakeDashboard{jobs: `[]`} + srv := dash.start(t) + handler, _ := newPollTestHandler(t, srv.URL) + handler.EndpointPollInterval = time.Hour + + stop := make(chan struct{}) + done := make(chan struct{}) + go func() { + handler.PollAdditionalEndpointsPeriodically(stop) + close(done) + }() + + // ShutdownChan stays open; the stop signal alone must end the loop. + close(stop) + g.Eventually(done).Should(BeClosed()) +} + +// TestPeriodicPollingCancelsInFlightRequestOnShutdown verifies the stop signal aborts a +// cycle mid-request: Run joins this goroutine before the final poll, so a blocked request +// must neither stall shutdown nor store anything after being canceled. +func TestPeriodicPollingCancelsInFlightRequestOnShutdown(t *testing.T) { + g := NewWithT(t) + + tmpRoot := t.TempDir() + t.Setenv("RAY_TMP_ROOT", tmpRoot) + g.Expect(os.MkdirAll(filepath.Join(tmpRoot, "session_1"), 0o755)).To(Succeed()) + g.Expect(os.Symlink(filepath.Join(tmpRoot, "session_1"), filepath.Join(tmpRoot, "session_latest"))).To(Succeed()) + + var entered sync.Once + enteredCh := make(chan struct{}) + release := make(chan struct{}) + defer close(release) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + entered.Do(func() { close(enteredCh) }) + <-release + })) + t.Cleanup(srv.Close) + + handler, writer := newPollTestHandler(t, srv.URL) + handler.EndpointPollInterval = time.Hour + + stop := make(chan struct{}) + done := make(chan struct{}) + go func() { + handler.PollAdditionalEndpointsPeriodically(stop) + close(done) + }() + + g.Eventually(enteredCh).Should(BeClosed()) + close(stop) + g.Eventually(done, "5s").Should(BeClosed()) + g.Expect(writtenKeys(writer)).To(BeEmpty()) +} + +// TestPeriodicPollTimeoutBuffersLateResult verifies a timed-out join does not share state +// and the periodic goroutine can still transfer its result and exit after a blocked write. +func TestPeriodicPollTimeoutBuffersLateResult(t *testing.T) { + g := NewWithT(t) + + tmpRoot := t.TempDir() + t.Setenv("RAY_TMP_ROOT", tmpRoot) + g.Expect(os.MkdirAll(filepath.Join(tmpRoot, "session_1"), 0o755)).To(Succeed()) + g.Expect(os.Symlink(filepath.Join(tmpRoot, "session_1"), filepath.Join(tmpRoot, "session_latest"))).To(Succeed()) + + dash := &fakeDashboard{jobs: `[]`} + srv := dash.start(t) + writer := &blockingStorageWriter{ + entered: make(chan struct{}), + release: make(chan struct{}), + } + handler, _ := newPollTestHandler(t, srv.URL) + handler.Writer = writer + handler.EndpointPollInterval = time.Hour + + stop := make(chan struct{}) + results := handler.startPeriodicEndpointPolling(stop) + g.Eventually(writer.entered).Should(BeClosed()) + close(stop) + + result, joined := waitForPeriodicPollResult(results, 0) + g.Expect(joined).To(BeFalse()) + g.Expect(result.state).To(BeNil()) + + close(writer.release) + g.Eventually(func() int { return len(results) }).Should(Equal(1)) + lateResult := <-results + g.Expect(lateResult.sessionName).To(Equal("session_1")) + g.Expect(lateResult.state).NotTo(BeNil()) +} + +// TestPollAllEndpointsStopsWhenContextExpires verifies a canceled context prevents new requests. +func TestPollAllEndpointsStopsWhenContextExpires(t *testing.T) { + g := NewWithT(t) + + dash := &fakeDashboard{jobs: `[{"job_id": "01000000", "status": "SUCCEEDED"}]`} + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + handler.pollAllEndpoints(ctx, "session_1", newDatasetPollState(), false) + + g.Expect(dash.requestsFor("/")).To(BeEmpty()) + g.Expect(writtenKeys(writer)).To(BeEmpty()) +} + +// TestPollSingleEndpointDoesNotStoreAfterCancellation verifies a response fetched just as +// shutdown starts remains retryable instead of being recorded as durably stored. +func TestPollSingleEndpointDoesNotStoreAfterCancellation(t *testing.T) { + for _, tc := range []struct { + name string + body string + }{ + {name: "non-empty response", body: `{"datasets": [{"dataset": "complete"}]}`}, + {name: "empty response", body: `{"datasets": []}`}, + } { + t.Run(tc.name, func(t *testing.T) { + g := NewWithT(t) + ctx, cancel := context.WithCancel(context.Background()) + handler, writer := newPollTestHandler(t, "http://dashboard") + handler.HttpClient = &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: &cancelOnEOFBody{ + Reader: strings.NewReader(tc.body), + cancel: cancel, + }, + Request: req, + }, nil + })} + + outcome := handler.pollSingleEndpoint(ctx, dataDatasetsEndpointPrefix+"01000000", "session_1", false) + + g.Expect(outcome).To(Equal(pollFailed)) + g.Expect(ctx.Err()).To(MatchError(context.Canceled)) + g.Expect(writtenKeys(writer)).To(BeEmpty()) + }) + } +} + +// TestPolledEndpointsAppendsConfiguredOnes verifies RAY_COLLECTOR_ADDITIONAL_ENDPOINTS adds +// to the built-in set, deduplicated. +func TestPolledEndpointsAppendsConfiguredOnes(t *testing.T) { + g := NewWithT(t) + + handler, _ := newPollTestHandler(t, "http://unused") + g.Expect(handler.polledEndpoints()).To(Equal(staticPolledEndpoints)) + + handler.AdditionalEndpoints = []string{ + trainRunsEndpoint, + serveApplicationsEndpoint, // already built in + trainRunsEndpoint, // repeated by the user + } + g.Expect(handler.polledEndpoints()).To(Equal([]string{ + serveApplicationsEndpoint, + placementGroupsEndpoint, + trainRunsEndpoint, + })) + + // The built-in list itself must not be mutated by the append. + g.Expect(staticPolledEndpoints).To(Equal([]string{ + serveApplicationsEndpoint, + placementGroupsEndpoint, + })) +} + +// TestPollAllEndpointsStoresConfiguredEndpoint verifies a configured endpoint is stored too. +func TestPollAllEndpointsStoresConfiguredEndpoint(t *testing.T) { + g := NewWithT(t) + + dash := &fakeDashboard{jobs: `[]`} + srv := dash.start(t) + handler, writer := newPollTestHandler(t, srv.URL) + handler.AdditionalEndpoints = []string{trainRunsEndpoint} + + handler.pollAllEndpoints(context.Background(), "session_1", newDatasetPollState(), false) + + g.Expect(dash.requestsFor(trainRunsEndpoint)).To(HaveLen(1)) + g.Expect(writtenKeys(writer)).To(ContainElement( + "cluster-dir/session_1/fetched_endpoints/restful__api__train__v2__runs__v1")) + g.Expect(writtenKeys(writer)).To(HaveLen(3)) +} + +// TestServeSnapshotFollowsLiveButSurvivesShutdown verifies a periodic poll mirrors the live +// cluster (empty included) while the final shutdown poll cannot erase a converged snapshot. +func TestServeSnapshotFollowsLiveButSurvivesShutdown(t *testing.T) { + g := NewWithT(t) + + var mu sync.Mutex + serveBody := `{"applications": {"app": {"status": "RUNNING"}}}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsEndpoint { + _, _ = w.Write([]byte(`[]`)) + return + } + mu.Lock() + defer mu.Unlock() + _, _ = w.Write([]byte(serveBody)) + })) + t.Cleanup(srv.Close) + + handler, writer := newPollTestHandler(t, srv.URL) + serveKey := "cluster-dir/session_1/fetched_endpoints/" + + utils.EndpointPathToStorageKey(serveApplicationsEndpoint) + storedServe := func() string { + writer.mu.Lock() + defer writer.mu.Unlock() + return writer.writtenFiles[serveKey] + } + setServeBody := func(body string) { + mu.Lock() + defer mu.Unlock() + serveBody = body + } + + handler.pollAllEndpoints(context.Background(), "session_1", newDatasetPollState(), false) + g.Expect(storedServe()).To(ContainSubstring("RUNNING")) + + // A final poll must not erase the converged snapshot with a dying dashboard's answer. + setServeBody(`{"applications": {}}`) + handler.pollAllEndpoints(context.Background(), "session_1", newDatasetPollState(), true) + g.Expect(storedServe()).To(ContainSubstring("RUNNING")) + + // A periodic poll mirrors the live cluster, empty included. + handler.pollAllEndpoints(context.Background(), "session_1", newDatasetPollState(), false) + g.Expect(storedServe()).To(Equal(`{"applications": {}}`)) +} + +func TestHasServeApplications(t *testing.T) { + g := NewWithT(t) + + g.Expect(hasServeApplications([]byte(`{"applications": {}}`))).To(BeFalse()) + g.Expect(hasServeApplications([]byte(`{"applications": {"a": {}}}`))).To(BeTrue()) + g.Expect(hasServeApplications([]byte(`{}`))).To(BeFalse()) + // An unexpected shape is stored rather than silently dropped. + g.Expect(hasServeApplications([]byte(`not json`))).To(BeTrue()) +} + +// TestIsEmptyPayloadOnlyGuardsKnownEndpoints verifies the guard never suppresses endpoints it +// does not understand. +func TestIsEmptyPayloadOnlyGuardsKnownEndpoints(t *testing.T) { + g := NewWithT(t) + + // Serve is guarded only on the final shutdown poll; datasets always. + g.Expect(isEmptyPayload(serveApplicationsEndpoint, []byte(`{"applications": {}}`), true)).To(BeTrue()) + g.Expect(isEmptyPayload(serveApplicationsEndpoint, []byte(`{"applications": {}}`), false)).To(BeFalse()) + g.Expect(isEmptyPayload(dataDatasetsEndpointPrefix+"01000000", []byte(`{"datasets": []}`), false)).To(BeTrue()) + g.Expect(isEmptyPayload(placementGroupsEndpoint, []byte(`{}`), true)).To(BeFalse()) + g.Expect(isEmptyPayload(trainRunsEndpoint, []byte(`{}`), true)).To(BeFalse()) +} + +func TestHasDatasets(t *testing.T) { + g := NewWithT(t) + + g.Expect(hasDatasets([]byte(`{"datasets": []}`))).To(BeFalse()) + g.Expect(hasDatasets([]byte(`{"datasets": [{"dataset": "a"}]}`))).To(BeTrue()) + g.Expect(hasDatasets([]byte(`{}`))).To(BeFalse()) + // An unexpected shape is stored rather than silently dropped. + g.Expect(hasDatasets([]byte(`not json`))).To(BeTrue()) +} + +// TestTerminalJobStatusesMatchRay guards against drift from Ray's JobStatus enum. +func TestTerminalJobStatusesMatchRay(t *testing.T) { + g := NewWithT(t) + + for _, status := range []string{"SUCCEEDED", "FAILED", "STOPPED"} { + g.Expect(terminalJobStatuses[status]).To(BeTrue(), "%s should be terminal", status) + } + for _, status := range []string{"PENDING", "RUNNING", ""} { + g.Expect(terminalJobStatuses[status]).To(BeFalse(), "%q should not be terminal", status) + } +} + +// TestFakeDashboardJobsShapeMatchesRay documents the /api/jobs/ fields the collector depends on. +func TestFakeDashboardJobsShapeMatchesRay(t *testing.T) { + g := NewWithT(t) + + var jobs []struct { + JobID string `json:"job_id"` + Status string `json:"status"` + } + body := `[{"job_id": "01000000", "submission_id": "raysubmit_x", "status": "SUCCEEDED", "type": "SUBMISSION"}]` + g.Expect(json.Unmarshal([]byte(body), &jobs)).To(Succeed()) + g.Expect(jobs).To(HaveLen(1)) + g.Expect(jobs[0].JobID).To(Equal("01000000")) + g.Expect(jobs[0].Status).To(Equal("SUCCEEDED")) +} diff --git a/historyserver/pkg/collector/types/types.go b/historyserver/pkg/collector/types/types.go index 9358841f38e..0f90cd08e21 100644 --- a/historyserver/pkg/collector/types/types.go +++ b/historyserver/pkg/collector/types/types.go @@ -24,6 +24,7 @@ type RayCollectorConfig struct { PushInterval time.Duration DashboardAddress string + // AdditionalEndpoints are polled on top of the collector's built-in set. AdditionalEndpoints []string EndpointPollInterval time.Duration diff --git a/historyserver/pkg/historyserver/router.go b/historyserver/pkg/historyserver/router.go index 28eb9046ffd..4797f6b2354 100644 --- a/historyserver/pkg/historyserver/router.go +++ b/historyserver/pkg/historyserver/router.go @@ -1121,8 +1121,9 @@ func (s *ServerHandler) getAdditionalEndpoint(req *restful.Request, resp *restfu clusterLogPathPrefix := s.getClusterLogPathPrefix(req) // Use the full request URI (path + query) for storage key lookup. - // The collector stores keys using the full endpoint URL from RAY_COLLECTOR_ADDITIONAL_ENDPOINTS, - // which may include query params (e.g., "/api/v0/placement_groups?detail=1&limit=10000"). + // The collector stores keys using the full endpoint URI it polls (see the + // collector's poll.go), which may include query params (e.g., + // "/api/v0/placement_groups?detail=1&limit=10000"). // RequestURI() includes query params when present, and equals URL.Path when absent. storageKey := utils.EndpointPathToStorageKey(req.Request.URL.RequestURI()) endpointPath := path.Join(sessionName, utils.RAY_SESSIONDIR_FETCHED_ENDPOINTS_NAME, storageKey) @@ -1173,6 +1174,13 @@ func emptyResponseForEndpoint(urlPath string) []byte { }) return data default: + // Nothing is stored for a job that never used Ray Data; match the live dashboard. + if strings.HasPrefix(trimmed, "/api/data/datasets/") { + data, _ := json.Marshal(map[string]interface{}{ + "datasets": []interface{}{}, + }) + return data + } return nil } } diff --git a/historyserver/pkg/storage/clusterlogs/clusterlogs.go b/historyserver/pkg/storage/clusterlogs/clusterlogs.go index 907aee72c4a..36aabc479e1 100644 --- a/historyserver/pkg/storage/clusterlogs/clusterlogs.go +++ b/historyserver/pkg/storage/clusterlogs/clusterlogs.go @@ -43,6 +43,12 @@ func SessionDir(rootDir, ownerKind, ownerName, namespace, clusterName, sessionNa return path.Join(cp, sessionName) } +// FetchedEndpointsDir returns the directory containing dashboard endpoint snapshots: +// //fetched_endpoints +func FetchedEndpointsDir(prefix, sessionName string) string { + return path.Join(prefix, sessionName, utils.RAY_SESSIONDIR_FETCHED_ENDPOINTS_NAME) +} + // NodeDir returns the path to a node's directory under a session: // // func NodeDir(rootDir, ownerKind, ownerName, namespace, clusterName, sessionName, nodeName string) string { diff --git a/historyserver/pkg/storage/clusterlogs/clusterlogs_test.go b/historyserver/pkg/storage/clusterlogs/clusterlogs_test.go index 53b8a9deaf2..7d14e08613b 100644 --- a/historyserver/pkg/storage/clusterlogs/clusterlogs_test.go +++ b/historyserver/pkg/storage/clusterlogs/clusterlogs_test.go @@ -24,6 +24,11 @@ func TestClusterLogsPaths(t *testing.T) { t.Errorf("SessionDir() = %q, want %q", got, wantSession) } + wantFetchedEndpoints := wantSession + "/fetched_endpoints" + if got := FetchedEndpointsDir(wantPrefix, session); got != wantFetchedEndpoints { + t.Errorf("FetchedEndpointsDir() = %q, want %q", got, wantFetchedEndpoints) + } + wantNode := wantSession + "/node-1" if got := NodeDir(rootDir, ownerKind, ownerName, ns, cluster, session, node); got != wantNode { t.Errorf("NodeDir() = %q, want %q", got, wantNode) diff --git a/historyserver/test/e2e/collector_test.go b/historyserver/test/e2e/collector_test.go index 26d24f810d0..7d37b2f6054 100644 --- a/historyserver/test/e2e/collector_test.go +++ b/historyserver/test/e2e/collector_test.go @@ -5,6 +5,8 @@ import ( "encoding/json" "fmt" "io" + "net/http" + "path" "path/filepath" "strings" "testing" @@ -80,6 +82,14 @@ func TestCollector(t *testing.T) { name: "Timezone: collector should fetch and store timezone endpoint data once on startup", testFunc: testCollectorStoresTimezone, }, + { + name: "Serve applications: collector should poll and store a converged Serve snapshot", + testFunc: testCollectorStoresServeApplications, + }, + { + name: "Ray Data datasets: collector should store per-job datasets only for the job that used Ray Data", + testFunc: testCollectorStoresDataDatasets, + }, } for _, tt := range tests { @@ -432,76 +442,44 @@ func testCollectorStoresTimezone(test Test, g *WithT, namespace *corev1.Namespac DeleteS3Bucket(test, g, s3Client) } -// testCollectorStoresPlacementGroups verifies that the Head collector periodically polls -// /api/v0/placement_groups from the Ray Dashboard and stores the result in S3. -// -// The placement_groups endpoint is configured via RAY_COLLECTOR_ADDITIONAL_ENDPOINTS in -// raycluster.yaml and polled by PollAdditionalEndpointsPeriodically. -// -// The test case follows these steps: -// 1. Prepare test environment by applying a Ray cluster with the collector -// 2. Submit a RayJob that creates a detached placement group (so the PG persists after the job) -// 3. Get the sessionID from the head pod to build the expected S3 key -// 4. Wait for the placement groups file to appear in S3 at {sessionName}/fetched_endpoints/restful__api__v0__placement_groups -// 5. Read the file and verify it contains valid JSON with a non-empty placement_groups list -// 6. Delete S3 bucket to ensure test isolation +// testCollectorStoresPlacementGroups verifies the collector stores the placement_groups +// snapshot; the RayJob creates a detached placement group so the PG outlives the job. +// The history-server replay of this endpoint is covered by testDeadClusterPlacementGroups. func testCollectorStoresPlacementGroups(test Test, g *WithT, namespace *corev1.Namespace, s3Client *s3.S3) { rayCluster := PrepareTestEnv(test, g, namespace, s3Client) - // Submit a RayJob that creates a detached placement group named "test_pg". - // The detached lifetime ensures the PG persists after the job exits, so the - // collector captures non-empty data when polling /api/v0/placement_groups. ApplyRayJobAndWaitForCompletion(test, g, namespace, rayCluster) sessionID := GetSessionIDFromHeadPod(test, g, rayCluster) - // The collector stores the endpoint with query params (as configured in RAY_COLLECTOR_ADDITIONAL_ENDPOINTS). + // Matches placementGroupsEndpoint in the collector's poll.go, query params included. storageKey := utils.EndpointPathToStorageKey("/api/v0/placement_groups?detail=1&limit=10000") sessionDir := clusterlogs.SessionDir("log", "", "", rayCluster.Namespace, rayCluster.Name, sessionID) pgKey := fmt.Sprintf("%s/%s/%s", sessionDir, utils.RAY_SESSIONDIR_FETCHED_ENDPOINTS_NAME, storageKey) LogWithTimestamp(test.T(), "Waiting for placement groups data to appear at S3 key: %s", pgKey) - - var pgBody []byte g.Eventually(func(gg Gomega) { - result, err := s3Client.GetObject(&s3.GetObjectInput{ - Bucket: aws.String(S3BucketName), - Key: aws.String(pgKey), - }) - gg.Expect(err).NotTo(HaveOccurred()) - defer result.Body.Close() - - body, err := io.ReadAll(result.Body) - gg.Expect(err).NotTo(HaveOccurred()) - gg.Expect(body).NotTo(BeEmpty(), "Placement groups file should not be empty") - - // Verify it is valid JSON with a non-empty placement groups list. - var response map[string]interface{} - err = json.Unmarshal(body, &response) - gg.Expect(err).NotTo(HaveOccurred(), "Placement groups response should be valid JSON") - - // The Ray State API v2 returns {"result": true, "msg": "", "data": {"result": {"total": N, "result": [...], ...}}}. - gg.Expect(response).To(HaveKey("result"), "Placement groups response should contain result field") - gg.Expect(response["result"]).To(BeTrue(), "result field should be true") - gg.Expect(response).To(HaveKey("data"), "Placement groups response should contain data field") - data, ok := response["data"].(map[string]interface{}) - gg.Expect(ok).To(BeTrue(), "data field should be a JSON object") - gg.Expect(data).To(HaveKey("result"), "data should contain result field") - resultObj, ok := data["result"].(map[string]interface{}) - gg.Expect(ok).To(BeTrue(), "data.result field should be a JSON object") - gg.Expect(resultObj).To(HaveKey("result"), "data.result should contain result field") - - pgList, ok := resultObj["result"].([]interface{}) - gg.Expect(ok).To(BeTrue(), "data.result.result should be a JSON array") - gg.Expect(pgList).NotTo(BeEmpty(), "placement groups list should not be empty (RayJob creates a detached PG)") - - pgBody = body + assertPlacementGroupsNonEmpty(gg, readS3Object(gg, s3Client, pgKey)) }, TestTimeoutMedium).Should(Succeed()) - LogWithTimestamp(test.T(), "Placement groups data stored successfully: %s", string(pgBody)) - DeleteS3Bucket(test, g, s3Client) } +// assertPlacementGroupsNonEmpty requires at least one PG in the State API envelope +// {"result": true, "data": {"result": {"result": [...]}}}; the RayJob creates a detached one. +func assertPlacementGroupsNonEmpty(g Gomega, body []byte) { + var response struct { + Result bool `json:"result"` + Data struct { + Result struct { + Result []json.RawMessage `json:"result"` + } `json:"result"` + } `json:"data"` + } + g.Expect(json.Unmarshal(body, &response)).To(Succeed(), "placement groups response should be valid JSON") + g.Expect(response.Result).To(BeTrue(), "state API result field should be true") + g.Expect(response.Data.Result.Result).NotTo(BeEmpty(), "placement groups list should not be empty") +} + // verifyS3SessionDirs verifies file contents in logs/, node_events/, and job_events/ directories under a session prefix in S3. // There are two phases of verification: // 1. Verify file contents in logs/ directory @@ -649,3 +627,186 @@ func assertAllEventTypesCovered(test Test, g Gomega, events []rayEvent) { g.Expect(foundEventTypes[string(eventType)]).To(BeTrue(), "Event type %s not found", eventType) } } + +// readS3Object reads an object body, failing the enclosing Eventually if it is absent. +func readS3Object(g Gomega, s3Client *s3.S3, key string) []byte { + result, err := s3Client.GetObject(&s3.GetObjectInput{ + Bucket: aws.String(S3BucketName), + Key: aws.String(key), + }) + g.Expect(err).NotTo(HaveOccurred()) + defer result.Body.Close() + + body, err := io.ReadAll(result.Body) + g.Expect(err).NotTo(HaveOccurred()) + return body +} + +// listFetchedEndpoints lists polled-endpoint objects by storage-key prefix. Listing beats +// constructing the key: the session name would have to come from an already-deleted head pod. +func listFetchedEndpoints(g Gomega, s3Client *s3.S3, clusterPrefix, storageKeyPrefix string) []string { + marker := "/" + utils.RAY_SESSIONDIR_FETCHED_ENDPOINTS_NAME + "/" + + var keys []string + err := s3Client.ListObjectsV2Pages(&s3.ListObjectsV2Input{ + Bucket: aws.String(S3BucketName), + Prefix: aws.String(clusterPrefix + "/"), + }, func(page *s3.ListObjectsV2Output, _ bool) bool { + for _, obj := range page.Contents { + key := aws.StringValue(obj.Key) + idx := strings.Index(key, marker) + if idx >= 0 && strings.HasPrefix(key[idx+len(marker):], storageKeyPrefix) { + keys = append(keys, key) + } + } + return true + }) + g.Expect(err).NotTo(HaveOccurred()) + return keys +} + +// enterClusterForOwner sets the session cookie via enter_cluster, which takes the owner's +// name and resolves the generated cluster name itself; a plain RayCluster is its own owner. +func enterClusterForOwner(test Test, g *WithT, client *http.Client, historyServerURL, namespace, ownerKind, ownerName, clusterName, session string) { + enterURL := fmt.Sprintf("%s/enter_cluster/%s/%s/%s/%s", historyServerURL, namespace, ownerKind, ownerName, session) + LogWithTimestamp(test.T(), "Setting cluster context: %s", enterURL) + + g.Eventually(func(gg Gomega) { + var result map[string]any + gg.Expect(json.Unmarshal(getHistoryServerJSON(gg, client, enterURL), &result)).To(Succeed()) + gg.Expect(result["result"]).To(Equal("success")) + gg.Expect(result["name"]).To(Equal(clusterName), "enter_cluster should resolve the owner to its generated cluster") + gg.Expect(result["session"]).To(Equal(session)) + }, TestTimeoutShort).Should(Succeed()) +} + +// getHistoryServerJSON GETs a history server URL and returns the body, requiring 200. +func getHistoryServerJSON(g Gomega, client *http.Client, url string) []byte { + resp, err := client.Get(url) + g.Expect(err).NotTo(HaveOccurred()) + defer resp.Body.Close() + g.Expect(resp.StatusCode).To(Equal(http.StatusOK)) + + body, err := io.ReadAll(resp.Body) + g.Expect(err).NotTo(HaveOccurred()) + return body +} + +// testCollectorStoresServeApplications verifies the collector stores a converged Serve +// snapshot, then replays it through the history server after the RayService is deleted. +// The round trip matters: collector and server derive the storage key from different +// inputs, and a mismatch surfaces as a valid-but-empty 200, not an error. +func testCollectorStoresServeApplications(test Test, g *WithT, namespace *corev1.Namespace, s3Client *s3.S3) { + rayService := ApplyRayServiceAndWaitForRunning(test, g, namespace) + clusterName := rayService.Status.ActiveServiceStatus.RayClusterName + clusterPrefix := clusterlogs.Prefix("log", utils.RayServiceKind, rayService.Name, namespace.Name, clusterName) + + // Matches serveApplicationsEndpoint in the collector's poll.go. + storageKey := utils.EndpointPathToStorageKey("/api/serve/applications/") + LogWithTimestamp(test.T(), "Waiting for a converged Serve snapshot under S3 prefix: %s", clusterPrefix) + + g.Eventually(func(gg Gomega) { + keys := listFetchedEndpoints(gg, s3Client, clusterPrefix, storageKey) + gg.Expect(keys).To(HaveLen(1), "the head collector stores exactly one Serve snapshot per session") + assertServeAppConverged(gg, readS3Object(gg, s3Client, keys[0])) + }, TestTimeoutMedium).Should(Succeed()) + + DeleteRayServiceAndWait(test, g, namespace.Name, rayService.Name, clusterName) + + ApplyHistoryServer(test, g, namespace, "") + historyServerURL := GetHistoryServerURL(test, g, namespace) + clusterInfo := getClusterFromList(test, g, historyServerURL, clusterName, namespace.Name) + g.Expect(clusterInfo.SessionName).NotTo(Equal(LiveSessionName), "Cluster should be a dead session after deletion") + + client := CreateHTTPClientWithCookieJar(g) + enterClusterForOwner(test, g, client, historyServerURL, namespace.Name, + utils.RayServiceKind, rayService.Name, clusterName, clusterInfo.SessionName) + + LogWithTimestamp(test.T(), "Replaying /api/serve/applications/ through the history server") + g.Eventually(func(gg Gomega) { + assertServeAppConverged(gg, getHistoryServerJSON(gg, client, historyServerURL+"/api/serve/applications/")) + }, TestTimeoutShort).Should(Succeed()) + + DeleteS3Bucket(test, g, s3Client) +} + +// assertServeAppConverged requires RUNNING/HEALTHY: a mid-deploy snapshot and the empty +// fallback are both valid JSON. +func assertServeAppConverged(g Gomega, body []byte) { + var response struct { + Applications map[string]struct { + Status string `json:"status"` + Deployments map[string]struct { + Status string `json:"status"` + } `json:"deployments"` + } `json:"applications"` + } + g.Expect(json.Unmarshal(body, &response)).To(Succeed(), "Serve response should be valid JSON") + + app, ok := response.Applications["history-e2e"] + g.Expect(ok).To(BeTrue(), "applications should contain history-e2e, got %v", response.Applications) + g.Expect(app.Status).To(Equal("RUNNING"), "history-e2e should have converged") + + deployment, ok := app.Deployments["NoOp"] + g.Expect(ok).To(BeTrue(), "history-e2e should contain the NoOp deployment") + g.Expect(deployment.Status).To(Equal("HEALTHY"), "NoOp replica should be healthy") +} + +// testCollectorStoresDataDatasets verifies per-job Ray Data snapshots: exactly one object +// (jobs without datasets are never stored), surviving a self-shutdown cluster, and served +// back by the history server for the URI the frontend requests. +func testCollectorStoresDataDatasets(test Test, g *WithT, namespace *corev1.Namespace, s3Client *s3.S3) { + rayJob := ApplyRayDataJobAndWaitForCompletion(test, g, namespace) + clusterName := rayJob.Status.RayClusterName + clusterPrefix := clusterlogs.Prefix("log", utils.RayJobKind, rayJob.Name, namespace.Name, clusterName) + + // "/api/data/datasets/" maps to "restful__api__data__datasets"; per-job keys append + // "__{job_id}", so this prefix matches every stored job. + storageKeyPrefix := utils.EndpointPathToStorageKey("/api/data/datasets/") + LogWithTimestamp(test.T(), "Waiting for a non-empty datasets object under S3 prefix: %s", clusterPrefix) + + var jobID string + g.Eventually(func(gg Gomega) { + keys := listFetchedEndpoints(gg, s3Client, clusterPrefix, storageKeyPrefix) + gg.Expect(keys).To(HaveLen(1), + "exactly one job used Ray Data, so exactly one datasets object should exist") + assertDatasetsNonEmpty(gg, readS3Object(gg, s3Client, keys[0])) + + // Ray assigns the job ID, so recover it from the key the collector chose. + jobID = strings.TrimPrefix(path.Base(keys[0]), storageKeyPrefix+"__") + gg.Expect(jobID).NotTo(BeEmpty()) + }, TestTimeoutMedium).Should(Succeed()) + + // shutdownAfterJobFinishes tears the cluster down on its own, which is what turns + // the session into a replayable one. + g.Eventually(func() error { + _, err := GetRayCluster(test, namespace.Name, clusterName) + return err + }, TestTimeoutMedium).Should(WithTransform(k8serrors.IsNotFound, BeTrue())) + + ApplyHistoryServer(test, g, namespace, "") + historyServerURL := GetHistoryServerURL(test, g, namespace) + clusterInfo := getClusterFromList(test, g, historyServerURL, clusterName, namespace.Name) + g.Expect(clusterInfo.SessionName).NotTo(Equal(LiveSessionName), "Cluster should be a dead session after shutdown") + + client := CreateHTTPClientWithCookieJar(g) + enterClusterForOwner(test, g, client, historyServerURL, namespace.Name, + utils.RayJobKind, rayJob.Name, clusterName, clusterInfo.SessionName) + + datasetsURL := fmt.Sprintf("%s/api/data/datasets/%s", historyServerURL, jobID) + LogWithTimestamp(test.T(), "Replaying %s through the history server", datasetsURL) + g.Eventually(func(gg Gomega) { + assertDatasetsNonEmpty(gg, getHistoryServerJSON(gg, client, datasetsURL)) + }, TestTimeoutShort).Should(Succeed()) + + DeleteS3Bucket(test, g, s3Client) +} + +// assertDatasetsNonEmpty requires real stats: empty coming back means the object was missing. +func assertDatasetsNonEmpty(g Gomega, body []byte) { + var response struct { + Datasets []json.RawMessage `json:"datasets"` + } + g.Expect(json.Unmarshal(body, &response)).To(Succeed(), "datasets response should be valid JSON") + g.Expect(response.Datasets).NotTo(BeEmpty(), "datasets must be non-empty") +} diff --git a/historyserver/test/support/rayjob.go b/historyserver/test/support/rayjob.go index f6d7d059c70..0ff13b9b8d1 100644 --- a/historyserver/test/support/rayjob.go +++ b/historyserver/test/support/rayjob.go @@ -11,6 +11,8 @@ import ( const ( rayJobManifestPath = "../../config/rayjob.yaml" + // Self-contained; the generated cluster name is only known from the RayJob status. + rayDataManifestPath = "../../config/ray-data.yaml" ) // ApplyRayJobAndWaitForCompletion applies a Ray job to the existing Ray cluster and waits for it to complete successfully. @@ -35,3 +37,27 @@ func ApplyRayJobAndWaitForCompletion(test Test, g *WithT, namespace *corev1.Name return rayJob } + +// ApplyRayDataJobAndWaitForCompletion applies the Ray Data RayJob and waits for success; +// the returned Status.RayClusterName is the only way to learn the generated cluster name. +func ApplyRayDataJobAndWaitForCompletion(test Test, g *WithT, namespace *corev1.Namespace) *rayv1.RayJob { + rayJobFromYaml := DeserializeRayJobYAML(test, rayDataManifestPath) + rayJobFromYaml.Namespace = namespace.Name + + rayJob, err := test.Client().Ray().RayV1(). + RayJobs(namespace.Name). + Create(test.Ctx(), rayJobFromYaml, metav1.CreateOptions{}) + g.Expect(err).NotTo(HaveOccurred()) + LogWithTimestamp(test.T(), "Created RayJob %s/%s successfully", rayJob.Namespace, rayJob.Name) + + LogWithTimestamp(test.T(), "Waiting for RayJob %s/%s to complete successfully", rayJob.Namespace, rayJob.Name) + g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium). + Should(WithTransform(RayJobStatus, Equal(rayv1.JobStatusSucceeded))) + + rayJob, err = GetRayJob(test, rayJob.Namespace, rayJob.Name) + g.Expect(err).NotTo(HaveOccurred()) + g.Expect(rayJob.Status.RayClusterName).NotTo(BeEmpty()) + LogWithTimestamp(test.T(), "RayJob %s/%s succeeded on cluster %s", rayJob.Namespace, rayJob.Name, rayJob.Status.RayClusterName) + + return rayJob +} diff --git a/historyserver/test/support/rayservice.go b/historyserver/test/support/rayservice.go new file mode 100644 index 00000000000..5adabb03e88 --- /dev/null +++ b/historyserver/test/support/rayservice.go @@ -0,0 +1,51 @@ +package support + +import ( + . "github.com/onsi/gomega" + corev1 "k8s.io/api/core/v1" + k8serrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1" + . "github.com/ray-project/kuberay/ray-operator/test/support" +) + +// Self-contained; the generated cluster name is only known from the RayService status. +const rayServiceManifestPath = "../../config/rayservice.yaml" + +// ApplyRayServiceAndWaitForRunning applies the RayService and waits until it is Running. +func ApplyRayServiceAndWaitForRunning(test Test, g *WithT, namespace *corev1.Namespace) *rayv1.RayService { + rayServiceFromYaml := DeserializeRayServiceYAML(test, rayServiceManifestPath) + rayServiceFromYaml.Namespace = namespace.Name + + rayService, err := test.Client().Ray().RayV1(). + RayServices(namespace.Name). + Create(test.Ctx(), rayServiceFromYaml, metav1.CreateOptions{}) + g.Expect(err).NotTo(HaveOccurred()) + LogWithTimestamp(test.T(), "Created RayService %s/%s successfully", rayService.Namespace, rayService.Name) + + LogWithTimestamp(test.T(), "Waiting for RayService %s/%s to be running", rayService.Namespace, rayService.Name) + g.Eventually(RayService(test, rayService.Namespace, rayService.Name), TestTimeoutMedium). + Should(WithTransform(RayServiceStatus, Equal(rayv1.Running))) + + rayService, err = GetRayService(test, rayService.Namespace, rayService.Name) + g.Expect(err).NotTo(HaveOccurred()) + g.Expect(rayService.Status.ActiveServiceStatus.RayClusterName).NotTo(BeEmpty()) + LogWithTimestamp(test.T(), "RayService %s/%s is running on cluster %s", + rayService.Namespace, rayService.Name, rayService.Status.ActiveServiceStatus.RayClusterName) + + return rayService +} + +// DeleteRayServiceAndWait deletes a RayService and waits until its RayCluster is gone. +func DeleteRayServiceAndWait(test Test, g *WithT, namespace, name, clusterName string) { + err := test.Client().Ray().RayV1().RayServices(namespace).Delete(test.Ctx(), name, metav1.DeleteOptions{}) + g.Expect(err).NotTo(HaveOccurred()) + LogWithTimestamp(test.T(), "Deleted RayService %s/%s", namespace, name) + + g.Eventually(func() error { + _, err := GetRayCluster(test, namespace, clusterName) + return err + }, TestTimeoutMedium).Should(WithTransform(k8serrors.IsNotFound, BeTrue())) + LogWithTimestamp(test.T(), "RayCluster %s/%s fully deleted", namespace, clusterName) +} diff --git a/historyserver/test/support/support.go b/historyserver/test/support/support.go index f9b43b442a7..ba2f4fb8a0e 100644 --- a/historyserver/test/support/support.go +++ b/historyserver/test/support/support.go @@ -1,7 +1,6 @@ package support import ( - "context" "fmt" "net/http" "net/http/cookiejar" @@ -38,11 +37,7 @@ func GetContainerStatusByName(pod *corev1.Pod, containerName string) (*corev1.Co } func PortForwardService(test Test, g *WithT, namespace, serviceName string, port int) { - ctx, cancel := context.WithCancel(context.Background()) - test.T().Cleanup(cancel) - - kubectlCmd := exec.CommandContext( - ctx, + kubectlCmd := exec.Command( "kubectl", "-n", namespace, "port-forward", @@ -51,6 +46,13 @@ func PortForwardService(test Test, g *WithT, namespace, serviceName string, port ) err := kubectlCmd.Start() g.Expect(err).NotTo(HaveOccurred()) + + // Kill and reap on cleanup: a leaked forward keeps the port, so the next test's + // forward cannot bind and silently talks to this test's deleted namespace. + test.T().Cleanup(func() { + _ = kubectlCmd.Process.Kill() + _ = kubectlCmd.Wait() + }) } // InstallGrafanaAndPrometheus installs Grafana and Prometheus in the cluster for testing.