mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
feat: add compute node detection for available GPUs/WebGPU
- Add ComputeNodeDetector to detect available compute resources - Detects CUDA GPUs, Metal GPUs, WebGPU connections, and CPU cores - Shows '1 node available (1 Metal GPU)' instead of confusing process count - Display node count in dev mode and SSE transport (not stdio to avoid protocol corruption) - Ready for WebGPU connections from Chrome (would show as 2 nodes)
This commit is contained in:
@@ -0,0 +1,179 @@
|
||||
"""Compute node detection and management for distributed processing."""
|
||||
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
from typing import List, Dict, Any
|
||||
|
||||
|
||||
class ComputeNodeDetector:
|
||||
"""Detect available compute nodes (GPUs, WebGPU, CPUs) for distributed work."""
|
||||
|
||||
@staticmethod
|
||||
def detect_local_gpus() -> List[Dict[str, Any]]:
|
||||
"""Detect local GPU devices."""
|
||||
gpus = []
|
||||
|
||||
# Try NVIDIA GPUs
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=2
|
||||
)
|
||||
if result.returncode == 0:
|
||||
for line in result.stdout.strip().split('\n'):
|
||||
if line:
|
||||
name, memory = line.split(', ')
|
||||
gpus.append({
|
||||
"type": "cuda",
|
||||
"name": name,
|
||||
"memory": memory,
|
||||
"id": f"cuda:{len(gpus)}"
|
||||
})
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
# Try Metal GPUs (macOS)
|
||||
if platform.system() == "Darwin":
|
||||
try:
|
||||
# Check for Metal support
|
||||
result = subprocess.run(
|
||||
["system_profiler", "SPDisplaysDataType"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=2
|
||||
)
|
||||
if result.returncode == 0 and "Metal" in result.stdout:
|
||||
# Parse GPU info from system_profiler
|
||||
lines = result.stdout.split('\n')
|
||||
for i, line in enumerate(lines):
|
||||
if 'Chipset Model:' in line:
|
||||
gpu_name = line.split(':')[1].strip()
|
||||
gpus.append({
|
||||
"type": "metal",
|
||||
"name": gpu_name,
|
||||
"memory": "Shared",
|
||||
"id": f"metal:{len(gpus)}"
|
||||
})
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
return gpus
|
||||
|
||||
@staticmethod
|
||||
def detect_webgpu_nodes() -> List[Dict[str, Any]]:
|
||||
"""Detect connected WebGPU nodes (from browsers)."""
|
||||
webgpu_nodes = []
|
||||
|
||||
# Check for WebGPU connections (would need actual WebSocket/server to track)
|
||||
# For now, check if a WebGPU server is running
|
||||
webgpu_port = os.environ.get("HANZO_WEBGPU_PORT", "8765")
|
||||
try:
|
||||
import socket
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
result = sock.connect_ex(('localhost', int(webgpu_port)))
|
||||
sock.close()
|
||||
if result == 0:
|
||||
webgpu_nodes.append({
|
||||
"type": "webgpu",
|
||||
"name": "Chrome WebGPU",
|
||||
"memory": "Browser",
|
||||
"id": "webgpu:0"
|
||||
})
|
||||
except:
|
||||
pass
|
||||
|
||||
return webgpu_nodes
|
||||
|
||||
@staticmethod
|
||||
def detect_cpu_nodes() -> List[Dict[str, Any]]:
|
||||
"""Detect CPU compute nodes."""
|
||||
import multiprocessing
|
||||
|
||||
return [{
|
||||
"type": "cpu",
|
||||
"name": f"{platform.processor() or 'CPU'}",
|
||||
"cores": multiprocessing.cpu_count(),
|
||||
"id": "cpu:0"
|
||||
}]
|
||||
|
||||
@classmethod
|
||||
def get_all_nodes(cls) -> List[Dict[str, Any]]:
|
||||
"""Get all available compute nodes."""
|
||||
nodes = []
|
||||
|
||||
# Detect GPUs
|
||||
gpus = cls.detect_local_gpus()
|
||||
nodes.extend(gpus)
|
||||
|
||||
# Detect WebGPU connections
|
||||
webgpu = cls.detect_webgpu_nodes()
|
||||
nodes.extend(webgpu)
|
||||
|
||||
# If no GPUs/WebGPU, add CPU as compute node
|
||||
if not nodes:
|
||||
nodes.extend(cls.detect_cpu_nodes())
|
||||
|
||||
return nodes
|
||||
|
||||
@classmethod
|
||||
def get_node_count(cls) -> int:
|
||||
"""Get total number of available compute nodes."""
|
||||
return len(cls.get_all_nodes())
|
||||
|
||||
@classmethod
|
||||
def get_node_summary(cls) -> str:
|
||||
"""Get a summary string of available nodes."""
|
||||
nodes = cls.get_all_nodes()
|
||||
if not nodes:
|
||||
return "No compute nodes available"
|
||||
|
||||
count = len(nodes)
|
||||
node_word = "node" if count == 1 else "nodes"
|
||||
|
||||
# Group by type
|
||||
types = {}
|
||||
for node in nodes:
|
||||
node_type = node["type"]
|
||||
if node_type not in types:
|
||||
types[node_type] = 0
|
||||
types[node_type] += 1
|
||||
|
||||
# Build summary
|
||||
parts = []
|
||||
for node_type, type_count in types.items():
|
||||
if node_type == "cuda":
|
||||
parts.append(f"{type_count} CUDA GPU{'s' if type_count > 1 else ''}")
|
||||
elif node_type == "metal":
|
||||
parts.append(f"{type_count} Metal GPU{'s' if type_count > 1 else ''}")
|
||||
elif node_type == "webgpu":
|
||||
parts.append(f"{type_count} WebGPU")
|
||||
elif node_type == "cpu":
|
||||
parts.append(f"{type_count} CPU")
|
||||
|
||||
type_str = ", ".join(parts)
|
||||
return f"{count} {node_word} available ({type_str})"
|
||||
|
||||
|
||||
def print_node_status():
|
||||
"""Print current node status."""
|
||||
detector = ComputeNodeDetector()
|
||||
nodes = detector.get_all_nodes()
|
||||
|
||||
print(f"\n🖥️ Compute Nodes: {len(nodes)}")
|
||||
for node in nodes:
|
||||
if node["type"] in ["cuda", "metal"]:
|
||||
print(f" • {node['id']}: {node['name']} ({node['memory']})")
|
||||
elif node["type"] == "webgpu":
|
||||
print(f" • {node['id']}: {node['name']}")
|
||||
elif node["type"] == "cpu":
|
||||
print(f" • {node['id']}: {node['name']} ({node['cores']} cores)")
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test the detector
|
||||
print_node_status()
|
||||
print(ComputeNodeDetector.get_node_summary())
|
||||
@@ -172,6 +172,17 @@ class DevServer:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"\n🚀 Starting Hanzo AI in development mode...")
|
||||
|
||||
# Show compute nodes
|
||||
try:
|
||||
from hanzo_mcp.compute_nodes import ComputeNodeDetector
|
||||
detector = ComputeNodeDetector()
|
||||
summary = detector.get_node_summary()
|
||||
logger.info(f"🖥️ {summary}")
|
||||
except Exception:
|
||||
# Silently ignore if compute node detection fails
|
||||
pass
|
||||
|
||||
logger.info(f"🔧 Hot reload enabled - watching for file changes")
|
||||
logger.info(f"📁 Project: {self.project_dir or 'current directory'}")
|
||||
logger.info(f"🌐 Transport: {transport}\n")
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""MCP server implementing Hanzo capabilities."""
|
||||
|
||||
import os
|
||||
import atexit
|
||||
import signal
|
||||
import logging
|
||||
@@ -219,6 +220,18 @@ class HanzoMCPServer:
|
||||
for path in allowed_paths_list:
|
||||
self.permission_manager.add_allowed_path(path)
|
||||
|
||||
# Show compute nodes only in non-stdio mode (to avoid corrupting protocol)
|
||||
if transport != "stdio" and not os.environ.get("HANZO_QUIET"):
|
||||
try:
|
||||
from hanzo_mcp.compute_nodes import ComputeNodeDetector
|
||||
detector = ComputeNodeDetector()
|
||||
summary = detector.get_node_summary()
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"🖥️ {summary}")
|
||||
except Exception:
|
||||
# Silently ignore if compute node detection fails
|
||||
pass
|
||||
|
||||
# Set up cleanup handlers before running
|
||||
self._setup_cleanup_handlers()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user