from collections import defaultdict from dataclasses import dataclass from enum import Enum from typing import Dict, List, Tuple import numpy as np from ray.core.generated.common_pb2 import TaskStatus from ray.util.state.common import TaskState TaskName = str class SchedulingPhase(str, Enum): """Phases of the Ray task scheduling lifecycle. See ``src/ray/protobuf/common.proto`` TaskStatus for the source of truth. Normal tasks:: PENDING_ARGS_AVAIL -> PENDING_NODE_ASSIGNMENT -> PENDING_OBJ_STORE_MEM_AVAIL (sub-state, metrics only) -> PENDING_ARGS_FETCH (sub-state, metrics only) -> SUBMITTED_TO_WORKER -> RUNNING -> FINISHED Actor tasks:: PENDING_ARGS_AVAIL -> PENDING_NODE_ASSIGNMENT -> PENDING_OBJ_STORE_MEM_AVAIL (sub-state, metrics only) -> PENDING_ARGS_FETCH (sub-state, metrics only) -> SUBMITTED_TO_WORKER -> PENDING_ACTOR_TASK_ARGS_FETCH -> PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY -> RUNNING -> FINISHED Each member maps to a (start, end) pair of ``TaskStatus`` values. Use ``phase.start_name`` / ``phase.end_name`` to get the string names that appear in task event dicts. """ # Fine-grained phases (consecutive state transitions) ARGS_WAIT_MS = "args_wait_ms" NODE_ASSIGNMENT_MS = "node_assignment_ms" OBJ_STORE_MEM_WAIT_MS = "obj_store_mem_wait_ms" ARGS_FETCH_MS = "args_fetch_ms" WORKER_DISPATCH_MS = "worker_dispatch_ms" ACTOR_TASK_ARGS_FETCH_MS = "actor_task_args_fetch_ms" ACTOR_TASK_ORDERING_MS = "actor_task_ordering_ms" EXECUTION_MS = "execution_ms" # Coarse aggregate phases SCHEDULING_MS = "scheduling_ms" WORKER_STARTUP_MS = "worker_startup_ms" TOTAL_OVERHEAD_MS = "total_overhead_ms" @property def boundaries(self) -> Tuple[int, int]: """Return the (start_status, end_status) TaskStatus pair for this phase.""" S = TaskStatus match self: # Fine-grained case SchedulingPhase.ARGS_WAIT_MS: return (S.PENDING_ARGS_AVAIL, S.PENDING_NODE_ASSIGNMENT) case SchedulingPhase.NODE_ASSIGNMENT_MS: return (S.PENDING_NODE_ASSIGNMENT, S.PENDING_OBJ_STORE_MEM_AVAIL) case SchedulingPhase.OBJ_STORE_MEM_WAIT_MS: return (S.PENDING_OBJ_STORE_MEM_AVAIL, S.PENDING_ARGS_FETCH) case SchedulingPhase.ARGS_FETCH_MS: return (S.PENDING_ARGS_FETCH, S.SUBMITTED_TO_WORKER) case SchedulingPhase.WORKER_DISPATCH_MS: return (S.SUBMITTED_TO_WORKER, S.PENDING_ACTOR_TASK_ARGS_FETCH) case SchedulingPhase.ACTOR_TASK_ARGS_FETCH_MS: return ( S.PENDING_ACTOR_TASK_ARGS_FETCH, S.PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY, ) case SchedulingPhase.ACTOR_TASK_ORDERING_MS: return (S.PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY, S.RUNNING) case SchedulingPhase.EXECUTION_MS: return (S.RUNNING, S.FINISHED) # Coarse aggregates case SchedulingPhase.SCHEDULING_MS: return (S.PENDING_NODE_ASSIGNMENT, S.SUBMITTED_TO_WORKER) case SchedulingPhase.WORKER_STARTUP_MS: return (S.SUBMITTED_TO_WORKER, S.RUNNING) case SchedulingPhase.TOTAL_OVERHEAD_MS: return (S.PENDING_ARGS_AVAIL, S.RUNNING) case _: raise ValueError(f"Unknown scheduling phase: {self}") @property def start_name(self) -> str: """TaskStatus string name for the start boundary.""" return TaskStatus.Name(self.boundaries[0]) @property def end_name(self) -> str: """TaskStatus string name for the end boundary.""" return TaskStatus.Name(self.boundaries[1]) @dataclass class PhaseStats: """Aggregated statistics for a single scheduling phase (all values in ms).""" p50: float p99: float mean: float max: float min: float @dataclass class SchedulingOverheadSummary: """Summary of scheduling overhead across matched tasks.""" num_tasks: int # Each key is a SchedulingPhase enum str. I avoided storing the enum directly # to make it more human readable. phases: Dict[str, PhaseStats] @classmethod def from_task_states(cls, tasks: List[TaskState]) -> "SchedulingOverheadSummary": """Build a summary from a list of TaskState objects. Computes per-phase scheduling stats by deriving durations from each task's event timestamps. """ phase_durations: Dict[SchedulingPhase, List[float]] = defaultdict(list) for t in tasks: if not t.events: continue ts: Dict[str, float] = {e["state"]: e["created_ms"] for e in t.events} for phase in SchedulingPhase: if phase.start_name in ts and phase.end_name in ts: phase_durations[phase].append( ts[phase.end_name] - ts[phase.start_name] ) phases: Dict[str, PhaseStats] = {} for phase, vals in phase_durations.items(): arr = np.array(vals) phases[phase.value] = PhaseStats( p50=float(np.percentile(arr, 50)), p99=float(np.percentile(arr, 99)), mean=float(np.mean(arr)), max=float(np.max(arr)), min=float(np.min(arr)), ) return cls(num_tasks=len(tasks), phases=phases) @dataclass class BucketedSchedulingOverhead: """Scheduling overhead for a single operator within a time bucket.""" start_ms: float end_ms: float summary: SchedulingOverheadSummary def collect_scheduling_overhead( operator_names: List[str], num_buckets: int = 4, ) -> Dict[TaskName, List[BucketedSchedulingOverhead]]: """Collect per-operator scheduling overhead from the Ray State API, bucketed into ``num_buckets`` equal time intervals. Queries tasks once (up to ``limit``), keeps tasks whose ``name`` starts with any entry in ``operator_names`` (longest prefix wins), then buckets by the matched logical operator name. Global time range comes from ``creation_time_ms`` of the kept tasks. """ from ray.util.state.api import list_tasks unfiltered: List[TaskState] = list_tasks( detail=True, limit=10_000, raise_on_missing_output=False, ) # task.name (Ray task display name) -> logical operator_name from operator_names task_name_to_operator: Dict[TaskName, str] = {} matched_tasks: List[TaskState] = [] for t in unfiltered: if not t.name: continue if t.name not in task_name_to_operator: found_op_name: bool = False for op_name in operator_names: if op_name in t.name: task_name_to_operator[t.name] = op_name found_op_name = True break if not found_op_name: continue matched_tasks.append(t) all_tasks = matched_tasks # 1. Sort all tasks by creation_time_ms. all_tasks = [t for t in all_tasks if t.creation_time_ms is not None] if not all_tasks: return {} all_tasks.sort(key=lambda t: t.creation_time_ms) min_ts = all_tasks[0].creation_time_ms max_ts = all_tasks[-1].creation_time_ms assert min_ts is not None and max_ts is not None bucket_width = (max_ts - min_ts) / num_buckets if num_buckets > 0 else 0 # 2. Build bucket boundaries. if bucket_width == 0: boundaries = [(min_ts, max_ts + 1)] else: boundaries = [] for i in range(num_buckets): lo = min_ts + i * bucket_width hi = lo + bucket_width boundaries.append((lo, hi)) # Pre-allocate per-bucket, per-operator task lists. # bucket_tasks[bucket_idx][operator_name] -> list of tasks bucket_tasks: List[Dict[TaskName, List[TaskState]]] = [ defaultdict(list) for _ in boundaries ] # 3. Single pass: assign each sorted task to its bucket. def _in_bucket(t: TaskState, bucket: Tuple[float, float]) -> bool: ts = t.creation_time_ms assert ts is not None return bucket[0] <= ts < bucket[1] bucket_idx = 0 for t in all_tasks: while bucket_idx < len(boundaries) - 1 and not _in_bucket( t=t, bucket=boundaries[bucket_idx] ): bucket_idx += 1 logical_op = task_name_to_operator[t.name] bucket_tasks[bucket_idx][logical_op].append(t) # Build the result dict: operator_name -> list of bucketed summaries. result: Dict[TaskName, List[BucketedSchedulingOverhead]] = defaultdict(list) for i, (lo, hi) in enumerate(boundaries): for op_name, tasks in bucket_tasks[i].items(): result[op_name].append( BucketedSchedulingOverhead( start_ms=lo, end_ms=hi, summary=SchedulingOverheadSummary.from_task_states(tasks), ) ) return dict(result)