"""Hang risk scoring based on timeout thresholds.
Analyzes collective operations to identify those that are approaching
timeout thresholds. This enables proactive detection of potential hangs
before they cause training failures.
Risk Levels:
- Normal: < 1% of timeout
- Elevated: 1-10% of timeout
- Warning: 10-50% of timeout
- Critical: > 50% of timeout
Default timeout reference: TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC = 600s
"""
from dataclasses import dataclass
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import Collective, CollectiveType
from wafer_core.lib.distributed_traces.models.trace_session import (
    HangRiskLevel,
    HangRiskReport,
    TraceSession,
)
# Default PyTorch NCCL timeout
DEFAULT_TIMEOUT_SEC = 600
DEFAULT_TIMEOUT_NS = DEFAULT_TIMEOUT_SEC * 1_000_000_000
@dataclass
class RiskThresholds:
    """Thresholds for risk level classification.
    Attributes:
        elevated_pct: Threshold for elevated risk (default 1%)
        warning_pct: Threshold for warning risk (default 10%)
        critical_pct: Threshold for critical risk (default 50%)
    """
    elevated_pct: float = 1.0
    warning_pct: float = 10.0
    critical_pct: float = 50.0
class HangRiskAnalyzer:
    """Analyzes hang risk for collective operations."""
    def __init__(
        self,
        session: TraceSession,
        timeout_ns: int = DEFAULT_TIMEOUT_NS,
        thresholds: Optional[RiskThresholds] = None,
    ) -> None:
        """Initialize the analyzer.
        Args:
            session: TraceSession to analyze
            timeout_ns: Timeout threshold in nanoseconds
            thresholds: Risk level thresholds
        """
        self.session = session
        self.timeout_ns = timeout_ns
        self.thresholds = thresholds or RiskThresholds()
    def analyze(self) -> HangRiskReport:
        """Perform hang risk analysis.
        Returns:
            HangRiskReport with risk assessment
        """
        collectives_by_risk: dict[HangRiskLevel, list[Collective]] = {
            level: [] for level in HangRiskLevel
        }
        highest_risk: Optional[Collective] = None
        highest_risk_pct = 0.0
        for timeline in self.session.timelines.values():
            for collective in timeline.collectives:
                risk_level, risk_pct = self._assess_risk(collective)
                collectives_by_risk[risk_level].append(collective)
                if risk_pct > highest_risk_pct:
                    highest_risk_pct = risk_pct
                    highest_risk = collective

        risk_distribution = {
            level: len(collectives) for level, collectives in collectives_by_risk.items()
        }
        # Generate recommendations
        recommendations = self._generate_recommendations(collectives_by_risk)
        return HangRiskReport(
            timeout_threshold_ns=self.timeout_ns,
            collectives_by_risk=collectives_by_risk,
            highest_risk_collective=highest_risk,
            risk_distribution=risk_distribution,
            recommendations=recommendations,
        )
    def _assess_risk(self, collective: Collective) -> tuple[HangRiskLevel, float]:
        """Assess hang risk for a single collective.
        Args:
            collective: Collective to assess
        Returns:
            Tuple of (HangRiskLevel, risk_percentage)
        """
        duration_pct = (collective.duration_ns / self.timeout_ns) * 100
        if duration_pct >= self.thresholds.critical_pct:
            return HangRiskLevel.CRITICAL, duration_pct
        elif duration_pct >= self.thresholds.warning_pct:
            return HangRiskLevel.WARNING, duration_pct
        elif duration_pct >= self.thresholds.elevated_pct:
            return HangRiskLevel.ELEVATED, duration_pct
        else:
            return HangRiskLevel.NORMAL, duration_pct
    def get_risk_score(self, collective: Collective) -> float:
        """Get risk score as percentage of timeout.
        Args:
            collective: Collective to score
        Returns:
            Risk score as percentage (0-100+)
        """
        return (collective.duration_ns / self.timeout_ns) * 100
    def _generate_recommendations(
        self,
        collectives_by_risk: dict[HangRiskLevel, list[Collective]],
    ) -> list[str]:
        """Generate recommendations based on risk analysis.
        Args:
            collectives_by_risk: Collectives grouped by risk level
        Returns:
            List of recommendation strings
        """
        recommendations = []
        critical_count = len(collectives_by_risk[HangRiskLevel.CRITICAL])
        warning_count = len(collectives_by_risk[HangRiskLevel.WARNING])
        elevated_count = len(collectives_by_risk[HangRiskLevel.ELEVATED])
        if critical_count > 0:
            recommendations.append(
                f"CRITICAL: {critical_count} collective(s) exceeded 50% of timeout threshold. "
                "Immediate investigation required."
            )
            # Analyze patterns in critical collectives
            critical_types = self._analyze_collective_patterns(
                collectives_by_risk[HangRiskLevel.CRITICAL]
            )
            if critical_types:
                recommendations.append(
                    f"  Most affected operations: {', '.join(critical_types)}"
                )
        if warning_count > 0:
            recommendations.append(
                f"WARNING: {warning_count} collective(s) between 10-50% of timeout. "
                "Consider investigating root cause."
            )
        if critical_count > 0 or warning_count > 0:
            recommendations.extend([
                "",
                "Potential causes and mitigations:",
                "  1. Network congestion - Check for competing traffic",
                "  2. GPU memory pressure - Reduce batch size or model size",
                "  3. Uneven workload - Check for data imbalance across ranks",
                "  4. Checkpoint I/O blocking - Move checkpointing off critical path",
                "",
                "Environment variable tuning:",
                f"  - Current timeout: {self.timeout_ns // 1_000_000_000}s",
                "  - Increase with: TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=1200",
                "  - Enable async error handling: TORCH_NCCL_ASYNC_ERROR_HANDLING=1",
            ])
        if elevated_count > 5:
            recommendations.append(
                f"Note: {elevated_count} collectives have elevated duration (1-10% of timeout). "
                "This may indicate network performance issues."
            )
        if not recommendations:
            recommendations.append(
                "All collectives are within normal duration thresholds."
            )
        return recommendations
    def _analyze_collective_patterns(
        self,
        collectives: list[Collective],
    ) -> list[str]:
        """Analyze patterns in problematic collectives.
        Args:
            collectives: List of collectives to analyze
        Returns:
            List of collective type names that appear most frequently
        """
        type_counts: dict[CollectiveType, int] = {}
        for c in collectives:
            type_counts.setdefault(c.collective_type, 0)
            type_counts[c.collective_type] += 1

        sorted_types = sorted(type_counts.items(), key=lambda x: x[1], reverse=True)
        return [t.value for t, _ in sorted_types[:3]]
def analyze_hang_risk(
    session: TraceSession,
    timeout_sec: float = DEFAULT_TIMEOUT_SEC,
    thresholds: Optional[RiskThresholds] = None,
) -> HangRiskReport:
    """Convenience function to analyze hang risk in a session.
    Args:
        session: TraceSession to analyze
        timeout_sec: Timeout threshold in seconds
        thresholds: Optional custom risk thresholds
    Returns:
        HangRiskReport
    """
    timeout_ns = int(timeout_sec * 1_000_000_000)
    analyzer = HangRiskAnalyzer(session, timeout_ns, thresholds)
    return analyzer.analyze()
