Skip to content

show_env

zeus.show_env

Collect information about the environment and display it.

  • Python version
  • Package availablility and versions: Zeus, PyTorch, JAX.
  • NVIDIA GPU availability: Number of GPUs and moels.
  • AMD GPU availability: Number of GPUs and models.
  • Intel RAPL availability: Number of CPUs and whether DRAM measurements are available.

show_env

show_env()

Collect information about the environment and display it.

Source code in zeus/show_env.py
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def show_env():
    """Collect information about the environment and display it."""
    print(SECTION_SEPARATOR)
    print(f"Python version: {platform.python_version()}\n")

    print(SECTION_SEPARATOR)
    package_availability = "\nPackage availability and versions:\n"
    package_availability += f"  Zeus: {zeus.__version__}\n"

    try:
        torch_available = framework.torch_is_available()
        torch_cuda_available = True
    except RuntimeError:
        torch_available = framework.torch_is_available(ensure_cuda=False)
        torch_cuda_available = False

    if torch_available and torch_cuda_available:
        torch = framework.MODULE_CACHE["torch"]
        package_availability += f"  PyTorch: {torch.__version__} (with CUDA support)\n"
    elif torch_available and not torch_cuda_available:
        torch = framework.MODULE_CACHE["torch"]
        package_availability += (
            f"  PyTorch: {torch.__version__} (without CUDA support)\n"
        )
    else:
        package_availability += "  PyTorch: not available\n"

    try:
        jax_available = framework.jax_is_available()
        jax_cuda_available = True
    except RuntimeError:
        jax_available = framework.jax_is_available(ensure_cuda=False)
        jax_cuda_available = False

    if jax_available and jax_cuda_available:
        jax = framework.MODULE_CACHE["jax"]
        package_availability += f"  JAX: {jax.__version__} (with CUDA support)\n"
    elif jax_available and not jax_cuda_available:
        jax = framework.MODULE_CACHE["jax"]
        package_availability += f"  JAX: {jax.__version__} (without CUDA support)\n"
    else:
        package_availability += "  JAX: not available\n"

    print(package_availability)

    print(SECTION_SEPARATOR)
    gpu_availability = "\nGPU availability:\n"
    try:
        gpus = get_gpus()
    except ZeusGPUInitError:
        gpus = EmptyGPUs()
    if len(gpus) > 0:
        for i in range(len(gpus)):
            gpu_availability += f"  GPU {i}: {gpus.getName(i)}\n"
    else:
        gpu_availability += "  No GPUs available.\n"
    print(gpu_availability)

    print(SECTION_SEPARATOR)
    cpu_availability = "\nCPU availability:\n"
    try:
        cpus = get_cpus()
    except ZeusCPUInitError:
        cpus = EmptyCPUs()
    if len(cpus) > 0:
        assert isinstance(cpus, RAPLCPUs)
        for i in range(len(cpus)):
            cpu_availability += f"  CPU {i}:\n    CPU measurements available ({cpus.cpus[i].rapl_file.path})\n"
            if cpus.supportsGetDramEnergyConsumption(i):
                dram = cpus.cpus[i].dram
                assert dram is not None
                cpu_availability += f"    DRAM measurements available ({dram.path})\n"
            else:
                cpu_availability += "    DRAM measurements unavailable\n"
    else:
        cpu_availability += "  No CPUs available.\n"
    print(cpu_availability)

    print(SECTION_SEPARATOR)