# colab-fleet

> Offload heavy CPU/ML work to **free Google Colab** from Claude Code (or any MCP client) — with a free-tier lock, auto-teardown, a job queue, and checkpoint-resume.

*English · [한국어](README.ko.md)*

Run the expensive stuff — ML training, hyperparameter sweeps, backtests — on a **free Google Colab runtime** instead of your own machine, and pull back only the artifacts. An **MCP client** (Claude Code, pi, …) submits a job with a single tool call (`gpu_submit(...)`) and the results are fetched automatically when it finishes.

You need no GPU and no heavy libraries locally. A **free Colab account** is enough.

---

## How is this different from the official `colab` CLI?

One line: Google's `colab` CLI is **a tool a human types**; colab-fleet is **a tool an AI calls**. Same engine (the official CLI), with guardrails and automation layered on top.

> Google did **not** ship an MCP server — only a command-line CLI. This project wraps it so an agent can drive it directly.

| | Official `colab` CLI (raw) | colab-fleet (this wrapper) |
|------|------------------------|------------------------|
| **Who drives it** | a human, in a terminal | an AI, calling `gpu_submit(...)` mid-conversation |
| **Billing accidents** | paid accelerators (A100, TPU) attach silently and bill you | **only free tiers (cpu/t4) allowed, paid ones hard-rejected** → zero spend |
| **Timeouts** | 30s default — an ML run gets cut off | per-job, long by default (1h) + checkpoint continuation |
| **Session cleanup** | forget once and it idles forever | **auto-stop on finish**, orphaned sessions reaped |
| **Concurrency** | unguarded (free tier allows 1 session → collisions) | serialized through a queue |
| **Data** | upload/download by hand | inputs uploaded → run → outputs **fetched automatically** |
| **Credentials** | your problem | tokens and emails **never leak into tool output** |

Analogy: the official CLI is a **manual transmission**; colab-fleet is **automatic + seatbelt + navigation** on the same engine (the free Colab runtime).

---

## Install (3 steps)

```bash
git clone https://github.com/yazzang-homelab/colab-fleet.git
cd colab-fleet
./install.sh              # creates a venv, installs the colab CLI + mcp
```

When `install.sh` finishes it prints the remaining two steps:

**1) Authenticate once with your own Google account** (opens a browser flow)

```bash
.venv/bin/colab sessions
```

Tokens are stored only under `~/.config/colab-cli/` and **never flow through this tool**.

**2) Register with your MCP client**

- Claude Code:
  ```bash
  claude mcp add -s user colab-fleet -- "$PWD/.venv/bin/python" "$PWD/server.py"
  ```
- pi / others (add to your `mcpServers` config):
  ```json
  "colab-fleet": { "command": "/abs/path/colab-fleet/.venv/bin/python", "args": ["/abs/path/colab-fleet/server.py"] }
  ```

Check it: call `gpu_doctor()` from the agent — auth ✅ means you're done.

**3) ⭐ (optional)** After install, `install.sh` — or, if you missed it, the first tool call (`gpu_doctor`/`gpu_submit`) — asks for a star **exactly once**. The agent only relays that message; **you** press the star. The tool never runs `gh repo star` on your behalf (i.e. it does not touch your GitHub account).

```bash
gh repo star yazzang-homelab/colab-fleet   # or hit ⭐ on the repo page
```

Nothing is gated behind it. To silence the message entirely, set `COLAB_FLEET_NO_STAR=1`.

> Requirements: Python 3.10+, git. A local GPU and heavy libraries are **not needed** — everything runs on Colab.

---

## Usage

```python
# 0) health check
gpu_doctor()

# 1) round-trip smoke test (no deps → ephemeral run)
gpu_submit(".../examples/selftest.py", accel="cpu", outputs="/content/selftest.json")
gpu_status(<id>); gpu_logs(<id>); gpu_fetch(<id>)

# 2) parallel hyperparameter sweep (deps + outputs → managed run)
gpu_submit(".../examples/sklearn_gridsearch.py", accel="cpu",
           deps="scikit-learn joblib", args="--folds 5",
           outputs="/content/result.json", timeout=1200)

# 3) train on your own CSV (inputs + deps + outputs → managed run)
gpu_submit(".../examples/train_on_csv.py", accel="cpu",
           deps="scikit-learn pandas joblib",
           inputs="/path/to/data.csv", args="--csv /content/data.csv",
           outputs="/content/model.pkl,/content/metrics.json", timeout=1800)

# real DL that needs a GPU (T4):
gpu_submit(".../my_torch_train.py", accel="t4", deps="torch ...", outputs="...")
```

Artifacts land in `~/.colab-fleet/jobs/job-<id>/` (configurable via env var).

### Tools

| Tool | Description |
|------|------|
| `gpu_submit(script, accel, deps, args, inputs, outputs, ckpt, timeout, retries, label)` | submit a job (async), returns a job id |
| `gpu_status(job_id=0)` | job detail; `0` returns recent jobs + active sessions |
| `gpu_logs(job_id, lines)` | tail logs per stage |
| `gpu_fetch(job_id)` | list artifact paths |
| `gpu_sessions()` | active Colab sessions |
| `gpu_stop(name="all")` | stop sessions (compute-unit safe) |
| `gpu_doctor()` | health, auth, and queue status |

---

## Guardrails

- **Free-tier lock**: only `accel ∈ {cpu, t4}`. A100/H100/L4/TPU are **hard-rejected** (`colabq.accel_flags`), so paid compute-unit spend is structurally zero. This also disarms the CLI footgun where an unrecognized GPU name silently falls back to A100.
- **Auto-teardown**: every job runs `colab stop` in a `try/finally`. The dispatcher also reaps orphaned sessions on shutdown.
- **Serial queue**: the free tier allows one session, so jobs are serialized with `flock`; concurrent submissions queue up.
- **No credential/PII leakage**: tool output contains no tokens or email addresses.

## Checkpoint-resume (long jobs)

If your script periodically saves to the `ckpt` path and loads it on startup, the dispatcher keeps the last checkpoint when a session hits the 12h/90min cap and re-uploads it on retry — so a job **finishes across multiple sessions**. Use `retries` to set the retry count.

## Environment variables (optional)

`COLAB_FLEET_AUTH` (oauth2/adc) · `COLAB_FLEET_TIMEOUT` · `COLAB_FLEET_ARTIFACTS` (artifact location) · `COLAB_FLEET_CONFIG` · `COLAB_FLEET_DB` · `COLAB_FLEET_BIN` (force the `colab` binary path) · `COLAB_FLEET_NO_STAR` (`=1` disables the one-time star message) · `COLAB_FLEET_STAR_MARKER` (file recording whether the message was shown; default `~/.config/colab-fleet/star-nudged`).

## Troubleshooting

| Symptom | Cause / fix |
|------|-----------|
| `gpu_doctor` auth ❌ | one-time login not done → run `.venv/bin/colab sessions` for the browser flow |
| 401/403 | missing scopes → log in again (or `COLAB_FLEET_AUTH=adc` + `gcloud auth application-default login`) |
| no GPU assigned (400 on new) | free T4 availability fluctuates → fall back to `accel="cpu"` |
| job stuck in `queued` | dispatcher not running → start it with `.venv/bin/python dispatch.py` |
| upload 500 / failure | `colab upload` uses the Jupyter API, which **500s above ~80MB per input**. Shrink with a subset or compression. Upload failures are surfaced as job failures (never hidden) |
| `done` but no artifacts | the remote script raised → success sentinel not printed → treated as failure. Check the stderr traceback via `gpu_logs` |

## Notes

- Google Colab's free runtime has a **fair-use policy**. This tool performs no circumvention (no multi-account rotation, no keep-alive bots) — just ordinary job offloading and cleanup.
- The free CPU runtime is roughly 2 vCPUs; the real win is **clean RAM and isolation from your dev machine**, not core count. If you need actual GPU acceleration, use `accel="t4"`.

## License

MIT. The engine it drives, [`google-colab-cli`](https://pypi.org/project/google-colab-cli/), is Apache-2.0 (separate).