from __future__ import annotations
import hashlib
import os
import platform
import subprocess
import sys
from pathlib import Path
from benchmarks.gpu_comparison.results import GPUInfo, SystemInfo
FRAMEWORK_DIR = Path(__file__).parent
REPO_ROOT = FRAMEWORK_DIR.parent.parent
BRAHE_EOP_FILE = REPO_ROOT / "data" / "eop" / "finals.all.iau2000.txt"
BRAHE_SPACE_WEATHER_FILE = REPO_ROOT / "data" / "space_weather" / "sw19571001.txt"
BRAHE_GRAVITY_FILE = REPO_ROOT / "data" / "gravity_models" / "EGM2008_360.gfc"
RESULTS_DIR = FRAMEWORK_DIR / "results"
def set_data_alignment_env() -> None:
os.environ.setdefault("BRAHE_EOP_FILE", str(BRAHE_EOP_FILE))
os.environ.setdefault("BRAHE_SPACE_WEATHER_FILE", str(BRAHE_SPACE_WEATHER_FILE))
os.environ.setdefault("BRAHE_GRAVITY_FILE", str(BRAHE_GRAVITY_FILE))
def sha256_of(path: Path) -> str:
h = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1 << 20), b""):
h.update(chunk)
return h.hexdigest()
def data_alignment_record() -> dict[str, str]:
return {
"eop_file": str(BRAHE_EOP_FILE),
"eop_sha256": sha256_of(BRAHE_EOP_FILE),
"space_weather_file": str(BRAHE_SPACE_WEATHER_FILE),
"space_weather_sha256": sha256_of(BRAHE_SPACE_WEATHER_FILE),
"gravity_model_file": str(BRAHE_GRAVITY_FILE),
"gravity_model_sha256": sha256_of(BRAHE_GRAVITY_FILE),
}
def _detect_cpu_model() -> str:
if sys.platform.startswith("linux"):
try:
for line in Path("/proc/cpuinfo").read_text().splitlines():
if line.startswith("model name"):
return line.split(":", 1)[1].strip()
except OSError:
pass
if sys.platform == "darwin":
try:
r = subprocess.run(
["sysctl", "-n", "machdep.cpu.brand_string"],
capture_output=True, text=True, timeout=5,
)
if r.returncode == 0:
return r.stdout.strip()
except (FileNotFoundError, subprocess.TimeoutExpired):
pass
return platform.processor() or "unknown"
def _detect_physical_cores() -> int:
try:
import psutil
n = psutil.cpu_count(logical=False)
if n:
return int(n)
except ImportError:
pass
return os.cpu_count() or 1
def _detect_ram_gb() -> int:
try:
import psutil
return int(psutil.virtual_memory().total / (1024 ** 3))
except ImportError:
pass
if sys.platform.startswith("linux"):
try:
for line in Path("/proc/meminfo").read_text().splitlines():
if line.startswith("MemTotal:"):
kb = int(line.split()[1])
return kb // (1024 * 1024)
except OSError:
pass
return 0
def _detect_cuda_version_from_smi() -> str:
try:
r = subprocess.run(
["nvidia-smi", "-q"], capture_output=True, text=True, timeout=10,
)
if r.returncode != 0:
return "unknown"
for line in r.stdout.splitlines():
line = line.strip()
if line.startswith("CUDA Version"):
return line.split(":", 1)[1].strip()
except (FileNotFoundError, subprocess.TimeoutExpired):
pass
return "unknown"
def _detect_gpus() -> list[GPUInfo]:
try:
r = subprocess.run(
[
"nvidia-smi",
"--query-gpu=index,name,memory.total,driver_version",
"--format=csv,noheader,nounits",
],
capture_output=True, text=True, timeout=10,
)
if r.returncode != 0:
return []
except (FileNotFoundError, subprocess.TimeoutExpired):
return []
cuda_version = _detect_cuda_version_from_smi()
gpus: list[GPUInfo] = []
for line in r.stdout.strip().splitlines():
parts = [p.strip() for p in line.split(",")]
if len(parts) < 4:
continue
try:
gpus.append(GPUInfo(
index=int(parts[0]),
model=parts[1],
memory_mb=int(parts[2]),
driver=parts[3],
cuda_runtime=cuda_version,
))
except ValueError:
continue
return gpus
def _detect_rust_version() -> str:
try:
r = subprocess.run(["rustc", "--version"], capture_output=True, text=True, timeout=5)
return r.stdout.strip() if r.returncode == 0 else "unknown"
except (FileNotFoundError, subprocess.TimeoutExpired):
return "unknown"
def _detect_jax_version() -> str:
try:
import jax
return jax.__version__
except ImportError:
return "not installed"
def _detect_brahe_version() -> str:
try:
from importlib.metadata import version
return version("brahe")
except Exception:
return "unknown"
def _detect_astrojax_version() -> str:
try:
from importlib.metadata import version
return version("astrojax")
except Exception:
return "unknown"
def _git_sha(repo_path: Path) -> str | None:
try:
r = subprocess.run(
["git", "-C", str(repo_path), "rev-parse", "--short", "HEAD"],
capture_output=True, text=True, timeout=5,
)
if r.returncode != 0:
return None
return r.stdout.strip()
except (FileNotFoundError, subprocess.TimeoutExpired):
return None
def _astrojax_repo_path() -> Path | None:
try:
import astrojax
p = Path(astrojax.__file__).resolve()
for parent in p.parents:
if (parent / ".git").exists():
return parent
except ImportError:
pass
return None
def _detect_rayon_threads() -> int:
val = os.environ.get("RAYON_NUM_THREADS")
if val and val.isdigit():
return int(val)
return os.cpu_count() or 1
def collect_system_info() -> SystemInfo:
astrojax_repo = _astrojax_repo_path()
return SystemInfo(
cpu_model=_detect_cpu_model(),
cpu_physical_cores=_detect_physical_cores(),
cpu_logical_cores=os.cpu_count() or 1,
ram_gb=_detect_ram_gb(),
os=f"{platform.system()} {platform.release()}",
python_version=platform.python_version(),
rust_version=_detect_rust_version(),
brahe_version=_detect_brahe_version(),
brahe_git_sha=_git_sha(REPO_ROOT),
astrojax_version=_detect_astrojax_version(),
astrojax_git_sha=_git_sha(astrojax_repo) if astrojax_repo else None,
jax_version=_detect_jax_version(),
gpus=_detect_gpus(),
rayon_threads=_detect_rayon_threads(),
)