Team dựng hệ thống trên GKE với hàng đợi và worker TPU Spot: có request thì bật chip, rảnh thì tắt.

• Phần cứng: 1 chip TPU v6e (ct6e-standard-1t, 32 GiB HBM, 176 GiB RAM host) trên GKE, node pool Spot ở us-central1, tự scale được về 0.

• Phần mềm: image torch-tpu (torch 2.14, torch_tpu 0.2.0) cộng diffusers bản main. TorchTPU là điểm then chốt — code PyTorch/diffusers gần như giữ nguyên, chỉ đổi device=”tpu”, không phải viết lại bằng XLA thuần.

• Model thử nghiệm: FLUX 3 Image Klein và FLUX.1-schnell (ảnh), Wan2.2-TI2V-5B (video không tiếng), MiniMax-H3 33B (video có tiếng).

• Thành phần chính: API FastAPI trên node pool thường (không cần TPU); hàng đợi jobs trong Firestore (queued → running → done/failed), worker nhận job bằng compare-and-swap để không trùng; kết quả lưu Cloud Storage và tải qua signed URL; đĩa PVC giữ weights và compile cache TorchTPU (model lớn, cache theo shape nên cần đĩa bền).

• Worker: mỗi worker là một Deployment 1 replica trên 1 chip, strategy Recreate để không có hai pod tranh một chip. Worker chính giữ Klein và Wan2.2 trong RAM host, chỉ đưa model đang dùng lên HBM — đổi giữa ảnh và video mất khoảng 12–15 giây, không phải compile lại. Worker MiniMax-H3 chạy riêng vì weights lớn.

• Spot rẻ hơn on-demand nhưng node có thể bị thu hồi bất cứ lúc nào. Khi worker khởi động lại, các job đang chạy dở được đưa lại hàng đợi.

• Cold start worker chính khoảng 8–26 phút (cấp node + nạp model + warmup); H3 khoảng 10–25 phút. Compile cache lưu trên PVC nên lần sau dùng lại. Worker tự tắt sau 15 phút rảnh (H3: 30 phút) để thu hồi node.

Điểm cần chú ý khi cài đặt là giữ nguyên bản torch và torch_tpu:

FROM us-docker.pkg.dev/ml-oss-artifacts-published/torch-tpu-public-release-container/torch-tpu:latest
RUN pip freeze | grep -iE ‘^(torch|torch[_-]tpu|torchvision|torchax|jax|libtpu)’ > /tmp/pins.txt \
&& pip install --no-cache-dir -c /tmp/pins.txt -r requirements-worker.txt

Tiếp theo là quyết định thành phần nào đặt trên TPU:

# FLUX.1-schnell: text encoder ở CPU, chỉ transformer + VAE lên TPU
pipe.transformer.to(”tpu”); pipe.vae.to(”tpu”)

Một số điểm team Cloud Ace rút ra khi dựng:

• 🔹 Compile theo hình dạng tensor. Mỗi độ dài prompt, độ phân giải hay số frame là một lần compile (30–110 giây), nên cần cố định tập kích thước và warmup sẵn.

• 🔹 Chỉ đưa từng phần cần thiết lên TPU. Không chuyển cả pipeline lên TPU một lần bằng pipeline.to(”tpu”), vì phần text encoder đang chạy trên CPU sẽ bị kéo lên theo. Cũng không dùng from_pipe() để tạo pipeline mới, vì hàm này đổi model về định dạng số lớn hơn (float32). Với FLUX, dung lượng model khi đó tăng lên khoảng 44 GB, vượt quá bộ nhớ 32 GB của một chip.