Files
2026-07-13 13:17:40 +08:00

305 lines
11 KiB
Python

# ABOUTME: Periodically samples Ray object-store state via ray.util.state.
# ABOUTME: Emits per-(node, operator) primary-bytes time series for spill diagnosis.
import csv
import os
import re
import threading
import time
import traceback
_MAPWORKER_RE = re.compile(r"^MapWorker\((.+)\)$")
def start(outdir, interval_s=5, fast_window_s=60, fast_interval_s=1):
"""Start object-store sampling in a background thread.
Outputs ``object_store_state.csv`` with one row per (tick, owner_node, operator)
showing primary-object counts and bytes plus a reference-type breakdown.
Args:
outdir: Shared storage directory for output files.
interval_s: Steady-state seconds between samples. Default 5s — fast enough
to catch short-lived primaries between produce and GC. ``list_objects()``
is ~100ms per call on this benchmark; 5s adds ~2% daemon overhead.
fast_window_s: Seconds at the start of the run during which the daemon
ticks every ``fast_interval_s`` instead of ``interval_s``. Resolves
placement-ramp events (e.g. GPU-node CPU saturation by ReadFiles) to
sub-5s precision. Set to 0 to disable.
fast_interval_s: Tick interval during the fast window. Default 1s.
"""
thread = threading.Thread(
target=_loop,
args=(outdir, interval_s, fast_window_s, fast_interval_s),
daemon=True,
)
thread.start()
return thread
def _loop(outdir, interval_s, fast_window_s=0, fast_interval_s=1):
os.makedirs(outdir, exist_ok=True)
csv_path = os.path.join(outdir, "object_store_state.csv")
actors_path = os.path.join(outdir, "actor_placement.csv")
plasma_path = os.path.join(outdir, "plasma_stats.csv")
obj_fields = [
"timestamp",
"owner_ip",
"operator",
"n_objects",
"bytes_total",
"bytes_pinned",
"bytes_local_ref",
"bytes_used_by_pending_task",
"bytes_other",
]
actor_fields = [
"timestamp",
"node_ip",
"operator",
"n_actors",
]
# Per-node Plasma stats from raylet GetNodeStats RPC. Distinguishes
# primary bytes (objects whose authoritative owner is this node) from
# total used bytes (= primary + secondary copies + framework overhead).
# Also exposes spill/restore totals per node for time-correlated analysis.
plasma_fields = [
"timestamp",
"node_ip",
"bytes_used",
"bytes_avail",
"bytes_primary",
"bytes_fallback",
"n_local_objects",
"spilled_bytes_total",
"spilled_objects_total",
"restored_bytes_total",
"restored_objects_total",
]
obj_f = open(csv_path, "w", newline="")
obj_writer = csv.DictWriter(obj_f, fieldnames=obj_fields)
obj_writer.writeheader()
obj_f.flush()
actor_f = open(actors_path, "w", newline="")
actor_writer = csv.DictWriter(actor_f, fieldnames=actor_fields)
actor_writer.writeheader()
actor_f.flush()
plasma_f = open(plasma_path, "w", newline="")
plasma_writer = csv.DictWriter(plasma_f, fieldnames=plasma_fields)
plasma_writer.writeheader()
plasma_f.flush()
if fast_window_s > 0:
print(
f"object_store_monitor: writing {csv_path}, {actors_path}, {plasma_path} "
f"every {fast_interval_s}s for {fast_window_s}s starting from first actor, "
f"then every {interval_s}s"
)
else:
print(
f"object_store_monitor: writing {csv_path}, {actors_path}, {plasma_path} "
f"every {interval_s}s"
)
# Anchor the fast window to "first time we observe any actors" so the 1s
# sampling lands on the actual ramp-up regardless of how long the job
# idle-waits before spawning actors. None until first non-zero count.
fast_window_started_at = None
while True:
try:
now = time.time()
actor_to_op, actor_placement = _snapshot_actors()
for row in actor_placement:
row["timestamp"] = now
actor_writer.writerow(row)
actor_f.flush()
plasma_rows = _snapshot_plasma_per_node()
for row in plasma_rows:
row["timestamp"] = now
plasma_writer.writerow(row)
plasma_f.flush()
obj_rows = _snapshot_objects(actor_to_op)
for row in obj_rows:
row["timestamp"] = now
obj_writer.writerow(row)
obj_f.flush()
n_actors = sum(r["n_actors"] for r in actor_placement)
if fast_window_started_at is None and n_actors > 0:
fast_window_started_at = time.time()
print(
f"object_store_monitor: first actor detected; fast window "
f"({fast_interval_s}s tick) active for next {fast_window_s}s"
)
plasma_used_total_gb = sum(r["bytes_used"] for r in plasma_rows) / 1e9
spill_total_gb = sum(r["spilled_bytes_total"] for r in plasma_rows) / 1e9
print(
f"object_store_monitor: tick={now:.0f} "
f"actors={n_actors} "
f"objects={sum(r['n_objects'] for r in obj_rows)} "
f"plasma_used={plasma_used_total_gb:.1f}GB "
f"spilled={spill_total_gb:.1f}GB"
)
except Exception as e:
print(f"object_store_monitor: WARN {e}")
traceback.print_exc()
in_fast_window = (
fast_window_started_at is not None
and (time.time() - fast_window_started_at) < fast_window_s
)
time.sleep(fast_interval_s if in_fast_window else interval_s)
def _snapshot_actors():
"""Return ((ip, pid) → operator_name) map and per-(node_ip, op) actor counts.
ActorState carries node_id (a hex string) but not the IP. We join against
list_nodes() to recover the IP — which is what list_objects() reports for
object owners and what we cross-reference against in the metrics extractor.
"""
from collections import defaultdict
from ray.util.state import list_actors, list_nodes
# raise_on_missing_output=False: allow partial results when the state
# API truncates due to data size. Without this the daemon dies on the
# first heavy snapshot.
nodes = list_nodes(limit=10000, raise_on_missing_output=False)
node_id_to_ip = {n["node_id"]: n.get("node_ip", "unknown") for n in nodes}
actors = list_actors(
filters=[("state", "=", "ALIVE")],
limit=10000,
raise_on_missing_output=False,
)
actor_to_op = {}
counts = defaultdict(int)
for a in actors:
cls = a.get("class_name", "") or ""
m = _MAPWORKER_RE.match(cls)
if not m:
continue
operator = m.group(1)
node_ip = node_id_to_ip.get(a.get("node_id"), "unknown")
if a.get("pid") not in (None, 0) and node_ip != "unknown":
actor_to_op[(node_ip, int(a["pid"]))] = operator
counts[(node_ip, operator)] += 1
placement = [
{"node_ip": ip, "operator": op, "n_actors": n} for (ip, op), n in counts.items()
]
return actor_to_op, placement
def _snapshot_objects(actor_to_op):
"""Aggregate live objects by (owner_ip, operator)."""
from collections import defaultdict
from ray.util.state import list_objects
# API server caps the limit at 10000 unless RAY_MAX_LIMIT_FROM_API_SERVER
# is set. For runs with > 10k objects, set that env var on the head node:
# RAY_MAX_LIMIT_FROM_API_SERVER=200000
# to avoid truncated samples.
#
# raise_on_missing_output=False: the state API also truncates when the
# response size (not just count) exceeds an internal RPC limit; with the
# default the daemon dies the moment objects get heavy. Partial data is
# preferable to no data here.
objs = list_objects(limit=10_000, raise_on_missing_output=False)
agg = defaultdict(
lambda: {
"n_objects": 0,
"bytes_total": 0,
"bytes_pinned": 0,
"bytes_local_ref": 0,
"bytes_used_by_pending_task": 0,
"bytes_other": 0,
}
)
for obj in objs:
# obj.type is WORKER / DRIVER / SPILL_WORKER / RESTORE_WORKER. We focus on
# WORKER since that's the productive workload.
if obj.get("type") != "WORKER":
continue
owner_ip = obj.get("ip") or "unknown"
pid = obj.get("pid")
operator = (
actor_to_op.get((owner_ip, int(pid))) if pid is not None else None
) or "_other"
size = int(obj.get("object_size", 0) or 0)
ref = obj.get("reference_type") or ""
bucket = agg[(owner_ip, operator)]
bucket["n_objects"] += 1
bucket["bytes_total"] += size
if ref == "PINNED_IN_MEMORY":
bucket["bytes_pinned"] += size
elif ref == "LOCAL_REFERENCE":
bucket["bytes_local_ref"] += size
elif ref == "USED_BY_PENDING_TASK":
bucket["bytes_used_by_pending_task"] += size
else:
bucket["bytes_other"] += size
rows = []
for (ip, op), b in agg.items():
rows.append({"owner_ip": ip, "operator": op, **b})
return rows
def _snapshot_plasma_per_node():
"""Per-node Plasma stats from each raylet's GetNodeStats RPC.
Returns a list of dicts, one per alive node, with the fields documented
in plasma_fields. The bytes_used vs bytes_primary split is the key
diagnostic — bytes_used - bytes_primary = bytes occupied by secondary
copies + framework overhead, which list_objects() can't see.
"""
import ray
from ray._private.internal_api import node_stats
rows = []
for node in ray.nodes():
if not node.get("Alive"):
continue
ip = node.get("NodeManagerAddress")
port = node.get("NodeManagerPort")
if not ip or not port:
continue
try:
reply = node_stats(
node_manager_address=ip,
node_manager_port=port,
include_memory_info=False,
)
except Exception:
# Best-effort — skip nodes whose raylet RPC times out.
continue
s = reply.store_stats
rows.append(
{
"node_ip": ip,
"bytes_used": int(s.object_store_bytes_used),
"bytes_avail": int(s.object_store_bytes_avail),
"bytes_primary": int(s.object_store_bytes_primary_copy),
"bytes_fallback": int(getattr(s, "object_store_bytes_fallback", 0)),
"n_local_objects": int(s.num_local_objects),
"spilled_bytes_total": int(s.spilled_bytes_total),
"spilled_objects_total": int(s.spilled_objects_total),
"restored_bytes_total": int(s.restored_bytes_total),
"restored_objects_total": int(s.restored_objects_total),
}
)
return rows