google-colab-cli 0.5.9__tar.gz → 0.6.0__tar.gz

google_colab_cli-0.6.0/.python-version

@@ -0,0 +1 @@
+3.12

{google_colab_cli-0.5.9 → google_colab_cli-0.6.0}/AGENTS.md

@@ -5,11 +5,11 @@
 - **Common**: `common.py` centralizes shared `State` (lazy-loading) and session resolution.
 - **Client**: `ColabClient` handles API interactions (assignment, unassignment).
 - **Auth**: `auth.py` exposes a single `get_credentials(config_path, provider)` facade that dispatches on the `AuthProvider` enum. Two providers are supported, selected via the global `--auth=oauth2|adc` flag (default `oauth2`):
-  - `oauth2`: public `google-auth-oauthlib` `InstalledAppFlow`, token cached at `~/.config/colab-cli/token.json`. Requires an explicit client OAuth config (`-c/--client-oauth-config`, default `~/.colab-cli-oauth-config.json`); the previously-bundled `oauth_config.json` resource fallback was removed (commit `20eb88e`).
-  - `adc`: Google Application Default Credentials via `google.auth.default()`. The CLI passes `scopes=PUBLIC_SCOPES` (which includes `colaboratory`) and re-applies via `creds.with_scopes()` for credential types that support it. **User credentials minted by `gcloud auth application-default login` ignore the `scopes=` kwarg AND raise `NotImplementedError` on `with_scopes`**: ADC users must explicitly re-authenticate with `gcloud auth application-default login --scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory`. `userinfo.email` is required by the session backend at `colab.research.google.com` (assign/unassign/sessions return 401 without it); `colaboratory` is required by the `RuntimeService` at `colab.pa.googleapis.com` (keep-alive returns 403 without it); `openid` and `cloud-platform` are mandated by `gcloud` itself, which rejects scope lists that omit `cloud-platform` with `Invalid value for [--scopes]`. Service-account / GCE / GKE / impersonated creds get the right scopes transparently via `with_scopes`.
+  - `oauth2`: public `google-auth-oauthlib` `InstalledAppFlow`, token cached at `~/.config/colab-cli/token.json`. Reads the client OAuth config from `-c/--client-oauth-config` (default `~/.colab-cli-oauth-config.json`), falling back to the **bundled** `src/colab_cli/oauth_config.json` resource (re-added in PR #41 / `9f44fe2`, 2026-05-29 — the earlier "removed in `20eb88e`" note was stale/incorrect; the file exists and `auth.py:_get_google_auth_credentials` loads it via `importlib.resources`). As of 2026-06-11 the flow is a **remote copy-paste flow**, not a localhost server: `_run_remote_flow` sets `redirect_uri=https://sdk.cloud.google.com/applicationdefaultauthcode.html` + `token_usage=remote`, prints the URL, and reads the pasted code via `input()`. NEVER revert to OOB (`urn:ietf:wg:oauth:2.0:oob`) — Google blocked it in 2022 ("OOB flow has been blocked"); the `sdk.cloud.google.com` redirect is registered only to the bundled cloud-SDK client (`764086051850-...`), so any other client id gets `redirect_uri_mismatch`. Server-side acceptance/rejection of these variants is verifiable GET-only by building the authorization URL and inspecting whether Google reaches sign-in vs. an OAuth error page (no resources allocated).
+  - `adc`: Google Application Default Credentials via `google.auth.default()`. The CLI passes `scopes=PUBLIC_SCOPES` (which includes `colaboratory`) and re-applies via `creds.with_scopes()` for credential types that support it. **User credentials minted by `gcloud auth application-default login` ignore the `scopes=` kwarg AND raise `NotImplementedError` on `with_scopes`**: ADC users must explicitly re-authenticate with `gcloud auth application-default login --scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory`. `userinfo.email` is required by the session backend at `colab.research.google.com` (assign/unassign/sessions/keep-alive return 401 without it); `colaboratory` is retained for forward compatibility and other Colab features (keep-alive no longer uses `colab.pa.googleapis.com` — see the Keep-alive note below); `openid` and `cloud-platform` are mandated by `gcloud` itself, which rejects scope lists that omit `cloud-platform` with `Invalid value for [--scopes]`. Service-account / GCE / GKE / impersonated creds get the right scopes transparently via `with_scopes`.
 - **Backend Hosts**: Two distinct backends with different requirements:
-  - `colab.research.google.com` (session backend / `tun/m/...`): accepts the `userinfo.email` scope.
-  - `colab.pa.googleapis.com` (`RuntimeService`, used by `KeepAliveAssignment`): requires (a) the `colaboratory` OAuth scope, AND (b) `X-Goog-Api-Client` header containing the substring `grpc-web`. Both are enforced server-side; missing either yields HTTP 400 / 403 with descriptive `google.rpc.DebugInfo` payloads. Always log `response_body` on failure for these RPCs to avoid silent debugging.
+  - `colab.research.google.com` (session backend / `tun/m/...`): accepts the `userinfo.email` scope. Handles assign, unassign, the contents API, **and keep-alive** (see below).
+  - **Keep-alive (2026-06-15, issue #14)**: keep-alive is a **Tunnel Frontend (TFE) HTTP ping** — `GET https://colab.research.google.com/tun/m/<endpoint>/keep-alive/` with header `X-Colab-Tunnel: Google`, authenticated by the user's own Gaia bearer token (same host/credential as `assign`). TFE records `LastActiveTime` before forwarding, refreshing the idle timer. The VM usually doesn't answer on this path, so the request commonly **read-times-out even on success** — `client.keep_alive_assignment` therefore catches `requests.exceptions.ReadTimeout` and treats it as success, while genuine HTTP errors (e.g. 404 for a deleted assignment) propagate. This mirrors the official `colab-vscode` extension's `sendKeepAlive` (`src/colab/client.ts`). **DO NOT revert to the `colab.pa.googleapis.com` `RuntimeService/KeepAliveAssignment` RPC**: that RPC requires the caller to be a `serviceusage` consumer of Colab's internal project `1014160490159`, which no ordinary user account is, so it returned HTTP 403 `USER_PROJECT_DENIED` for every external user (issue #14) and silently idle-pruned their sessions within minutes. The browser only made that RPC work by riding the user's `google.com` cookie through an internal cookie-proxy (`colab.clients6.google.com`), which a headless bearer-token CLI cannot use. Verified live 2026-06-15 with a third-party account (the RPC 403'd; the tunnel ping succeeded and kept the VM alive).
 - **Runtime**: `ColabRuntime` wraps `jupyter-kernel-client` for execution.
 - **State**:
   - `StateStore` persists session metadata in `~/.config/colab-cli/sessions.json`.
@@ -60,7 +60,7 @@
 15. **Two distinct auths — never confuse them**: This codebase has two unrelated authentication concerns: (a) **CLI-to-Colab-control-plane**: how `Client` authenticates HTTP requests to `colab.research.google.com` and `colab.pa.googleapis.com`. Driven by the global `--auth={oauth2,adc}` flag and `auth.py:get_credentials`. The OAuth2 path is bootstrapped automatically by `google_auth_oauthlib`'s `InstalledAppFlow` on first invocation when `~/.config/colab-cli/token.json` doesn't exist; **no separate command is needed** — any `--auth=oauth2` invocation (e.g. `colab --auth=oauth2 sessions`) triggers the browser consent flow. (b) **VM-side credentials**: how the `colab auth` subcommand injects user GCP credentials *into* the running Colab kernel so notebook code (e.g. `gcloud`, BigQuery client) can make authenticated calls from inside the VM. This uses the `USE_AUTH_EPHEM='0'` gcloud-fallback path executed on the kernel via `ColabRuntime`. **NEVER tell a user to "run `colab auth`" as a prerequisite for fixing CLI-side auth issues** — they're orthogonal layers and the suggestion misleads.
 16. **Detached daemons inherit nothing useful**: When spawning a detached child process via `subprocess.Popen` (e.g. `spawn_keep_alive`), the child does NOT inherit the parent's parsed Typer flags. It re-parses argv from scratch, so any global flag the parent saw via `--auth=adc` or `--config /tmp/foo.json` MUST be re-emitted as part of the child's command line. Forgetting this causes silent fallback to Typer defaults: in 2026-04-30 this manifested as the keep-alive daemon (a) using OAUTH2 instead of the parent's ADC, and (b) reading from `~/.config/colab-cli/sessions.json` instead of the parent's `--config` path. Always propagate every relevant global flag in `cmd = [sys.executable, "-m", ...]` BEFORE the subcommand name (Typer requires global flags before the subcommand).
 17. **Persist-before-spawn for daemons that read shared state**: When a detached child reads from a state file the parent owns (e.g. `state.store.get(session_name)`), the parent MUST persist BEFORE spawning. Otherwise the child can race ahead of the parent's `add()` call and observe an empty store. Symptom in 2026-04-30: keep-alive daemon exits immediately with `keep_alive_stopped reason=session_not_found iters=1 duration=0.0s`. Fix: call `state.store.add(s)` once before `spawn_keep_alive`, and again after (to capture the PID).
-18. **`colab.pa.googleapis.com` rejects ADC user creds without `X-Goog-User-Project`**: When calling `colab.pa.googleapis.com` from ADC user credentials minted by `gcloud auth application-default login`, the bearer token carries the user's gcloud quota project. The colab API key sent alongside it belongs to a different project (Colab, `1014160490159`). The backend enforces project-match between the two and returns HTTP 400 `CONSUMER_INVALID` ("The API Key and the authentication credential are from different projects."). Fix: send `X-Goog-User-Project: 1014160490159` to pin the consumer project to Colab's. Any signed-in user has implicit access to the colab project via the public web client, so this works for ordinary user accounts. Service-account / GCE / GKE / impersonated creds don't hit this because their bearer token IS owned by a project that matches.
+18. **[SUPERSEDED 2026-06-15 — keep-alive no longer calls `colab.pa.googleapis.com`; see issue #14 / the Keep-alive note above] `colab.pa.googleapis.com` rejects ADC user creds without `X-Goog-User-Project`**: When calling `colab.pa.googleapis.com` from ADC user credentials minted by `gcloud auth application-default login`, the bearer token carries the user's gcloud quota project. The colab API key sent alongside it belongs to a different project (Colab, `1014160490159`). The backend enforces project-match between the two and returns HTTP 400 `CONSUMER_INVALID` ("The API Key and the authentication credential are from different projects."). The old fix was to send `X-Goog-User-Project: 1014160490159` to pin the consumer project — but that in turn required `serviceusage.serviceUsageConsumer` on project `1014160490159`, which ordinary users lack, yielding HTTP 403 `USER_PROJECT_DENIED` (the issue #14 root cause). Both failure modes are now moot because keep-alive uses the TFE tunnel ping on `colab.research.google.com` instead. Retained here as institutional knowledge: if any future code path must call `colab.pa.googleapis.com` with a bearer token, it will face this same project-entitlement wall for non-internal accounts.
 19. **Pydantic validation requires a schema**: `_issue_request` accepts an optional `schema=` and historically called `TypeAdapter(schema).validate_python(...)` unconditionally. When `schema=None` (a caller that doesn't care about the response body — e.g. fire-and-forget RPCs like `KeepAliveAssignment` that return `[]`), this raises `pydantic.ValidationError: Input should be None`. Always guard with `if schema is None: return` after the empty-body short-circuit.
 20. **Suggest the branch-diff review command after committing**: The user reviews changes with `git diff main..<branch-name>` (full cumulative diff against `main`, not just the latest commit). After landing one or more commits on a feature branch, ALWAYS suggest the exact command — e.g. "Review with `git diff main..sort-help-commands`" — instead of `git show <sha>` (which only shows a single commit and misses context when a branch has multiple commits). Encoded 2026-05-05 after suggesting `git show 9d9c7da` for a branch the user wanted to review holistically.
 21. **Verify research-tool claims with primary sources**: Research tools and AI assistants can be confidently wrong, especially about edge cases or features outside their training corpus. When such a tool says "X is not used / not parsed / doesn't exist", treat it as a hypothesis to verify, not a fact. Always cross-check against the primary source (the actual code or config) — and when a tool names files, check whether the indirection chain it describes actually exists. The cost of believing the tool when it's wrong is shipping a non-functional feature; the cost of double-checking is small. Encoded 2026-05-05 after a `colab url` first-cut shipped the wrong URL format because of unverified output.

google_colab_cli-0.6.0/CHANGELOG.md

@@ -0,0 +1,50 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+The package version is derived from the git tag via `hatch-vcs`; each release
+below corresponds to a tag of the same name.
+
+## [0.6.0] - 2026-06-16
+
+### Changed
+
+- **auth:** OAuth2 login now uses a remote copy-paste flow instead of a
+  localhost callback server. The CLI prints an authorization URL with
+  `redirect_uri=https://sdk.cloud.google.com/applicationdefaultauthcode.html`
+  and `token_usage=remote`, then reads the pasted code from stdin. This works
+  in headless/remote environments where a browser cannot reach a local
+  callback port. (#54)
+
+### Added
+
+- **display output:** Rich rendering for `display_data` output via a shared
+  `render_display_data()` helper. HTML is converted with `html2text` and
+  rendered as Markdown, following a `text/markdown > text/html > text/plain`
+  priority; `text/plain` is wrapped with `Text.from_ansi` to preserve embedded
+  ANSI escapes. Applied consistently across `exec`, `console`/`repl`, and
+  automation call sites. (#58)
+
+### Fixed
+
+- **keep-alive:** Replace the `RuntimeService/KeepAliveAssignment` RPC on
+  `colab.pa.googleapis.com` with a Tunnel Frontend (TFE) HTTP ping
+  (`GET /tun/m/<endpoint>/keep-alive/` with `X-Colab-Tunnel: Google`) on
+  `colab.research.google.com`, authenticated by the user's own bearer token.
+  The old RPC required `serviceusage` consumer access to Colab's internal
+  project and returned HTTP 403 `USER_PROJECT_DENIED` for every external user,
+  causing their sessions to be idle-pruned within minutes. The TFE ping needs
+  no project entitlement; because the VM often does not answer on this path, a
+  `ReadTimeout` is treated as success while genuine HTTP errors propagate.
+  (#14, #61)
+
+### Removed
+
+- Dead grpc-web client-registry / API-key code path and the now-irrelevant
+  `colaboratory`-scope / `pa.googleapis.com` pre-flight remediation messaging,
+  superseded by the TFE keep-alive ping. (#61)
+
+[0.6.0]: https://github.com/googlecolab/google-colab-cli/compare/v0.5.11...v0.6.0

google_colab_cli-0.6.0/COLAB_SKILL.md

@@ -0,0 +1,79 @@
+# Skill: Colab Session Operator
+
+Operate Google Colab environments via the `colab` CLI: provision GPU/TPU sessions, run Python/shell on the VM, sync files, and capture work as notebooks.
+
+## When to activate
+- Creating or managing TPU/GPU sessions.
+- Running Python or shell on a remote Colab VM.
+- Syncing files between local and remote.
+- Automating environment setup (packages, auth, Drive).
+- Exporting session history as a Jupyter notebook.
+
+## Mental model (read this first)
+- **A session == a live Jupyter kernel on a rented VM.** `colab new` allocates a billable VM; `colab stop` releases it. Nothing reclaims it automatically except a 24h keep-alive cap, so an unstopped session burns compute units indefinitely.
+- **Kernel state PERSISTS across `colab exec` / `colab repl` calls in the same session.** Each invocation reattaches to the *same* kernel (the kernel ID is cached in local state) and only closes the websocket on exit — it does **not** shut the kernel down. So imports, variables, and defined functions survive between separate `colab exec` commands. Build up state incrementally; don't re-import everything each call. (`colab stop` and `colab restart-kernel` are what actually reset it.)
+- **Default working directory is `/content`.** Every `exec`/`repl`/`run` `cd`s there first; prefer absolute paths (`/content/...`) for file work. For `colab ls/rm/upload/download`, absolute `/content/...` paths work and the default `ls` path is `content` (VM root).
+- **`colab` is fire-and-forget.** Each command authenticates, does one thing, and exits. A detached background daemon (spawned by `colab new`) handles keep-alive; you don't manage it.
+
+## Authentication (the #1 thing that blocks agents)
+- The global flag is `--auth={adc,oauth2}` and the **default is `adc`** (Application Default Credentials). It must come *before* the subcommand: `colab --auth=adc new -s x`.
+- **ADC setup** (most reliable for headless/agent use). The Colab backends need a specific scope set, so re-mint ADC with all four scopes:
+  ```bash
+  gcloud auth application-default login \
+    --scopes=openid,\
+  https://www.googleapis.com/auth/cloud-platform,\
+  https://www.googleapis.com/auth/userinfo.email,\
+  https://www.googleapis.com/auth/colaboratory
+  ```
+  Why all four: `userinfo.email` (session backend `colab.research.google.com`, else 401), `colaboratory` (RuntimeService `colab.pa.googleapis.com` keep-alive, else 403), `openid`+`cloud-platform` (mandated by gcloud itself; it rejects scope lists missing `cloud-platform`).
+- **oauth2 setup**: `colab --auth=oauth2 <anything>` triggers a browser consent flow on first use (token cached at `~/.config/colab-cli/token.json`). Requires a client config at `~/.colab-cli-oauth-config.json` (or `-c PATH`). The browser step means it usually needs a human; prefer ADC for agents.
+- **Verify auth in one shot**: `colab sessions` (read-only, lists server assignments) or `colab whoami` (hidden debug command: prints the active email, scopes, audience, and expiry). When any call 403s against `colab.pa.googleapis.com`, the cause is almost always a missing scope — `colab whoami` shows it instantly.
+- **`colab new` pre-flights the keep-alive RPC** right after allocating. If your token lacks the `colaboratory` scope it unassigns the fresh VM (so you don't leak a billable assignment) and prints the exact remediation. Follow that message rather than retrying blindly.
+- **Do NOT confuse `colab auth` with CLI authentication.** `colab auth` injects *VM-side* GCP credentials into the running kernel (so notebook code can call BigQuery/GCS); it is orthogonal to how the CLI itself authenticates. Never suggest "run `colab auth`" to fix a CLI 401/403 — that's a scope/identity problem fixed via the `gcloud` command above.
+
+## Workflow
+
+### Provision
+- `colab new -s <name>` (CPU). Add `--gpu A100` or `--tpu v6e1` for accelerators. **Always pass `-s <name>`** — an omitted name is auto-generated as a random 6-hex string, which makes later commands ambiguous.
+- Supported `--gpu`: `T4`, `L4`, `G4`, `H100`, `A100`. Supported `--tpu`: `v5e1`, `v6e1`.
+- **Gotcha**: an unrecognized `--gpu` value silently falls back to **A100** (which then usually fails the next step). A `400` on `colab new` with an accelerator means no quota/entitlement for it on this account — fall back to `--gpu T4` or omit the flag for CPU.
+- Accelerator availability is tier-gated; most accounts can only get CPU. Don't assume a GPU/TPU will allocate.
+
+### Execute
+- **Preferred**: `colab exec -s <name> -f <script.py>` runs a local script on the remote VM (read locally, sent to the kernel — no manual upload needed).
+- **Piped code**: `echo "print(1)" | colab exec -s <name>` or `cat script.py | colab exec -s <name>`.
+- **Notebooks**: `colab exec -s <name> -f nb.ipynb` runs each code cell and writes results to `<basename>_output.ipynb` next to the input. A `# @title Foo` first line labels the cell in progress output.
+- **Plots/images**: PNG/JPEG outputs are intercepted. Use `--output-image <path>` on `exec`/`repl` to save to a known location (otherwise a temp path is printed). Inline terminal-image escapes are auto-suppressed when stdout isn't a TTY, so piped/captured output stays clean.
+- **Shell**: `echo "cmd" | colab console -s <name>` for batch shell. Console wraps bash in tmux, so even piped output contains terminal-control bytes — filter with `grep -a` for a specific line. `exec` is faster when you don't need a real shell.
+- **Never run `colab repl`, `colab console`, `colab auth`, or `colab drivemount` interactively from an agent** — they expect a TTY and will hang. `repl`/`console` accept piped stdin and exit on EOF; `auth`/`drivemount` genuinely require a human at the terminal.
+
+### Ephemeral one-shot jobs (`colab run`)
+- `colab run [--gpu T4] [--tpu v6e1] [--keep] [-s NAME] script.py [args...]` = `new` + `exec` + `stop` in one command. It provisions a fresh VM, runs the script with `sys.argv` and `__name__ == "__main__"` set like native `python script.py args`, then tears the VM down (unless `--keep`).
+- **Exit codes propagate**: an uncaught exception or `sys.exit(N)` in the script makes `colab run` exit non-zero (CPython semantics: `sys.exit()`/`sys.exit(0)` → 0, `sys.exit(N)` → N, `sys.exit("msg")` → 1).
+- **Stream separation**: `colab run` writes its own `[colab] ...` chatter to **stderr** and the script's output to **stdout** — so `colab run job.py > out.txt` captures only the script's stdout. (`colab exec` streams the script's stdout/stderr live to your stdout/stderr.)
+- Works as a shebang: `#!/usr/bin/env -S colab run --gpu T4` makes a `chmod +x`'d `.py` a self-contained "rent a GPU, run, clean up" script. After editing CLI behavior, reinstall before testing shebangs — they resolve `colab` via `$PATH`, not the editable install.
+- A nonexistent script path exits non-zero **before** allocating a VM (no wasted compute).
+
+### Automate
+- `colab auth -s <name>` — VM-side GCP creds, needed before in-VM GCS/BigQuery calls (interactive; not agent-runnable).
+- `colab drivemount -s <name> [PATH]` — mounts Drive at `/content/drive` by default (interactive; not agent-runnable).
+- `colab install -s <name> pkg1 pkg2` — installs via `uv pip install --system`, falling back to `pip`. Also `colab install -s <name> -r requirements.txt`.
+
+### Inspect & report
+- `colab help` (or `colab help <cmd>`) lists/explains commands; the listing is alphabetical.
+- `colab sessions` lists server-side assignments and auto-prunes stale local entries. Orphans with no local record show as `[?]`.
+- `colab status [-s <name>]` shows hardware, IDLE/BUSY, and last execution.
+- `colab log -s <name> [-n 20] [-t TYPE]` shows recent structured events; invaluable when a task fails (keep-alive errors carry the raw `response_body`).
+- `colab log -s <name> -o summary.ipynb` exports the session as a notebook (also `.md`, `.txt`, `.jsonl` by suffix).
+- `colab url -s <name>` prints a browser URL that attaches the Colab web UI to your existing CLI session instead of allocating a new VM (add `--open` to launch it).
+- `colab skill` / `colab readme` print this skill and the README (handy for self-discovery).
+
+## Safety
+- **Always `colab stop -s <name>` when done** — idle VMs burn compute units. `colab run` (without `--keep`) self-cleans even if the script errors.
+- Local state lives in `~/.config/colab-cli/sessions.json` (settings in `settings.json`, history in `history/*.jsonl`). Don't edit by hand.
+- **Isolate parallel/agent runs** with the global `--config <path>` flag to point session state at a scratch file (e.g. `colab --config /tmp/agent.json new -s job`). The keep-alive daemon inherits `--auth` and `--config` automatically.
+
+## Recovery
+- "Session not found" / 404 / 401 on exec: the backend pruned the VM. `colab exec`/`repl` detect this and clean up local state automatically — run `colab sessions` and re-create with `colab new`.
+- Execution timeout or wedged kernel: `colab restart-kernel -s <name>` (keeps the VM, resets the kernel), or `colab stop` then `colab new`.
+- Keep-alive daemon died (`colab log` shows `keep_alive_stopped reason=consecutive_4xx_errors`): almost always the missing `colaboratory` scope — re-auth per the Authentication section.

{google_colab_cli-0.5.9 → google_colab_cli-0.6.0}/PKG-INFO

@@ -1,12 +1,26 @@
 Metadata-Version: 2.4
 Name: google-colab-cli
-Version: 0.5.9
+Version: 0.6.0
 Summary: CLI for interacting with Colab.
+Project-URL: Homepage, https://github.com/googlecolab/google-colab-cli
+Project-URL: Repository, https://github.com/googlecolab/google-colab-cli
+Project-URL: Issues, https://github.com/googlecolab/google-colab-cli/issues
+License-Expression: Apache-2.0
 License-File: LICENSE
-Requires-Python: >=3.13
+Keywords: cli,colab,google-colab,jupyter,notebook
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Software Development
+Requires-Python: >=3.12
 Requires-Dist: click>=8.0
+Requires-Dist: filelock>=3.29.2
 Requires-Dist: google-auth-oauthlib>=1.3.0
 Requires-Dist: google-auth>=2.49.1
+Requires-Dist: html2text>=2024.2.26
 Requires-Dist: jupyter-kernel-client
 Requires-Dist: nbformat>=5.10.4
 Requires-Dist: packaging>=24.0
@@ -26,6 +40,14 @@ A command-line interface for Google Colab. Provision high-performance CPU, GPU,
 
 Designed to support seamless developer productivity, headless automation, and AI agent integrations.
 
+[Demo](https://github.com/user-attachments/assets/656226a9-af13-4fdb-8eda-d7de747336a2)
+
+> [!NOTE]
+> **Platform support:** the Colab CLI currently supports **Linux and macOS** only. Windows is not supported at this time.
+
+> [!TIP]
+> Looking for in-notebook, interactive agent-assisted coding instead of a terminal workflow? See the [Colab MCP Server](https://github.com/googlecolab/colab-mcp).
+
 ---
 
 ## Key Features

{google_colab_cli-0.5.9 → google_colab_cli-0.6.0}/README.md

@@ -4,6 +4,14 @@ A command-line interface for Google Colab. Provision high-performance CPU, GPU,
 
 Designed to support seamless developer productivity, headless automation, and AI agent integrations.
 
+[Demo](https://github.com/user-attachments/assets/656226a9-af13-4fdb-8eda-d7de747336a2)
+
+> [!NOTE]
+> **Platform support:** the Colab CLI currently supports **Linux and macOS** only. Windows is not supported at this time.
+
+> [!TIP]
+> Looking for in-notebook, interactive agent-assisted coding instead of a terminal workflow? See the [Colab MCP Server](https://github.com/googlecolab/colab-mcp).
+
 ---
 
 ## Key Features

{google_colab_cli-0.5.9 → google_colab_cli-0.6.0}/docs/01_session_management.md

@@ -1,3 +1,9 @@
+---
+log:
+2026-06-15: Switched the keep-alive daemon from the `colab.pa.googleapis.com` `RuntimeService/KeepAliveAssignment` RPC to a Tunnel Frontend HTTP ping (`GET /tun/m/<endpoint>/keep-alive/` with `X-Colab-Tunnel: Google`) on `colab.research.google.com`. The RPC required `serviceusage` consumer access to Colab's internal project `1014160490159`, which ordinary user accounts lack, so every external user hit HTTP 403 `USER_PROJECT_DENIED` and their CLI sessions were idle-pruned within minutes (issue #14). Reproduced live with a third-party account; verified the tunnel ping is accepted by the same bearer-token credential that already works for `assign`. A `ReadTimeout` on the ping is treated as success (TFE records activity before forwarding to the often-non-responding VM). Generalized the pre-flight remediation messaging away from the now-irrelevant `colaboratory`/`pa.googleapis.com` framing, and removed the dead grpc-web client-registry/API-key code.
+2026-06-10: Replaced the POSIX-only `fcntl.flock` file locking in `_LockedFileStore` with the cross-platform `filelock` library (reported broken on Windows). Reads use `ReadWriteLock.read_lock()` (shared) and writes use `write_lock()` (exclusive), preserving the original `LOCK_SH`/`LOCK_EX` semantics. The lock is constructed with `is_singleton=False` so two `StateStore` instances for the same path in one process don't collapse into a single reentrant lock (which would raise `RuntimeError` on multi-threaded write contention). Added shared-read, cross-process exclusion, and multi-thread/multi-process regression tests.
+---
+
 # Design: Session Management (`new`, `status`, `stop`, `sessions`)
 
 ## Overview
@@ -65,15 +71,16 @@ The CLI maps user flags to these backend parameters:
 ### 5. Keep-Alive Protocol
 To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 minutes), the CLI implements a background keep-alive mechanism.
 - **Daemon Process**: Since the CLI is a fire-and-forget tool, `colab new` spawns a detached background process running a hidden `keep-alive` command.
-- **RPC**: Every 60 seconds, the daemon calls `google.internal.colab.v1.RuntimeService/KeepAliveAssignment` at `colab.pa.googleapis.com`. The wire format is grpc-web JSON: `Content-Type: application/json+protobuf`, body `["<endpoint>"]` (positional protojson), `X-Goog-Api-Client: grpc-web/0.1`, `x-user-agent: grpc-web-javascript/0.1`. **Both** the `colaboratory` OAuth scope (see `04_automation_and_utility.md`) and the `grpc-web` substring in `X-Goog-Api-Client` are server-enforced; missing either yields a descriptive 403/400 response.
-- **Pre-flight (`colab new`, OAuth2/ADC only)**: Immediately after a successful `assign`, the CLI invokes `keep_alive_assignment` once synchronously. If the response is 403 with a `SCOPE_NOT_PERMITTED` body, it unassigns the new VM (to avoid leaking a billable assignment) and prints a per-provider remediation message before exiting non-zero. Other errors are tolerated — the daemon will retry and surface them via the structured event log.
+- **Tunnel ping**: Every 60 seconds, the daemon issues `GET https://colab.research.google.com/tun/m/<endpoint>/keep-alive/` with the header `X-Colab-Tunnel: Google`, authenticated with the user's own Gaia bearer token (the same credential and host used for `/tun/m/assign`). The Tunnel Frontend (TFE) records `LastActiveTime` before forwarding the request, which refreshes the idle timer. This matches the official `colab-vscode` extension's `sendKeepAlive`. TFE notes the activity on arrival and then forwards to the VM, which often does not answer on this path — so the request commonly read-times-out even though the keep-alive succeeded; a `ReadTimeout` is therefore treated as success, while genuine HTTP errors (e.g. 404 for a deleted assignment) propagate.
+  - **Why not the RuntimeService RPC**: The previous implementation called `google.internal.colab.v1.RuntimeService/KeepAliveAssignment` at `colab.pa.googleapis.com` with `X-Goog-User-Project: 1014160490159`. That path requires the caller to be a `serviceusage` consumer of Colab's internal project `1014160490159`, which no ordinary user account is — so it returned HTTP 403 `USER_PROJECT_DENIED` for every external user, causing CLI sessions to be idle-pruned within minutes (issue #14). Dropping the header instead produced HTTP 400 `CONSUMER_INVALID` (public API-key project ≠ bearer-token quota project). The browser only succeeds because it rides the user's `google.com` cookie through an internal cookie-proxy (`colab.clients6.google.com`), which a headless bearer-token client cannot use. The TFE tunnel ping needs no project entitlement and works for any account that can assign a VM.
+- **Pre-flight (`colab new`, OAuth2/ADC only)**: Immediately after a successful `assign`, the CLI invokes `keep_alive_assignment` once synchronously. If the response is 403 with a `SCOPE_NOT_PERMITTED` body, it unassigns the new VM (to avoid leaking a billable assignment) and prints a per-provider remediation message before exiting non-zero. Other errors are tolerated — the daemon will retry and surface them via the structured event log. (Because keep-alive now uses the same backend/credential as `assign`, a scope failure at this stage is rare — assignment would normally have failed first.)
 - **Structured logging**: The daemon emits `keep_alive_started` (with `pid`, `endpoint`), one `keep_alive_error` per failed iteration (with `status_code`, `error_type`, truncated `error`, `response_body`, `iteration`, `consecutive_4xx`), and `keep_alive_stopped` (with `reason`, `iterations`, `duration_seconds`, optional `last_error`, optional `expected_endpoint`/`actual_endpoint`). All three are rendered specially by `colab log` so users get diagnostic context without parsing JSONL by hand.
 - **Termination**:
     - **Explicit**: `colab stop` terminates the daemon using its stored PID.
     - **Implicit**: If a session is pruned (e.g., during `sync_sessions`), its daemon is also terminated.
     - **Safety Fallback**: The daemon automatically terminates after 24 hours to prevent permanent zombie processes.
     - **State Check**: The daemon periodically verifies that its session still exists in the local state store; if missing, it exits.
-    - **Repeated 4xx**: After two consecutive 4xx responses, the daemon exits with `reason=consecutive_4xx_errors`. The pre-flight in `colab new` now catches the most common cause (missing `colaboratory` scope) before it reaches this branch.
+    - **Repeated 4xx**: After two consecutive 4xx responses, the daemon exits with `reason=consecutive_4xx_errors`. With the TFE tunnel ping, a normal read-timeout is not counted as a 4xx (it is treated as success), so this branch is now reached only by genuine HTTP errors such as a 404 for a deleted/expired assignment.
 
 ## TODO / Future Work
 - **Backend Sync**: Implement a way to reconcile the local `sessions.json` with the output of `colab sessions`.
@@ -88,6 +95,12 @@ To prevent Colab VMs from being deleted due to idle timeouts (standard is ~90 mi
 - Use `requests` for robust HTTP interactions and `pydantic` for schema validation.
 - Handle authentication headers (likely `Authorization: Bearer <token>` or cookies).
 
+### State Persistence & File Locking
+- **Stores**: `StateStore` (`sessions.json`) and `SettingsStore` (`settings.json`) both derive from `_LockedFileStore`, which guards concurrent access across independent `colab` invocations and the detached keep-alive daemon.
+- **Cross-platform locking**: Locking uses the [`filelock`](https://pypi.org/project/filelock/) library rather than `fcntl.flock`. `fcntl` is POSIX-only and is unavailable on Windows, so the original implementation crashed on import there. `filelock` provides the same advisory cross-process locking on Linux, macOS, and Windows.
+- **Shared vs. exclusive**: Each store owns a `filelock.ReadWriteLock` bound to a sidecar file (`<path>.lock`). Reads acquire `read_lock()` (shared — multiple concurrent readers allowed) and writes acquire `write_lock()` (exclusive). This preserves the `LOCK_SH`/`LOCK_EX` distinction of the previous `fcntl` implementation.
+- **`is_singleton=False`**: The `ReadWriteLock` is created with `is_singleton=False`. With `filelock`'s default (`True`), two `ReadWriteLock` objects for the same path *within a single process* are deduplicated into one reentrant lock; its reentrancy guard then raises `RuntimeError` when two threads each construct their own `StateStore` and contend for the write lock. Disabling the singleton registry makes each store's lock independent so they serialize via the underlying file lock instead.
+
 ## Testing Strategy
 TDD is mandatory for all session management features.
 
@@ -100,3 +113,8 @@ TDD is mandatory for all session management features.
 ### 2. State Store Validation
 - **Test Case**: Verify `StateStore` correctly handles file locking and multiple concurrent reads/writes.
 - **Test Case**: Verify `--config` override correctly directs all operations to the specified file path.
+- **Test Case (cross-platform locking)**: Verify the store locks via `filelock.ReadWriteLock` on the `<path>.lock` sidecar and does not import the POSIX-only `fcntl`.
+- **Test Case (shared/exclusive semantics)**: Verify reads go through `read_lock()` and writes through `write_lock()`.
+- **Test Case (cross-process exclusion)**: Hold the write lock from a separate process and confirm the store's in-process write blocks until release.
+- **Test Case (concurrent readers)**: Hold a read lock from a separate process and confirm the store can still complete a read concurrently.
+- **Test Case (multi-thread regression)**: Two `StateStore` instances writing from different threads must serialize without raising `RuntimeError` (guards the `is_singleton=False` choice).

{google_colab_cli-0.5.9 → google_colab_cli-0.6.0}/docs/02_execution_and_interactive.md

@@ -3,6 +3,8 @@ log:
 2026-05-07: Fixed `colab console` piped-stdin handling. Previously a piped invocation (e.g. `echo 'cmd' | colab console -s s`) sent the command and then hung indefinitely because the previous EOF handler emitted a bare `\x04` (Ctrl-D), which the remote `tmux`-wrapped bash treats as a literal character rather than a session terminator. The new handler sends `exit\n` (which bash actually exits on) and then closes the websocket from the client side after a short grace period (`PIPED_EOF_GRACE_SECONDS = 0.5s`) so any tail output (bash `logout`, tmux `[exited]`) makes it back to the user. TTY mode is unchanged: real-terminal EOF is left to the remote shell. Verified live: `echo 'echo HELLO' | colab console -s s` now exits in ~1.2s instead of hanging.
 
 2026-05-07: Fixed `print_kitty` (used by `colab exec --output-image` and any image-producing exec) to no-op when `sys.stdout.isatty()` is false. The Kitty Graphics Protocol escape sequence is meaningless when stdout is a file or pipe and was visually corrupting captured output (a multi-KB base64 PNG blob would land in log files, grep targets, or showboat captures). Image bytes are still saved to disk via `handle_image`'s file-write path; only the inline-render attempt is suppressed.
+
+2026-06-04: Bumped the default `--timeout` for `colab exec` from 10s to 30s (and the matching `colab run` default) so brief silent tasks are less likely to hit a premature `TimeoutError`. Explicit `--timeout` overrides are unaffected.
 ---
 
 # Design: Execution and Interactive Interaction (`repl`, `exec`, `console`)
@@ -26,7 +28,7 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel
     - If file path is local: Read content, send as code.
     - If file path is remote: Execute `!python <path>`.
 - **Multi-Modal Output**: Handle `display_data` messages (e.g., `image/png`, `text/html`). For the CLI, we'll save images to temporary files and print their paths, or if the terminal supports it (e.g., iTerm2), inline them.
-- **Timeout Configuration**: Exposes a `--timeout` flag (default 10s) to allow long-running silent tasks (like model compilation or data downloading) to execute without being prematurely killed.
+- **Timeout Configuration**: Exposes a `--timeout` flag (default 30s) to allow long-running silent tasks (like model compilation or data downloading) to execute without being prematurely killed.
 
 ### 3. Console (`colab console`)
 - **Implementation**: Connects directly to the backend terminal endpoint (`/colab/tty`) via WebSockets using `websocket-client`.

{google_colab_cli-0.5.9 → google_colab_cli-0.6.0}/docs/04_automation_and_utility.md

@@ -1,5 +1,6 @@
 ---
 log:
+2026-06-11: Replaced the `oauth2` provider's `run_local_server()` (localhost redirect) with a remote copy-paste flow (`_run_remote_flow` in `auth.py`). The CLI now prints an authorization URL built with `redirect_uri=https://sdk.cloud.google.com/applicationdefaultauthcode.html` and `token_usage=remote`, then reads the pasted authorization code via `input()` and exchanges it with `flow.fetch_token(code=...)`. This is the same flow `gcloud auth application-default login` uses and works identically in local and remote/headless/container environments, removing the heuristic of whether to auto-open a browser. Confirmed server-side acceptance with a live GET-only check against the bundled cloud-SDK client (`764086051850-...`); the OOB redirect and a non-bundled client id were both verified to be rejected (`OOB flow has been blocked` / `redirect_uri_mismatch`). Unit tests in `tests/test_auth.py` assert no localhost server is started, the redirect URI + `token_usage=remote` are set, and the pasted code is exchanged.
 2026-06-01: Enabled `colab update --install` self-update on macOS in addition to Linux. Refactored platform check logic to keep the implementation DRY and updated both tests and documentation. Also, on these platforms, an additional message is shown recommending `colab update --install` to upgrade in place, positioned above the standard `pip`/`uv` installation command.
 2026-05-29: Added default OAuth2 client config (`oauth_config.json`) as a bundled package resource and restored fallback loading logic in `get_credentials()`. The CLI now falls back to using these default credentials when no explicit local config is found. Added `integration/repro_bundled_oauth` integration test.
 2026-05-27: Refactored `colab README` and `colab AGENT` to bundle `README.md` and `AGENTS.md` via Hatchling's `force-include` and read them using `importlib.resources` instead of `importlib.metadata`. `colab AGENT` now correctly prints `AGENTS.md`.
@@ -23,11 +24,25 @@ managing local state, or inspecting the environment.
 The CLI supports two authentication strategies for talking to the Colab
 backend, selected via the global `--auth=<provider>` flag:
 
-1.  **`oauth2`** (default): Standard public InstalledAppFlow via
-    `google-auth-oauthlib`. Opens a browser for consent, caches the refresh
-    token at `~/.config/colab-cli/token.json`. If no local config is provided
-    via `-c/--client-oauth-config` or found at `~/.colab-cli-oauth-config.json`,
-    it falls back to a bundled `oauth_config.json` containing default OAuth credentials.
+1.  **`oauth2`** (default): Public `InstalledAppFlow` via
+    `google-auth-oauthlib`, but run with a **remote copy-paste flow** rather
+    than a localhost server. The CLI prints an authorization URL (with
+    `token_usage=remote`) using the registered HTTPS landing page
+    `https://sdk.cloud.google.com/applicationdefaultauthcode.html`; the user
+    signs in, copies the code Google displays, and pastes it back at the
+    prompt. The refresh token is cached at `~/.config/colab-cli/token.json`.
+    This is the same mechanism `gcloud auth application-default login` uses,
+    and it behaves identically on local, remote, headless, and container
+    hosts (no auto-opened browser, no bound port). We deliberately do **not**
+    use `run_local_server()` (environment-dependent) or the out-of-band (OOB)
+    redirect `urn:ietf:wg:oauth:2.0:oob` (blocked by Google in 2022 — see
+    `_run_remote_flow` / `REMOTE_REDIRECT_URI` in `auth.py`). The
+    `sdk.cloud.google.com` redirect is registered to the cloud-SDK OAuth
+    client (`764086051850-...`), which is also the client shipped in the
+    bundled `oauth_config.json`; reusing it with any other client id yields
+    `redirect_uri_mismatch`. If no local config is provided via
+    `-c/--client-oauth-config` or found at `~/.colab-cli-oauth-config.json`,
+    it falls back to that bundled `oauth_config.json`.
 2.  **`adc`**: Application Default Credentials via `google.auth.default()`.
     Honors the standard ADC discovery chain
     (`GOOGLE_APPLICATION_CREDENTIALS`, `gcloud auth application-default
@@ -41,18 +56,18 @@ allowing the core `Client` to remain authentication-agnostic — it only sees a
 
 ### Required Scopes
 
-The CLI talks to two distinct backends, each with different scope demands:
+The CLI talks to the Colab session backend at `colab.research.google.com`
+for assignment, unassignment, the contents API, **and keep-alive** (the TFE
+tunnel ping — see `01_session_management.md`). The `userinfo.email` scope is
+sufficient for this host.
 
--   `colab.research.google.com` (session assignment / unassignment /
-    contents API): the `userinfo.email` scope is sufficient.
--   `colab.pa.googleapis.com` (`RuntimeService`, used by
-    `KeepAliveAssignment`): **requires** the
-    `https://www.googleapis.com/auth/colaboratory` scope. Without it, every
-    request returns HTTP 403 with body `[7,"Request had insufficient
-    authentication scopes.",...]` and a `DebugInfo` mentioning
-    `SCOPE_NOT_PERMITTED`. (The frontend additionally requires
-    `X-Goog-Api-Client` to contain `grpc-web` — see
-    `01_session_management.md` §5.)
+> Historical note: keep-alive previously used the `RuntimeService`
+> (`KeepAliveAssignment`) at `colab.pa.googleapis.com`, which required the
+> `https://www.googleapis.com/auth/colaboratory` scope **and** the caller to
+> be a `serviceusage` consumer of Colab's internal project `1014160490159`.
+> The latter is impossible for ordinary user accounts, which made keep-alive
+> fail with HTTP 403 `USER_PROJECT_DENIED` for all external users (issue #14).
+> Keep-alive no longer touches `colab.pa.googleapis.com`.
 
 How each provider supplies the scope:
 
@@ -76,12 +91,11 @@ How each provider supplies the scope:
     ```
 
     `userinfo.email` is required for the session backend at
-    `colab.research.google.com` (otherwise assign/unassign/sessions return
-    HTTP 401); `colaboratory` is required for the `RuntimeService` at
-    `colab.pa.googleapis.com` (otherwise keep-alive returns HTTP 403);
-    `openid` and `cloud-platform` are mandated by `gcloud` itself
-    (`gcloud auth application-default login` rejects scope lists that
-    omit `cloud-platform` with `Invalid value for [--scopes]`).
+    `colab.research.google.com` (otherwise assign/unassign/sessions/keep-alive
+    return HTTP 401); `colaboratory` is retained for forward compatibility and
+    other Colab features; `openid` and `cloud-platform` are mandated by
+    `gcloud` itself (`gcloud auth application-default login` rejects scope
+    lists that omit `cloud-platform` with `Invalid value for [--scopes]`).
 
 `colab new` performs a one-shot keep-alive pre-flight after `assign`
 succeeds so missing-scope failures surface immediately (with per-provider

{google_colab_cli-0.5.9 → google_colab_cli-0.6.0}/docs/05_run_command.md

@@ -2,6 +2,7 @@
 log:
 2026-05-12: Initial design and implementation of `colab run <script.py> [args...]`. Combines `colab new` + `colab exec` + `colab stop` into a single fire-and-forget invocation so a Python file can use `#!/usr/bin/env -S colab run` as a shebang line and execute on a freshly-allocated Colab VM. Adds `--keep` (skip auto-stop), `--gpu` / `--tpu` (passthrough to session creation), `-s/--session` (name the ephemeral session), and propagates the script's exit status (non-zero on any uncaught exception in the kernel). The script's `sys.argv` is re-set inside the kernel to mirror native `python script.py arg1 arg2` semantics, and `__name__` is set to `"__main__"`.
 2026-05-12: Native CPython exit-code semantics for `sys.exit()` / `raise SystemExit(...)` from the script body. The Colab kernel reports a `SystemExit` as `output_type=='error'`, which under the previous logic would have (a) printed the IPython traceback (`An exception has occurred, use %tb...`) and (b) flagged the run as a failure regardless of the integer exit code. Now: `sys.exit()` / `sys.exit(0)` exit 0 silently; `sys.exit(N)` exits N; `sys.exit('msg')` exits 1 (matching CPython). The IPython "To exit: use 'exit', 'quit', or Ctrl-D." UserWarning is filtered via the prelude. Encoded after running `examples/gpu_hello.py` end-to-end and seeing the noisy `SystemExit: 0` traceback at the end of an otherwise-successful GPU run.
+2026-06-04: Bumped the default value of the `--timeout` flag from 10.0s to 30.0s so short-but-silent tasks aren't prematurely killed out of the box. Mirrors the same change for `colab exec`.
 ---
 
 # Design: `colab run` — Shebang-Compatible One-Shot Execution
@@ -29,7 +30,7 @@ colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...
 | `--gpu` | str | None | Same set as `colab new --gpu` (T4, L4, G4, H100, A100). |
 | `--tpu` | str | None | Same set as `colab new --tpu` (v5e1, v6e1). |
 | `--keep` | bool | False | Do **not** stop the session after the script finishes. |
-| `--timeout` | float | 10.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. |
+| `--timeout` | float | 30.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. |
 
 ### Shebang usage
 With `--keep` and `--gpu` baked into the shebang line, an entire one-file workload becomes:

google_colab_cli-0.6.0/examples/finetune_run.py

@@ -0,0 +1,153 @@
+# -----------------------------------------------------------------------------
+# Setup: HF_TOKEN is required
+# -----------------------------------------------------------------------------
+# This script downloads google/gemma-3-1b-it, which is a gated model. You need
+# a Hugging Face access token. Get one at https://huggingface.co/settings/tokens
+# and accept the model license at https://huggingface.co/google/gemma-3-1b-it
+#
+# 1. Set HF_TOKEN in your local shell, persistently:
+#
+#       echo 'export HF_TOKEN=hf_yourTokenHere' | tee -a ~/.zshrc ~/.bashrc
+#       source ~/.zshrc                    # or open a new terminal
+#       echo $HF_TOKEN                     # verify — should print your token
+#
+# 2. Pipe the local env var into the colab kernel before running this script:
+#
+#       echo "import os; os.environ['HF_TOKEN'] = '$HF_TOKEN'" | colab exec
+#
+# 3. Verify the kernel received it:
+#
+#       echo 'import os; print(bool(os.environ.get("HF_TOKEN")))' | colab exec
+#       # → should print: True
+#
+# 4. Run this script:
+#
+#       colab exec -f finetune_run.py
+#
+# Note: HF_TOKEN lives in the colab kernel for the lifetime of the session.
+# If you `colab stop` or the session expires, you'll need to re-pipe it (step 2).
+# -----------------------------------------------------------------------------
+
+import os
+
+os.system("pip install -q -U 'bitsandbytes>=0.46.1'")
+
+import torch
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model
+from trl import SFTConfig, SFTTrainer
+
+MODEL_ID = "google/gemma-3-1b-it"
+NUM_SAMPLES = 200  # demo size; bump to 5000+ for a real run
+MAX_STEPS = 60  # demo cap; set to -1 for full-epoch training
+
+# -------- Data --------
+# philschmid/gretel-synthetic-text-to-sql has sql_prompt, sql_context, sql.
+# We hand SFTTrainer a "messages" column and let it apply the chat template.
+print("Loading dataset...")
+dataset = load_dataset("philschmid/gretel-synthetic-text-to-sql", split="train").select(
+    range(NUM_SAMPLES)
+)
+
+
+def to_messages(example):
+    user_msg = (
+        "You are a SQL expert. Given the schema, write a SQL query that "
+        "answers the question. Reply with only the SQL.\n\n"
+        f"Schema:\n{example['sql_context']}\n\n"
+        f"Question:\n{example['sql_prompt']}"
+    )
+    return {
+        "messages": [
+            {"role": "user", "content": user_msg},
+            {"role": "assistant", "content": example["sql"]},
+        ]
+    }
+
+
+dataset = dataset.map(to_messages, remove_columns=dataset.column_names)
+
+# -------- Model (4-bit QLoRA, bf16 throughout) --------
+# Everything is bf16 — matches Gemma's natural dtype, matches TRL's default
+# T4 (Turing) has no hardware bf16, so this is slower than fp16 would be (~2x)
+print(f"Loading {MODEL_ID} in 4-bit...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+    ),
+    device_map="auto",
+)
+
+model = get_peft_model(
+    model,
+    LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules="all-linear",
+        task_type="CAUSAL_LM",
+    ),
+)
+# Required for QLoRA backward: makes the embedding output require grad so that
+# gradients can flow into the LoRA params attached to layers downstream of the
+# frozen 4-bit base.
+model.enable_input_require_grads()
+model.print_trainable_parameters()
+
+# -------- Train --------
+# All other knobs use SFTConfig defaults (which include bf16=True,
+# gradient_checkpointing=True, logging_steps=10). The overrides below are just
+# the demo cap, batch sizing that fits T4 VRAM, and silencing wandb/tensorboard.
+print("Training...")
+trainer = SFTTrainer(
+    model=model,
+    train_dataset=dataset,
+    processing_class=tokenizer,
+    args=SFTConfig(
+        output_dir="./results",
+        max_steps=MAX_STEPS,
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=2,
+        # Standard QLoRA LR. SFTConfig defaults to 2e-5, which is too low for
+        # LoRA adapters to learn anything meaningful in 60 steps.
+        learning_rate=2e-4,
+        # Compute loss only on the assistant's SQL, not on the schema/question.
+        assistant_only_loss=True,
+        # Off so KV cache works during the inference step at the end.
+        gradient_checkpointing=False,
+        report_to="none",
+    ),
+)
+trainer.train()
+
+# -------- Save --------
+out_dir = "./gemma-3-1b-qlora-adapter"
+trainer.model.save_pretrained(out_dir)
+tokenizer.save_pretrained(out_dir)
+print(f"Saved adapter to {out_dir}")
+
+# -------- Inference check --------
+sample = dataset[0]
+prompt = tokenizer.apply_chat_template(
+    sample["messages"][:1],  # just the user turn
+    tokenize=False,
+    add_generation_prompt=True,
+)
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+with torch.no_grad():
+    out_ids = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=False,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+generated = tokenizer.decode(
+    out_ids[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
+)
+print(f"\nGold:  {sample['messages'][1]['content']}")
+print(f"Model: {generated.strip()}")