longbow-quarrel/Dockerfile.tpu at main · 23skdu/longbow-quarrel · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
# =============================================================================
# Longbow-Quarrel - Google TPU Build
# =============================================================================
# This Dockerfile builds Quarrel with Google TPU/XLA acceleration.
# Uses XLA (Accelerated Linear Algebra) for TPU execution.
#
# Build:    docker build -f Dockerfile.tpu -t longbow-quarrel:tpu .
# Run:      docker run --privileged -v $(pwd)/models:/data longbow-quarrel:tpu --model /data/model.gguf
# =============================================================================

# -----------------------------------------------------------------------------
# Build Stage: TPU Compilation
# -----------------------------------------------------------------------------
FROM ubuntu:22.04 AS tpu-builder

# Install build dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    golang-go \
    git \
    build-essential \
    pkg-config \
    wget \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Install Go 1.26 for XLA compatibility
RUN wget -q https://go.dev/dl/go1.26.3.linux-amd64.tar.gz && \
    rm -rf /usr/local/go && \
    tar -C /usr/local -xzf go1.26.3.linux-amd64.tar.gz && \
    rm go1.26.3.linux-amd64.tar.gz

ENV PATH=/usr/local/go/bin:$PATH \
    GO111MODULE=on

# Install JAX with TPU support (includes libtpu)
# JAX provides the TPU XLA backend
RUN pip install --no-cache-dir jax[tpu]>=0.4.20 -f https://storage.googleapis.com/jax-releases/jax_lib.html

# XLA/TPU environment setup
ENV LIBTPU_INIT_ARGS="force_authoritative=true" \
    TPU_NAME=localtpu \
    XLA_FLAGS="--xla_tpu_enable_latency_hiding_communication=true"

# Copy go mod files
WORKDIR /app
COPY go.mod go.sum ./
RUN go mod download

# Copy source code
COPY . .

# Build with TPU support
# Note: Uses custom tpu build tag or libtpu integration
RUN CGO_ENABLED=1 GOOS=linux GOARCH=amd64 go build -tags tpu \
    -o quarrel-tpu ./cmd/quarrel

# -----------------------------------------------------------------------------
# Runtime Stage: TPU Runtime
# -----------------------------------------------------------------------------
FROM ubuntu:22.04

# Install TPU runtime dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3-pip \
    libtpu0 \
    && rm -rf /var/lib/apt/lists/*

# Install JAX runtime (libtpu)
RUN pip3 install --no-cache-dir jax[tpu]>=0.4.20 -f https://storage.googleapis.com/jax-releases/jax_lib.html

# XLA/TPU environment setup
ENV LIBTPU_INIT_ARGS="force_authoritative=true" \
    TPU_NAME=localtpu \
    XLA_FLAGS="--xla_tpu_enable_latency_hiding_communication=true" \
    LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/local/lib

WORKDIR /app

COPY --from=tpu-builder /app/quarrel-tpu ./quarrel

CMD ["./quarrel", "--help"]

# =============================================================================
# Alternative: Cloud TPU Build
# =============================================================================
# For Google Cloud TPU v5e/v5p:
#
#   FROM gcr.io/deeplearning-platform-release/base:latest
#   RUN pip install jax[tpu]
#   COPY . .
#   RUN CGO_ENABLED=1 go build -tags tpu -o quarrel-tpu ./cmd/quarrel
#
# =============================================================================
# Build Instructions
# =============================================================================
# Local TPU:
#   docker build -f Dockerfile.tpu -t longbow-quarrel:tpu .
#   docker run --privileged -v $(pwd)/models:/data longbow-quarrel:tpu --model /data/model.gguf
#
# Cloud TPU:
#   gcloud compute tpus zones us-central1-a list --accelerator-type v5e-256
#   docker build -f Dockerfile.tpu -t longbow-quarrel:tpu .
#   gcloud compute ssh <tpu-vm> --zone us-central1-a
#   docker run --network=host -v $(pwd)/models:/data longbow-quarrel:tpu --model /data/model.gguf
#
# =============================================================================
# TPU Requirements
# =============================================================================
# - Cloud TPU: v5e, v5p, or v4
# - Local: USB TPU or PCIe TPU
# - XLA: Latest stable
# - Python: 3.10+
#
# Run with TPU:
#   docker run --privileged -v /path/to/models:/data longbow-quarrel:tpu \
#       --model /data/smollm2.gguf --prompt "Hello"
#
# =============================================================================