"""AMD rocprofv3 and ROCm Systems Profiler collection setup.
Sets up environment variables and commands for collecting RCCL traces:
- rocprofv3: API-level tracing with RCCL support
- ROCm Systems Profiler (omnitrace): Full system profiling
Reference:
- rocprofv3: https://rocm.docs.amd.com/projects/rocprofiler-sdk/en/develop/how-to/using-rocprofv3.html
- RCCL: https://rocm.docs.amd.com/projects/rccl/en/develop/how-to/rccl-usage-tips.html
"""
import os
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
@dataclass
class Rocprofv3Config:
    """Configuration for rocprofv3 collection.
    Attributes:
        output_dir: Directory for output files
        output_format: Output format (csv, json)
        enable_rccl_trace: Enable RCCL API tracing
        enable_hip_trace: Enable HIP API tracing
        enable_hsa_trace: Enable HSA API tracing
        enable_kernel_trace: Enable kernel timing
    """
    output_dir: str = "./rocprofv3_traces"
    output_format: str = "csv"
    enable_rccl_trace: bool = True
    enable_hip_trace: bool = True
    enable_hsa_trace: bool = False
    enable_kernel_trace: bool = True
@dataclass
class RocsysConfig:
    """Configuration for ROCm Systems Profiler collection.
    Attributes:
        output_dir: Directory for output files
        enable_rcclp: Enable RCCL profiling
        enable_hip: Enable HIP tracing
        perfetto_output: Generate Perfetto-compatible output
    """
    output_dir: str = "./rocsys_traces"
    enable_rcclp: bool = True
    enable_hip: bool = True
    perfetto_output: bool = True
@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 AMDCollector:
    """Collects RCCL traces on AMD GPUs."""
    def __init__(
        self,
        rocprofv3_config: Optional[Rocprofv3Config] = None,
        rocsys_config: Optional[RocsysConfig] = None,
    ) -> None:
        """Initialize the collector.
        Args:
            rocprofv3_config: rocprofv3 configuration
            rocsys_config: ROCm Systems Profiler configuration
        """
        self.rocprofv3_config = rocprofv3_config or Rocprofv3Config()
        self.rocsys_config = rocsys_config or RocsysConfig()
    def get_env_vars(
        self,
        enable_rccl_debug: bool = False,
    ) -> dict[str, str]:
        """Get environment variables for trace collection.
        Args:
            enable_rccl_debug: Enable RCCL debug logging
        Returns:
            Dictionary of environment variables to set
        """
        env = {}
        # ROCm Systems Profiler (omnitrace) variables
        if self.rocsys_config.enable_rcclp:
            env["ROCPROFSYS_USE_RCCLP"] = "ON"
        if self.rocsys_config.enable_hip:
            env["ROCPROFSYS_USE_HIP"] = "ON"
        # RCCL debug variables
        if enable_rccl_debug:
            env["NCCL_DEBUG"] = "INFO"
            env["NCCL_DEBUG_SUBSYS"] = "ALL"
        # Trace output directory
        env["ROCPROFSYS_OUTPUT_PATH"] = str(
            Path(self.rocsys_config.output_dir).absolute()
        )
        return env
    def build_rocprofv3_command(
        self,
        target_command: list[str],
    ) -> list[str]:
        """Build rocprofv3 profiling command.
        Args:
            target_command: Command to profile
        Returns:
            Full rocprofv3 command with options
        """
        cmd = ["rocprofv3"]
        # Output directory
        output_dir = Path(self.rocprofv3_config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        cmd.extend(["-d", str(output_dir.absolute())])
        # Output format
        cmd.extend(["--output-format", self.rocprofv3_config.output_format])
        # Trace options
        if self.rocprofv3_config.enable_rccl_trace:
            cmd.append("--rccl-trace")
        if self.rocprofv3_config.enable_hip_trace:
            cmd.append("--hip-trace")
        if self.rocprofv3_config.enable_hsa_trace:
            cmd.append("--hsa-trace")
        if self.rocprofv3_config.enable_kernel_trace:
            cmd.append("--kernel-trace")
        cmd.append("--")
        cmd.extend(target_command)
        return cmd
    def build_omnitrace_command(
        self,
        target_command: list[str],
    ) -> list[str]:
        """Build omnitrace (ROCm Systems Profiler) command.
        Args:
            target_command: Command to profile
        Returns:
            Full omnitrace command with options
        """
        cmd = ["omnitrace-run"]
        # Output directory
        output_dir = Path(self.rocsys_config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        cmd.extend(["-o", str(output_dir.absolute())])
        # Perfetto output
        if self.rocsys_config.perfetto_output:
            cmd.append("--perfetto-output")
        cmd.append("--")
        cmd.extend(target_command)
        return cmd
    def collect_with_rocprofv3(
        self,
        target_command: list[str],
        working_dir: Optional[str] = None,
    ) -> CollectionResult:
        """Run collection with rocprofv3.
        Args:
            target_command: Command to run with tracing
            working_dir: Working directory for command
        Returns:
            CollectionResult with output files
        """
        rocprof_cmd = self.build_rocprofv3_command(target_command)
        try:
            result = subprocess.run(
                rocprof_cmd,
                cwd=working_dir,
                capture_output=True,
                text=True,
            )

            output_dir = Path(self.rocprofv3_config.output_dir)
            output_files = []
            if output_dir.exists():
                for ext in ["csv", "json"]:
                    output_files.extend([str(f) for f in output_dir.glob(f"*.{ext}")])
            return CollectionResult(
                success=result.returncode == 0,
                output_files=output_files,
                command=rocprof_cmd,
                stdout=result.stdout,
                stderr=result.stderr,
                error=result.stderr if result.returncode != 0 else None,
            )
        except FileNotFoundError:
            return CollectionResult(
                success=False,
                command=rocprof_cmd,
                error="rocprofv3 not found. Please install ROCm profiler.",
            )
        except Exception as e:
            return CollectionResult(
                success=False,
                command=rocprof_cmd,
                error=str(e),
            )
    def collect_with_omnitrace(
        self,
        target_command: list[str],
        working_dir: Optional[str] = None,
    ) -> CollectionResult:
        """Run collection with ROCm Systems Profiler (omnitrace).
        Args:
            target_command: Command to profile
            working_dir: Working directory for command
        Returns:
            CollectionResult with output files
        """
        env = {**os.environ, **self.get_env_vars()}
        omnitrace_cmd = self.build_omnitrace_command(target_command)
        try:
            result = subprocess.run(
                omnitrace_cmd,
                env=env,
                cwd=working_dir,
                capture_output=True,
                text=True,
            )

            output_dir = Path(self.rocsys_config.output_dir)
            output_files = []
            if output_dir.exists():
                for ext in ["proto", "perfetto-trace", "json", "rocpd"]:
                    output_files.extend([str(f) for f in output_dir.glob(f"*.{ext}")])
            return CollectionResult(
                success=result.returncode == 0,
                output_files=output_files,
                command=omnitrace_cmd,
                stdout=result.stdout,
                stderr=result.stderr,
                error=result.stderr if result.returncode != 0 else None,
            )
        except FileNotFoundError:
            return CollectionResult(
                success=False,
                command=omnitrace_cmd,
                error="omnitrace-run not found. Please install ROCm Systems Profiler.",
            )
        except Exception as e:
            return CollectionResult(
                success=False,
                command=omnitrace_cmd,
                error=str(e),
            )
def get_amd_collection_env(
    output_dir: str = "./traces",
    enable_rccl: bool = True,
    enable_hip: bool = True,
) -> dict[str, str]:
    """Convenience function to get AMD collection environment variables.
    Args:
        output_dir: Output directory for traces
        enable_rccl: Enable RCCL profiling
        enable_hip: Enable HIP tracing
    Returns:
        Dictionary of environment variables
    """
    config = RocsysConfig(
        output_dir=output_dir,
        enable_rcclp=enable_rccl,
        enable_hip=enable_hip,
    )
    collector = AMDCollector(rocsys_config=config)
    return collector.get_env_vars()
