"""Wafer-managed trace collection setup.

This module provides functionality to set up environment variables
and configuration for collecting distributed traces:
- NVIDIA: NCCL Inspector and Nsight Systems
- AMD: rocprofv3 and ROCm Systems Profiler
- PyTorch: torch.profiler configuration
"""

from wafer_core.lib.distributed_traces.collection.amd_collector import (
    AMDCollector,
    get_amd_collection_env,
)
from wafer_core.lib.distributed_traces.collection.nvidia_collector import (
    NVIDIACollector,
    get_nvidia_collection_env,
)
from wafer_core.lib.distributed_traces.collection.pytorch_collector import (
    PyTorchCollector,
    get_pytorch_collection_env,
)

__all__ = [
    # NVIDIA
    "NVIDIACollector",
    "get_nvidia_collection_env",
    # AMD
    "AMDCollector",
    "get_amd_collection_env",
    # PyTorch
    "PyTorchCollector",
    "get_pytorch_collection_env",
]
