ROCm HW support

packages/node/meshnet_node/hardware.py

@@ -123,6 +123,24 @@ def _detect_nvidia_smi_gpu_memory() -> dict | None:
     return None
 
 
+def _detect_torch_cuda_inventory(torch_module) -> dict | None:
+    """Return torch-visible CUDA/HIP GPU metadata without running kernels."""
+    try:
+        if not torch_module.cuda.is_available() or torch_module.cuda.device_count() < 1:
+            return None
+        idx = torch_module.cuda.current_device()
+        name = torch_module.cuda.get_device_name(idx)
+        props = torch_module.cuda.get_device_properties(idx)
+        vram_mb = int(props.total_memory // (1024 * 1024))
+        gpu = {"gpu_name": name, "vram_mb": max(0, vram_mb)}
+        gcn_arch = getattr(props, "gcnArchName", None)
+        if gcn_arch:
+            gpu["gcn_arch"] = str(gcn_arch)
+        return gpu
+    except Exception:
+        return None
+
+
 def _torch_cuda_is_executable(torch_module) -> bool:
     """Return True only if this Python process can execute a CUDA tensor op."""
     try:
@@ -139,7 +157,7 @@ def _torch_cuda_is_executable(torch_module) -> bool:
 def _gpu_inventory_profile(gpu: dict | None, ram_mb: int) -> dict | None:
     if gpu is None:
         return None
-    return {
+    profile = {
         "device": "cpu",
         "gpu_name": gpu["gpu_name"],
         "vram_mb": gpu["vram_mb"],
@@ -148,20 +166,26 @@ def _gpu_inventory_profile(gpu: dict | None, ram_mb: int) -> dict | None:
         "ram_mb": ram_mb,
         "cuda_available": False,
     }
+    if gpu.get("gcn_arch"):
+        profile["gcn_arch"] = gpu["gcn_arch"]
+    return profile
 
 
 def detect_hardware() -> dict:
     """Detect GPU model and available VRAM. Returns hardware profile dict."""
     ram_mb = _detect_ram_mb()
+    torch_gpu: dict | None = None
     try:
         import torch  # type: ignore[import]
+
+        torch_gpu = _detect_torch_cuda_inventory(torch)
         if _torch_cuda_is_executable(torch):
-            idx = torch.cuda.current_device()
-            name = torch.cuda.get_device_name(idx)
-            props = torch.cuda.get_device_properties(idx)
-            vram_mb = props.total_memory // (1024 * 1024)
+            if torch_gpu is None:
+                torch_gpu = _detect_torch_cuda_inventory(torch)
+            name = torch_gpu["gpu_name"] if torch_gpu is not None else "CUDA GPU"
+            vram_mb = torch_gpu["vram_mb"] if torch_gpu is not None else 0
             shared_vram_mb = max(0, ram_mb // 2)
-            return {
+            profile = {
                 "device": "cuda",
                 "gpu_name": name,
                 "vram_mb": vram_mb,
@@ -170,9 +194,16 @@ def detect_hardware() -> dict:
                 "ram_mb": ram_mb,
                 "cuda_available": True,
             }
+            if torch_gpu is not None and torch_gpu.get("gcn_arch"):
+                profile["gcn_arch"] = torch_gpu["gcn_arch"]
+            return profile
     except ImportError:
         pass
 
+    torch_inventory = _gpu_inventory_profile(torch_gpu, ram_mb)
+    if torch_inventory is not None:
+        return torch_inventory
+
     nvidia_gpu = _gpu_inventory_profile(_detect_nvidia_smi_gpu_memory(), ram_mb)
     if nvidia_gpu is not None:
         return nvidia_gpu