diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..75c1be0 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,5 @@ +build/ +.git/ +tests/ +docs/ +.env diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..229c393 --- /dev/null +++ b/.env.example @@ -0,0 +1,17 @@ +# Tailscale auth key — only needed on first run once state is persisted to the volume. +# Generate one at: https://login.tailscale.com/admin/settings/keys +TS_AUTHKEY=tskey-auth-xxxx + +# Device name on the tailnet. Aperture reaches tsheadroom at http://..ts.net/ +TS_HOSTNAME=tsheadroom + +# headroom-ai variant to install at build time. +# base — tool-output compression only (SmartCrusher); low memory, no ML model. +# ml — tool + text/prose compression (Kompress ML); ~600 MB/worker resident. +# Changing this requires a rebuild: docker compose up -d --build +HEADROOM_VARIANT=ml + +# Number of persistent Python workers. Each holds a resident ML model copy when using ml. +# ~600 MB/worker — e.g. POOL_SIZE=4 uses ~2.4 GB. Size deliberately. +# Tip: start with POOL_SIZE=1 on a fresh volume to let the model download, then raise it. +POOL_SIZE=4 diff --git a/.gitignore b/.gitignore index 05a8b93..7b73ae9 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,9 @@ __pycache__/ /tsheadroom.config.json *.config.json.tmp +# Local Docker config dir (bind-mounted into the container); copy of the example. +/config/ + # tsnet state — contains tailscaled.state (the node's private key). Never commit. # `t/` is the conventional local -state-dir; also guard the state files by name. /t/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..8134d27 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,27 @@ +# Stage 1: build the Go binary +FROM golang:1.26 AS builder +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY . . +RUN go build -o /build/tsheadroom . + +# Stage 2: Python runtime with headroom-ai +FROM python:3.13-slim +ARG HEADROOM_VARIANT=base + +RUN python -m venv /venv + +# Install the chosen headroom-ai variant. +# python:3.13-slim ships prebuilt wheels so no Rust toolchain is needed. +# 'ml' adds Kompress (~600 MB ML model downloaded on first use, cached in a volume). +RUN if [ "$HEADROOM_VARIANT" = "ml" ]; then \ + /venv/bin/pip install --no-cache-dir 'headroom-ai[ml]'; \ + else \ + /venv/bin/pip install --no-cache-dir 'headroom-ai'; \ + fi + +COPY --from=builder /build/tsheadroom /app/tsheadroom +COPY worker.py /app/worker.py + +ENTRYPOINT ["/app/tsheadroom"] diff --git a/README.md b/README.md index 80bebe7..439b309 100644 --- a/README.md +++ b/README.md @@ -167,6 +167,39 @@ WantedBy=multi-user.target `StateDirectory=tsheadroom` provisions `/var/lib/tsheadroom` with the right ownership. After the device has authenticated once, you can drop the `TS_AUTHKEY` line. +### Run it with Docker + +If you'd rather not set up a Go toolchain, a virtualenv, and a `systemd` unit by hand, the repo ships a `Dockerfile` and `docker-compose.yml` that build and run tsheadroom for you. You can skip the [Install](#install) and [Run](#run) steps above; deploys and updates become a single `docker compose up -d --build`. + +```shell +git clone https://github.com/tailscale/tsheadroom.git +cd tsheadroom + +cp .env.example .env # set TS_AUTHKEY, variant, pool size +mkdir -p config && cp tsheadroom.config.example.json config/config.json + +docker compose up -d --build +``` + +This builds the binary, installs `headroom-ai`, joins your tailnet as `tsheadroom`, and serves the hook on `:80` — the same result as the manual install, in one command. Settings live in `.env` (what `.env.example` ships with): + +| Variable | `.env.example` | Description | +|---|---|---| +| `TS_AUTHKEY` | — | Tailscale [auth key](https://tailscale.com/docs/features/access-control/auth-keys); only needed until the device authenticates once. | +| `TS_HOSTNAME` | `tsheadroom` | Device name on the tailnet. | +| `HEADROOM_VARIANT` | `ml` | `base` for tool-output compression, `ml` for text + tool compression (installs `headroom-ai[ml]`). See [Choose tool-output or text compression](#choose-tool-output-or-text-compression). Changing it needs a rebuild. | +| `POOL_SIZE` | `4` | Number of Python workers; each holds a resident copy of the ML model under `ml`. | + +The device identity and the HuggingFace model cache persist in named volumes (`tsheadroom-state`, `tsheadroom-hf-cache`), so restarts never re-authenticate the device or re-download the ~600 MB model. Compression knobs live in the bind-mounted `config/config.json`: edit it and restart, or change them live with `PUT /config` (see [Tune compression](#tune-compression-runtime-config)). + +Update to a newer version with: + +```shell +git pull && docker compose up -d --build +``` + +> **Migrating from a manual or `systemd` install?** Before the first `up`, copy your existing `-state-dir` into the `tsheadroom-state` volume (and `~/.cache/huggingface` into `tsheadroom-hf-cache`) so the container keeps the same device identity and skips the model download. Stop the old service first — the two can't share one tailnet identity. + ### Flags | Flag | Default | Description | diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..784b3e5 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,37 @@ +services: + tsheadroom: + build: + context: . + args: + HEADROOM_VARIANT: ${HEADROOM_VARIANT:-base} + command: + - -python + - /venv/bin/python + - -worker + - /app/worker.py + - -hostname + - ${TS_HOSTNAME:-tsheadroom} + - -pool-size + - "${POOL_SIZE:-4}" + - -state-dir + - /var/lib/tsheadroom + - -config + - /etc/tsheadroom/config.json + - -v + environment: + TS_AUTHKEY: ${TS_AUTHKEY:-} + volumes: + # Tailnet device identity. Must persist or the device re-authenticates as a new node. + - tsheadroom-state:/var/lib/tsheadroom + # HuggingFace ML model cache (~600 MB). Shared by all workers; survives restarts. + - tsheadroom-hf-cache:/root/.cache/huggingface + # Compression knobs. A directory (not a single file) so the binary's atomic + # PUT /config write can rename within it; edit ./config/config.json on the + # host and restart, or use the PUT /config API. Seed it once: + # mkdir -p config && cp tsheadroom.config.example.json config/config.json + - ./config:/etc/tsheadroom + restart: unless-stopped + +volumes: + tsheadroom-state: + tsheadroom-hf-cache: diff --git a/tsheadroom.config.example.json b/tsheadroom.config.example.json new file mode 100644 index 0000000..f6da067 --- /dev/null +++ b/tsheadroom.config.example.json @@ -0,0 +1,9 @@ +{ + "compress_user_messages": false, + "compress_system_messages": true, + "protect_recent": 4, + "protect_analysis_context": true, + "target_ratio": null, + "min_tokens_to_compress": 250, + "kompress_model": null +}