Environment check script

check_environment.py

#! /usr/bin/env python

"""
Script for checking various elements of the runtime environment and print a large amount of diagnostic information.

This is meant to be used for debugging environments used with MONAI, but doesn't directly need MONAI itself. It will
print information about the environment including trying to get installed packages, test PyTorch with CUDA, and then
have MONAI print its debugging information if no errors encountered. If MONAI is not installed this script should still
work and produce useful information.
"""

import os
import sys
import platform
import multiprocessing
import subprocess
import shutil
import argparse
import getpass

USER = getpass.getuser()


def print_platform():
    """
    Print basic platform information.
    """
    print(platform.platform())
    uname = list(platform.uname())
    uname[1] = "<host>"  # remove node/hostname
    print("uname:", uname)
    print("CPU:", platform.processor(), "Count:", multiprocessing.cpu_count())
    print("Python:", sys.executable, platform.python_implementation(), platform.python_version())


def print_environment_vars():
    """
    Print all environment variables other than a few known pointless ones.
    """
    print("Environment:")
    for k, v in os.environ.items():
        if k not in ("LS_COLORS", "PS1", "PS2"):
            print(f"  {k}:", v.replace(USER, "<user>"))


def print_environment():
    """
    Print the installed environment using `conda` or `pip`, fail if neither are present.
    """
    try:
        cmd = ("conda", "env", "export") if shutil.which("conda") else ("pip", "list")

        result = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
        print(result.decode())
        return True
    except Exception as e:
        print(f"Exception encountered getting environment with conda/pip: {e}", file=sys.stderr)
        return False


def check_torch():
    """
    Check PyTorch is installed and a tensor can be created.
    """
    try:
        import torch

        t = torch.rand(2, 3) * 5
        print("PyTorch:", torch.__version__, torch.__path__)
        print("Test tensor:", t.flatten())
        return True
    except ImportError:
        print("PyTorch not installed", file=sys.stderr)
        return False


def check_torch_cuda():
    """
    Check CUDA capability in PyTorch by moving a tensor to each available device.
    """
    import torch

    dcount = torch.cuda.device_count()
    print("PyTorch GPU Count:", dcount)

    if dcount == 0:
        print("PyTorch has 0 device count", file=sys.stderr)
        return False

    try:
        for d in range(dcount):
            print(f"  {torch.cuda.get_device_properties(d)}")
            t = torch.rand(2, 3).to(torch.device(f"cuda:{d}")) * 5
            print("Test tensor:", t.flatten())
        return True
    except Exception as e:
        print(f"PyTorch encountered exception creating GPU tensor on device {d}: {e}", file=sys.stderr)
        return False


def check_monai():
    """
    Check MONAI by importing it then printing its debug info.
    """
    try:
        import monai

        monai.config.deviceconfig.print_debug_info()
        return True
    except ImportError:
        print("MONAI not installed", file=sys.stderr)
        return False


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog="check_environment", description="Script for checking MONAI environment")
    parser.add_argument("--env", default=False, action="store_true", help="Print environment info")
    parser.add_argument("--envvars", default=False, action="store_true", help="Include environment variables")
    parser.add_argument("--monai", default=False, action="store_true", help="Print MONAI info")
    args = parser.parse_args()

    if args.env:
        print("=" * 10, "Checking Environment", "=" * 10)
        print_platform()
        if args.envvars:
            print_environment_vars()
        print_environment()

    print("=" * 10, "Checking PyTorch", "=" * 10)

    if not check_torch():
        print("Exiting early, no valid PyTorch install found.", file=sys.stderr)
        sys.exit(1)

    if not check_torch_cuda():
        print("Exiting early, PyTorch encountered CUDA error.", file=sys.stderr)
        sys.exit(1)

    if args.monai:
        print("=" * 10, "Checking MONAI", "=" * 10)
        check_monai()