| """ |
| GPU Tools for SPARKNET |
| Tools for GPU monitoring and management |
| """ |
|
|
| from typing import Optional |
| from loguru import logger |
| from .base_tool import BaseTool, ToolResult |
| from ..utils.gpu_manager import get_gpu_manager |
|
|
|
|
| class GPUMonitorTool(BaseTool): |
| """Tool for monitoring GPU status.""" |
|
|
| def __init__(self): |
| super().__init__( |
| name="gpu_monitor", |
| description="Monitor GPU status, memory usage, and utilization", |
| ) |
| self.add_parameter("gpu_id", "int", "Specific GPU ID to monitor (optional)", required=False, default=None) |
| self.gpu_manager = get_gpu_manager() |
|
|
| async def execute(self, gpu_id: Optional[int] = None, **kwargs) -> ToolResult: |
| """ |
| Monitor GPU status. |
| |
| Args: |
| gpu_id: Specific GPU ID or None for all GPUs |
| |
| Returns: |
| ToolResult with GPU information |
| """ |
| try: |
| if gpu_id is not None: |
| |
| info = self.gpu_manager.get_gpu_info(gpu_id) |
|
|
| if "error" in info: |
| return ToolResult( |
| success=False, |
| output=None, |
| error=info["error"], |
| ) |
|
|
| output = self._format_gpu_info(info) |
|
|
| return ToolResult( |
| success=True, |
| output=output, |
| metadata=info, |
| ) |
| else: |
| |
| all_info = self.gpu_manager.get_all_gpu_info() |
|
|
| output_lines = [] |
| for info in all_info: |
| if "error" not in info: |
| output_lines.append(self._format_gpu_info(info)) |
|
|
| output = "\n\n".join(output_lines) |
|
|
| return ToolResult( |
| success=True, |
| output=output, |
| metadata={"gpus": all_info}, |
| ) |
|
|
| except Exception as e: |
| logger.error(f"GPU monitoring error: {e}") |
| return ToolResult( |
| success=False, |
| output=None, |
| error=f"Monitoring error: {str(e)}", |
| ) |
|
|
| def _format_gpu_info(self, info: dict) -> str: |
| """Format GPU info for display.""" |
| return ( |
| f"GPU {info['gpu_id']}: {info['name']}\n" |
| f" Memory: {info['memory_used'] / 1024**3:.2f} GB / {info['memory_total'] / 1024**3:.2f} GB " |
| f"({info['memory_percent']:.1f}% used)\n" |
| f" Free Memory: {info['memory_free'] / 1024**3:.2f} GB\n" |
| f" GPU Utilization: {info['gpu_utilization']}%\n" |
| f" Temperature: {info['temperature']}°C" |
| ) |
|
|
|
|
| class GPUSelectTool(BaseTool): |
| """Tool for selecting best available GPU.""" |
|
|
| def __init__(self): |
| super().__init__( |
| name="gpu_select", |
| description="Select the best available GPU based on free memory", |
| ) |
| self.add_parameter("min_memory_gb", "float", "Minimum required memory in GB", required=False, default=8.0) |
| self.gpu_manager = get_gpu_manager() |
|
|
| async def execute(self, min_memory_gb: float = 8.0, **kwargs) -> ToolResult: |
| """ |
| Select best GPU. |
| |
| Args: |
| min_memory_gb: Minimum required memory |
| |
| Returns: |
| ToolResult with selected GPU ID |
| """ |
| try: |
| gpu_id = self.gpu_manager.select_best_gpu(min_memory_gb) |
|
|
| if gpu_id is None: |
| return ToolResult( |
| success=False, |
| output=None, |
| error=f"No GPU found with {min_memory_gb} GB free memory", |
| ) |
|
|
| info = self.gpu_manager.get_gpu_info(gpu_id) |
|
|
| output = ( |
| f"Selected GPU {gpu_id}: {info['name']}\n" |
| f"Free Memory: {info['memory_free'] / 1024**3:.2f} GB" |
| ) |
|
|
| return ToolResult( |
| success=True, |
| output=output, |
| metadata={ |
| "gpu_id": gpu_id, |
| "gpu_info": info, |
| }, |
| ) |
|
|
| except Exception as e: |
| logger.error(f"GPU selection error: {e}") |
| return ToolResult( |
| success=False, |
| output=None, |
| error=f"Selection error: {str(e)}", |
| ) |
|
|
