wehub-resource-sync
267 lines
9.1 KiB
Python
267 lines
9.1 KiB
Python
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)
|