colab_execute
Execute Python code on Google Colab GPU/TPU runtimes, with support for multiple accelerators, background jobs, and Google Drive file integration.
Instructions
Execute Python code on a Google Colab GPU/TPU runtime.
Primary tool for running GPU/TPU-accelerated Python code (ML training, inference, CUDA operations). Allocates hardware, runs the code, and returns structured JSON with per-cell output, errors, and stderr.
After execution:
If background=True, poll with colab_poll(job_id) for results.
If OOM error occurs, reduce batch_size or upgrade to A100/H100.
If import error, add
pip install <pkg>before your main code.To persist outputs to Google Drive, use the drive_save parameter.
To pre-load data from Drive, use the drive_fetch parameter.
Common issues:
QUOTA_EXCEEDED: Colab rate limit hit. Wait a few minutes or switch to a different accelerator type.
CUDA_ERROR: Driver mismatch or GPU init failure. Retry, or try T4.
Only ONE background job can run at a time (Colab limitation).
Args: code: Python code to execute on the Colab runtime. accelerator: Hardware accelerator type. Default: "T4". GPU types: "T4" - NVIDIA Tesla T4 (16 GB, free-tier) "L4" - NVIDIA L4 (24 GB, Colab Pro) "A100" - NVIDIA A100 (40 GB, Colab Pro/Pro+) "H100" - NVIDIA H100 (80 GB, Colab Pro+) "G4" - NVIDIA G4 (Colab Pro+) TPU types: "V5E1" - TPU v5e-1 (Colab Pro+) "V6E1" - TPU v6e-1 (Colab Pro+) high_memory: Enable high-memory runtime (more RAM). Default: False. timeout: Max execution time in seconds (10-3600). Default: 300. background: Run in background (non-blocking). Default: False. When True, returns immediately with a job_id that can be polled via colab_poll. Incompatible with drive_fetch/drive_save. drive_fetch: JSON mapping Drive paths to Colab paths. Files are downloaded from Google Drive BEFORE your code runs. Example: '{"colab_data/train.csv": "/content/train.csv"}' Requires prior colab_drive_upload to place files on Drive. drive_save: JSON mapping Colab paths to Drive paths. Files are uploaded to Google Drive AFTER your code finishes, using a freshly obtained token (safe for long-running tasks). Example: '{"/content/model.pt": "results/model.pt"}'
Input Schema
| Name | Required | Description | Default |
|---|---|---|---|
| code | Yes | ||
| timeout | No | ||
| background | No | ||
| drive_save | No | ||
| accelerator | No | T4 | |
| drive_fetch | No | ||
| high_memory | No |
Output Schema
| Name | Required | Description | Default |
|---|---|---|---|
| result | Yes |