@hyperspaceng/neural-pods 0.60.0

package/README.md

@@ -0,0 +1,511 @@
+# pi
+
+Deploy and manage LLMs on GPU pods with automatic vLLM configuration for agentic workloads.
+
+## Installation
+
+```bash
+npm install -g @mariozechner/pi
+```
+
+## What is pi?
+
+`pi` simplifies running large language models on remote GPU pods. It automatically:
+- Sets up vLLM on fresh Ubuntu pods
+- Configures tool calling for agentic models (Qwen, GPT-OSS, GLM, etc.)
+- Manages multiple models on the same pod with "smart" GPU allocation
+- Provides OpenAI-compatible API endpoints for each model
+- Includes an interactive agent with file system tools for testing
+
+## Quick Start
+
+```bash
+# Set required environment variables
+export HF_TOKEN=your_huggingface_token      # Get from https://huggingface.co/settings/tokens
+export PI_API_KEY=your_api_key              # Any string you want for API authentication
+
+# Setup a DataCrunch pod with NFS storage (models path auto-extracted)
+pi pods setup dc1 "ssh root@1.2.3.4" \
+  --mount "sudo mount -t nfs -o nconnect=16 nfs.fin-02.datacrunch.io:/your-pseudo /mnt/hf-models"
+
+# Start a model (automatic configuration for known models)
+pi start Qwen/Qwen2.5-Coder-32B-Instruct --name qwen
+
+# Send a single message to the model
+pi agent qwen "What is the Fibonacci sequence?"
+
+# Interactive chat mode with file system tools
+pi agent qwen -i
+
+# Use with any OpenAI-compatible client
+export OPENAI_BASE_URL='http://1.2.3.4:8001/v1'
+export OPENAI_API_KEY=$PI_API_KEY
+```
+
+## Prerequisites
+
+- Node.js 18+
+- HuggingFace token (for model downloads)
+- GPU pod with:
+  - Ubuntu 22.04 or 24.04
+  - SSH root access
+  - NVIDIA drivers installed
+  - Persistent storage for models
+
+## Supported Providers
+
+### Primary Support
+
+**DataCrunch** - Best for shared model storage
+- NFS volumes sharable across multiple pods in same region
+- Models download once, use everywhere
+- Ideal for teams or multiple experiments
+
+**RunPod** - Good persistent storage
+- Network volumes persist independently
+- Cannot share between running pods simultaneously
+- Good for single-pod workflows
+
+### Also Works With
+- Vast.ai (volumes locked to specific machine)
+- Prime Intellect (no persistent storage)
+- AWS EC2 (with EFS setup)
+- Any Ubuntu machine with NVIDIA GPUs, CUDA driver, and SSH
+
+## Commands
+
+### Pod Management
+
+```bash
+pi pods setup <name> "<ssh>" [options]        # Setup new pod
+  --mount "<mount_command>"                   # Run mount command during setup
+  --models-path <path>                        # Override extracted path (optional)
+  --vllm release|nightly|gpt-oss              # vLLM version (default: release)
+
+pi pods                                       # List all configured pods
+pi pods active <name>                         # Switch active pod
+pi pods remove <name>                         # Remove pod from local config
+pi shell [<name>]                             # SSH into pod
+pi ssh [<name>] "<command>"                   # Run command on pod
+```
+
+**Note**: When using `--mount`, the models path is automatically extracted from the mount command's target directory. You only need `--models-path` if not using `--mount` or to override the extracted path.
+
+#### vLLM Version Options
+
+- `release` (default): Stable vLLM release, recommended for most users
+- `nightly`: Latest vLLM features, needed for newest models like GLM-4.5
+- `gpt-oss`: Special build for OpenAI's GPT-OSS models only
+
+### Model Management
+
+```bash
+pi start <model> --name <name> [options]  # Start a model
+  --memory <percent>      # GPU memory: 30%, 50%, 90% (default: 90%)
+  --context <size>        # Context window: 4k, 8k, 16k, 32k, 64k, 128k
+  --gpus <count>          # Number of GPUs to use (predefined models only)
+  --pod <name>            # Target specific pod (overrides active)
+  --vllm <args...>        # Pass custom args directly to vLLM
+
+pi stop [<name>]          # Stop model (or all if no name given)
+pi list                   # List running models with status
+pi logs <name>            # Stream model logs (tail -f)
+```
+
+### Agent & Chat Interface
+
+```bash
+pi agent <name> "<message>"               # Single message to model
+pi agent <name> "<msg1>" "<msg2>"         # Multiple messages in sequence
+pi agent <name> -i                        # Interactive chat mode
+pi agent <name> -i -c                     # Continue previous session
+
+# Standalone OpenAI-compatible agent (works with any API)
+pi-agent --base-url http://localhost:8000/v1 --model llama-3.1 "Hello"
+pi-agent --api-key sk-... "What is 2+2?"  # Uses OpenAI by default
+pi-agent --json "What is 2+2?"            # Output event stream as JSONL
+pi-agent -i                                # Interactive mode
+```
+
+The agent includes tools for file operations (read, list, bash, glob, rg) to test agentic capabilities, particularly useful for code navigation and analysis tasks.
+
+## Predefined Model Configurations
+
+`pi` includes predefined configurations for popular agentic models, so you do not have to specify `--vllm` arguments manually. `pi` will also check if the model you selected can actually run on your pod with respect to the number of GPUs and available VRAM. Run `pi start` without additional arguments to see a list of predefined models that can run on the active pod.
+
+### Qwen Models
+```bash
+# Qwen2.5-Coder-32B - Excellent coding model, fits on single H100/H200
+pi start Qwen/Qwen2.5-Coder-32B-Instruct --name qwen
+
+# Qwen3-Coder-30B - Advanced reasoning with tool use
+pi start Qwen/Qwen3-Coder-30B-A3B-Instruct --name qwen3
+
+# Qwen3-Coder-480B - State-of-the-art on 8xH200 (data-parallel mode)
+pi start Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8 --name qwen-480b
+```
+
+### GPT-OSS Models
+```bash
+# Requires special vLLM build during setup
+pi pods setup gpt-pod "ssh root@1.2.3.4" --models-path /workspace --vllm gpt-oss
+
+# GPT-OSS-20B - Fits on 16GB+ VRAM
+pi start openai/gpt-oss-20b --name gpt20
+
+# GPT-OSS-120B - Needs 60GB+ VRAM
+pi start openai/gpt-oss-120b --name gpt120
+```
+
+### GLM Models
+```bash
+# GLM-4.5 - Requires 8-16 GPUs, includes thinking mode
+pi start zai-org/GLM-4.5 --name glm
+
+# GLM-4.5-Air - Smaller version, 1-2 GPUs
+pi start zai-org/GLM-4.5-Air --name glm-air
+```
+
+### Custom Models with --vllm
+
+For models not in the predefined list, use `--vllm` to pass arguments directly to vLLM:
+
+```bash
+# DeepSeek with custom settings
+pi start deepseek-ai/DeepSeek-V3 --name deepseek --vllm \
+  --tensor-parallel-size 4 --trust-remote-code
+
+# Mistral with pipeline parallelism
+pi start mistralai/Mixtral-8x22B-Instruct-v0.1 --name mixtral --vllm \
+  --tensor-parallel-size 8 --pipeline-parallel-size 2
+
+# Any model with specific tool parser
+pi start some/model --name mymodel --vllm \
+  --tool-call-parser hermes --enable-auto-tool-choice
+```
+
+## DataCrunch Setup
+
+DataCrunch offers the best experience with shared NFS storage across pods:
+
+### 1. Create Shared Filesystem (SFS)
+- Go to DataCrunch dashboard → Storage → Create SFS
+- Choose size and datacenter
+- Note the mount command (e.g., `sudo mount -t nfs -o nconnect=16 nfs.fin-02.datacrunch.io:/hf-models-fin02-8ac1bab7 /mnt/hf-models-fin02`)
+
+### 2. Create GPU Instance
+- Create instance in same datacenter as SFS
+- Share the SFS with the instance
+- Get SSH command from dashboard
+
+### 3. Setup with pi
+```bash
+# Get mount command from DataCrunch dashboard
+pi pods setup dc1 "ssh root@instance.datacrunch.io" \
+  --mount "sudo mount -t nfs -o nconnect=16 nfs.fin-02.datacrunch.io:/your-pseudo /mnt/hf-models"
+
+# Models automatically stored in /mnt/hf-models (extracted from mount command)
+```
+
+### 4. Benefits
+- Models persist across instance restarts
+- Share models between multiple instances in same datacenter
+- Download once, use everywhere
+- Pay only for storage, not compute time during downloads
+
+## RunPod Setup
+
+RunPod offers good persistent storage with network volumes:
+
+### 1. Create Network Volume (optional)
+- Go to RunPod dashboard → Storage → Create Network Volume
+- Choose size and region
+
+### 2. Create GPU Pod
+- Select "Network Volume" during pod creation (if using)
+- Attach your volume to `/runpod-volume`
+- Get SSH command from pod details
+
+### 3. Setup with pi
+```bash
+# With network volume
+pi pods setup runpod "ssh root@pod.runpod.io" --models-path /runpod-volume
+
+# Or use workspace (persists with pod but not shareable)
+pi pods setup runpod "ssh root@pod.runpod.io" --models-path /workspace
+```
+
+
+## Multi-GPU Support
+
+### Automatic GPU Assignment
+When running multiple models, pi automatically assigns them to different GPUs:
+```bash
+pi start model1 --name m1  # Auto-assigns to GPU 0
+pi start model2 --name m2  # Auto-assigns to GPU 1
+pi start model3 --name m3  # Auto-assigns to GPU 2
+```
+
+### Specify GPU Count for Predefined Models
+For predefined models with multiple configurations, use `--gpus` to control GPU usage:
+```bash
+# Run Qwen on 1 GPU instead of all available
+pi start Qwen/Qwen2.5-Coder-32B-Instruct --name qwen --gpus 1
+
+# Run GLM-4.5 on 8 GPUs (if it has an 8-GPU config)
+pi start zai-org/GLM-4.5 --name glm --gpus 8
+```
+
+If the model doesn't have a configuration for the requested GPU count, you'll see available options.
+
+### Tensor Parallelism for Large Models
+For models that don't fit on a single GPU:
+```bash
+# Use all available GPUs
+pi start meta-llama/Llama-3.1-70B-Instruct --name llama70b --vllm \
+  --tensor-parallel-size 4
+
+# Specific GPU count
+pi start Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8 --name qwen480 --vllm \
+  --data-parallel-size 8 --enable-expert-parallel
+```
+
+## API Integration
+
+All models expose OpenAI-compatible endpoints:
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+    base_url="http://your-pod-ip:8001/v1",
+    api_key="your-pi-api-key"
+)
+
+# Chat completion with tool calling
+response = client.chat.completions.create(
+    model="Qwen/Qwen2.5-Coder-32B-Instruct",
+    messages=[
+        {"role": "user", "content": "Write a Python function to calculate fibonacci"}
+    ],
+    tools=[{
+        "type": "function",
+        "function": {
+            "name": "execute_code",
+            "description": "Execute Python code",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "code": {"type": "string"}
+                },
+                "required": ["code"]
+            }
+        }
+    }],
+    tool_choice="auto"
+)
+```
+
+## Standalone Agent CLI
+
+`pi` includes a standalone OpenAI-compatible agent that can work with any API:
+
+```bash
+# Install globally to get pi-agent command
+npm install -g @mariozechner/pi
+
+# Use with OpenAI
+pi-agent --api-key sk-... "What is machine learning?"
+
+# Use with local vLLM
+pi-agent --base-url http://localhost:8000/v1 \
+         --model meta-llama/Llama-3.1-8B-Instruct \
+         --api-key dummy \
+         "Explain quantum computing"
+
+# Interactive mode
+pi-agent -i
+
+# Continue previous session
+pi-agent --continue "Follow up question"
+
+# Custom system prompt
+pi-agent --system-prompt "You are a Python expert" "Write a web scraper"
+
+# Use responses API (for GPT-OSS models)
+pi-agent --api responses --model openai/gpt-oss-20b "Hello"
+```
+
+The agent supports:
+- Session persistence across conversations
+- Interactive TUI mode with syntax highlighting
+- File system tools (read, list, bash, glob, rg) for code navigation
+- Both Chat Completions and Responses API formats
+- Custom system prompts
+
+## Tool Calling Support
+
+`pi` automatically configures appropriate tool calling parsers for known models:
+
+- **Qwen models**: `hermes` parser (Qwen3-Coder uses `qwen3_coder`)
+- **GLM models**: `glm4_moe` parser with reasoning support
+- **GPT-OSS models**: Uses `/v1/responses` endpoint, as tool calling (function calling in OpenAI parlance) is currently a [WIP with the `v1/chat/completions` endpoint](https://docs.vllm.ai/projects/recipes/en/latest/OpenAI/GPT-OSS.html#tool-use).
+- **Custom models**: Specify with `--vllm --tool-call-parser <parser> --enable-auto-tool-choice`
+
+To disable tool calling:
+```bash
+pi start model --name mymodel --vllm --disable-tool-call-parser
+```
+
+## Memory and Context Management
+
+### GPU Memory Allocation
+Controls how much GPU memory vLLM pre-allocates:
+- `--memory 30%`: High concurrency, limited context
+- `--memory 50%`: Balanced (default)
+- `--memory 90%`: Maximum context, low concurrency
+
+### Context Window
+Sets maximum input + output tokens:
+- `--context 4k`: 4,096 tokens total
+- `--context 32k`: 32,768 tokens total
+- `--context 128k`: 131,072 tokens total
+
+Example for coding workload:
+```bash
+# Large context for code analysis, moderate concurrency
+pi start Qwen/Qwen2.5-Coder-32B-Instruct --name coder \
+  --context 64k --memory 70%
+```
+
+**Note**: When using `--vllm`, the `--memory`, `--context`, and `--gpus` parameters are ignored. You'll see a warning if you try to use them together.
+
+## Session Persistence
+
+The interactive agent mode (`-i`) saves sessions for each project directory:
+
+```bash
+# Start new session
+pi agent qwen -i
+
+# Continue previous session (maintains chat history)
+pi agent qwen -i -c
+```
+
+Sessions are stored in `~/.pi/sessions/` organized by project path and include:
+- Complete conversation history
+- Tool call results
+- Token usage statistics
+
+## Architecture & Event System
+
+The agent uses a unified event-based architecture where all interactions flow through `AgentEvent` types. This enables:
+- Consistent UI rendering across console and TUI modes
+- Session recording and replay
+- Clean separation between API calls and UI updates
+- JSON output mode for programmatic integration
+
+Events are automatically converted to the appropriate API format (Chat Completions or Responses) based on the model type.
+
+### JSON Output Mode
+
+Use `--json` flag to output the event stream as JSONL (JSON Lines) for programmatic consumption:
+```bash
+pi-agent --api-key sk-... --json "What is 2+2?"
+```
+
+Each line is a complete JSON object representing an event:
+```jsonl
+{"type":"user_message","text":"What is 2+2?"}
+{"type":"assistant_start"}
+{"type":"assistant_message","text":"2 + 2 = 4"}
+{"type":"token_usage","inputTokens":10,"outputTokens":5,"totalTokens":15,"cacheReadTokens":0,"cacheWriteTokens":0}
+```
+
+## Troubleshooting
+
+### OOM (Out of Memory) Errors
+- Reduce `--memory` percentage
+- Use smaller model or quantized version (FP8)
+- Reduce `--context` size
+
+### Model Won't Start
+```bash
+# Check GPU usage
+pi ssh "nvidia-smi"
+
+# Check if port is in use
+pi list
+
+# Force stop all models
+pi stop
+```
+
+### Tool Calling Issues
+- Not all models support tool calling reliably
+- Try different parser: `--vllm --tool-call-parser mistral`
+- Or disable: `--vllm --disable-tool-call-parser`
+
+### Access Denied for Models
+Some models (Llama, Mistral) require HuggingFace access approval. Visit the model page and click "Request access".
+
+### vLLM Build Issues
+If using `--vllm nightly` fails, try:
+- Use `--vllm release` for stable version
+- Check CUDA compatibility with `pi ssh "nvidia-smi"`
+
+### Agent Not Finding Messages
+If the agent shows configuration instead of your message, ensure quotes around messages with special characters:
+```bash
+# Good
+pi agent qwen "What is this file about?"
+
+# Bad (shell might interpret special chars)
+pi agent qwen What is this file about?
+```
+
+## Advanced Usage
+
+### Working with Multiple Pods
+```bash
+# Override active pod for any command
+pi start model --name test --pod dev-pod
+pi list --pod prod-pod
+pi stop test --pod dev-pod
+```
+
+### Custom vLLM Arguments
+```bash
+# Pass any vLLM argument after --vllm
+pi start model --name custom --vllm \
+  --quantization awq \
+  --enable-prefix-caching \
+  --max-num-seqs 256 \
+  --gpu-memory-utilization 0.95
+```
+
+### Monitoring
+```bash
+# Watch GPU utilization
+pi ssh "watch -n 1 nvidia-smi"
+
+# Check model downloads
+pi ssh "du -sh ~/.cache/huggingface/hub/*"
+
+# View all logs
+pi ssh "ls -la ~/.vllm_logs/"
+
+# Check agent session history
+ls -la ~/.pi/sessions/
+```
+
+## Environment Variables
+
+- `HF_TOKEN` - HuggingFace token for model downloads
+- `PI_API_KEY` - API key for vLLM endpoints
+- `PI_CONFIG_DIR` - Config directory (default: `~/.pi`)
+- `OPENAI_API_KEY` - Used by `pi-agent` when no `--api-key` provided
+
+## License
+
+MIT

package/package.json

@@ -0,0 +1,40 @@
+{
+	"name": "@hyperspaceng/neural-pods",
+	"version": "0.60.0",
+	"description": "CLI tool for managing vLLM deployments on GPU pods",
+	"type": "module",
+	"bin": {
+		"pi-pods": "dist/cli.js"
+	},
+	"scripts": {
+		"clean": "shx rm -rf dist",
+		"build": "tsgo -p tsconfig.build.json && shx chmod +x dist/cli.js && shx cp src/models.json dist/ && shx cp -r scripts dist/",
+		"prepublishOnly": "npm run clean && npm run build"
+	},
+	"files": [
+		"dist",
+		"scripts"
+	],
+	"keywords": [
+		"llm",
+		"vllm",
+		"gpu",
+		"ai",
+		"cli"
+	],
+	"author": "Hyperspace Technologies <hyperspace@hyperspace.ng>",
+	"license": "MIT",
+	"repository": {
+		"type": "git",
+		"url": "git+https://github.com/badlogic/pi-mono.git",
+		"directory": "packages/pods"
+	},
+	"engines": {
+		"node": ">=20.0.0"
+	},
+	"dependencies": {
+		"@hyperspaceng/neural-agent-core": "^0.60.0",
+		"chalk": "^5.5.0"
+	},
+	"devDependencies": {}
+}

package/scripts/model_run.sh

@@ -0,0 +1,83 @@
+#!/usr/bin/env bash
+# Model runner script - runs sequentially, killed by pi stop
+set -euo pipefail
+
+# These values are replaced before upload by pi CLI
+MODEL_ID="{{MODEL_ID}}"
+NAME="{{NAME}}"
+PORT="{{PORT}}"
+VLLM_ARGS="{{VLLM_ARGS}}"
+
+# Trap to ensure cleanup on exit and kill any child processes
+cleanup() {
+    local exit_code=$?
+    echo "Model runner exiting with code $exit_code"
+    # Kill any child processes
+    pkill -P $$ 2>/dev/null || true
+    exit $exit_code
+}
+trap cleanup EXIT TERM INT
+
+# Force colored output even when not a TTY
+export FORCE_COLOR=1
+export PYTHONUNBUFFERED=1
+export TERM=xterm-256color
+export RICH_FORCE_TERMINAL=1
+export CLICOLOR_FORCE=1
+
+# Source virtual environment
+source /root/venv/bin/activate
+
+echo "========================================="
+echo "Model Run: $NAME"
+echo "Model ID: $MODEL_ID"
+echo "Port: $PORT"
+if [ -n "$VLLM_ARGS" ]; then
+    echo "vLLM Args: $VLLM_ARGS"
+fi
+echo "========================================="
+echo ""
+
+# Download model (with color progress bars)
+echo "Downloading model (will skip if cached)..."
+HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$MODEL_ID"
+
+if [ $? -ne 0 ]; then
+    echo "❌ ERROR: Failed to download model" >&2
+    exit 1
+fi
+
+echo ""
+echo "✅ Model download complete"
+echo ""
+
+# Build vLLM command
+VLLM_CMD="vllm serve '$MODEL_ID' --port $PORT --api-key '$PI_API_KEY'"
+if [ -n "$VLLM_ARGS" ]; then
+    VLLM_CMD="$VLLM_CMD $VLLM_ARGS"
+fi
+
+echo "Starting vLLM server..."
+echo "Command: $VLLM_CMD"
+echo "========================================="
+echo ""
+
+# Run vLLM in background so we can monitor it
+echo "Starting vLLM process..."
+bash -c "$VLLM_CMD" &
+VLLM_PID=$!
+
+# Monitor the vLLM process
+echo "Monitoring vLLM process (PID: $VLLM_PID)..."
+wait $VLLM_PID
+VLLM_EXIT_CODE=$?
+
+if [ $VLLM_EXIT_CODE -ne 0 ]; then
+    echo "❌ ERROR: vLLM exited with code $VLLM_EXIT_CODE" >&2
+    # Make sure to exit the script command too
+    kill -TERM $$ 2>/dev/null || true
+    exit $VLLM_EXIT_CODE
+fi
+
+echo "✅ vLLM exited normally"
+exit 0

package/scripts/pod_setup.sh

@@ -0,0 +1,336 @@
+#!/usr/bin/env bash
+# GPU pod bootstrap for vLLM deployment
+set -euo pipefail
+
+# Parse arguments passed from pi CLI
+MOUNT_COMMAND=""
+MODELS_PATH=""
+HF_TOKEN=""
+PI_API_KEY=""
+VLLM_VERSION="release"  # Default to release
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --mount)
+            MOUNT_COMMAND="$2"
+            shift 2
+            ;;
+        --models-path)
+            MODELS_PATH="$2"
+            shift 2
+            ;;
+        --hf-token)
+            HF_TOKEN="$2"
+            shift 2
+            ;;
+        --vllm-api-key)
+            PI_API_KEY="$2"
+            shift 2
+            ;;
+        --vllm)
+            VLLM_VERSION="$2"
+            shift 2
+            ;;
+        *)
+            echo "ERROR: Unknown option: $1" >&2
+            exit 1
+            ;;
+    esac
+done
+
+# Validate required parameters
+if [ -z "$HF_TOKEN" ]; then
+    echo "ERROR: HF_TOKEN is required" >&2
+    exit 1
+fi
+
+if [ -z "$PI_API_KEY" ]; then
+    echo "ERROR: PI_API_KEY is required" >&2
+    exit 1
+fi
+
+if [ -z "$MODELS_PATH" ]; then
+    echo "ERROR: MODELS_PATH is required" >&2
+    exit 1
+fi
+
+echo "=== Starting pod setup ==="
+
+# Install system dependencies
+apt update -y
+apt install -y python3-pip python3-venv git build-essential cmake ninja-build curl wget lsb-release htop pkg-config
+
+# --- Install matching CUDA toolkit -------------------------------------------
+echo "Checking CUDA driver version..."
+DRIVER_CUDA_VERSION=$(nvidia-smi | grep "CUDA Version" | awk '{print $9}')
+echo "Driver supports CUDA: $DRIVER_CUDA_VERSION"
+
+# Check if nvcc exists and its version
+if command -v nvcc &> /dev/null; then
+    NVCC_VERSION=$(nvcc --version | grep "release" | awk '{print $6}' | cut -d, -f1)
+    echo "Current nvcc version: $NVCC_VERSION"
+else
+    NVCC_VERSION="none"
+    echo "nvcc not found"
+fi
+
+# Install CUDA toolkit matching driver version if needed
+if [[ "$NVCC_VERSION" != "$DRIVER_CUDA_VERSION" ]]; then
+    echo "Installing CUDA Toolkit $DRIVER_CUDA_VERSION to match driver..."
+
+    # Detect Ubuntu version
+    UBUNTU_VERSION=$(lsb_release -rs)
+    UBUNTU_CODENAME=$(lsb_release -cs)
+
+    echo "Detected Ubuntu $UBUNTU_VERSION ($UBUNTU_CODENAME)"
+
+    # Map Ubuntu version to NVIDIA repo path
+    if [[ "$UBUNTU_VERSION" == "24.04" ]]; then
+        REPO_PATH="ubuntu2404"
+    elif [[ "$UBUNTU_VERSION" == "22.04" ]]; then
+        REPO_PATH="ubuntu2204"
+    elif [[ "$UBUNTU_VERSION" == "20.04" ]]; then
+        REPO_PATH="ubuntu2004"
+    else
+        echo "Warning: Unsupported Ubuntu version $UBUNTU_VERSION, trying ubuntu2204"
+        REPO_PATH="ubuntu2204"
+    fi
+
+    # Add NVIDIA package repositories
+    wget https://developer.download.nvidia.com/compute/cuda/repos/${REPO_PATH}/x86_64/cuda-keyring_1.1-1_all.deb
+    dpkg -i cuda-keyring_1.1-1_all.deb
+    rm cuda-keyring_1.1-1_all.deb
+    apt-get update
+
+    # Install specific CUDA toolkit version
+    # Convert version format (12.9 -> 12-9)
+    CUDA_VERSION_APT=$(echo $DRIVER_CUDA_VERSION | sed 's/\./-/')
+    echo "Installing cuda-toolkit-${CUDA_VERSION_APT}..."
+    apt-get install -y cuda-toolkit-${CUDA_VERSION_APT}
+
+    # Add CUDA to PATH
+    export PATH=/usr/local/cuda-${DRIVER_CUDA_VERSION}/bin:$PATH
+    export LD_LIBRARY_PATH=/usr/local/cuda-${DRIVER_CUDA_VERSION}/lib64:${LD_LIBRARY_PATH:-}
+
+    # Verify installation
+    nvcc --version
+else
+    echo "CUDA toolkit $NVCC_VERSION matches driver version"
+    export PATH=/usr/local/cuda-${DRIVER_CUDA_VERSION}/bin:$PATH
+    export LD_LIBRARY_PATH=/usr/local/cuda-${DRIVER_CUDA_VERSION}/lib64:${LD_LIBRARY_PATH:-}
+fi
+
+# --- Install uv (fast Python package manager) --------------------------------
+curl -LsSf https://astral.sh/uv/install.sh | sh
+export PATH="$HOME/.local/bin:$PATH"
+
+# --- Install Python 3.12 if not available ------------------------------------
+if ! command -v python3.12 &> /dev/null; then
+    echo "Python 3.12 not found. Installing via uv..."
+    uv python install 3.12
+fi
+
+# --- Clean up existing environments and caches -------------------------------
+echo "Cleaning up existing environments and caches..."
+
+# Remove existing venv for a clean installation
+VENV="$HOME/venv"
+if [ -d "$VENV" ]; then
+    echo "Removing existing virtual environment..."
+    rm -rf "$VENV"
+fi
+
+# Remove uv cache to ensure fresh installs
+if [ -d "$HOME/.cache/uv" ]; then
+    echo "Clearing uv cache..."
+    rm -rf "$HOME/.cache/uv"
+fi
+
+# Remove vLLM cache to avoid conflicts
+if [ -d "$HOME/.cache/vllm" ]; then
+    echo "Clearing vLLM cache..."
+    rm -rf "$HOME/.cache/vllm"
+fi
+
+# --- Create and activate venv ------------------------------------------------
+echo "Creating fresh virtual environment..."
+uv venv --python 3.12 --seed "$VENV"
+source "$VENV/bin/activate"
+
+# --- Install PyTorch and vLLM ------------------------------------------------
+echo "Installing vLLM and dependencies (version: $VLLM_VERSION)..."
+case "$VLLM_VERSION" in
+    release)
+        echo "Installing vLLM release with PyTorch..."
+        # Install vLLM with automatic PyTorch backend selection
+        # vLLM will automatically install the correct PyTorch version
+        uv pip install vllm>=0.10.0 --torch-backend=auto || {
+            echo "ERROR: Failed to install vLLM"
+            exit 1
+        }
+        ;;
+    nightly)
+        echo "Installing vLLM nightly with PyTorch..."
+        echo "This will install the latest nightly build of vLLM..."
+
+        # Install vLLM nightly with PyTorch
+        uv pip install -U vllm \
+            --torch-backend=auto \
+            --extra-index-url https://wheels.vllm.ai/nightly || {
+            echo "ERROR: Failed to install vLLM nightly"
+            exit 1
+        }
+
+        echo "vLLM nightly successfully installed!"
+        ;;
+    gpt-oss)
+        echo "Installing GPT-OSS special build with PyTorch nightly..."
+        echo "WARNING: This build is ONLY for GPT-OSS models!"
+        echo "Installing PyTorch nightly and cutting-edge dependencies..."
+
+        # Convert CUDA version format for PyTorch (12.4 -> cu124)
+        PYTORCH_CUDA="cu$(echo $DRIVER_CUDA_VERSION | sed 's/\.//')"
+        echo "Using PyTorch nightly with ${PYTORCH_CUDA} (driver supports ${DRIVER_CUDA_VERSION})"
+
+        # The GPT-OSS build will pull PyTorch nightly and other dependencies
+        # via the extra index URLs. We don't pre-install torch here to avoid conflicts.
+        uv pip install --pre vllm==0.10.1+gptoss \
+            --extra-index-url https://wheels.vllm.ai/gpt-oss/ \
+            --extra-index-url https://download.pytorch.org/whl/nightly/${PYTORCH_CUDA} \
+            --index-strategy unsafe-best-match || {
+            echo "ERROR: Failed to install GPT-OSS vLLM build"
+            echo "This automatically installs PyTorch nightly with ${PYTORCH_CUDA}, Triton nightly, and other dependencies"
+            exit 1
+        }
+
+        # Install gpt-oss library for tool support
+        uv pip install gpt-oss || {
+            echo "WARNING: Failed to install gpt-oss library (needed for tool use)"
+        }
+        ;;
+    *)
+        echo "ERROR: Unknown vLLM version: $VLLM_VERSION"
+        exit 1
+        ;;
+esac
+
+# --- Install additional packages ---------------------------------------------
+echo "Installing additional packages..."
+# Note: tensorrt removed temporarily due to CUDA 13.0 compatibility issues
+# TensorRT still depends on deprecated nvidia-cuda-runtime-cu13 package
+uv pip install huggingface-hub psutil hf_transfer
+
+# --- FlashInfer installation (optional, improves performance) ----------------
+echo "Attempting FlashInfer installation (optional)..."
+if uv pip install flashinfer-python; then
+    echo "FlashInfer installed successfully"
+else
+    echo "FlashInfer not available, using Flash Attention instead"
+fi
+
+# --- Mount storage if provided -----------------------------------------------
+if [ -n "$MOUNT_COMMAND" ]; then
+    echo "Setting up mount..."
+
+    # Create mount point directory if it doesn't exist
+    mkdir -p "$MODELS_PATH"
+
+    # Execute the mount command
+    eval "$MOUNT_COMMAND" || {
+        echo "WARNING: Mount command failed, continuing without mount"
+    }
+
+    # Verify mount succeeded (optional, may not always be a mount point)
+    if mountpoint -q "$MODELS_PATH" 2>/dev/null; then
+        echo "Storage successfully mounted at $MODELS_PATH"
+    else
+        echo "Note: $MODELS_PATH is not a mount point (might be local storage)"
+    fi
+fi
+
+# --- Model storage setup ------------------------------------------------------
+echo ""
+echo "=== Setting up model storage ==="
+echo "Storage path: $MODELS_PATH"
+
+# Check if the path exists and is writable
+if [ ! -d "$MODELS_PATH" ]; then
+    echo "Creating model storage directory: $MODELS_PATH"
+    mkdir -p "$MODELS_PATH"
+fi
+
+if [ ! -w "$MODELS_PATH" ]; then
+    echo "ERROR: Model storage path is not writable: $MODELS_PATH"
+    echo "Please check permissions"
+    exit 1
+fi
+
+# Create the huggingface cache directory structure in the models path
+mkdir -p "${MODELS_PATH}/huggingface/hub"
+
+# Remove any existing cache directory or symlink
+if [ -e ~/.cache/huggingface ] || [ -L ~/.cache/huggingface ]; then
+    echo "Removing existing ~/.cache/huggingface..."
+    rm -rf ~/.cache/huggingface 2>/dev/null || true
+fi
+
+# Create parent directory if needed
+mkdir -p ~/.cache
+
+# Create symlink from ~/.cache/huggingface to the models path
+ln -s "${MODELS_PATH}/huggingface" ~/.cache/huggingface
+echo "Created symlink: ~/.cache/huggingface -> ${MODELS_PATH}/huggingface"
+
+# Verify the symlink works
+if [ -d ~/.cache/huggingface/hub ]; then
+    echo "✓ Model storage configured successfully"
+
+    # Check available space
+    AVAILABLE_SPACE=$(df -h "$MODELS_PATH" | awk 'NR==2 {print $4}')
+    echo "Available space: $AVAILABLE_SPACE"
+else
+    echo "ERROR: Could not verify model storage setup"
+    echo "The symlink was created but the target directory is not accessible"
+    exit 1
+fi
+
+# --- Configure environment ----------------------------------------------------
+mkdir -p ~/.config/vllm
+touch ~/.config/vllm/do_not_track
+
+# Write environment to .bashrc for persistence
+cat >> ~/.bashrc << EOF
+
+# Pi vLLM environment
+[ -d "\$HOME/venv" ] && source "\$HOME/venv/bin/activate"
+export PATH="/usr/local/cuda-${DRIVER_CUDA_VERSION}/bin:\$HOME/.local/bin:\$PATH"
+export LD_LIBRARY_PATH="/usr/local/cuda-${DRIVER_CUDA_VERSION}/lib64:\${LD_LIBRARY_PATH:-}"
+export HF_TOKEN="${HF_TOKEN}"
+export PI_API_KEY="${PI_API_KEY}"
+export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN}"
+export HF_HUB_ENABLE_HF_TRANSFER=1
+export VLLM_NO_USAGE_STATS=1
+export VLLM_DO_NOT_TRACK=1
+export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+EOF
+
+# Create log directory for vLLM
+mkdir -p ~/.vllm_logs
+
+# --- Output GPU info for pi CLI to parse -------------------------------------
+echo ""
+echo "===GPU_INFO_START==="
+nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader | while IFS=, read -r id name memory; do
+    # Trim whitespace
+    id=$(echo "$id" | xargs)
+    name=$(echo "$name" | xargs)
+    memory=$(echo "$memory" | xargs)
+    echo "{\"id\": $id, \"name\": \"$name\", \"memory\": \"$memory\"}"
+done
+echo "===GPU_INFO_END==="
+
+echo ""
+echo "=== Setup complete ==="
+echo "Pod is ready for vLLM deployments"
+echo "Models will be cached at: $MODELS_PATH"