"""Straggler detection and scoring.
A "straggler" is a rank that consistently takes longer than others for
collective operations, causing all other ranks to wait. This module
detects stragglers and quantifies their impact on training throughput.
Straggler Score = (rank_duration - median_duration) * (num_ranks - 1)
This represents the total wait time imposed on other ranks by this
rank being slow.
"""
from dataclasses import dataclass
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import CollectiveType
from wafer_core.lib.distributed_traces.models.trace_session import (
    CollectiveMatch,
    StragglerReport,
    TraceSession,
)
@dataclass
class RankStragglerStats:
    """Per-rank straggler statistics.
    Attributes:
        rank: The rank number
        total_impact_ns: Total straggler impact in nanoseconds
        num_straggler_events: Number of times this rank was the straggler
        avg_delay_ns: Average delay when this rank was straggler
        max_delay_ns: Maximum delay caused
        primary_collective_type: Collective type where straggling is most common
    """
    rank: int
    total_impact_ns: int = 0
    num_straggler_events: int = 0
    avg_delay_ns: int = 0
    max_delay_ns: int = 0
    primary_collective_type: Optional[CollectiveType] = None
class StragglerAnalyzer:
    """Detects and analyzes straggler behavior in distributed training."""
    def __init__(
        self,
        session: TraceSession,
        matches: list[CollectiveMatch],
        threshold_std_devs: float = 2.0,
    ) -> None:
        """Initialize the analyzer.
        Args:
            session: TraceSession being analyzed
            matches: Pre-computed collective matches across ranks
            threshold_std_devs: Standard deviations above mean to flag straggler
        """
        self.session = session
        self.matches = matches
        self.threshold_std_devs = threshold_std_devs
    def analyze(self) -> StragglerReport:
        """Perform straggler analysis.
        Returns:
            StragglerReport with per-rank scores and flagged stragglers
        """
        rank_stats: dict[int, RankStragglerStats] = {
            rank: RankStragglerStats(rank=rank) for rank in self.session.ranks
        }

        type_counts: dict[int, dict[CollectiveType, int]] = {
            rank: {} for rank in self.session.ranks
        }
        top_impacts: list[tuple[CollectiveMatch, int, int]] = []  # (match, rank, impact)
        for match in self.matches:
            if match.num_ranks < 2:
                continue

            straggler_rank = match.straggler_rank
            impact = match.straggler_impact_ns
            if impact <= 0:
                continue

            stats = rank_stats[straggler_rank]
            stats.total_impact_ns += impact
            stats.num_straggler_events += 1
            if impact > stats.max_delay_ns:
                stats.max_delay_ns = impact

            type_counts[straggler_rank].setdefault(match.collective_type, 0)
            type_counts[straggler_rank][match.collective_type] += 1

            top_impacts.append((match, straggler_rank, impact))
        # Finalize stats
        for rank, stats in rank_stats.items():
            if stats.num_straggler_events > 0:
                stats.avg_delay_ns = stats.total_impact_ns // stats.num_straggler_events

            if type_counts[rank]:
                stats.primary_collective_type = max(
                    type_counts[rank].items(),
                    key=lambda x: x[1],
                )[0]
        rank_scores = {rank: float(stats.total_impact_ns) for rank, stats in rank_stats.items()}

        flagged_ranks = self._identify_flagged_ranks(rank_scores)
        top_impacts.sort(key=lambda x: x[2], reverse=True)
        top_collectives = [(m, r) for m, r, _ in top_impacts[:10]]
        # Total impact
        total_impact = sum(stats.total_impact_ns for stats in rank_stats.values())
        is_consistent = len(flagged_ranks) == 1 and flagged_ranks[0] in [
            r for r, s in rank_stats.items() if s.num_straggler_events > len(self.matches) * 0.5
        ]
        return StragglerReport(
            rank_scores=rank_scores,
            flagged_ranks=flagged_ranks,
            top_collectives=top_collectives,
            total_straggler_impact_ns=total_impact,
            is_consistent=is_consistent,
        )
    def _identify_flagged_ranks(self, rank_scores: dict[int, float]) -> list[int]:
        """Identify ranks that exceed the straggler threshold.
        Args:
            rank_scores: Per-rank cumulative straggler scores
        Returns:
            List of rank numbers that are flagged as stragglers
        """
        if not rank_scores:
            return []
        scores = list(rank_scores.values())
        if len(scores) < 2:
            return []
        mean_score = sum(scores) / len(scores)
        variance = sum((s - mean_score) ** 2 for s in scores) / len(scores)
        std_dev = variance ** 0.5
        if std_dev == 0:
            return []
        # Flag ranks exceeding threshold
        threshold = mean_score + self.threshold_std_devs * std_dev
        flagged = [rank for rank, score in rank_scores.items() if score > threshold]
        return sorted(flagged)
    def get_straggler_summary(self) -> dict[str, object]:
        """Get a summary of straggler behavior.
        Returns:
            Dictionary with summary statistics
        """
        report = self.analyze()
        if not report.rank_scores:
            return {
                "has_stragglers": False,
                "num_flagged_ranks": 0,
                "total_impact_ms": 0,
            }
        return {
            "has_stragglers": len(report.flagged_ranks) > 0,
            "num_flagged_ranks": len(report.flagged_ranks),
            "flagged_ranks": report.flagged_ranks,
            "total_impact_ms": report.total_straggler_impact_ns / 1_000_000,
            "is_consistent": report.is_consistent,
            "worst_rank": max(report.rank_scores.items(), key=lambda x: x[1])[0]
            if report.rank_scores
            else None,
        }
def detect_stragglers(
    session: TraceSession,
    matches: Optional[list[CollectiveMatch]] = None,
    threshold_std_devs: float = 2.0,
) -> StragglerReport:
    """Convenience function to detect stragglers in a session.
    Args:
        session: TraceSession to analyze
        matches: Pre-computed matches (computed if not provided)
        threshold_std_devs: Standard deviations for flagging
    Returns:
        StragglerReport
    """
    if matches is None:
        from wafer_core.lib.distributed_traces.correlation import match_collectives_across_ranks
        matches = match_collectives_across_ranks(session)
    analyzer = StragglerAnalyzer(session, matches, threshold_std_devs)
    return analyzer.analyze()
