diffusers/tools

1 2# Unlike the rest of the PyTorch this file must be python2 compliant.3# This script outputs relevant system environment info4# Run it with `python collect_env.py`.5import datetime6import locale7import re8import subprocess9import sys10import os11from collections import namedtuple12 13 14try:15    import torch16    TORCH_AVAILABLE = True17except (ImportError, NameError, AttributeError, OSError):18    TORCH_AVAILABLE = False19 20# System Environment Information21SystemEnv = namedtuple('SystemEnv', [22    'torch_version',23    'is_debug_build',24    'cuda_compiled_version',25    'gcc_version',26    'clang_version',27    'cmake_version',28    'os',29    'libc_version',30    'python_version',31    'python_platform',32    'is_cuda_available',33    'cuda_runtime_version',34    'cuda_module_loading',35    'nvidia_driver_version',36    'nvidia_gpu_models',37    'cudnn_version',38    'pip_version',  # 'pip' or 'pip3'39    'pip_packages',40    'conda_packages',41    'hip_compiled_version',42    'hip_runtime_version',43    'miopen_runtime_version',44    'caching_allocator_config',45    'is_xnnpack_available',46    'cpu_info',47])48 49 50def run(command):51    """Returns (return-code, stdout, stderr)"""52    shell = True if type(command) is str else False53    p = subprocess.Popen(command, stdout=subprocess.PIPE,54                         stderr=subprocess.PIPE, shell=shell)55    raw_output, raw_err = p.communicate()56    rc = p.returncode57    if get_platform() == 'win32':58        enc = 'oem'59    else:60        enc = locale.getpreferredencoding()61    output = raw_output.decode(enc)62    err = raw_err.decode(enc)63    return rc, output.strip(), err.strip()64 65 66def run_and_read_all(run_lambda, command):67    """Runs command using run_lambda; reads and returns entire output if rc is 0"""68    rc, out, _ = run_lambda(command)69    if rc != 0:70        return None71    return out72 73 74def run_and_parse_first_match(run_lambda, command, regex):75    """Runs command using run_lambda, returns the first regex match if it exists"""76    rc, out, _ = run_lambda(command)77    if rc != 0:78        return None79    match = re.search(regex, out)80    if match is None:81        return None82    return match.group(1)83 84def run_and_return_first_line(run_lambda, command):85    """Runs command using run_lambda and returns first line if output is not empty"""86    rc, out, _ = run_lambda(command)87    if rc != 0:88        return None89    return out.split('\n')[0]90 91 92def get_conda_packages(run_lambda):93    conda = os.environ.get('CONDA_EXE', 'conda')94    out = run_and_read_all(run_lambda, "{} list".format(conda))95    if out is None:96        return out97 98    return "\n".join(99        line100        for line in out.splitlines()101        if not line.startswith("#")102        and any(103            name in line104            for name in {105                "torch",106                "numpy",107                "cudatoolkit",108                "soumith",109                "mkl",110                "magma",111                "triton",112            }113        )114    )115 116def get_gcc_version(run_lambda):117    return run_and_parse_first_match(run_lambda, 'gcc --version', r'gcc (.*)')118 119def get_clang_version(run_lambda):120    return run_and_parse_first_match(run_lambda, 'clang --version', r'clang version (.*)')121 122 123def get_cmake_version(run_lambda):124    return run_and_parse_first_match(run_lambda, 'cmake --version', r'cmake (.*)')125 126 127def get_nvidia_driver_version(run_lambda):128    if get_platform() == 'darwin':129        cmd = 'kextstat | grep -i cuda'130        return run_and_parse_first_match(run_lambda, cmd,131                                         r'com[.]nvidia[.]CUDA [(](.*?)[)]')132    smi = get_nvidia_smi()133    return run_and_parse_first_match(run_lambda, smi, r'Driver Version: (.*?) ')134 135 136def get_gpu_info(run_lambda):137    if get_platform() == 'darwin' or (TORCH_AVAILABLE and hasattr(torch.version, 'hip') and torch.version.hip is not None):138        if TORCH_AVAILABLE and torch.cuda.is_available():139            return torch.cuda.get_device_name(None)140        return None141    smi = get_nvidia_smi()142    uuid_regex = re.compile(r' \(UUID: .+?\)')143    rc, out, _ = run_lambda(smi + ' -L')144    if rc != 0:145        return None146    # Anonymize GPUs by removing their UUID147    return re.sub(uuid_regex, '', out)148 149 150def get_running_cuda_version(run_lambda):151    return run_and_parse_first_match(run_lambda, 'nvcc --version', r'release .+ V(.*)')152 153 154def get_cudnn_version(run_lambda):155    """This will return a list of libcudnn.so; it's hard to tell which one is being used"""156    if get_platform() == 'win32':157        system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')158        cuda_path = os.environ.get('CUDA_PATH', "%CUDA_PATH%")159        where_cmd = os.path.join(system_root, 'System32', 'where')160        cudnn_cmd = '{} /R "{}\\bin" cudnn*.dll'.format(where_cmd, cuda_path)161    elif get_platform() == 'darwin':162        # CUDA libraries and drivers can be found in /usr/local/cuda/. See163        # https://docs.nvidia.com/cuda/cuda-installation-guide-mac-os-x/index.html#install164        # https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#installmac165        # Use CUDNN_LIBRARY when cudnn library is installed elsewhere.166        cudnn_cmd = 'ls /usr/local/cuda/lib/libcudnn*'167    else:168        cudnn_cmd = 'ldconfig -p | grep libcudnn | rev | cut -d" " -f1 | rev'169    rc, out, _ = run_lambda(cudnn_cmd)170    # find will return 1 if there are permission errors or if not found171    if len(out) == 0 or (rc != 1 and rc != 0):172        l = os.environ.get('CUDNN_LIBRARY')173        if l is not None and os.path.isfile(l):174            return os.path.realpath(l)175        return None176    files_set = set()177    for fn in out.split('\n'):178        fn = os.path.realpath(fn)  # eliminate symbolic links179        if os.path.isfile(fn):180            files_set.add(fn)181    if not files_set:182        return None183    # Alphabetize the result because the order is non-deterministic otherwise184    files = sorted(files_set)185    if len(files) == 1:186        return files[0]187    result = '\n'.join(files)188    return 'Probably one of the following:\n{}'.format(result)189 190 191def get_nvidia_smi():192    # Note: nvidia-smi is currently available only on Windows and Linux193    smi = 'nvidia-smi'194    if get_platform() == 'win32':195        system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')196        program_files_root = os.environ.get('PROGRAMFILES', 'C:\\Program Files')197        legacy_path = os.path.join(program_files_root, 'NVIDIA Corporation', 'NVSMI', smi)198        new_path = os.path.join(system_root, 'System32', smi)199        smis = [new_path, legacy_path]200        for candidate_smi in smis:201            if os.path.exists(candidate_smi):202                smi = '"{}"'.format(candidate_smi)203                break204    return smi205 206 207# example outputs of CPU infos208#  * linux209#    Architecture:            x86_64210#      CPU op-mode(s):        32-bit, 64-bit211#      Address sizes:         46 bits physical, 48 bits virtual212#      Byte Order:            Little Endian213#    CPU(s):                  128214#      On-line CPU(s) list:   0-127215#    Vendor ID:               GenuineIntel216#      Model name:            Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz217#        CPU family:          6218#        Model:               106219#        Thread(s) per core:  2220#        Core(s) per socket:  32221#        Socket(s):           2222#        Stepping:            6223#        BogoMIPS:            5799.78224#        Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr225#                             sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl226#                             xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16227#                             pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand228#                             hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced229#                             fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap230#                             avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1231#                             xsaves wbnoinvd ida arat avx512vbmi pku ospke avx512_vbmi2 gfni vaes vpclmulqdq232#                             avx512_vnni avx512_bitalg tme avx512_vpopcntdq rdpid md_clear flush_l1d arch_capabilities233#    Virtualization features:234#      Hypervisor vendor:     KVM235#      Virtualization type:   full236#    Caches (sum of all):237#      L1d:                   3 MiB (64 instances)238#      L1i:                   2 MiB (64 instances)239#      L2:                    80 MiB (64 instances)240#      L3:                    108 MiB (2 instances)241#    NUMA:242#      NUMA node(s):          2243#      NUMA node0 CPU(s):     0-31,64-95244#      NUMA node1 CPU(s):     32-63,96-127245#    Vulnerabilities:246#      Itlb multihit:         Not affected247#      L1tf:                  Not affected248#      Mds:                   Not affected249#      Meltdown:              Not affected250#      Mmio stale data:       Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown251#      Retbleed:              Not affected252#      Spec store bypass:     Mitigation; Speculative Store Bypass disabled via prctl and seccomp253#      Spectre v1:            Mitigation; usercopy/swapgs barriers and __user pointer sanitization254#      Spectre v2:            Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence255#      Srbds:                 Not affected256#      Tsx async abort:       Not affected257#  * win32258#    Architecture=9259#    CurrentClockSpeed=2900260#    DeviceID=CPU0261#    Family=179262#    L2CacheSize=40960263#    L2CacheSpeed=264#    Manufacturer=GenuineIntel265#    MaxClockSpeed=2900266#    Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz267#    ProcessorType=3268#    Revision=27142269#270#    Architecture=9271#    CurrentClockSpeed=2900272#    DeviceID=CPU1273#    Family=179274#    L2CacheSize=40960275#    L2CacheSpeed=276#    Manufacturer=GenuineIntel277#    MaxClockSpeed=2900278#    Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz279#    ProcessorType=3280#    Revision=27142281 282def get_cpu_info(run_lambda):283    rc, out, err = 0, '', ''284    if get_platform() == 'linux':285        rc, out, err = run_lambda('lscpu')286    elif get_platform() == 'win32':287        rc, out, err = run_lambda('wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID,\288        CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE')289    elif get_platform() == 'darwin':290        rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")291    cpu_info = 'None'292    if rc == 0:293        cpu_info = out294    else:295        cpu_info = err296    return cpu_info297 298 299def get_platform():300    if sys.platform.startswith('linux'):301        return 'linux'302    elif sys.platform.startswith('win32'):303        return 'win32'304    elif sys.platform.startswith('cygwin'):305        return 'cygwin'306    elif sys.platform.startswith('darwin'):307        return 'darwin'308    else:309        return sys.platform310 311 312def get_mac_version(run_lambda):313    return run_and_parse_first_match(run_lambda, 'sw_vers -productVersion', r'(.*)')314 315 316def get_windows_version(run_lambda):317    system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')318    wmic_cmd = os.path.join(system_root, 'System32', 'Wbem', 'wmic')319    findstr_cmd = os.path.join(system_root, 'System32', 'findstr')320    return run_and_read_all(run_lambda, '{} os get Caption | {} /v Caption'.format(wmic_cmd, findstr_cmd))321 322 323def get_lsb_version(run_lambda):324    return run_and_parse_first_match(run_lambda, 'lsb_release -a', r'Description:\t(.*)')325 326 327def check_release_file(run_lambda):328    return run_and_parse_first_match(run_lambda, 'cat /etc/*-release',329                                     r'PRETTY_NAME="(.*)"')330 331 332def get_os(run_lambda):333    from platform import machine334    platform = get_platform()335 336    if platform == 'win32' or platform == 'cygwin':337        return get_windows_version(run_lambda)338 339    if platform == 'darwin':340        version = get_mac_version(run_lambda)341        if version is None:342            return None343        return 'macOS {} ({})'.format(version, machine())344 345    if platform == 'linux':346        # Ubuntu/Debian based347        desc = get_lsb_version(run_lambda)348        if desc is not None:349            return '{} ({})'.format(desc, machine())350 351        # Try reading /etc/*-release352        desc = check_release_file(run_lambda)353        if desc is not None:354            return '{} ({})'.format(desc, machine())355 356        return '{} ({})'.format(platform, machine())357 358    # Unknown platform359    return platform360 361 362def get_python_platform():363    import platform364    return platform.platform()365 366 367def get_libc_version():368    import platform369    if get_platform() != 'linux':370        return 'N/A'371    return '-'.join(platform.libc_ver())372 373 374def get_pip_packages(run_lambda):375    """Returns `pip list` output. Note: will also find conda-installed pytorch376    and numpy packages."""377    # People generally have `pip` as `pip` or `pip3`378    # But here it is invoked as `python -mpip`379    def run_with_pip(pip):380        out = run_and_read_all(run_lambda, pip + ["list", "--format=freeze"])381        return "\n".join(382            line383            for line in out.splitlines()384            if any(385                name in line386                for name in {387                    "torch",388                    "numpy",389                    "mypy",390                    "flake8",391                    "triton",392                }393            )394        )395 396    pip_version = 'pip3' if sys.version[0] == '3' else 'pip'397    out = run_with_pip([sys.executable, '-mpip'])398 399    return pip_version, out400 401 402def get_cachingallocator_config():403    ca_config = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '')404    return ca_config405 406 407def get_cuda_module_loading_config():408    if TORCH_AVAILABLE and torch.cuda.is_available():409        torch.cuda.init()410        config = os.environ.get('CUDA_MODULE_LOADING', '')411        return config412    else:413        return "N/A"414 415 416def is_xnnpack_available():417    if TORCH_AVAILABLE:418        import torch.backends.xnnpack419        return str(torch.backends.xnnpack.enabled)  # type: ignore[attr-defined]420    else:421        return "N/A"422 423def get_env_info():424    run_lambda = run425    pip_version, pip_list_output = get_pip_packages(run_lambda)426 427    if TORCH_AVAILABLE:428        version_str = torch.__version__429        debug_mode_str = str(torch.version.debug)430        cuda_available_str = str(torch.cuda.is_available())431        cuda_version_str = torch.version.cuda432        if not hasattr(torch.version, 'hip') or torch.version.hip is None:  # cuda version433            hip_compiled_version = hip_runtime_version = miopen_runtime_version = 'N/A'434        else:  # HIP version435            def get_version_or_na(cfg, prefix):436                _lst = [s.rsplit(None, 1)[-1] for s in cfg if prefix in s]437                return _lst[0] if _lst else 'N/A'438 439            cfg = torch._C._show_config().split('\n')440            hip_runtime_version = get_version_or_na(cfg, 'HIP Runtime')441            miopen_runtime_version = get_version_or_na(cfg, 'MIOpen')442            cuda_version_str = 'N/A'443            hip_compiled_version = torch.version.hip444    else:445        version_str = debug_mode_str = cuda_available_str = cuda_version_str = 'N/A'446        hip_compiled_version = hip_runtime_version = miopen_runtime_version = 'N/A'447 448    sys_version = sys.version.replace("\n", " ")449 450    return SystemEnv(451        torch_version=version_str,452        is_debug_build=debug_mode_str,453        python_version='{} ({}-bit runtime)'.format(sys_version, sys.maxsize.bit_length() + 1),454        python_platform=get_python_platform(),455        is_cuda_available=cuda_available_str,456        cuda_compiled_version=cuda_version_str,457        cuda_runtime_version=get_running_cuda_version(run_lambda),458        cuda_module_loading=get_cuda_module_loading_config(),459        nvidia_gpu_models=get_gpu_info(run_lambda),460        nvidia_driver_version=get_nvidia_driver_version(run_lambda),461        cudnn_version=get_cudnn_version(run_lambda),462        hip_compiled_version=hip_compiled_version,463        hip_runtime_version=hip_runtime_version,464        miopen_runtime_version=miopen_runtime_version,465        pip_version=pip_version,466        pip_packages=pip_list_output,467        conda_packages=get_conda_packages(run_lambda),468        os=get_os(run_lambda),469        libc_version=get_libc_version(),470        gcc_version=get_gcc_version(run_lambda),471        clang_version=get_clang_version(run_lambda),472        cmake_version=get_cmake_version(run_lambda),473        caching_allocator_config=get_cachingallocator_config(),474        is_xnnpack_available=is_xnnpack_available(),475        cpu_info=get_cpu_info(run_lambda),476    )477 478env_info_fmt = """479PyTorch version: {torch_version}480Is debug build: {is_debug_build}481CUDA used to build PyTorch: {cuda_compiled_version}482ROCM used to build PyTorch: {hip_compiled_version}483 484OS: {os}485GCC version: {gcc_version}486Clang version: {clang_version}487CMake version: {cmake_version}488Libc version: {libc_version}489 490Python version: {python_version}491Python platform: {python_platform}492Is CUDA available: {is_cuda_available}493CUDA runtime version: {cuda_runtime_version}494CUDA_MODULE_LOADING set to: {cuda_module_loading}495GPU models and configuration: {nvidia_gpu_models}496Nvidia driver version: {nvidia_driver_version}497cuDNN version: {cudnn_version}498HIP runtime version: {hip_runtime_version}499MIOpen runtime version: {miopen_runtime_version}500Is XNNPACK available: {is_xnnpack_available}501 502CPU:503{cpu_info}504 505Versions of relevant libraries:506{pip_packages}507{conda_packages}508""".strip()509 510 511def pretty_str(envinfo):512    def replace_nones(dct, replacement='Could not collect'):513        for key in dct.keys():514            if dct[key] is not None:515                continue516            dct[key] = replacement517        return dct518 519    def replace_bools(dct, true='Yes', false='No'):520        for key in dct.keys():521            if dct[key] is True:522                dct[key] = true523            elif dct[key] is False:524                dct[key] = false525        return dct526 527    def prepend(text, tag='[prepend]'):528        lines = text.split('\n')529        updated_lines = [tag + line for line in lines]530        return '\n'.join(updated_lines)531 532    def replace_if_empty(text, replacement='No relevant packages'):533        if text is not None and len(text) == 0:534            return replacement535        return text536 537    def maybe_start_on_next_line(string):538        # If `string` is multiline, prepend a \n to it.539        if string is not None and len(string.split('\n')) > 1:540            return '\n{}\n'.format(string)541        return string542 543    mutable_dict = envinfo._asdict()544 545    # If nvidia_gpu_models is multiline, start on the next line546    mutable_dict['nvidia_gpu_models'] = \547        maybe_start_on_next_line(envinfo.nvidia_gpu_models)548 549    # If the machine doesn't have CUDA, report some fields as 'No CUDA'550    dynamic_cuda_fields = [551        'cuda_runtime_version',552        'nvidia_gpu_models',553        'nvidia_driver_version',554    ]555    all_cuda_fields = dynamic_cuda_fields + ['cudnn_version']556    all_dynamic_cuda_fields_missing = all(557        mutable_dict[field] is None for field in dynamic_cuda_fields)558    if TORCH_AVAILABLE and not torch.cuda.is_available() and all_dynamic_cuda_fields_missing:559        for field in all_cuda_fields:560            mutable_dict[field] = 'No CUDA'561        if envinfo.cuda_compiled_version is None:562            mutable_dict['cuda_compiled_version'] = 'None'563 564    # Replace True with Yes, False with No565    mutable_dict = replace_bools(mutable_dict)566 567    # Replace all None objects with 'Could not collect'568    mutable_dict = replace_nones(mutable_dict)569 570    # If either of these are '', replace with 'No relevant packages'571    mutable_dict['pip_packages'] = replace_if_empty(mutable_dict['pip_packages'])572    mutable_dict['conda_packages'] = replace_if_empty(mutable_dict['conda_packages'])573 574    # Tag conda and pip packages with a prefix575    # If they were previously None, they'll show up as ie '[conda] Could not collect'576    if mutable_dict['pip_packages']:577        mutable_dict['pip_packages'] = prepend(mutable_dict['pip_packages'],578                                               '[{}] '.format(envinfo.pip_version))579    if mutable_dict['conda_packages']:580        mutable_dict['conda_packages'] = prepend(mutable_dict['conda_packages'],581                                                 '[conda] ')582    mutable_dict['cpu_info'] = envinfo.cpu_info583    return env_info_fmt.format(**mutable_dict)584 585 586def get_pretty_env_info():587    return pretty_str(get_env_info())588 589 590def main():591    print("Collecting environment information...")592    output = get_pretty_env_info()593    print(output)594 595    if TORCH_AVAILABLE and hasattr(torch, 'utils') and hasattr(torch.utils, '_crash_handler'):596        minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR597        if sys.platform == "linux" and os.path.exists(minidump_dir):598            dumps = [os.path.join(minidump_dir, dump) for dump in os.listdir(minidump_dir)]599            latest = max(dumps, key=os.path.getctime)600            ctime = os.path.getctime(latest)601            creation_time = datetime.datetime.fromtimestamp(ctime).strftime('%Y-%m-%d %H:%M:%S')602            msg = "\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time) + \603                  "if this is related to your bug please include it when you file a report ***"604            print(msg, file=sys.stderr)605 606 607 608if __name__ == '__main__':609    main()610