scripts/check_environment.py

#!/usr/bin/env python3
"""Report core OneScience/DTK and BoltzGen dependency versions."""

from __future__ import annotations

import importlib
import sys
from importlib.metadata import PackageNotFoundError, version
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "model"))


def package_version(distribution: str) -> str:
    try:
        return version(distribution)
    except PackageNotFoundError:
        return "NOT INSTALLED"


def main() -> int:
    print("python:", sys.version.split()[0])
    try:
        import torch

        print("torch:", torch.__version__)
        print("HIP:", torch.version.hip)
        print("DCU available:", torch.cuda.is_available())
        print("device count:", torch.cuda.device_count())
        if torch.cuda.is_available() and torch.cuda.device_count():
            print("device 0:", torch.cuda.get_device_name(0))
    except Exception as exc:  # noqa: BLE001
        print("torch check failed:", repr(exc))
        return 1

    for distribution in (
        "pytorch-lightning", "torchmetrics", "hydra-core", "omegaconf",
        "biotite", "gemmi", "mashumaro", "einx", "hydride", "pydssp",
        "logomaker", "frozendict",
    ):
        print(f"{distribution}:", package_version(distribution))

    try:
        importlib.import_module("boltzgen.cli.boltzgen")
        print("boltzgen import: OK")
    except Exception as exc:  # noqa: BLE001
        print("boltzgen import failed:", repr(exc))
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())