Enhance error handling and logging in pipeline and stage processing. Added detailed logging for prompt loading failures, file writing issues, and stage shutdown processes. Introduced a new logging utility for orchestrator metrics and streamlined stats handling in the PipelinedOmniLLM class.
This commit is contained in:
tzhouam committed
1 parent
3c2ad1bf12
commit
c5631ca753
6 files changed
+657
-345
No files matched your search
@@ -162,6 +162,7 @@ def main():
|
||||
args.prompts = loaded
|
||||
print(f"[Info] Loaded {len(args.prompts)} prompts from {args.pt_prompts}")
|
||||
except Exception as e:
|
||||
print(f"[Error] Failed to load pt-prompts: {e}")
|
||||
raise
|
||||
|
||||
if args.prompts is None:
|
||||
@@ -233,8 +234,8 @@ def main():
|
||||
try:
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
f.writelines(lines)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"[Warn] Failed writing text file {out_txt}: {e}")
|
||||
print(f"Request ID: {request_id}, Text saved to {out_txt}")
|
||||
elif stage_outputs.final_output_type == "audio":
|
||||
for output in stage_outputs.request_output:
|
||||
@@ -255,8 +256,8 @@ def main():
|
||||
try:
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
f.writelines(lines)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"[Warn] Failed writing text file {out_txt}: {e}")
|
||||
print(f"Request ID: {request_id}, Saved audio to {output_wav} and text to {out_txt}")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,548 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl
|
||||
|
||||
|
||||
def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
|
||||
try:
|
||||
if not log_file:
|
||||
return
|
||||
try:
|
||||
if os.path.exists(log_file):
|
||||
os.remove(log_file)
|
||||
except Exception:
|
||||
pass
|
||||
# Per-stage logs and stats
|
||||
for sid in range(num_stages):
|
||||
try:
|
||||
p = f"{log_file}.stage{sid}.log"
|
||||
if os.path.exists(p):
|
||||
os.remove(p)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
p = f"{log_file}.stage{sid}.stats.jsonl"
|
||||
if os.path.exists(p):
|
||||
os.remove(p)
|
||||
except Exception:
|
||||
pass
|
||||
# Orchestrator stats files
|
||||
try:
|
||||
p = f"{log_file}.orchestrator.stats.jsonl"
|
||||
if os.path.exists(p):
|
||||
os.remove(p)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
p = f"{log_file}.overall.stats.jsonl"
|
||||
if os.path.exists(p):
|
||||
os.remove(p)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None:
|
||||
try:
|
||||
if not log_file:
|
||||
return
|
||||
has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
|
||||
if not has_file_handler:
|
||||
fh = logging.FileHandler(log_file)
|
||||
fh.setLevel(logging.DEBUG)
|
||||
fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
|
||||
logger.addHandler(fh)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optional[str], Optional[str]]:
|
||||
stats_file: Optional[str] = None
|
||||
overall_file: Optional[str] = None
|
||||
try:
|
||||
if enable_stats and log_file:
|
||||
stats_file = f"{log_file}.orchestrator.stats.jsonl"
|
||||
overall_file = f"{log_file}.overall.stats.jsonl"
|
||||
except Exception:
|
||||
stats_file = None
|
||||
overall_file = None
|
||||
return stats_file, overall_file
|
||||
|
||||
|
||||
def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
|
||||
if not path:
|
||||
return
|
||||
try:
|
||||
_append_jsonl(path, record) # type: ignore[arg-type]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def log_transfer_tx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
|
||||
size_bytes: int, tx_time_ms: float, used_shm: bool) -> None:
|
||||
_safe_append_jsonl(stats_file, {
|
||||
"type": "transfer_stats",
|
||||
"from_stage": from_stage,
|
||||
"to_stage": to_stage,
|
||||
"request_id": request_id,
|
||||
"size_bytes": int(size_bytes),
|
||||
"tx_time_ms": float(tx_time_ms),
|
||||
"tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
|
||||
"used_shm": bool(used_shm),
|
||||
})
|
||||
|
||||
|
||||
def log_transfer_rx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
|
||||
rx_bytes: int, rx_decode_time_ms: float, in_flight_time_ms: float) -> None:
|
||||
_safe_append_jsonl(stats_file, {
|
||||
"type": "transfer_rx_stats",
|
||||
"from_stage": from_stage,
|
||||
"to_stage": to_stage,
|
||||
"request_id": request_id,
|
||||
"rx_bytes": int(rx_bytes),
|
||||
"rx_decode_time_ms": float(rx_decode_time_ms),
|
||||
"in_flight_time_ms": float(in_flight_time_ms),
|
||||
"rx_time_per_kb_ms": (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0,
|
||||
})
|
||||
|
||||
|
||||
def log_transfer_total(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
|
||||
size_bytes: int, tx_time_ms: float, in_flight_time_ms: float,
|
||||
rx_decode_time_ms: float, total_time_ms: float) -> None:
|
||||
_safe_append_jsonl(stats_file, {
|
||||
"type": "transfer_total_stats",
|
||||
"from_stage": from_stage,
|
||||
"to_stage": to_stage,
|
||||
"request_id": request_id,
|
||||
"size_bytes": int(size_bytes),
|
||||
"tx_time_ms": float(tx_time_ms),
|
||||
"in_flight_time_ms": float(in_flight_time_ms),
|
||||
"rx_decode_time_ms": float(rx_decode_time_ms),
|
||||
"total_time_ms": float(total_time_ms),
|
||||
"total_time_per_kb_ms": float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0,
|
||||
})
|
||||
|
||||
|
||||
def log_orchestrator_e2e(stats_file: Optional[str], request_id: Any, final_stage_id: int,
|
||||
e2e_time_ms: float, num_tokens_out: int, e2e_time_per_token_ms: float) -> None:
|
||||
_safe_append_jsonl(stats_file, {
|
||||
"type": "orchestrator_request_e2e",
|
||||
"request_id": request_id,
|
||||
"final_stage_id": final_stage_id,
|
||||
"e2e_time_ms": float(e2e_time_ms),
|
||||
"num_tokens_out": int(num_tokens_out),
|
||||
"e2e_time_per_token_ms": float(e2e_time_per_token_ms),
|
||||
})
|
||||
|
||||
|
||||
def log_orchestrator_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None:
|
||||
_safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
|
||||
|
||||
|
||||
def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None:
|
||||
_safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
|
||||
|
||||
|
||||
def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None:
|
||||
_safe_append_jsonl(overall_stats_file, record)
|
||||
|
||||
|
||||
def log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, batch_size: int,
|
||||
num_tokens_out: int, stage_gen_time_ms: float, tokens_per_s: float,
|
||||
rx_transfer_bytes: int, rx_decode_time_ms: float, rx_mbps: float) -> None:
|
||||
_safe_append_jsonl(stats_file, {
|
||||
"type": "stage_request_stats",
|
||||
"stage_id": stage_id,
|
||||
"request_id": request_id,
|
||||
"batch_size": int(batch_size),
|
||||
"num_tokens_out": int(num_tokens_out),
|
||||
"stage_gen_time_ms": float(stage_gen_time_ms),
|
||||
"tokens_per_s": float(tokens_per_s),
|
||||
"rx_transfer_bytes": int(rx_transfer_bytes),
|
||||
"rx_decode_time_ms": float(rx_decode_time_ms),
|
||||
"rx_mbps": float(rx_mbps),
|
||||
})
|
||||
|
||||
|
||||
def log_stage_running_avg(stats_file: Optional[str], stage_id: int, total_tokens: int,
|
||||
total_gen_time_ms: float, avg_tokens_per_s: float) -> None:
|
||||
_safe_append_jsonl(stats_file, {
|
||||
"type": "stage_running_avg",
|
||||
"stage_id": stage_id,
|
||||
"total_tokens": int(total_tokens),
|
||||
"total_gen_time_ms": float(total_gen_time_ms),
|
||||
"avg_tokens_per_s": float(avg_tokens_per_s),
|
||||
})
|
||||
|
||||
|
||||
def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: int,
|
||||
batch_gen_time_ms: float, request_ids: List[Any]) -> None:
|
||||
_safe_append_jsonl(stats_file, {
|
||||
"type": "stage_batch_stats",
|
||||
"stage_id": stage_id,
|
||||
"batch_size": int(batch_size),
|
||||
"batch_gen_time_ms": float(batch_gen_time_ms),
|
||||
"request_ids": list(request_ids),
|
||||
})
|
||||
|
||||
|
||||
# ----------------- Aggregation helpers for orchestrator -----------------
|
||||
|
||||
def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
|
||||
stage_req_counts: List[int],
|
||||
stage_total_time_ms: List[float],
|
||||
stage_total_tokens: List[int],
|
||||
stage_id: int,
|
||||
req_id: Any,
|
||||
metrics: Dict[str, Any]) -> None:
|
||||
try:
|
||||
stage_req_counts[stage_id] += 1
|
||||
stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
|
||||
rid_int = int(req_id)
|
||||
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
pr_stages = pr["stages"] # type: ignore[index]
|
||||
pr_stages[stage_id] = {
|
||||
"stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
|
||||
"num_tokens_out": int(metrics.get("num_tokens_out", 0)),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
|
||||
transfer_agg: Dict[tuple[int, int], Dict[str, float]],
|
||||
per_request: Dict[int, Dict[str, Any]],
|
||||
stage_id: int,
|
||||
req_id: Any,
|
||||
rx_bytes: float,
|
||||
rx_ms: float,
|
||||
in_flight_ms: float) -> Optional[tuple[int, float, float]]:
|
||||
try:
|
||||
# Update RX aggregates for (stage_id-1 -> stage_id)
|
||||
if stage_id > 0:
|
||||
key = (stage_id - 1, stage_id)
|
||||
agg = transfer_agg.get(key)
|
||||
if agg is None:
|
||||
agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
|
||||
"sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
|
||||
"sum_total_ms": 0.0, "total_count": 0.0}
|
||||
transfer_agg[key] = agg
|
||||
agg["sum_rx_bytes"] += float(rx_bytes)
|
||||
agg["sum_rx_ms"] += float(rx_ms)
|
||||
agg["rx_count"] += 1.0
|
||||
|
||||
# Try combine with sender-side timing if present
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
except Exception:
|
||||
rid_int = req_id # best effort
|
||||
s = transfer_edge_req.get((stage_id - 1, stage_id, rid_int))
|
||||
if s is None:
|
||||
return None
|
||||
tx_ms = float(s.get("tx_ms", 0.0))
|
||||
size_b = float(s.get("size_bytes", rx_bytes))
|
||||
total_ms = tx_ms + float(in_flight_ms) + float(rx_ms)
|
||||
agg["sum_total_ms"] += total_ms
|
||||
agg["total_count"] += 1.0
|
||||
# accumulate per-request transfer totals
|
||||
try:
|
||||
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index]
|
||||
pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes) # type: ignore[index]
|
||||
except Exception:
|
||||
pass
|
||||
return int(size_b), float(tx_ms), float(total_ms)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, float]],
|
||||
transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
|
||||
from_stage: int,
|
||||
to_stage: int,
|
||||
req_id: Any,
|
||||
size_bytes: int,
|
||||
tx_ms: float) -> None:
|
||||
try:
|
||||
key = (from_stage, to_stage)
|
||||
agg = transfer_agg.get(key)
|
||||
if agg is None:
|
||||
agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
|
||||
"sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
|
||||
"sum_total_ms": 0.0, "total_count": 0.0}
|
||||
transfer_agg[key] = agg
|
||||
agg["sum_bytes"] += float(size_bytes)
|
||||
agg["sum_ms"] += float(tx_ms)
|
||||
agg["count"] += 1.0
|
||||
# Store sender-side timing for per-request combination
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
except Exception:
|
||||
rid_int = req_id
|
||||
transfer_edge_req[(from_stage, to_stage, rid_int)] = {
|
||||
"tx_ms": float(tx_ms),
|
||||
"size_bytes": float(size_bytes),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
|
||||
total = 0
|
||||
for _ro in engine_outputs:
|
||||
try:
|
||||
outs = getattr(_ro, "outputs", None)
|
||||
if outs and len(outs) > 0:
|
||||
tokens = getattr(outs[0], "token_ids", None)
|
||||
if tokens is not None:
|
||||
total += len(tokens)
|
||||
except Exception:
|
||||
pass
|
||||
return total
|
||||
|
||||
|
||||
def build_stage_summary(stage_req_counts: List[int],
|
||||
stage_total_tokens: List[int],
|
||||
stage_total_time_ms: List[float]) -> List[Dict[str, Any]]:
|
||||
summary: List[Dict[str, Any]] = []
|
||||
for sid in range(len(stage_req_counts)):
|
||||
reqs = stage_req_counts[sid]
|
||||
tokens = stage_total_tokens[sid]
|
||||
total_ms = float(stage_total_time_ms[sid])
|
||||
avg_req = (total_ms / reqs) if reqs > 0 else 0.0
|
||||
avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
|
||||
summary.append({
|
||||
"stage_id": sid,
|
||||
"requests": int(reqs),
|
||||
"tokens": int(tokens),
|
||||
"total_time_ms": total_ms,
|
||||
"avg_time_per_request_ms": avg_req,
|
||||
"avg_tokens_per_s": avg_tok,
|
||||
})
|
||||
return summary
|
||||
|
||||
|
||||
def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]) -> List[Dict[str, Any]]:
|
||||
summary: List[Dict[str, Any]] = []
|
||||
for (src, dst), agg in transfer_agg.items():
|
||||
sum_bytes = float(agg.get("sum_bytes", 0.0))
|
||||
sum_ms = float(agg.get("sum_ms", 0.0))
|
||||
samples = int(agg.get("count", 0.0))
|
||||
tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
|
||||
sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
|
||||
sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
|
||||
samples_rx = int(agg.get("rx_count", 0.0))
|
||||
rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
|
||||
sum_total_ms = float(agg.get("sum_total_ms", 0.0))
|
||||
samples_total = int(agg.get("total_count", 0.0))
|
||||
total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
|
||||
summary.append({
|
||||
"from_stage": src,
|
||||
"to_stage": dst,
|
||||
"samples": samples,
|
||||
"total_bytes": int(sum_bytes),
|
||||
"total_time_ms": sum_ms,
|
||||
"tx_mbps": tx_mbps,
|
||||
"rx_samples": samples_rx,
|
||||
"rx_total_bytes": int(sum_rx_bytes),
|
||||
"rx_total_time_ms": sum_rx_ms,
|
||||
"rx_mbps": rx_mbps,
|
||||
"total_samples": samples_total,
|
||||
"total_transfer_time_ms": sum_total_ms,
|
||||
"total_mbps": total_mbps,
|
||||
})
|
||||
return summary
|
||||
|
||||
|
||||
class OrchestratorMetrics:
|
||||
def __init__(self, num_stages: int, enable_stats: bool,
|
||||
stats_file: Optional[str], overall_stats_file: Optional[str],
|
||||
wall_start_ts: float) -> None:
|
||||
self.num_stages = int(num_stages)
|
||||
self.enable_stats = bool(enable_stats)
|
||||
self.stats_file = stats_file
|
||||
self.overall_stats_file = overall_stats_file
|
||||
self.stage_total_time_ms: List[float] = [0.0 for _ in range(self.num_stages)]
|
||||
self.stage_total_tokens: List[int] = [0 for _ in range(self.num_stages)]
|
||||
self.stage_req_counts: List[int] = [0 for _ in range(self.num_stages)]
|
||||
self.transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {}
|
||||
self.transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {}
|
||||
self.e2e_total_ms: float = 0.0
|
||||
self.e2e_total_tokens: int = 0
|
||||
self.e2e_count: int = 0
|
||||
self.e2e_done: set[int] = set()
|
||||
self.per_request: Dict[int, Dict[str, Any]] = {}
|
||||
self.sum_per_request_transfer_ms: float = 0.0
|
||||
self.wall_start_ts: float = float(wall_start_ts)
|
||||
self.last_finish_ts: float = float(wall_start_ts)
|
||||
self.stage_seen_batches: Dict[int, set] = {sid: set() for sid in range(self.num_stages)}
|
||||
self.stage_first_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
|
||||
self.stage_last_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
|
||||
|
||||
def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: Dict[str, Any]) -> None:
|
||||
record_stage_metrics(
|
||||
self.per_request,
|
||||
self.stage_req_counts,
|
||||
self.stage_total_time_ms,
|
||||
self.stage_total_tokens,
|
||||
stage_id,
|
||||
req_id,
|
||||
metrics,
|
||||
)
|
||||
try:
|
||||
batch_id_raw = metrics.get("batch_id", None)
|
||||
if batch_id_raw is not None:
|
||||
batch_id = int(batch_id_raw)
|
||||
if batch_id not in self.stage_seen_batches[stage_id]:
|
||||
self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0))
|
||||
self.stage_seen_batches[stage_id].add(batch_id)
|
||||
except Exception:
|
||||
pass
|
||||
rx_b = float(metrics.get("rx_transfer_bytes", 0.0))
|
||||
rx_ms = float(metrics.get("rx_decode_time_ms", 0.0))
|
||||
in_flight_ms = float(metrics.get("rx_in_flight_time_ms", 0.0))
|
||||
combined = aggregate_rx_and_maybe_total(
|
||||
self.transfer_edge_req,
|
||||
self.transfer_agg,
|
||||
self.per_request,
|
||||
stage_id,
|
||||
req_id,
|
||||
rx_b,
|
||||
rx_ms,
|
||||
in_flight_ms,
|
||||
)
|
||||
if self.enable_stats and self.stats_file and stage_id > 0:
|
||||
log_transfer_rx(
|
||||
self.stats_file,
|
||||
stage_id - 1,
|
||||
stage_id,
|
||||
req_id,
|
||||
int(rx_b),
|
||||
float(rx_ms),
|
||||
float(in_flight_ms),
|
||||
)
|
||||
if combined is not None:
|
||||
size_b_c, tx_ms_c, total_ms_c = combined
|
||||
log_transfer_total(
|
||||
self.stats_file,
|
||||
stage_id - 1,
|
||||
stage_id,
|
||||
req_id,
|
||||
int(size_b_c),
|
||||
float(tx_ms_c),
|
||||
float(in_flight_ms),
|
||||
float(rx_ms),
|
||||
float(total_ms_c),
|
||||
)
|
||||
|
||||
def on_forward(self, from_stage: int, to_stage: int, req_id: Any,
|
||||
size_bytes: int, tx_ms: float, used_shm: bool) -> None:
|
||||
# Mark first input time for the destination stage if not set
|
||||
if self.stage_first_ts[to_stage] is None:
|
||||
self.stage_first_ts[to_stage] = time.time()
|
||||
if self.enable_stats and self.stats_file:
|
||||
log_transfer_tx(
|
||||
self.stats_file,
|
||||
from_stage,
|
||||
to_stage,
|
||||
req_id,
|
||||
int(size_bytes),
|
||||
float(tx_ms),
|
||||
bool(used_shm),
|
||||
)
|
||||
record_sender_transfer_agg(
|
||||
self.transfer_agg,
|
||||
self.transfer_edge_req,
|
||||
from_stage,
|
||||
to_stage,
|
||||
req_id,
|
||||
int(size_bytes),
|
||||
float(tx_ms),
|
||||
)
|
||||
|
||||
def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: List[Any],
|
||||
req_start_ts: float) -> None:
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
except Exception:
|
||||
rid_int = req_id
|
||||
_t0 = float(req_start_ts)
|
||||
_t1 = time.time()
|
||||
# Update last output time for this stage
|
||||
prev_last = self.stage_last_ts[stage_id]
|
||||
self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1)
|
||||
self.last_finish_ts = max(self.last_finish_ts, _t1)
|
||||
e2e_ms = (_t1 - _t0) * 1000.0
|
||||
num_tokens = count_tokens_from_outputs(engine_outputs)
|
||||
self.e2e_total_ms += e2e_ms
|
||||
self.e2e_total_tokens += int(num_tokens)
|
||||
self.e2e_count += 1
|
||||
self.e2e_done.add(rid_int)
|
||||
pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
per_req_record = {
|
||||
"type": "overall_request",
|
||||
"request_id": rid_int,
|
||||
"e2e_time_ms": e2e_ms,
|
||||
"num_tokens_out": int(num_tokens),
|
||||
"transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
|
||||
"transfers_total_bytes": int(pr.get("transfers_bytes", 0)),
|
||||
"stages": pr.get("stages", {}),
|
||||
}
|
||||
self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
|
||||
if self.enable_stats and self.overall_stats_file:
|
||||
log_overall_record(self.overall_stats_file, per_req_record)
|
||||
if self.enable_stats and self.stats_file:
|
||||
e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
|
||||
log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
|
||||
|
||||
def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]:
|
||||
# Compute stage summary using wall time between first input and last output per stage
|
||||
stage_summary: List[Dict[str, Any]] = []
|
||||
for sid in range(self.num_stages):
|
||||
first_ts = self.stage_first_ts[sid]
|
||||
last_ts = self.stage_last_ts[sid]
|
||||
total_ms = (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
|
||||
reqs = self.stage_req_counts[sid]
|
||||
tokens = self.stage_total_tokens[sid]
|
||||
avg_req = (total_ms / reqs) if reqs > 0 else 0.0
|
||||
avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
|
||||
stage_summary.append({
|
||||
"stage_id": sid,
|
||||
"requests": int(reqs),
|
||||
"tokens": int(tokens),
|
||||
"total_time_ms": float(total_ms),
|
||||
"avg_time_per_request_ms": float(avg_req),
|
||||
"avg_tokens_per_s": float(avg_tok),
|
||||
})
|
||||
transfer_summary = build_transfer_summary(self.transfer_agg)
|
||||
e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
|
||||
e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
|
||||
wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
|
||||
summary: Dict[str, Any] = {
|
||||
"e2e_requests": int(self.e2e_count),
|
||||
"e2e_total_time_ms": float(wall_time_ms),
|
||||
"e2e_sum_time_ms": float(self.e2e_total_ms),
|
||||
"e2e_total_tokens": int(self.e2e_total_tokens),
|
||||
"e2e_avg_time_per_request_ms": e2e_avg_req,
|
||||
"e2e_avg_tokens_per_s": e2e_avg_tok,
|
||||
"wall_time_ms": wall_time_ms,
|
||||
"final_stage_id": final_stage_id_for_e2e,
|
||||
"stages": stage_summary,
|
||||
"transfers": transfer_summary,
|
||||
}
|
||||
if self.enable_stats and self.stats_file:
|
||||
log_orchestrator_summary(self.stats_file, summary)
|
||||
if self.enable_stats and self.overall_stats_file:
|
||||
log_overall_summary(self.overall_stats_file, summary)
|
||||
return summary
|
||||
|
||||
|
||||
@@ -116,18 +116,18 @@ class OmniStage:
|
||||
if self._in_q is not None:
|
||||
try:
|
||||
self._in_q.put_nowait(None)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
|
||||
if self._proc is not None:
|
||||
try:
|
||||
self._proc.join(timeout=5)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True)
|
||||
if self._proc.is_alive():
|
||||
try:
|
||||
self._proc.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
|
||||
|
||||
def submit(self, payload: Dict[str, Any]) -> None:
|
||||
assert self._in_q is not None
|
||||
@@ -185,6 +185,11 @@ def _stage_worker(
|
||||
"""Stage worker entry: device setup, LLM init, batching, SHM IPC."""
|
||||
import logging as _logging
|
||||
from vllm_omni.entrypoints.omni_llm import OmniStageLLM # noqa: WPS433
|
||||
from vllm_omni.entrypoints.log_utils import ( # noqa: WPS433
|
||||
log_stage_request_stats,
|
||||
log_stage_running_avg,
|
||||
log_stage_batch_stats,
|
||||
)
|
||||
import queue as _queue
|
||||
import time as _time
|
||||
from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl, serialize_obj as _ser
|
||||
@@ -217,6 +222,8 @@ def _stage_worker(
|
||||
# Aggregates for running average
|
||||
_agg_total_tokens = 0
|
||||
_agg_total_gen_time_ms = 0.0
|
||||
# Monotonic batch id per stage process for orchestrator dedup on time aggregation
|
||||
_batch_seq = 0
|
||||
|
||||
# Device mapping
|
||||
try:
|
||||
@@ -294,6 +301,7 @@ def _stage_worker(
|
||||
print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True)
|
||||
print("--------------------------------", flush=True)
|
||||
try:
|
||||
_batch_seq += 1
|
||||
gen_outputs: List[Any] = []
|
||||
_gen_t0 = _time.time()
|
||||
for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False):
|
||||
@@ -355,24 +363,9 @@ def _stage_worker(
|
||||
pass
|
||||
|
||||
if _stats_file:
|
||||
try:
|
||||
_avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
|
||||
_append_jsonl(_stats_file, {
|
||||
"type": "stage_running_avg",
|
||||
"stage_id": stage_id,
|
||||
"total_tokens": int(_agg_total_tokens),
|
||||
"total_gen_time_ms": _agg_total_gen_time_ms,
|
||||
"avg_tokens_per_s": _avg_tokens_per_s,
|
||||
})
|
||||
_append_jsonl(_stats_file, {
|
||||
"type": "stage_batch_stats",
|
||||
"stage_id": stage_id,
|
||||
"batch_size": len(batch_tasks),
|
||||
"batch_gen_time_ms": _gen_ms,
|
||||
"request_ids": list(batch_request_ids),
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
_avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
|
||||
log_stage_running_avg(_stats_file, stage_id, int(_agg_total_tokens), float(_agg_total_gen_time_ms), float(_avg_tokens_per_s))
|
||||
log_stage_batch_stats(_stats_file, stage_id, len(batch_tasks), float(_gen_ms), list(batch_request_ids))
|
||||
|
||||
# Emit per-request results
|
||||
for rid in batch_request_ids:
|
||||
@@ -382,10 +375,29 @@ def _stage_worker(
|
||||
_metrics = {
|
||||
"num_tokens_out": int(_count_tokens(r_outputs)),
|
||||
"stage_gen_time_ms": _gen_ms,
|
||||
"batch_id": int(_batch_seq),
|
||||
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
|
||||
"rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
|
||||
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
|
||||
}
|
||||
if _stats_file:
|
||||
_num_tokens = int(_metrics["num_tokens_out"]) # type: ignore[index]
|
||||
_tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0
|
||||
_rx_b = int(_metrics["rx_transfer_bytes"]) # type: ignore[index]
|
||||
_rx_ms = float(_metrics["rx_decode_time_ms"]) # type: ignore[index]
|
||||
_rx_mbps = (float(_rx_b) * 8.0) / (max(float(_rx_ms), 1e-6) * 1000.0) if _rx_b > 0 else 0.0
|
||||
log_stage_request_stats(
|
||||
_stats_file,
|
||||
stage_id,
|
||||
rid,
|
||||
len(batch_tasks),
|
||||
_num_tokens,
|
||||
float(_gen_ms),
|
||||
float(_tokens_per_s),
|
||||
_rx_b,
|
||||
_rx_ms,
|
||||
float(_rx_mbps),
|
||||
)
|
||||
if use_shm:
|
||||
out_q.put({
|
||||
"request_id": rid,
|
||||
|
||||
@@ -39,7 +39,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
|
||||
"[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
|
||||
stage_id, devices, selected_physical,
|
||||
)
|
||||
except Exception:
|
||||
except Exception as e:
|
||||
logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
|
||||
selected_physical = None
|
||||
elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
|
||||
logical_idx = max(0, int(devices))
|
||||
@@ -49,7 +50,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
|
||||
mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
|
||||
if 0 <= logical_idx < len(mapping):
|
||||
selected_physical = mapping[logical_idx]
|
||||
except Exception:
|
||||
except Exception as e:
|
||||
logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
|
||||
selected_physical = None
|
||||
if selected_physical is None:
|
||||
selected_physical = int(logical_idx)
|
||||
@@ -70,8 +72,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
|
||||
if torch.cuda.is_available():
|
||||
try:
|
||||
torch.cuda.set_device(0)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True)
|
||||
num = torch.cuda.device_count()
|
||||
info = []
|
||||
for i in range(num):
|
||||
@@ -84,8 +86,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
|
||||
"free": int(free),
|
||||
})
|
||||
logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
|
||||
|
||||
@@ -107,8 +109,8 @@ def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
|
||||
meta = {"name": shm.name, "size": len(payload)}
|
||||
try:
|
||||
shm.close()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.debug("Ensure parent dir failed for %s: %s", path, e)
|
||||
return meta
|
||||
|
||||
|
||||
|
||||
@@ -16,6 +16,24 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm_omni.entrypoints.omni_llm import OmniLLM, OmniStageLLM
|
||||
from vllm_omni.entrypoints.omni_stage import OmniStage
|
||||
from vllm_omni.entrypoints.utils import load_stage_configs_from_model
|
||||
from vllm_omni.entrypoints.log_utils import (
|
||||
remove_old_logs,
|
||||
configure_orchestrator_logger,
|
||||
init_stats_paths,
|
||||
record_stage_metrics,
|
||||
aggregate_rx_and_maybe_total,
|
||||
record_sender_transfer_agg,
|
||||
count_tokens_from_outputs,
|
||||
build_stage_summary,
|
||||
build_transfer_summary,
|
||||
log_transfer_tx,
|
||||
log_transfer_rx,
|
||||
log_transfer_total,
|
||||
log_orchestrator_e2e,
|
||||
log_orchestrator_summary,
|
||||
log_overall_summary,
|
||||
OrchestratorMetrics,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.entrypoints.pipeline_utils import (
|
||||
maybe_load_from_ipc as _load,
|
||||
@@ -65,32 +83,10 @@ class PipelinedOmniLLM(OmniLLM):
|
||||
# Optional file handler for orchestrator
|
||||
self._log_file = log_file
|
||||
if self._log_file:
|
||||
try:
|
||||
# Avoid duplicate handlers
|
||||
has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
|
||||
if not has_file_handler:
|
||||
fh = logging.FileHandler(self._log_file)
|
||||
fh.setLevel(logging.DEBUG)
|
||||
fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
|
||||
logger.addHandler(fh)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
except Exception:
|
||||
pass
|
||||
remove_old_logs(self._log_file, len(self.stage_list))
|
||||
configure_orchestrator_logger(logger, self._log_file)
|
||||
|
||||
# Orchestrator stats JSONL file
|
||||
self._stats_file: Optional[str] = None
|
||||
if self._enable_stats and self._log_file:
|
||||
try:
|
||||
self._stats_file = f"{self._log_file}.orchestrator.stats.jsonl"
|
||||
except Exception:
|
||||
self._stats_file = None
|
||||
# Overall stats JSONL file (per-request + summary)
|
||||
self._overall_stats_file: Optional[str] = None
|
||||
if self._enable_stats and self._log_file:
|
||||
try:
|
||||
self._overall_stats_file = f"{self._log_file}.overall.stats.jsonl"
|
||||
except Exception:
|
||||
self._overall_stats_file = None
|
||||
self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
|
||||
|
||||
self._init_sleep_seconds = max(0, int(init_sleep_seconds))
|
||||
self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
|
||||
@@ -122,19 +118,19 @@ class PipelinedOmniLLM(OmniLLM):
|
||||
for q in self._stage_in_queues:
|
||||
try:
|
||||
q.put_nowait(None)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e)
|
||||
for stage in self.stage_list:
|
||||
try:
|
||||
stage.stop_stage_worker()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
|
||||
|
||||
def __del__(self) -> None: # best-effort
|
||||
try:
|
||||
self.close()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
@@ -178,25 +174,16 @@ class PipelinedOmniLLM(OmniLLM):
|
||||
final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid)
|
||||
if final_stage_id_for_e2e < 0:
|
||||
final_stage_id_for_e2e = len(self.stage_list) - 1
|
||||
except Exception:
|
||||
except Exception as e:
|
||||
logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True)
|
||||
final_stage_id_for_e2e = len(self.stage_list) - 1
|
||||
# In-memory aggregators for this generate() call
|
||||
stage_total_time_ms: List[float] = [0.0 for _ in range(num_stages)]
|
||||
stage_total_tokens: List[int] = [0 for _ in range(num_stages)]
|
||||
stage_req_counts: List[int] = [0 for _ in range(num_stages)]
|
||||
transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {}
|
||||
# Per-edge per-request sender timing to combine with receiver timing later
|
||||
transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {}
|
||||
e2e_total_ms: float = 0.0
|
||||
e2e_total_tokens: int = 0
|
||||
e2e_count: int = 0
|
||||
e2e_done: set[int] = set()
|
||||
# Per-request overall aggregation
|
||||
per_request: Dict[int, Dict[str, Any]] = {}
|
||||
sum_per_request_transfer_ms: float = 0.0
|
||||
# Metrics/aggregation helper
|
||||
metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts)
|
||||
|
||||
# Seed stage-0 queue with all requests
|
||||
logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
|
||||
# Mark first input time for stage-0
|
||||
metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
|
||||
|
||||
for req_id, prompt in request_id_to_prompt.items():
|
||||
sp0: SamplingParams = sampling_params_list[0] # type: ignore[index]
|
||||
@@ -235,94 +222,14 @@ class PipelinedOmniLLM(OmniLLM):
|
||||
continue
|
||||
|
||||
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
|
||||
# Aggregate per-request metrics from stage worker if present
|
||||
# Mark last output time for this stage whenever we receive outputs
|
||||
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
|
||||
try:
|
||||
_m = result.get("metrics")
|
||||
if _m is not None:
|
||||
stage_req_counts[stage_id] += 1
|
||||
stage_total_time_ms[stage_id] += float(_m.get("stage_gen_time_ms", 0.0))
|
||||
stage_total_tokens[stage_id] += int(_m.get("num_tokens_out", 0))
|
||||
# record per-request stage metrics
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
pr_stages = pr["stages"] # type: ignore[index]
|
||||
pr_stages[stage_id] = {
|
||||
"stage_gen_time_ms": float(_m.get("stage_gen_time_ms", 0.0)),
|
||||
"num_tokens_out": int(_m.get("num_tokens_out", 0)),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
# Also aggregate receiver-side transfer decode time for (stage_id-1 -> stage_id)
|
||||
try:
|
||||
if stage_id > 0:
|
||||
key = (stage_id - 1, stage_id)
|
||||
agg = transfer_agg.get(key)
|
||||
if agg is None:
|
||||
agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
|
||||
"sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
|
||||
"sum_total_ms": 0.0, "total_count": 0.0}
|
||||
transfer_agg[key] = agg
|
||||
rx_b = float(_m.get("rx_transfer_bytes", 0.0))
|
||||
rx_ms = float(_m.get("rx_decode_time_ms", 0.0))
|
||||
in_flight_ms = float(_m.get("rx_in_flight_time_ms", 0.0))
|
||||
agg["sum_rx_bytes"] += rx_b
|
||||
agg["sum_rx_ms"] += rx_ms
|
||||
agg["rx_count"] += 1.0
|
||||
# If we have sender-side timing stored, compute combined (encode+enqueue + decode)
|
||||
try:
|
||||
edge_req = (stage_id - 1, stage_id, int(req_id))
|
||||
s = transfer_edge_req.get(edge_req)
|
||||
if s is not None:
|
||||
total_ms = float(s.get("tx_ms", 0.0)) + in_flight_ms + rx_ms
|
||||
agg["sum_total_ms"] += total_ms
|
||||
agg["total_count"] += 1.0
|
||||
# accumulate per-request transfer totals
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index]
|
||||
pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_b) # type: ignore[index]
|
||||
except Exception:
|
||||
pass
|
||||
if getattr(self, "_stats_file", None):
|
||||
try:
|
||||
size_b = float(s.get("size_bytes", rx_b))
|
||||
_append_jsonl(self._stats_file, { # type: ignore[arg-type]
|
||||
"type": "transfer_total_stats",
|
||||
"from_stage": stage_id - 1,
|
||||
"to_stage": stage_id,
|
||||
"request_id": req_id,
|
||||
"size_bytes": int(size_b),
|
||||
"tx_time_ms": float(s.get("tx_ms", 0.0)),
|
||||
"in_flight_time_ms": in_flight_ms,
|
||||
"rx_decode_time_ms": rx_ms,
|
||||
"total_time_ms": total_ms,
|
||||
"total_time_per_kb_ms": total_ms / max(size_b / 1024.0, 1e-6) if size_b > 0 else 0.0,
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
# Emit per-request RX stats to JSONL
|
||||
if getattr(self, "_stats_file", None):
|
||||
try:
|
||||
_append_jsonl(self._stats_file, { # type: ignore[arg-type]
|
||||
"type": "transfer_rx_stats",
|
||||
"from_stage": stage_id - 1,
|
||||
"to_stage": stage_id,
|
||||
"request_id": req_id,
|
||||
"rx_bytes": int(rx_b),
|
||||
"rx_decode_time_ms": rx_ms,
|
||||
"in_flight_time_ms": in_flight_ms,
|
||||
"rx_time_per_kb_ms": rx_ms / max(rx_b / 1024.0, 1e-6) if rx_b > 0 else 0.0,
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
metrics.on_stage_metrics(stage_id, req_id, _m)
|
||||
except Exception as e:
|
||||
logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e)
|
||||
logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id)
|
||||
stage.set_engine_outputs(engine_outputs)
|
||||
|
||||
@@ -338,67 +245,11 @@ class PipelinedOmniLLM(OmniLLM):
|
||||
|
||||
# End-to-end timing and time-per-token for final output (only once per request at the designated final stage)
|
||||
try:
|
||||
if stage_id == final_stage_id_for_e2e and req_id not in e2e_done:
|
||||
_t0 = _req_start_ts.get(req_id)
|
||||
if _t0 is not None:
|
||||
_t1 = time.time()
|
||||
_last_finish_ts = max(_last_finish_ts, _t1)
|
||||
_e2e_ms = (_t1 - _t0) * 1000.0
|
||||
# Count tokens from final stage outputs
|
||||
def _count_tokens(_ros: List[Any]) -> int: # type: ignore[name-defined]
|
||||
total = 0
|
||||
for _ro in _ros:
|
||||
try:
|
||||
outs = getattr(_ro, "outputs", None)
|
||||
if outs and len(outs) > 0:
|
||||
tokens = getattr(outs[0], "token_ids", None)
|
||||
if tokens is not None:
|
||||
total += len(tokens)
|
||||
except Exception:
|
||||
pass
|
||||
return total
|
||||
_num_tokens = _count_tokens(engine_outputs)
|
||||
_time_per_token_ms = (_e2e_ms / _num_tokens) if _num_tokens > 0 else 0.0
|
||||
# Update E2E aggregators
|
||||
e2e_total_ms += _e2e_ms
|
||||
e2e_total_tokens += int(_num_tokens)
|
||||
e2e_count += 1
|
||||
e2e_done.add(req_id)
|
||||
# Write per-request overall record
|
||||
try:
|
||||
rid_int = int(req_id)
|
||||
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
|
||||
per_req_record = {
|
||||
"type": "overall_request",
|
||||
"request_id": rid_int,
|
||||
"e2e_time_ms": _e2e_ms,
|
||||
"num_tokens_out": int(_num_tokens),
|
||||
"transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
|
||||
"transfers_total_bytes": int(pr.get("transfers_bytes", 0)),
|
||||
"stages": pr.get("stages", {}),
|
||||
}
|
||||
sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
|
||||
if getattr(self, "_overall_stats_file", None):
|
||||
try:
|
||||
_append_jsonl(self._overall_stats_file, per_req_record) # type: ignore[arg-type]
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
if getattr(self, "_stats_file", None):
|
||||
try:
|
||||
_append_jsonl(self._stats_file, { # type: ignore[arg-type]
|
||||
"type": "orchestrator_request_e2e",
|
||||
"request_id": req_id,
|
||||
"final_stage_id": stage_id,
|
||||
"e2e_time_ms": _e2e_ms,
|
||||
"num_tokens_out": int(_num_tokens),
|
||||
"e2e_time_per_token_ms": _time_per_token_ms,
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
|
||||
if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
|
||||
metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts))
|
||||
except Exception as e:
|
||||
logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e)
|
||||
|
||||
next_stage_id = stage_id + 1
|
||||
if next_stage_id < num_stages:
|
||||
@@ -427,40 +278,9 @@ class PipelinedOmniLLM(OmniLLM):
|
||||
self.stage_list[next_stage_id].submit(ipc_payload)
|
||||
t1 = time.time()
|
||||
tx_ms = (t1 - t0) * 1000.0
|
||||
if self._enable_stats and getattr(self, "_stats_file", None):
|
||||
try:
|
||||
_append_jsonl(self._stats_file, {
|
||||
"type": "transfer_stats",
|
||||
"from_stage": stage_id,
|
||||
"to_stage": next_stage_id,
|
||||
"request_id": req_id,
|
||||
"size_bytes": int(size_bytes),
|
||||
"tx_time_ms": tx_ms,
|
||||
"tx_mbps": (float(size_bytes) * 8.0) / (max(tx_ms, 1e-6) * 1000.0),
|
||||
"used_shm": bool("engine_inputs_shm" in ipc_payload),
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
# Update in-memory transfer aggregator
|
||||
try:
|
||||
key = (stage_id, next_stage_id)
|
||||
agg = transfer_agg.get(key)
|
||||
if agg is None:
|
||||
agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
|
||||
"sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
|
||||
"sum_total_ms": 0.0, "total_count": 0.0}
|
||||
transfer_agg[key] = agg
|
||||
agg["sum_bytes"] += float(size_bytes)
|
||||
agg["sum_ms"] += float(tx_ms)
|
||||
agg["count"] += 1.0
|
||||
# Store sender-side timing for per-request combination
|
||||
transfer_edge_req[(stage_id, next_stage_id, int(req_id))] = {
|
||||
"tx_ms": float(tx_ms),
|
||||
"size_bytes": float(size_bytes),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
metrics.on_forward(stage_id, next_stage_id, req_id, int(size_bytes), float(tx_ms), bool("engine_inputs_shm" in ipc_payload))
|
||||
except Exception as e:
|
||||
logger.warning("[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", req_id, e)
|
||||
self.stage_list[next_stage_id].submit({
|
||||
"request_id": req_id,
|
||||
"engine_inputs": next_inputs,
|
||||
@@ -478,81 +298,10 @@ class PipelinedOmniLLM(OmniLLM):
|
||||
|
||||
# Summarize and print stats
|
||||
try:
|
||||
stage_summary: List[Dict[str, Any]] = []
|
||||
for sid in range(num_stages):
|
||||
reqs = stage_req_counts[sid]
|
||||
tokens = stage_total_tokens[sid]
|
||||
total_ms = float(stage_total_time_ms[sid])
|
||||
avg_req = (total_ms / reqs) if reqs > 0 else 0.0
|
||||
avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
|
||||
stage_summary.append({
|
||||
"stage_id": sid,
|
||||
"requests": int(reqs),
|
||||
"tokens": int(tokens),
|
||||
"total_time_ms": total_ms,
|
||||
"avg_time_per_request_ms": avg_req,
|
||||
"avg_tokens_per_s": avg_tok,
|
||||
})
|
||||
|
||||
transfer_summary: List[Dict[str, Any]] = []
|
||||
for (src, dst), agg in transfer_agg.items():
|
||||
sum_bytes = float(agg.get("sum_bytes", 0.0))
|
||||
sum_ms = float(agg.get("sum_ms", 0.0))
|
||||
samples = int(agg.get("count", 0.0))
|
||||
tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
|
||||
sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
|
||||
sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
|
||||
samples_rx = int(agg.get("rx_count", 0.0))
|
||||
rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
|
||||
sum_total_ms = float(agg.get("sum_total_ms", 0.0))
|
||||
samples_total = int(agg.get("total_count", 0.0))
|
||||
total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
|
||||
transfer_summary.append({
|
||||
"from_stage": src,
|
||||
"to_stage": dst,
|
||||
"samples": samples,
|
||||
"total_bytes": int(sum_bytes),
|
||||
"total_time_ms": sum_ms,
|
||||
"tx_mbps": tx_mbps,
|
||||
"rx_samples": samples_rx,
|
||||
"rx_total_bytes": int(sum_rx_bytes),
|
||||
"rx_total_time_ms": sum_rx_ms,
|
||||
"rx_mbps": rx_mbps,
|
||||
"total_samples": samples_total,
|
||||
"total_transfer_time_ms": sum_total_ms,
|
||||
"total_mbps": total_mbps,
|
||||
})
|
||||
|
||||
e2e_avg_req = (e2e_total_ms / e2e_count) if e2e_count > 0 else 0.0
|
||||
e2e_avg_tok = (e2e_total_tokens * 1000.0 / e2e_total_ms) if e2e_total_ms > 0 else 0.0
|
||||
wall_time_ms = max(0.0, (_last_finish_ts - _wall_start_ts) * 1000.0)
|
||||
summary: Dict[str, Any] = {
|
||||
"e2e_requests": int(e2e_count),
|
||||
# 按你的需求：e2e_total_time_ms 直接使用墙钟时间（非各请求之和）
|
||||
"e2e_total_time_ms": float(wall_time_ms),
|
||||
# 额外保留各请求 E2E 之和，供需要时参考
|
||||
"e2e_sum_time_ms": float(e2e_total_ms),
|
||||
"e2e_total_tokens": int(e2e_total_tokens),
|
||||
"e2e_avg_time_per_request_ms": e2e_avg_req,
|
||||
"e2e_avg_tokens_per_s": e2e_avg_tok,
|
||||
"wall_time_ms": wall_time_ms,
|
||||
"final_stage_id": final_stage_id_for_e2e,
|
||||
"stages": stage_summary,
|
||||
"transfers": transfer_summary,
|
||||
}
|
||||
summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
|
||||
logger.info("[Summary] %s", summary)
|
||||
if self._enable_stats and getattr(self, "_stats_file", None):
|
||||
try:
|
||||
_append_jsonl(self._stats_file, {"type": "orchestrator_summary", **summary}) # type: ignore[arg-type]
|
||||
except Exception:
|
||||
pass
|
||||
if self._enable_stats and getattr(self, "_overall_stats_file", None):
|
||||
try:
|
||||
_append_jsonl(self._overall_stats_file, {"type": "overall_summary", **summary}) # type: ignore[arg-type]
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
|
||||
|
||||
return final_outputs
|
||||
|
||||
|
||||
@@ -126,7 +126,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
src = pe_cpu[
|
||||
num_computed_tokens : num_computed_tokens + overlay_len
|
||||
].to(dtype=self.dtype, device=self.device, non_blocking=True)
|
||||
start_offset = int(self.query_start_loc_cpu[req_index])
|
||||
start_offset = int(self.query_start_loc.cpu[req_index])
|
||||
self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
|
||||
src
|
||||
)
|
||||
@@ -381,7 +381,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
# Pass per-request additional information map for this step (no concat)
|
||||
if per_req_additional_information:
|
||||
model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
|
||||
# 始终传递每请求的运行期 additional_information（持久化在 request state）
|
||||
# Always pass per-request runtime additional_information (persisted in request state)
|
||||
try:
|
||||
per_req_runtime_info = []
|
||||
for req_id in self.input_batch.req_ids:
|
||||
@@ -396,10 +396,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
per_req_runtime_info
|
||||
)
|
||||
model_kwargs_extra["request_ids"] = self.input_batch.req_ids
|
||||
# 传递每个请求在本步扁平序列中的 token 区间，便于模型按请求映射 decode/prefill
|
||||
# Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
|
||||
req_token_spans = []
|
||||
for req_index in range(len(self.input_batch.req_ids)):
|
||||
start_offset = int(self.query_start_loc_cpu[req_index])
|
||||
start_offset = int(self.query_start_loc.cpu[req_index])
|
||||
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
||||
req_token_spans.append((start_offset, start_offset + sched_tokens))
|
||||
model_kwargs_extra["request_token_spans"] = req_token_spans
|
||||
@@ -429,9 +429,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
|
||||
hidden_states
|
||||
)
|
||||
# 模型侧可能返回 per-request 的 additional_information 更新（与模型无关的通用通道）。
|
||||
# 约定：multimodal_outputs["additional_information_update"] 为按 batch 顺序的 list[dict]，
|
||||
# runner 将其合并到对应 request 的 additional_information_cpu，供后续 decode 使用。
|
||||
# The model side may return per-request additional_information updates (model-agnostic channel).
|
||||
# Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
|
||||
# the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
|
||||
try:
|
||||
if (
|
||||
isinstance(multimodal_outputs, dict)
|
||||
@@ -440,7 +440,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
or "additional_information_update_by_req_id" in multimodal_outputs
|
||||
)
|
||||
):
|
||||
# 方案 A：按 batch 顺序的 list[dict]
|
||||
# Option A: list[dict] in batch order
|
||||
updates_list = multimodal_outputs.get(
|
||||
"additional_information_update"
|
||||
)
|
||||
@@ -450,7 +450,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
continue
|
||||
req_id = self.input_batch.req_ids[idx]
|
||||
self._merge_additional_information_update(req_id, upd)
|
||||
# 方案 B：按 req_id 的 dict[str, dict]
|
||||
# Option B: dict[str, dict] keyed by req_id
|
||||
updates_map = multimodal_outputs.get(
|
||||
"additional_information_update_by_req_id"
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user