()
| 12 | return torch.cuda.is_available() |
| 13 | |
| 14 | def get_env_info(): |
| 15 | info = { |
| 16 | "`dataflow-mm` version": __version__, |
| 17 | "Platform": platform.platform(), |
| 18 | "Python version": platform.python_version(), |
| 19 | "PyTorch version": torch.__version__, |
| 20 | "Torchvision version": torch.__version__, |
| 21 | } |
| 22 | if is_torch_cuda_available(): |
| 23 | info["PyTorch version"] += " (GPU)" |
| 24 | info["GPU type"] = torch.cuda.get_device_name() |
| 25 | info["GPU number"] = torch.cuda.device_count() |
| 26 | info["GPU memory"] = f"{torch.cuda.mem_get_info()[1] / (1024**3):.2f}GB" |
| 27 | |
| 28 | |
| 29 | try: |
| 30 | import deepspeed # type: ignore |
| 31 | |
| 32 | info["DeepSpeed version"] = deepspeed.__version__ |
| 33 | except Exception: |
| 34 | pass |
| 35 | |
| 36 | try: |
| 37 | import bitsandbytes # type: ignore |
| 38 | |
| 39 | info["Bitsandbytes version"] = bitsandbytes.__version__ |
| 40 | except Exception: |
| 41 | pass |
| 42 | |
| 43 | try: |
| 44 | import vllm |
| 45 | |
| 46 | info["vLLM version"] = vllm.__version__ |
| 47 | except Exception: |
| 48 | pass |
| 49 | try: |
| 50 | import subprocess |
| 51 | # get the dir of imdlbenco package |
| 52 | imdlbenco_dir = os.path.dirname(os.path.abspath(__file__)) |
| 53 | # move to this dir and get the git commit hash in a subprocess |
| 54 | # but don't change the current working directory |
| 55 | os.chdir(imdlbenco_dir) |
| 56 | commit_info = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True) |
| 57 | commit_hash = commit_info.stdout.strip() |
| 58 | info["Git commit"] = commit_hash |
| 59 | except Exception: |
| 60 | pass |
| 61 | |
| 62 | print(Fore.BLUE + "=" * os.get_terminal_size().columns + Style.RESET_ALL) |
| 63 | print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n") |
| 64 | print(Fore.BLUE + "=" * os.get_terminal_size().columns + Style.RESET_ALL) |
| 65 | |
| 66 | def cli_env(): |
| 67 | get_env_info() |
no test coverage detected