device_info.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Auto-detect GPU host label, hardware, and coarse location for engine registration."""
 
from __future__ import annotations
 
import json
import os
import platform
import re
import socket
import urllib.request
from typing import Any
 
 
def _fetch_aws_placement() -> str | None:
    try:
        req = urllib.request.Request(
            "http://169.254.169.254/latest/meta-data/placement/availability-zone",
            headers={"Metadata-Flavor": "Amazon"},
        )
        with urllib.request.urlopen(req, timeout=0.4) as resp:
            text = resp.read().decode().strip()
            return text or None
    except Exception:
        return None
 
 
def _sanitize_label(text: str, *, max_len: int = 64) -> str:
    cleaned = re.sub(r"[^\w.\-]+", "-", text.strip().lower()).strip("-")
    return (cleaned or "gpu-host")[:max_len]
 
 
def detect_gpu_info() -> dict[str, Any]:
    try:
        from aquin.compute.device import probe_backend
 
        probe = probe_backend()
        info: dict[str, Any] = {
            "cuda_available": probe.get("cuda_available", False),
            "mps_available": probe.get("mps_available", False),
            "backend": probe.get("backend"),
            "selected_device": probe.get("selected"),
            "rocm": probe.get("rocm", False),
            "accelerator_available": probe.get("accelerator_available", False),
            "gpus": probe.get("gpus") or [],
        }
        return info
    except Exception:
        return {
            "cuda_available": False,
            "mps_available": False,
            "backend": "cpu",
            "selected_device": "cpu",
            "rocm": False,
            "accelerator_available": False,
            "gpus": [],
        }
 
 
def detect_location() -> dict[str, Any]:
    hostname = socket.gethostname()
    aws_az = _fetch_aws_placement()
    region = (
        os.environ.get("AWS_REGION")
        or os.environ.get("AWS_DEFAULT_REGION")
        or os.environ.get("GCP_REGION")
        or os.environ.get("MODAL_REGION")
        or os.environ.get("RUNPOD_POD_HOSTNAME")
        or aws_az
    )
 
    provider: str | None = None
    if os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or aws_az:
        provider = "aws"
    elif os.environ.get("MODAL_TASK_ID") or os.environ.get("MODAL_REGION"):
        provider = "modal"
    elif os.environ.get("RUNPOD_POD_ID"):
        provider = "runpod"
 
    parts = [p for p in (region, hostname) if p]
    return {
        "hostname": hostname,
        "region": region,
        "provider": provider,
        "platform": platform.system(),
        "location_label": " · ".join(parts) if parts else hostname,
    }
 
 
def detect_device_label() -> str:
    hostname = socket.gethostname()
    gpu = detect_gpu_info()
    gpus = gpu.get("gpus") or []
    if gpus:
        short_gpu = re.sub(r"\s+", "-", str(gpus[0].get("name", "")).lower())
        short_gpu = re.sub(r"[^a-z0-9\-]", "", short_gpu)[:24]
        if short_gpu:
            return _sanitize_label(f"{hostname}-{short_gpu}")
    if gpu.get("mps_available"):
        return _sanitize_label(f"{hostname}-metal")
    return _sanitize_label(hostname)
 
 
def build_engine_profile() -> dict[str, Any]:
    gpu = detect_gpu_info()
    location = detect_location()
    device = detect_device_label()
    return {
        "device": device,
        "gpu": gpu,
        "location": location,
        "gpu_info": json.dumps({**gpu, **location}, separators=(",", ":")),
    }
 
 
def format_engine_profile_summary(profile: dict[str, Any]) -> str:
    loc = profile.get("location") or {}
    gpu = profile.get("gpu") or {}
    gpus = gpu.get("gpus") or []
    backend = gpu.get("backend") or "cpu"
    selected = gpu.get("selected_device") or "cpu"
 
    if gpus:
        g = gpus[0]
        gpu_line = f"{g.get('name', 'GPU')} · {g.get('vram_gb', '?')} GB ({backend})"
        if len(gpus) > 1:
            gpu_line += f" (+{len(gpus) - 1} more)"
    elif gpu.get("mps_available") or selected == "mps":
        gpu_line = "Apple Metal (MPS)"
    elif gpu.get("cuda_available"):
        gpu_line = f"CUDA/ROCm ({backend})"
    else:
        gpu_line = "no GPU accelerator"
 
    loc_line = loc.get("location_label") or loc.get("hostname") or "unknown host"
    return f"{loc_line}  ·  {gpu_line}"