mcp_call_gpu_available
Check GPU availability and performance in conda environments for PyTorch or TensorFlow, verifying Metal acceleration setup and providing benchmark comparisons.
Instructions
Check if GPU is available in torch for a specific conda environment.
Input: torch or tensorflow
if framework is not provided, it will default to torch.
Returns a detailed dictionary with the following information:
- "torch_version": PyTorch version string
- "python_version": Python version string
- "platform": Platform information string
- "processor": Processor type
- "architecture": CPU architecture
- "mps_available": True if MPS (Metal Performance Shaders) is available
- "mps_built": True if PyTorch was built with MPS support
- "mps_functional": True if MPS is functional, False otherwise
- "benchmarks": A list of benchmark results for different matrix sizes, each containing:
- "size": Matrix size used for benchmark
- "cpu_time": Time taken on CPU (seconds)
- "mps_time": Time taken on MPS (seconds)
- "speedup": Ratio of CPU time to MPS time (higher means MPS is faster)
This helps determine if GPU acceleration via Apple's Metal is properly configured
and functioning, with performance benchmarks for comparison.Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| env_name | Yes | ||
| framework | No | torch |
Implementation Reference
- src/mcp_server_my_mac/server.py:80-109 (handler)Primary handler for the 'mcp_call_gpu_available' tool. Decorated with @mcp.tool for registration. Dispatches to specific framework helpers based on input.
@mcp.tool(name="mcp_call_gpu_available") async def mcp_call_gpu_available(env_name: str, framework: str = "torch") -> dict: """ Check if GPU is available in torch for a specific conda environment. Input: torch or tensorflow if framework is not provided, it will default to torch. Returns a detailed dictionary with the following information: - "torch_version": PyTorch version string - "python_version": Python version string - "platform": Platform information string - "processor": Processor type - "architecture": CPU architecture - "mps_available": True if MPS (Metal Performance Shaders) is available - "mps_built": True if PyTorch was built with MPS support - "mps_functional": True if MPS is functional, False otherwise - "benchmarks": A list of benchmark results for different matrix sizes, each containing: - "size": Matrix size used for benchmark - "cpu_time": Time taken on CPU (seconds) - "mps_time": Time taken on MPS (seconds) - "speedup": Ratio of CPU time to MPS time (higher means MPS is faster) This helps determine if GPU acceleration via Apple's Metal is properly configured and functioning, with performance benchmarks for comparison. """ if framework == "torch": return load_gpu_available_mac_torch(env_name) elif framework == "tensorflow": return load_gpu_available_mac_tensorflow_benchmarks(env_name) return {"error": "Framework not supported"} - Helper function that executes a PyTorch script in the target conda environment to check MPS (Apple GPU) availability, functionality, and performance benchmarks via matrix multiplication timings.
def load_gpu_available_mac_torch(env_name: str) -> dict: """Get detailed information about PyTorch and MPS capabilities on Mac in the specified conda environment.""" conda_executable = find_conda_executable() if not conda_executable: return {"error": "Conda executable not found"} # Create a temporary Python script with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as f: f.write( "import torch\n" "import platform\n" "import sys\n" "import json\n" "import time\n" "\n" "info = {\n" " 'torch_version': torch.__version__,\n" " 'python_version': sys.version.split()[0],\n" " 'platform': platform.platform(),\n" " 'processor': platform.processor(),\n" " 'architecture': platform.machine(),\n" " 'mps_available': torch.backends.mps.is_available(),\n" " 'mps_built': torch.backends.mps.is_built(),\n" " 'benchmarks': []\n" "}\n" "\n" "# Try to get device information if MPS is available\n" "if torch.backends.mps.is_available():\n" " try:\n" " # Test MPS with a small tensor operation\n" " device = torch.device('mps')\n" " x = torch.ones(10, 10, device=device)\n" " y = x + x\n" " info['mps_functional'] = bool((y == 2).all().item())\n" " \n" " # GPU vs CPU benchmark\n" " matrix_sizes = [5000] # Test with different sizes\n" " \n" " for size in matrix_sizes:\n" " benchmark = {'size': size}\n" " \n" " # CPU benchmark\n" " a_cpu = torch.randn(size, size)\n" " b_cpu = torch.randn(size, size)\n" " start = time.time()\n" " c_cpu = torch.matmul(a_cpu, b_cpu)\n" " cpu_time = time.time() - start\n" " benchmark['cpu_time'] = cpu_time\n" " \n" " # MPS benchmark\n" " a_mps = torch.randn(size, size, device=device)\n" " b_mps = torch.randn(size, size, device=device)\n" " torch.mps.synchronize()\n" " start = time.time()\n" " c_mps = torch.matmul(a_mps, b_mps)\n" " torch.mps.synchronize()\n" " mps_time = time.time() - start\n" " benchmark['mps_time'] = mps_time\n" " \n" " # Calculate speedup\n" " benchmark['speedup'] = cpu_time / mps_time if mps_time > 0 else 0\n" " info['benchmarks'].append(benchmark)\n" " \n" " except Exception as e:\n" " info['mps_error'] = str(e)\n" " info['mps_functional'] = False\n" "else:\n" " # Get reason why MPS is not available\n" " info['mps_not_available_reason'] = 'PyTorch not built with MPS support' " "if not torch.backends.mps.is_built() else 'Hardware/OS not supported'\n" "\n" "print(json.dumps(info))" ) script_path = f.name try: # Create a clean environment without venv variables clean_env = os.environ.copy() # Remove virtual environment variables that might interfere for var in list(clean_env.keys()): if var.startswith("VIRTUAL_ENV") or var.startswith("PYTHONHOME"): clean_env.pop(var, None) # Update PATH to remove .venv entries if "PATH" in clean_env: path_parts = clean_env["PATH"].split(os.pathsep) clean_path = os.pathsep.join([p for p in path_parts if ".venv" not in p]) clean_env["PATH"] = clean_path # Execute with clean environment command = f"{conda_executable} run -n {env_name} python {script_path}" logging.debug(f"Executing: {command}") result = subprocess.run( command, shell=True, capture_output=True, text=True, env=clean_env, # Use the clean environment ) if result.returncode == 0: try: gpu_info = json.loads(result.stdout.strip()) logging.info(f"PyTorch MPS check: {gpu_info}") return gpu_info except json.JSONDecodeError: error_msg = f"Failed to parse PyTorch output: {result.stdout}" logging.error(error_msg) return {"error": error_msg, "raw_output": result.stdout} else: error_msg = f"Failed to check PyTorch MPS availability: {result.stderr}" logging.error(error_msg) return {"error": error_msg, "returncode": result.returncode} finally: # Clean up the temporary file os.remove(script_path) - Helper function for TensorFlow GPU (MPS) availability check and benchmarks in the specified conda environment, analogous to the PyTorch version.
def load_gpu_available_mac_tensorflow_benchmarks(env_name: str) -> dict: """Get detailed information about TensorFlow and MPS capabilities on Mac in the specified conda environment.""" conda_executable = find_conda_executable() if not conda_executable: return {"error": "Conda executable not found"} # Create a temporary Python script with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as f: f.write( "import tensorflow as tf\n" "import platform\n" "import sys\n" "import json\n" "import time\n" "import numpy as np\n" "\n" "info = {\n" " 'tf_version': tf.__version__,\n" " 'python_version': sys.version.split()[0],\n" " 'platform': platform.platform(),\n" " 'processor': platform.processor(),\n" " 'architecture': platform.machine(),\n" " 'gpu_devices': tf.config.list_physical_devices('GPU'),\n" " 'benchmarks': []\n" "}\n" "\n" "try:\n" " # Check if MPS is available\n" " info['mps_available'] = len(tf.config.list_physical_devices('GPU')) > 0\n" " \n" " if info['mps_available']:\n" " # Test with matrix multiplication benchmark\n" " matrix_sizes = [5000] # Test with different sizes\n" " \n" " for size in matrix_sizes:\n" " benchmark = {'size': size}\n" " \n" " # CPU benchmark\n" " with tf.device('/CPU:0'):\n" " a_cpu = tf.random.normal([size, size])\n" " b_cpu = tf.random.normal([size, size])\n" " start = time.time()\n" " c_cpu = tf.matmul(a_cpu, b_cpu)\n" " _ = c_cpu.numpy() # Force execution\n" " cpu_time = time.time() - start\n" " benchmark['cpu_time'] = cpu_time\n" " \n" " # GPU/MPS benchmark\n" " with tf.device('/GPU:0'):\n" " a_gpu = tf.random.normal([size, size])\n" " b_gpu = tf.random.normal([size, size])\n" " # Warmup run\n" " _ = tf.matmul(a_gpu, b_gpu).numpy()\n" " \n" " start = time.time()\n" " c_gpu = tf.matmul(a_gpu, b_gpu)\n" " _ = c_gpu.numpy() # Force execution\n" " gpu_time = time.time() - start\n" " benchmark['gpu_time'] = gpu_time\n" " \n" " # Calculate speedup\n" " benchmark['speedup'] = cpu_time / gpu_time if gpu_time > 0 else 0\n" " info['benchmarks'].append(benchmark)\n" " \n" " info['mps_functional'] = True\n" "except Exception as e:\n" " info['error'] = str(e)\n" " info['mps_functional'] = False\n" "\n" "print(json.dumps(info))" ) script_path = f.name try: # Create a clean environment without venv variables clean_env = os.environ.copy() # Remove virtual environment variables that might interfere for var in list(clean_env.keys()): if var.startswith("VIRTUAL_ENV") or var.startswith("PYTHONHOME"): clean_env.pop(var, None) # Update PATH to remove .venv entries if "PATH" in clean_env: path_parts = clean_env["PATH"].split(os.pathsep) clean_path = os.pathsep.join([p for p in path_parts if ".venv" not in p]) clean_env["PATH"] = clean_path # Execute with clean environment command = f"{conda_executable} run -n {env_name} python {script_path}" logging.debug(f"Executing: {command}") result = subprocess.run( command, shell=True, capture_output=True, text=True, env=clean_env, ) if result.returncode == 0: try: gpu_info = json.loads(result.stdout.strip()) logging.info(f"TensorFlow MPS check: {gpu_info}") return gpu_info except json.JSONDecodeError: error_msg = f"Failed to parse TensorFlow output: {result.stdout}" logging.error(error_msg) return {"error": error_msg, "raw_output": result.stdout} else: error_msg = f"Failed to check TensorFlow MPS availability: {result.stderr}" logging.error(error_msg) return {"error": error_msg, "returncode": result.returncode} finally: # Clean up the temporary file os.remove(script_path)