"""NVIDIA NCCL Inspector and Nsight Systems collection setup.
Sets up environment variables for collecting NCCL communication traces:
- NCCL Inspector: Lightweight always-on tracing
- Nsight Systems: Full system profiling with NCCL events
Reference:
- NCCL Inspector: https://developer.nvidia.com/blog/enhancing-communication-observability-of-ai-workloads-with-nccl-inspector/
- NCCL env vars: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html
"""
import os
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
@dataclass
class NCCLInspectorConfig:
    """Configuration for NCCL Inspector collection.
    Attributes:
        output_dir: Directory for output files
        output_format: Output format (jsonl, log)
        verbosity: Verbosity level (0-3)
        enable_timestamps: Whether to include timestamps
    """
    output_dir: str = "./nccl_inspector_traces"
    output_format: str = "jsonl"
    verbosity: int = 1
    enable_timestamps: bool = True
@dataclass
class NsysConfig:
    """Configuration for Nsight Systems collection.
    Attributes:
        output_path: Output file path (without extension)
        trace_options: What to trace (cuda, nccl, nvtx, osrt)
        stats: Whether to generate statistics
        cuda_graph_trace: Whether to trace CUDA graphs
        capture_range: Capture range specification
    """
    output_path: str = "./nsys_trace"
    trace_options: list[str] = field(default_factory=lambda: ["cuda", "nccl", "nvtx"])
    stats: bool = True
    cuda_graph_trace: bool = False
    capture_range: Optional[str] = None  # e.g., "cudaProfilerApi"
@dataclass
class CollectionResult:
    """Result of trace collection.
    Attributes:
        success: Whether collection completed successfully
        output_files: List of generated trace files
        command: Command that was executed
        stdout: Standard output
        stderr: Standard error
        error: Error message if failed
    """
    success: bool
    output_files: list[str] = field(default_factory=list)
    command: Optional[list[str]] = None
    stdout: Optional[str] = None
    stderr: Optional[str] = None
    error: Optional[str] = None
class NVIDIACollector:
    """Collects NCCL traces on NVIDIA GPUs."""
    def __init__(
        self,
        nccl_config: Optional[NCCLInspectorConfig] = None,
        nsys_config: Optional[NsysConfig] = None,
    ) -> None:
        """Initialize the collector.
        Args:
            nccl_config: NCCL Inspector configuration
            nsys_config: Nsight Systems configuration
        """
        self.nccl_config = nccl_config or NCCLInspectorConfig()
        self.nsys_config = nsys_config or NsysConfig()
    def get_env_vars(
        self,
        enable_nccl_inspector: bool = True,
        enable_pytorch_profiler: bool = True,
    ) -> dict[str, str]:
        """Get environment variables for trace collection.
        Args:
            enable_nccl_inspector: Enable NCCL Inspector
            enable_pytorch_profiler: Enable PyTorch profiler env vars
        Returns:
            Dictionary of environment variables to set
        """
        env = {}
        if enable_nccl_inspector:
            output_dir = Path(self.nccl_config.output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)
            # NCCL Inspector environment variables
            env["NCCL_PROFILER_PLUGIN"] = "inspector"
            env["NCCL_INSPECTOR_OUTPUT_DIR"] = str(output_dir.absolute())
            env["NCCL_INSPECTOR_FORMAT"] = self.nccl_config.output_format
            env["NCCL_INSPECTOR_VERBOSITY"] = str(self.nccl_config.verbosity)
            if self.nccl_config.enable_timestamps:
                env["NCCL_INSPECTOR_TIMESTAMPS"] = "1"
        if enable_pytorch_profiler:
            # PyTorch trace output
            env["TORCH_TRACE_OUTPUT_DIR"] = str(
                Path(self.nccl_config.output_dir).absolute()
            )
        return env
    def build_nsys_command(
        self,
        target_command: list[str],
    ) -> list[str]:
        """Build Nsight Systems profiling command.
        Args:
            target_command: Command to profile
        Returns:
            Full nsys command with options
        """
        cmd = ["nsys", "profile"]
        # Output file
        cmd.extend(["-o", self.nsys_config.output_path])
        # Trace options
        cmd.extend(["-t", ",".join(self.nsys_config.trace_options)])
        # Stats
        if self.nsys_config.stats:
            cmd.append("--stats=true")
        # CUDA graph trace
        if self.nsys_config.cuda_graph_trace:
            cmd.append("--cuda-graph-trace=node")
        # Capture range
        if self.nsys_config.capture_range:
            cmd.extend(["--capture-range", self.nsys_config.capture_range])
        cmd.append("--")
        cmd.extend(target_command)
        return cmd
    def collect_with_inspector(
        self,
        target_command: list[str],
        working_dir: Optional[str] = None,
    ) -> CollectionResult:
        """Run collection with NCCL Inspector.
        Args:
            target_command: Command to run with tracing
            working_dir: Working directory for command
        Returns:
            CollectionResult with output files
        """
        env = {**os.environ, **self.get_env_vars()}
        try:
            result = subprocess.run(
                target_command,
                env=env,
                cwd=working_dir,
                capture_output=True,
                text=True,
            )

            output_dir = Path(self.nccl_config.output_dir)
            output_files = []
            if output_dir.exists():
                output_files = [str(f) for f in output_dir.glob("*.jsonl")]
                output_files.extend([str(f) for f in output_dir.glob("*.log")])
            return CollectionResult(
                success=result.returncode == 0,
                output_files=output_files,
                command=target_command,
                stdout=result.stdout,
                stderr=result.stderr,
                error=result.stderr if result.returncode != 0 else None,
            )
        except Exception as e:
            return CollectionResult(
                success=False,
                command=target_command,
                error=str(e),
            )
    def collect_with_nsys(
        self,
        target_command: list[str],
        working_dir: Optional[str] = None,
    ) -> CollectionResult:
        """Run collection with Nsight Systems.
        Args:
            target_command: Command to profile
            working_dir: Working directory for command
        Returns:
            CollectionResult with output files
        """
        nsys_cmd = self.build_nsys_command(target_command)
        try:
            result = subprocess.run(
                nsys_cmd,
                cwd=working_dir,
                capture_output=True,
                text=True,
            )

            output_files = []
            output_path = Path(self.nsys_config.output_path)
            for ext in [".nsys-rep", ".sqlite", ".qdrep"]:
                candidate = output_path.with_suffix(ext)
                if candidate.exists():
                    output_files.append(str(candidate))
            return CollectionResult(
                success=result.returncode == 0,
                output_files=output_files,
                command=nsys_cmd,
                stdout=result.stdout,
                stderr=result.stderr,
                error=result.stderr if result.returncode != 0 else None,
            )
        except FileNotFoundError:
            return CollectionResult(
                success=False,
                command=nsys_cmd,
                error="nsys not found. Please install Nsight Systems.",
            )
        except Exception as e:
            return CollectionResult(
                success=False,
                command=nsys_cmd,
                error=str(e),
            )
def get_nvidia_collection_env(
    output_dir: str = "./traces",
    enable_inspector: bool = True,
    enable_pytorch: bool = True,
) -> dict[str, str]:
    """Convenience function to get NVIDIA collection environment variables.
    Args:
        output_dir: Output directory for traces
        enable_inspector: Enable NCCL Inspector
        enable_pytorch: Enable PyTorch profiler variables
    Returns:
        Dictionary of environment variables
    """
    config = NCCLInspectorConfig(output_dir=output_dir)
    collector = NVIDIACollector(nccl_config=config)
    return collector.get_env_vars(
        enable_nccl_inspector=enable_inspector,
        enable_pytorch_profiler=enable_pytorch,
    )
