#!/usr/bin/env python3

from __future__ import annotations

import importlib

import json

import platform

import sys

import argparse

from pathlib import Path

def parse_args() -> argparse.Namespace:

parser = argparse.ArgumentParser(description="Inspect Python, CUDA, and training package availability.")

return parser.parse_args()

def _module_info(name: str) -> dict[str, str]:

try:

module = importlib.import_module(name)

except Exception as exc: # pragma: no cover - diagnostics only

return {"status": "error", "error": repr(exc)}

version = getattr(module, "__version__", "unknown")

location = getattr(module, "__file__", "unknown")

return {"status": "ok", "version": str(version), "location": str(location)}

def main() -> int:

parse_args()

info: dict[str, object] = {

"python": {

"executable": sys.executable,

"version": sys.version,

"platform": platform.platform(),

},

"modules": {},

}

import torch # noqa: PLC0415

cuda: dict[str, object] = {

"torch_version": torch.__version__,

"torch_cuda": torch.version.cuda,

"cuda_available": torch.cuda.is_available(),

}

if torch.cuda.is_available():

cuda.update(

{

"device_name": torch.cuda.get_device_name(0),

"device_capability": list(torch.cuda.get_device_capability(0)),

"bf16_supported": torch.cuda.is_bf16_supported(),

"flash_sdp_enabled": torch.backends.cuda.flash_sdp_enabled(),

"mem_efficient_sdp_enabled": torch.backends.cuda.mem_efficient_sdp_enabled(),

"math_sdp_enabled": torch.backends.cuda.math_sdp_enabled(),

}

)

info["cuda"] = cuda

# Import Unsloth first so its import-time patches apply before TRL / Transformers.

module_names = [

"unsloth",

"transformers",

"trl",

"xformers",

"torchvision",

"flash_attn",

]

for name in module_names:

info["modules"][name] = _module_info(name)

try:

from transformers.utils import is_flash_attn_2_available # noqa: PLC0415

info["transformers_flash_attn_2_available"] = bool(is_flash_attn_2_available())

except Exception as exc: # pragma: no cover - diagnostics only

info["transformers_flash_attn_2_available"] = {"status": "error", "error": repr(exc)}

print(json.dumps(info, indent=2))

return 0

if __name__ == "__main__":

raise SystemExit(main())