File size: 2,077 Bytes
e020674 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 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 |
import os
import platform
from colorama import init, Fore, Style
from dataflow import __version__
import torch
def is_torch_cuda_available():
"""
Check if CUDA is available for PyTorch.
"""
return torch.cuda.is_available()
def get_env_info():
info = {
"`dataflow` version": __version__,
"Platform": platform.platform(),
"Python version": platform.python_version(),
"PyTorch version": torch.__version__,
"Torchvision version": torch.__version__,
}
if is_torch_cuda_available():
info["PyTorch version"] += " (GPU)"
info["GPU type"] = torch.cuda.get_device_name()
info["GPU number"] = torch.cuda.device_count()
info["GPU memory"] = f"{torch.cuda.mem_get_info()[1] / (1024**3):.2f}GB"
try:
import deepspeed # type: ignore
info["DeepSpeed version"] = deepspeed.__version__
except Exception:
pass
try:
import bitsandbytes # type: ignore
info["Bitsandbytes version"] = bitsandbytes.__version__
except Exception:
pass
try:
import vllm
info["vLLM version"] = vllm.__version__
except Exception:
pass
try:
import subprocess
# get the dir of imdlbenco package
imdlbenco_dir = os.path.dirname(os.path.abspath(__file__))
# move to this dir and get the git commit hash in a subprocess
# but don't change the current working directory
os.chdir(imdlbenco_dir)
commit_info = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True)
commit_hash = commit_info.stdout.strip()
info["Git commit"] = commit_hash
except Exception:
pass
print(Fore.BLUE + "=" * os.get_terminal_size().columns + Style.RESET_ALL)
print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n")
print(Fore.BLUE + "=" * os.get_terminal_size().columns + Style.RESET_ALL)
def cli_env():
get_env_info()
if __name__ == "__main__":
print(get_env_info()) |