Managing GPU Pods with @pi/pods

Relevant source files

The following files were used as context for generating this wiki page:

• packages/pods/src/cli.ts
• packages/pods/src/commands/pods.ts
• packages/pods/src/commands/models.ts
• README.md
• packages/coding-agent/src/main.ts

The @mariozechner/pi package, aliased as pi, is a command-line interface (CLI) tool designed to streamline the management of vLLM (Very Large Language Model) deployments on remote GPU pods. It provides a simple yet powerful abstraction for setting up GPU environments, starting and stopping language models, and interacting with them. This tool is a key component of the pi-mono monorepo, enabling developers to easily manage the lifecycle of LLM services required by other packages like the @mariozechner/pi-coding-agent.

The system operates by connecting to remote machines via SSH, running setup and model execution scripts, and maintaining a local configuration file to track the state of pods and the models running on them. This allows for multi-pod management, efficient GPU allocation, and seamless integration with AI agents.

Sources: README.md, packages/pods/src/cli.ts

Command Line Interface (CLI)

The pi CLI is the primary interface for all pod and model management tasks. The main commands are routed through packages/pods/src/cli.ts, which parses arguments and calls the appropriate functions.

Sources: packages/pods/src/cli.ts:13-28

Command	Description
pi pods setup <name> ...	Configures a new GPU pod.
pi pods	Lists all configured pods and indicates the active one.
pi pods active <name>	Sets the specified pod as the active one for subsequent commands.
pi pods remove <name>	Removes a pod's configuration from the local machine.
pi shell [<name>]	Opens an interactive SSH shell on the active or specified pod.
pi ssh [<name>] "<cmd>"	Executes a given command on the pod via SSH.
pi start <model> --name <name>	Starts a new vLLM model instance on the active pod.
pi stop [<name>]	Stops a specific model or all models if no name is provided.
pi list	Lists all models currently running on the active pod.
pi logs <name>	Streams the logs for a specific running model.
pi agent <name> ...	Interacts with a deployed model using the @pi/coding-agent.

All model management commands (start, stop, list, logs, agent) support a --pod <name> flag to override the currently active pod.

Sources: packages/pods/src/cli.ts:13-43, 149-156

Pod Management

Pod management revolves around setting up, selecting, and accessing remote GPU machines.

Pod Setup (pi pods setup)

The setup process automates the configuration of a new GPU pod for running vLLM. It requires an SSH command for access and a mount path for models.

The following diagram illustrates the pod setup flow initiated by the pi pods setup command.

This diagram shows the sequence of operations during pod setup, including local validations and remote execution steps.

Process:

1. Validation: The command first checks for the HF_TOKEN and PI_API_KEY environment variables, which are required for model downloads and API access.
2. SSH Connection Test: It verifies that the pod is accessible via the provided SSH command.
3. Script Transfer: The local pod_setup.sh script is copied to the pod's /tmp directory using scpFile.
4. Remote Execution: The setup script is executed on the pod via sshExecStream. This script handles Docker/NVIDIA toolkit installation, Python environment setup, and vLLM installation. The --vllm flag allows specifying release, nightly, or a special gpt-oss build.
5. GPU Detection: After setup, nvidia-smi is run on the pod to query and record the available GPUs, including their names and memory.
6. Configuration: The new pod's details (SSH command, GPU info, models path) are saved to the local configuration file and the pod is set as the active one.

Sources: packages/pods/src/commands/pods.ts:48-189

Listing and Switching Pods

Users can manage multiple pod configurations.

• pi pods: Lists all configured pods, their GPU hardware, vLLM version, and SSH connection string. The currently active pod is marked with a green asterisk (*).
• pi pods active <name>: Switches the active pod by updating the active key in the configuration file. This determines the default target for all model management commands.
• pi pods remove <name>: Deletes a pod's configuration from the local file. This does not affect the remote machine itself.

Sources: packages/pods/src/commands/pods.ts:15-46, 194-219

Model Management

Once a pod is set up, you can manage the lifecycle of vLLM model deployments on it.

Starting a Model (pi start)

The pi start command is the core function for deploying a model. It automates GPU allocation, port assignment, and process management.

The sequence diagram below shows the interaction between the user, the CLI, and the remote pod when starting a model.

Process Details:

1. Pod Selection: The command targets the active pod unless overridden by --pod <name>.
2. Port and GPU Allocation:
   • It finds the next available port starting from 8001.
   • It calls selectGPUs() to allocate GPUs in a round-robin fashion, prioritizing the least-used ones. The number of GPUs is determined by the model's configuration or the --gpus flag.
3. Configuration:
   • For known models (defined in models.json), it retrieves optimized vllm arguments.
   • For unknown models, it defaults to a single GPU.
   • Overrides like --memory and --context are applied to the vllm arguments.
   • If --vllm is used, all other configuration is ignored, and the provided arguments are passed directly to vLLM.
4. Execution:
   • A model_run.sh script is customized with the model ID, port, and vLLM arguments.
   • This script is uploaded to the pod.
   • The script is executed in the background using setsid to ensure it persists after the SSH session closes. A wrapper script ensures the process is logged to ~/.vllm_logs/<name>.log.
5. State Management: The model's name, PID, port, and allocated GPUs are saved to the pod's configuration.
6. Log Streaming: The CLI tails the model's log file, streaming the output until "Application startup complete" is detected or an error occurs. If startup fails (e.g., out of memory), the model's entry is removed from the configuration.

Sources: packages/pods/src/commands/models.ts:58-301

GPU Selection Logic

The selectGPUs function implements a simple but effective load-balancing strategy for model placement.

This diagram shows how GPUs are selected based on current usage to balance load across the pod.

This logic ensures that new models are placed on the GPUs with the fewest active models, distributing the load evenly.

Sources: packages/pods/src/commands/models.ts:39-56

Stopping and Listing Models

• pi stop <name>: Stops a specific model by killing its process tree (pkill -P <pid>) on the remote pod and removing it from the local configuration.
• pi stop: Stops all models running on the active pod.
• pi list: Lists all running models on the pod, showing their name, port, GPU allocation, and PID. It also performs a health check by verifying that the process is still running on the pod.
• pi logs <name>: Streams the log file for a running model using tail -f over SSH, allowing for real-time monitoring.

Sources: packages/pods/src/commands/models.ts:303-441

Interacting with Deployed Models

Once a model is running, it exposes an OpenAI-compatible API endpoint. The pi start command provides the necessary connection details. The primary way to interact with these models within the monorepo is through the pi agent command, which is an alias for the @mariozechner/pi-coding-agent.

# Example output from 'pi start'
✓ Model started successfully!

Connection Details:
──────────────────────────────────────────────────
Base URL:    http://<host>:<port>/v1
Model:       <model-id>
API Key:     <your-pi-api-key>
──────────────────────────────────────────────────

# Example usage provided
Chat with model:  pi agent <name> "Your message"
Interactive mode: pi agent <name> -i

The pi agent command from the @pi/pods package is a wrapper that invokes the main pi-coding-agent with the correct API endpoint and key for the specified deployed model. This creates a seamless workflow from deployment to interaction.

Sources: packages/pods/src/commands/models.ts:276-300, packages/pods/src/cli.ts:289-307, packages/coding-agent/src/main.ts:1000-1012

Summary

The @pi/pods package provides a comprehensive toolkit for managing the entire lifecycle of vLLM deployments. By abstracting away SSH commands, process management, and configuration, it allows developers to treat remote GPU resources as a simple, manageable pool for running LLMs. Its integration with the @pi/coding-agent further solidifies its role as a foundational infrastructure tool within the pi-mono ecosystem, bridging the gap between development and deployment of AI models.