pkg/serve.py

#!/usr/bin/env python
"""dixie-flatline: onegraph-spec7 substrate + PCK-04 lm_head vocabulary pruning.

Stack:
  - blake's onegraph (315.12 TPS public best): one CUDA-graph replay of K=7
    width-1 drafter iterations, ping-pong slots=3, fused sparse argmax.
  - PCK-04: pruned lm_head (K≈32k rows) scattered back to full-vocab at
    compute_logits time, via serve_patch_pck04.py meta-path hook.
  - SMP-02: slim greedy sampler fast path + lastchance prewarm.
  - PLE scale-fold / fast path.

serve.py is structurally identical to field-artifacts/onegraph-spec7-v0/serve.py;
only deltas are:
  1. WEIGHTS_BUCKET / LOCAL_MODEL_DIR / LOCAL_DRAFTER_DIR defaults updated for
     the int4-pck04-16k checkpoint.
  2. setup_pck04_path() injects submission-pck04/ into PYTHONPATH so the worker
     process imports serve_patch_pck04.py via sitecustomize (meta-path finder).
  3. main() calls setup_pck04_path() after setup_sitecustomize_path().

kenyan-duma composition delta (vs hayai-agent osoi-v0 serve.py, byte-identical
otherwise): ensure_drafter() gains an optional DRAFTER_BUCKET env branch
(hf buckets sync, same mechanism as ensure_weights) for serving retrained
drafter checkpoints, and logs the sha256 of the drafter model.safetensors it
actually loads so the run record proves WHICH drafter served (stale-dir trap).
With DRAFTER_BUCKET unset, behavior is identical to hayai's original. Greedy
spec decode emits the target's argmax regardless of drafter proposals, so
emitted tokens are governed by the target checkpoint alone.
"""

from __future__ import annotations

import glob
import http.server
import json
import os
import pathlib
import shutil
import signal
import subprocess
import sys
import sysconfig
import threading
import time
import urllib.error
import urllib.request
from collections.abc import Callable


WEIGHTS_BUCKET = os.environ.get(
    "WEIGHTS_BUCKET",
    "hf://buckets/gemma-challenge/gemma-dixie-flatline/weights/int4-pck04-16k",
)
LOCAL_MODEL_DIR = os.environ.get("LOCAL_MODEL_DIR", "/tmp/int4-pck04-16k")
DRAFTER_REPO = os.environ.get(
    "DRAFTER_REPO", "google/gemma-4-E4B-it-qat-q4_0-unquantized-assistant"
)
LOCAL_DRAFTER_DIR = os.environ.get("LOCAL_DRAFTER_DIR", "/tmp/qat-assistant")
# Optional bucket override for the drafter (takes precedence over DRAFTER_REPO);
# used to serve retrained drafter checkpoints. The loaded-file sha256 is logged
# either way so the run record proves WHICH drafter actually served.
DRAFTER_BUCKET = os.environ.get("DRAFTER_BUCKET")
DRAFTER_SHA256 = os.environ.get("DRAFTER_SHA256", "").lower()
CENTROID_TOP_K = int(os.environ.get("CENTROID_TOP_K", "64"))
JINJA2_VERSION = "3.1.6"
MARKUPSAFE_VERSION = "3.0.3"

TCMALLOC_CANDIDATES = [
    "/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4",
    "/usr/lib/libtcmalloc_minimal.so.4",
    "/usr/lib64/libtcmalloc_minimal.so.4",
]

Patcher = Callable[[str, pathlib.Path], tuple[str, bool]]


def replace_required(
    source: str,
    *,
    model_path: pathlib.Path,
    label: str,
    old: str,
    new: str,
    marker: str,
) -> tuple[str, bool]:
    """Apply one idempotent source replacement and fail on source drift."""
    if marker in source:
        return source, False
    old_count = source.count(old)
    if old_count != 1:
        raise RuntimeError(
            f"{label} patch pattern count is {old_count} in {model_path}; "
            "refusing to run a silent no-op baseline."
        )
    return source.replace(old, new, 1), True


PLE_TEXT_FAST_PATH_OLD = """        per_layer_inputs_mask = torch.logical_and(
            input_ids >= 0,
            input_ids < self.vocab_size_per_layer_input,
        )
        per_layer_inputs_tokens = torch.where(
            per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)
        )
        per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens)
"""

PLE_TEXT_FAST_PATH_NEW = """        # Challenge fast path: harness text token IDs are valid PLE IDs.
        # Multimodal serving still maps multimodal positions to token 0 before
        # this call in gemma4_mm.py, so the multimodal PLE contract is retained.
        per_layer_embeds = self.embed_tokens_per_layer(input_ids)
"""

PLE_RUNTIME_SCALE_OLD = (
    "        per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer\n"
)
PLE_RUNTIME_SCALE_NEW = (
    "        # PLE scale-fold: embed_scale_per_layer is folded into "
    "embedding weights after load.\n"
)

PLE_GATE_SCRATCH_OLD = """            gate = self.per_layer_input_gate(hidden_states)
            gate = torch.nn.functional.gelu(gate, approximate="tanh")
            gated_per_layer = gate * per_layer_input
            per_layer_contribution = self.per_layer_projection(gated_per_layer)
"""

PLE_GATE_SCRATCH_NEW = """            gate = self.per_layer_input_gate(hidden_states)
            gate = torch.nn.functional.gelu(gate, approximate="tanh")
            # PLE scratch reuse: in-place gate multiply when dtype-preserving.
            if gate.dtype == per_layer_input.dtype:
                gate.mul_(per_layer_input)
                gated_per_layer = gate
            else:
                gated_per_layer = gate * per_layer_input
            per_layer_contribution = self.per_layer_projection(gated_per_layer)
"""

PLE_COMBINE_SCRATCH_OLD = """        if per_layer_inputs is None:
            return per_layer_projection
        return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
"""

PLE_COMBINE_SCRATCH_NEW = """        if per_layer_inputs is None:
            return per_layer_projection
        # PLE scratch reuse: in-place projection add when dtype-preserving.
        if per_layer_projection.dtype == per_layer_inputs.dtype:
            per_layer_projection.add_(per_layer_inputs)
            return per_layer_projection * self.per_layer_input_scale
        return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
"""

SELF_DECODER_FOLD_ANCHOR = """    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids) * self.normalizer

    def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
"""
SELF_DECODER_FOLD_METHOD = """    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids) * self.normalizer

    @torch.inference_mode()
    def fold_per_layer_embed_scale(self) -> None:
        if self.embed_tokens_per_layer is None or self.embed_scale_per_layer is None:
            return
        if getattr(self.embed_tokens_per_layer, "_ple_embed_scale_folded", False):
            return
        if self.hidden_size_per_layer_input != 256:
            raise RuntimeError(
                "PLE scale-fold expected hidden_size_per_layer_input=256, "
                f"got {self.hidden_size_per_layer_input}"
            )
        if self.embed_scale_per_layer.numel() != 1:
            raise RuntimeError("PLE scale-fold expects scalar embed_scale_per_layer")

        scale = float(self.embed_scale_per_layer.item())
        expected_scale = float(self.hidden_size_per_layer_input ** 0.5)
        if scale != expected_scale:
            raise RuntimeError(
                f"PLE scale-fold expected scale {expected_scale}, got {scale}"
            )

        embedding = self.embed_tokens_per_layer
        if hasattr(embedding, "weight_scale"):
            target = embedding.weight_scale
            folded_name = "weight_scale"
        elif hasattr(embedding, "weight"):
            target = embedding.weight
            folded_name = "weight"
        else:
            raise RuntimeError(
                "PLE scale-fold found no weight_scale or weight on "
                "embed_tokens_per_layer"
            )

        if target.dtype != torch.bfloat16:
            raise RuntimeError(
                f"PLE scale-fold expects bf16 {folded_name}, got {target.dtype}"
            )
        if target.device.type != "cuda":
            raise RuntimeError(
                f"PLE scale-fold expects CUDA {folded_name}, got {target.device}"
            )

        target.data.mul_(scale)
        embedding._ple_embed_scale_folded = True
        logger.info("Folded Gemma4 PLE embed scale %s into %s", scale, folded_name)

    def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
"""

MODEL_DELEGATE_OLD = '''    def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
        """Get per-layer embeddings from embed_tokens_per_layer.

        Returns:
            Per-layer embeddings (num_tokens, num_layers,
            hidden_size_per_layer_input)
        """
        return self.self_decoder.get_per_layer_inputs(input_ids)

    def project_per_layer_inputs(
'''
MODEL_DELEGATE_NEW = '''    def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
        """Get per-layer embeddings from embed_tokens_per_layer.

        Returns:
            Per-layer embeddings (num_tokens, num_layers,
            hidden_size_per_layer_input)
        """
        return self.self_decoder.get_per_layer_inputs(input_ids)

    def fold_per_layer_embed_scale(self) -> None:
        self.self_decoder.fold_per_layer_embed_scale()

    def project_per_layer_inputs(
'''

LOADER_IMPORT_OLD = "import inspect\n"
LOADER_IMPORT_NEW = "import inspect\nimport os\n"
LOADER_HOOK_OLD = """    if model_config.quantization == "torchao":
        set_torchao_reload_attrs(model, model_config)
"""
LOADER_HOOK_NEW = """    if model_config.quantization == "torchao":
        set_torchao_reload_attrs(model, model_config)

    if os.environ.get("PLE_FOLD_EMBED_SCALE") == "1":
        fold_target_model = os.environ.get("PLE_FOLD_TARGET_MODEL")
        current_model = getattr(model_config, "model", None)
        if fold_target_model and current_model != fold_target_model:
            logger.info(
                "Skipping Gemma4 PLE embed_scale_per_layer fold for "
                "non-target model %s",
                current_model,
            )
        else:
            candidates = [
                model,
                getattr(model, "model", None),
                getattr(getattr(model, "language_model", None), "model", None),
            ]
            fold_applied = False
            for candidate in candidates:
                folder = getattr(candidate, "fold_per_layer_embed_scale", None)
                if folder is None:
                    continue
                logger.info("Folding Gemma4 PLE embed_scale_per_layer")
                folder()
                decoder = getattr(candidate, "self_decoder", None)
                embedding = getattr(decoder, "embed_tokens_per_layer", None)
                fold_applied = bool(
                    getattr(embedding, "_ple_embed_scale_folded", False)
                )
                if not fold_applied:
                    raise RuntimeError(
                        "PLE_FOLD_EMBED_SCALE=1 but fold_per_layer_embed_scale "
                        "did not mark embed_tokens_per_layer as folded"
                    )
                break
            if not fold_applied:
                raise RuntimeError(
                    "PLE_FOLD_EMBED_SCALE=1 but no target model candidate "
                    "exposed fold_per_layer_embed_scale"
                )
"""


def build_ple_text_fast_source(
    source: str, model_path: pathlib.Path
) -> tuple[str, bool]:
    """Build Gemma4 source with the exact PLE valid-token fast path applied.

    Args:
        source: Current text of vLLM's Gemma4 model file.
        model_path: Path included in failure messages for actionable startup errors.

    Returns:
        A pair of patched source text and whether the text changed.

    Raises:
        RuntimeError: If neither the original nor patched source block is present.
    """
    return replace_required(
        source,
        model_path=model_path,
        label="PLE fast path",
        old=PLE_TEXT_FAST_PATH_OLD,
        new=PLE_TEXT_FAST_PATH_NEW,
        marker="Challenge fast path: harness text token IDs are valid PLE IDs.",
    )


def patch_gemma4_source(source: str, model_path: pathlib.Path) -> tuple[str, bool]:
    changed_any = False
    for label, old, new, marker in (
        (
            "PLE valid-token fast path",
            PLE_TEXT_FAST_PATH_OLD,
            PLE_TEXT_FAST_PATH_NEW,
            "Challenge fast path: harness text token IDs are valid PLE IDs.",
        ),
        (
            "PLE scale-fold method",
            SELF_DECODER_FOLD_ANCHOR,
            SELF_DECODER_FOLD_METHOD,
            "def fold_per_layer_embed_scale",
        ),
        (
            "PLE runtime scale multiply",
            PLE_RUNTIME_SCALE_OLD,
            PLE_RUNTIME_SCALE_NEW,
            "PLE scale-fold: embed_scale_per_layer is folded into embedding weights",
        ),
        (
            "PLE gate scratch reuse",
            PLE_GATE_SCRATCH_OLD,
            PLE_GATE_SCRATCH_NEW,
            "PLE scratch reuse: in-place gate multiply",
        ),
        (
            "PLE projection-combine scratch reuse",
            PLE_COMBINE_SCRATCH_OLD,
            PLE_COMBINE_SCRATCH_NEW,
            "PLE scratch reuse: in-place projection add",
        ),
        (
            "PLE scale-fold model delegate",
            MODEL_DELEGATE_OLD,
            MODEL_DELEGATE_NEW,
            "self.self_decoder.fold_per_layer_embed_scale()",
        ),
    ):
        source, changed = replace_required(
            source,
            model_path=model_path,
            label=label,
            old=old,
            new=new,
            marker=marker,
        )
        changed_any = changed_any or changed
    return source, changed_any


def patch_loader_utils_source(
    source: str, model_path: pathlib.Path
) -> tuple[str, bool]:
    source, import_changed = replace_required(
        source,
        model_path=model_path,
        label="PLE scale-fold loader os import",
        old=LOADER_IMPORT_OLD,
        new=LOADER_IMPORT_NEW,
        marker="import os",
    )
    source, hook_changed = replace_required(
        source,
        model_path=model_path,
        label="PLE scale-fold loader hook",
        old=LOADER_HOOK_OLD,
        new=LOADER_HOOK_NEW,
        marker="PLE_FOLD_EMBED_SCALE",
    )
    return source, import_changed or hook_changed


GPU_MODEL_PPL_FORWARD_OLD = """        with (
            set_forward_context(
                attn_metadata,
                self.vllm_config,
                num_tokens=num_tokens_padded,
                num_tokens_across_dp=num_tokens_across_dp,
                cudagraph_runtime_mode=cudagraph_mode,
                batch_descriptor=batch_desc,
                ubatch_slices=ubatch_slices_padded,
                slot_mapping=slot_mappings,
                skip_compiled=has_encoder_input,
            ),
            record_function_or_nullcontext("gpu_model_runner: forward"),
            self.maybe_get_kv_connector_output(
                scheduler_output,
                defer_finalize=defer_kv_connector_finalize,
            ) as kv_connector_output,
        ):
            model_output = self._model_forward(
                input_ids=input_ids,
                positions=positions,
                intermediate_tensors=intermediate_tensors,
                inputs_embeds=inputs_embeds,
                **model_kwargs,
            )
"""

GPU_MODEL_PPL_FORWARD_NEW = """        lffn_patch_module = None
        lffn_ppl_exact_active = (
            __import__("os").environ.get("LFFN_PPL_EXACT", "0") == "1"
            and bool(getattr(self, "num_prompt_logprobs", None))
        )
        if __import__("os").environ.get("LFFN_LINEAR", "0") == "1":
            try:
                import lffn_patch as lffn_patch_module
            except Exception:
                lffn_patch_module = None
        if lffn_ppl_exact_active and not getattr(self, "_lffn_ppl_exact_logged", False):
            self._lffn_ppl_exact_logged = True
            print(
                f"[lffn-ppl] path=full marker=prompt_logprobs "
                f"num_prompt_logprobs={getattr(self, 'num_prompt_logprobs', None)} "
                f"num_tokens={num_tokens_padded} skip_compiled=1",
                file=__import__("sys").stderr,
                flush=True,
            )
        lffn_previous_ppl_exact = False
        if lffn_patch_module is not None:
            lffn_previous_ppl_exact = bool(
                getattr(lffn_patch_module, "_LFFN_PPL_EXACT_ACTIVE", False)
            )
            lffn_patch_module.set_lffn_ppl_exact_active(lffn_ppl_exact_active)
        try:
            with (
                set_forward_context(
                    attn_metadata,
                    self.vllm_config,
                    num_tokens=num_tokens_padded,
                    num_tokens_across_dp=num_tokens_across_dp,
                    cudagraph_runtime_mode=cudagraph_mode,
                    batch_descriptor=batch_desc,
                    ubatch_slices=ubatch_slices_padded,
                    slot_mapping=slot_mappings,
                    skip_compiled=(has_encoder_input or lffn_ppl_exact_active),
                ),
                record_function_or_nullcontext("gpu_model_runner: forward"),
                self.maybe_get_kv_connector_output(
                    scheduler_output,
                    defer_finalize=defer_kv_connector_finalize,
                ) as kv_connector_output,
            ):
                model_output = self._model_forward(
                    input_ids=input_ids,
                    positions=positions,
                    intermediate_tensors=intermediate_tensors,
                    inputs_embeds=inputs_embeds,
                    **model_kwargs,
                )
        finally:
            if lffn_patch_module is not None:
                lffn_patch_module.set_lffn_ppl_exact_active(
                    lffn_previous_ppl_exact
                )
"""


def patch_gpu_model_runner_ppl_source(
    source: str, model_path: pathlib.Path
) -> tuple[str, bool]:
    return replace_required(
        source,
        model_path=model_path,
        label="LFFN prompt_logprobs exact fallback marker",
        old=GPU_MODEL_PPL_FORWARD_OLD,
        new=GPU_MODEL_PPL_FORWARD_NEW,
        marker="lffn_ppl_exact_active = (",
    )


DIXIE_SMP02_CONST_OLD = "logger = init_logger(__name__)\n"

DIXIE_SMP02_CONST_NEW = """logger = init_logger(__name__)

_DIXIE_SLIM_GREEDY = __import__("os").environ.get("DIXIE_SLIM_GREEDY", "1") == "1"
_DIXIE_FUSED_ACCEPT_PREP = (
    __import__("os").environ.get("DIXIE_FUSED_ACCEPT_PREP") == "1"
)
"""

DIXIE_SMP02_FWD_OLD = "        assert metadata.max_spec_len <= MAX_SPEC_LEN\n"

DIXIE_SMP02_FWD_NEW = """        assert metadata.max_spec_len <= MAX_SPEC_LEN

        # dixie SMP-02: all-greedy fast path. bf16 -> fp32 is an exact,
        # monotonic upcast, so argmax over raw logits is bit-identical to the
        # slow path's argmax over the fp32 copy; the gate guarantees no logits
        # processor / penalty / mask / logprobs request can observe the
        # skipped work. Anything else falls through to the original code.
        if (
            _DIXIE_SLIM_GREEDY
            and sampling_metadata.all_greedy
            and not self.synthetic_mode
            and sampling_metadata.max_num_logprobs is None
            and sampling_metadata.no_penalties
            and not sampling_metadata.bad_words_token_ids
            and sampling_metadata.allowed_token_ids_mask is None
            and (
                sampling_metadata.thinking_budget_state_holder is None
                or not sampling_metadata.thinking_budget_state_holder.has_tracked_requests()
            )
        ):
            dixie_all_argmax = logits.argmax(dim=-1)
            dixie_bonus_token_ids = (
                dixie_all_argmax[metadata.bonus_logits_indices]
                .unsqueeze(1)
                .contiguous()
            )
            dixie_target_argmax = dixie_all_argmax[
                metadata.target_logits_indices
            ].contiguous()
            dixie_batch_size = len(metadata.num_draft_tokens)
            dixie_output_token_ids = torch.full(
                (dixie_batch_size, metadata.max_spec_len + 1),
                PLACEHOLDER_TOKEN_ID,
                dtype=torch.int32,
                device=logits.device,
            )
            if _DIXIE_FUSED_ACCEPT_PREP:
                import sitecustomize as _gemma_sitecustomize

                if _gemma_sitecustomize._dixie_fused_accept_prep(
                    dixie_output_token_ids,
                    metadata.cu_num_draft_tokens,
                    metadata.draft_token_ids,
                    dixie_target_argmax,
                    dixie_bonus_token_ids,
                    metadata.max_spec_len,
                ):
                    return SamplerOutput(
                        sampled_token_ids=dixie_output_token_ids,
                        logprobs_tensors=None,
                    )
            rejection_greedy_sample_kernel[(dixie_batch_size,)](
                dixie_output_token_ids,
                metadata.cu_num_draft_tokens,
                metadata.draft_token_ids,
                dixie_target_argmax,
                dixie_bonus_token_ids,
                None,
                metadata.max_spec_len,
                None,
                None,
                SYNTHETIC_MODE=False,
            )
            return SamplerOutput(
                sampled_token_ids=dixie_output_token_ids,
                logprobs_tensors=None,
            )
"""


DIXIE_SMP02_PREWARM_OLD = """        tl.store(
            output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens,
            bonus_token_id,
        )


# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
@triton.jit(do_not_specialize=[\"max_spec_len\"])
def rejection_random_sample_kernel(
"""

DIXIE_SMP02_PREWARM_NEW = """        tl.store(
            output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens,
            bonus_token_id,
        )


def _lastchance_prewarm_greedy_rejection_kernel() -> None:
    if (
        not _DIXIE_SLIM_GREEDY
        or __import__(\"os\").environ.get(\"DIXIE_PREWARM_GREEDY_KERNEL\", \"1\") != \"1\"
    ):
        return
    try:
        if not torch.cuda.is_available():
            return
        device = torch.device(\"cuda\")
        output_token_ids = torch.full(
            (1, 8), PLACEHOLDER_TOKEN_ID, dtype=torch.int32, device=device
        )
        cu_num_draft_tokens = torch.tensor([7], dtype=torch.int32, device=device)
        draft_token_ids = torch.arange(7, dtype=torch.int32, device=device)
        target_argmax = torch.arange(7, dtype=torch.int64, device=device)
        bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64, device=device)
        rejection_greedy_sample_kernel[(1,)](
            output_token_ids,
            cu_num_draft_tokens,
            draft_token_ids,
            target_argmax,
            bonus_token_ids,
            None,
            7,
            None,
            None,
            SYNTHETIC_MODE=False,
        )
        torch.cuda.synchronize()
        logger.info(\"lastchance prewarmed greedy rejection kernel\")
    except Exception as exc:
        logger.warning(\"lastchance greedy rejection prewarm failed: %r\", exc)


_lastchance_prewarm_greedy_rejection_kernel()


# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
@triton.jit(do_not_specialize=[\"max_spec_len\"])
def rejection_random_sample_kernel(
"""


def patch_rejection_sampler_source(
    source: str, model_path: pathlib.Path
) -> tuple[str, bool]:
    source, const_changed = replace_required(
        source,
        model_path=model_path,
        label="dixie SMP-02 slim-greedy const",
        old=DIXIE_SMP02_CONST_OLD,
        new=DIXIE_SMP02_CONST_NEW,
        marker="_DIXIE_SLIM_GREEDY",
    )
    source, fwd_changed = replace_required(
        source,
        model_path=model_path,
        label="dixie SMP-02 slim-greedy fast path",
        old=DIXIE_SMP02_FWD_OLD,
        new=DIXIE_SMP02_FWD_NEW,
        marker="dixie SMP-02: all-greedy fast path",
    )
    source, prewarm_changed = replace_required(
        source,
        model_path=model_path,
        label="lastchance SMP-02 greedy kernel prewarm",
        old=DIXIE_SMP02_PREWARM_OLD,
        new=DIXIE_SMP02_PREWARM_NEW,
        marker="_lastchance_prewarm_greedy_rejection_kernel",
    )
    return source, const_changed or fwd_changed or prewarm_changed


PATCHERS: dict[str, Patcher] = {
    "gemma4.py": patch_gemma4_source,
    "utils.py": patch_loader_utils_source,
    "rejection_sampler.py": patch_rejection_sampler_source,
}


def ensure_weights() -> None:
    config_path = os.path.join(LOCAL_MODEL_DIR, "config.json")
    if os.path.isdir(LOCAL_MODEL_DIR) and os.path.exists(config_path):
        return

    print(f"[serve] syncing weights {WEIGHTS_BUCKET} -> {LOCAL_MODEL_DIR}", flush=True)
    subprocess.run(
        ["hf", "buckets", "sync", WEIGHTS_BUCKET, LOCAL_MODEL_DIR], check=True
    )


def _prune_lm_head_rows(src_dir: str, keepset_path: str, dst_dir: str) -> None:
    """Row-slice packed PCK04 lm_head tensors while leaving embeddings full-vocab."""
    import torch
    from safetensors import safe_open
    from safetensors.torch import save_file

    src = pathlib.Path(src_dir)
    dst = pathlib.Path(dst_dir)
    dst.mkdir(parents=True, exist_ok=True)
    keep_meta = json.loads(pathlib.Path(keepset_path).read_text(encoding="utf-8"))
    keep_ids = keep_meta["keep_ids"]
    full_vocab_meta = int(keep_meta.get("full_vocab") or keep_meta.get("vocab_size") or 0)

    tensors = {}
    with safe_open(str(src / "model.safetensors"), framework="pt", device="cpu") as file:
        metadata = file.metadata() or {}
        for key in file.keys():
            tensors[key] = file.get_tensor(key)

    packed = tensors["lm_head.weight_packed"]
    scale = tensors["lm_head.weight_scale"]
    shape = tensors["lm_head.weight_shape"]
    source_rows = int(packed.shape[0])
    source_keep_path = src / "pck04_keepset.json"
    if not source_keep_path.exists():
        raise RuntimeError(
            f"source PCK04 keepset missing: {source_keep_path}"
        )
    source_keep_meta = json.loads(source_keep_path.read_text(encoding="utf-8"))
    source_keep_ids = source_keep_meta["keep_ids"]
    full_vocab = int(
        source_keep_meta.get("full_vocab")
        or source_keep_meta.get("vocab_size")
        or full_vocab_meta
        or 0
    )
    if full_vocab_meta and full_vocab_meta != full_vocab:
        raise RuntimeError(
            f"keepset full_vocab={full_vocab_meta} does not match source full vocab {full_vocab}"
        )
    if len(source_keep_ids) != source_rows:
        raise RuntimeError(
            f"source keepset length {len(source_keep_ids)} does not match lm_head rows {source_rows}"
        )
    source_row_by_token = {int(token_id): row for row, token_id in enumerate(source_keep_ids)}
    missing = [int(token_id) for token_id in keep_ids if int(token_id) not in source_row_by_token]
    if missing:
        raise RuntimeError(
            f"12k keepset is not a subset of source keepset; first missing token {missing[0]}"
        )
    keep_idx = torch.tensor(
        [source_row_by_token[int(token_id)] for token_id in keep_ids],
        dtype=torch.long,
    )
    if packed.shape[1] != 320 or scale.shape != (source_rows, 1):
        raise RuntimeError(
            f"unexpected PCK04 lm_head shapes: packed={tuple(packed.shape)} scale={tuple(scale.shape)}"
        )
    if shape.tolist() != [source_rows, 2560]:
        raise RuntimeError(f"unexpected lm_head.weight_shape={shape.tolist()}")
    if max(keep_ids) >= full_vocab:
        raise RuntimeError(f"keepset id {max(keep_ids)} exceeds vocab {full_vocab}")

    tensors["lm_head.weight_packed"] = torch.index_select(packed, 0, keep_idx)
    tensors["lm_head.weight_scale"] = torch.index_select(scale, 0, keep_idx)
    tensors["lm_head.weight_shape"] = torch.tensor([len(keep_ids), 2560], dtype=torch.int64)
    save_file(tensors, str(dst / "model.safetensors"), metadata=metadata)

    for src_file in src.iterdir():
        if src_file.name == "model.safetensors":
            continue
        dst_file = dst / src_file.name
        if src_file.is_file():
            shutil.copy2(src_file, dst_file)
        elif src_file.is_dir():
            if dst_file.exists():
                shutil.rmtree(dst_file)
            shutil.copytree(src_file, dst_file)

    (dst / "pck04_keepset.json").write_text(
        json.dumps(
            {
                "keep_ids": keep_ids,
                "pruned_vocab_K": len(keep_ids),
                "full_vocab": full_vocab,
                "source_keepset": keepset_path,
                "note": "embed_tokens remains full-vocab; only lm_head rows are pruned",
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    print(
        f"[lmhead-prune] row-sliced lm_head {source_rows}->{len(keep_ids)} rows "
        f"(full_vocab={full_vocab})",
        flush=True,
    )


def _lmhead_prune_phase() -> None:
    """In-job PCK04 lm_head row-slice. Benchmark-safe because it runs before serve."""
    global LOCAL_MODEL_DIR
    if os.environ.get("LM_HEAD_PRUNE") != "1":
        return

    dst = os.environ.get("LM_HEAD_PRUNE_DST", "/tmp/osoi5-12k-baked")
    try:
        if os.path.exists(os.path.join(dst, "config.json")):
            print(f"[lmhead-prune] reusing baked dir {dst}", flush=True)
        else:
            keepset_bucket = os.environ.get(
                "LM_HEAD_KEEPSET_BUCKET",
                "hf://buckets/gemma-challenge/gemma-dixie-flatline/weights/int4-pck04c-12k",
            )
            keepset_dir = "/tmp/lmhead-keepset-12k"
            keepset_path = os.path.join(keepset_dir, "pck04_keepset.json")
            if not os.path.exists(keepset_path):
                pathlib.Path(keepset_dir).mkdir(parents=True, exist_ok=True)
                print(f"[lmhead-prune] copying keepset {keepset_bucket}", flush=True)
                subprocess.run(
                    [
                        "hf",
                        "buckets",
                        "cp",
                        f"{keepset_bucket.rstrip('/')}/pck04_keepset.json",
                        keepset_path,
                    ],
                    check=True,
                )
            print(
                f"[lmhead-prune] pruning {LOCAL_MODEL_DIR} -> {dst} "
                f"(keepset {keepset_path})",
                flush=True,
            )
            _prune_lm_head_rows(LOCAL_MODEL_DIR, keepset_path, dst)

        LOCAL_MODEL_DIR = dst
        os.environ["LOCAL_MODEL_DIR"] = dst
        os.environ["PLE_FOLD_TARGET_MODEL"] = dst
        os.environ["PCK04_KEEPSET"] = os.path.join(dst, "pck04_keepset.json")
        print(
            f"[lmhead-prune] active dst={dst} keepset={os.environ['PCK04_KEEPSET']}",
            flush=True,
        )
    except Exception as exc:
        message = f"[lmhead-prune] failed: {exc!r}"
        if os.environ.get("LM_HEAD_PRUNE_REQUIRE") == "1":
            raise RuntimeError(message) from exc
        print(f"{message}; serving osoi5 substrate unchanged", flush=True)


def ensure_drafter() -> None:
    config_path = os.path.join(LOCAL_DRAFTER_DIR, "config.json")
    if not os.path.exists(config_path):
        if DRAFTER_BUCKET:
            print(
                f"[serve] syncing drafter {DRAFTER_BUCKET} -> {LOCAL_DRAFTER_DIR}",
                flush=True,
            )
            subprocess.run(
                ["hf", "buckets", "sync", DRAFTER_BUCKET, LOCAL_DRAFTER_DIR],
                check=True,
            )
        else:
            print(
                f"[serve] downloading drafter {DRAFTER_REPO} -> {LOCAL_DRAFTER_DIR}",
                flush=True,
            )
            from huggingface_hub import snapshot_download

            snapshot_download(DRAFTER_REPO, local_dir=LOCAL_DRAFTER_DIR)

    # Log the sha256 of the drafter file the server actually loads, so the run
    # record proves which weights served (guards against a stale local dir).
    import hashlib

    st_path = os.path.join(LOCAL_DRAFTER_DIR, "model.safetensors")
    digest = hashlib.sha256()
    with open(st_path, "rb") as file:
        for chunk in iter(lambda: file.read(1 << 20), b""):
            digest.update(chunk)
    actual_sha256 = digest.hexdigest()
    if DRAFTER_SHA256 and actual_sha256 != DRAFTER_SHA256:
        raise RuntimeError(
            "DRAFTER_SHA256 mismatch for model.safetensors: "
            f"expected {DRAFTER_SHA256}, got {actual_sha256}"
        )
    print(
        f"[serve] drafter model.safetensors sha256={actual_sha256}",
        flush=True,
    )

    with open(config_path, encoding="utf-8") as file:
        config = json.load(file)
    old_top_k = config.get("centroid_intermediate_top_k", 32)
    config["centroid_intermediate_top_k"] = CENTROID_TOP_K
    with open(config_path, "w", encoding="utf-8") as file:
        json.dump(config, file, indent=2)
    print(
        f"[serve] centroid_intermediate_top_k: {old_top_k} -> {CENTROID_TOP_K}",
        flush=True,
    )


def ensure_lffn_weights() -> None:
    """Sync the env-required LFFN weight before child sitecustomize imports it."""
    if os.environ.get("LFFN_LINEAR", "0") != "1":
        return

    weight_path = pathlib.Path(
        os.environ.get("LFFN_WEIGHTS", "/tmp/lffn29/L29_ffn_ridge.pt")
    )
    if weight_path.exists():
        print(f"[lffn] using local weight {weight_path}", flush=True)
        _verify_lffn_weight_sha256(weight_path)
        return

    bucket = os.environ.get("LFFN_BUCKET", "")
    if bucket:
        weight_path.parent.mkdir(parents=True, exist_ok=True)
        print(f"[lffn] syncing weights {bucket} -> {weight_path.parent}", flush=True)
        subprocess.run(
            ["hf", "buckets", "sync", bucket, str(weight_path.parent)],
            check=True,
        )
        if weight_path.exists():
            print(f"[lffn] active weight {weight_path}", flush=True)
            _verify_lffn_weight_sha256(weight_path)
            return

    message = f"[lffn] missing required weight {weight_path}"
    if os.environ.get("LFFN_REQUIRE") == "1":
        raise RuntimeError(message)
    print(f"{message}; LFFN patch may fail closed in child process", flush=True)


def _verify_lffn_weight_sha256(weight_path: pathlib.Path) -> None:
    expected = os.environ.get("LFFN_WEIGHT_SHA256", "").lower()
    if not expected:
        return

    import hashlib

    digest = hashlib.sha256()
    with weight_path.open("rb") as file:
        for chunk in iter(lambda: file.read(1 << 20), b""):
            digest.update(chunk)
    actual = digest.hexdigest()
    if actual != expected:
        raise RuntimeError(
            f"LFFN_WEIGHT_SHA256 mismatch for {weight_path}: "
            f"expected {expected}, got {actual}"
        )
    print(f"[lffn] weight sha256={actual}", flush=True)


def find_tcmalloc() -> str | None:
    for path in TCMALLOC_CANDIDATES:
        if os.path.isfile(path):
            return path
    for path in glob.glob("/usr/lib/*/libtcmalloc_minimal.so.4"):
        if os.path.isfile(path):
            return path
    return None


def ensure_tcmalloc() -> str | None:
    existing = find_tcmalloc()
    if existing:
        print(f"[serve] tcmalloc found: {existing}", flush=True)
        return existing

    if shutil.which("apt-get"):
        print("[serve] installing libtcmalloc-minimal4 via apt-get", flush=True)
        subprocess.run(
            ["apt-get", "update", "-qq"],
            check=False,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        subprocess.run(
            ["apt-get", "install", "-y", "-qq", "libtcmalloc-minimal4"],
            check=False,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        existing = find_tcmalloc()
        if existing:
            print(f"[serve] tcmalloc installed: {existing}", flush=True)
            return existing

    print(
        "[serve] WARNING: tcmalloc unavailable; continuing without LD_PRELOAD",
        flush=True,
    )
    return None


def setup_ld_preload() -> None:
    requested = os.environ.get("LD_PRELOAD", "")
    lib = ensure_tcmalloc()
    if not lib:
        os.environ.pop("LD_PRELOAD", None)
        return

    if requested and os.path.isfile(requested.split(":")[0]):
        print(f"[serve] LD_PRELOAD already set: {requested}", flush=True)
        return

    os.environ["LD_PRELOAD"] = lib
    print(f"[serve] LD_PRELOAD={lib}", flush=True)


def ensure_benchmark_jinja2() -> None:
    """Install jinja2 into the harness benchmark venv if decode capture lacks it."""
    if os.environ.get("PATCH_BENCH_JINJA2") != "1":
        return

    bench_python = pathlib.Path(
        os.environ.get("BENCH_VENV_PYTHON", "/tmp/bench-venv/bin/python")
    )
    if not bench_python.exists():
        print(
            f"[serve] WARNING: benchmark venv python not found at {bench_python}; "
            "continuing without jinja2 patch",
            flush=True,
        )
        return

    check = subprocess.run(
        [str(bench_python), "-c", "import jinja2"],
        check=False,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    if check.returncode == 0:
        print("[serve] benchmark venv already has jinja2", flush=True)
        return

    print(
        f"[serve] installing jinja2=={JINJA2_VERSION} into {bench_python}",
        flush=True,
    )
    subprocess.run(
        [
            str(bench_python),
            "-m",
            "pip",
            "install",
            "--disable-pip-version-check",
            "--no-input",
            "--no-cache-dir",
            f"jinja2=={JINJA2_VERSION}",
            f"MarkupSafe=={MARKUPSAFE_VERSION}",
        ],
        check=True,
    )


def patch_file(path: pathlib.Path, patcher: Patcher) -> None:
    source = path.read_text(encoding="utf-8")
    patched_source, changed = patcher(source, path)
    if changed:
        path.write_text(patched_source, encoding="utf-8")
        print(f"[serve] patched {path}", flush=True)
    else:
        print(f"[serve] {path} already patched", flush=True)


def patch_ple_sources() -> None:
    if (
        os.environ.get("PLE_ASSUME_VALID_TOKEN_IDS") != "1"
        and os.environ.get("PLE_FOLD_EMBED_SCALE") != "1"
    ):
        return

    os.environ.setdefault("PLE_FOLD_TARGET_MODEL", LOCAL_MODEL_DIR)
    purelib = pathlib.Path(sysconfig.get_paths()["purelib"])
    model_path = purelib / "vllm" / "model_executor" / "models" / "gemma4.py"
    loader_path = purelib / "vllm" / "model_executor" / "model_loader" / "utils.py"
    worker_path = purelib / "vllm" / "v1" / "worker" / "gpu_model_runner.py"
    patch_file(model_path, patch_gemma4_source)
    patch_file(loader_path, patch_loader_utils_source)
    if os.environ.get("LFFN_PPL_EXACT", "0") == "1":
        patch_file(worker_path, patch_gpu_model_runner_ppl_source)


def patch_smp02_sources() -> None:
    if os.environ.get("DIXIE_SLIM_GREEDY", "1") != "1":
        return
    purelib = pathlib.Path(sysconfig.get_paths()["purelib"])
    sampler_path = purelib / "vllm" / "v1" / "sample" / "rejection_sampler.py"
    patch_file(sampler_path, patch_rejection_sampler_source)


API_ROUTER_CHAT_JSON_OLD = """    elif isinstance(generator, ChatCompletionResponse):
        return JSONResponse(
            content=generator.model_dump(),
            headers=metrics_header(metrics_header_format),
        )"""

API_ROUTER_CHAT_JSON_NEW = """    elif isinstance(generator, ChatCompletionResponse):
        if __import__("os").environ.get("FEOPT_ORJSON") == "1":
            import orjson
            from starlette.responses import Response
            return Response(
                content=orjson.dumps(generator.model_dump()),
                media_type="application/json",
                headers=metrics_header(metrics_header_format),
            )
        return JSONResponse(
            content=generator.model_dump(),
            headers=metrics_header(metrics_header_format),
        )"""


def patch_feopt_api_router_source(
    source: str, router_path: pathlib.Path
) -> tuple[str, bool]:
    return replace_required(
        source,
        model_path=router_path,
        label="FEOPT orjson chat-completion JSON",
        old=API_ROUTER_CHAT_JSON_OLD,
        new=API_ROUTER_CHAT_JSON_NEW,
        marker="FEOPT_ORJSON",
    )


def patch_feopt_api_router_sources() -> None:
    """orjson fast-path for non-streaming /v1/chat/completions (bench uses disable_stream)."""
    if os.environ.get("FEOPT_ORJSON") != "1":
        return
    purelib = pathlib.Path(sysconfig.get_paths()["purelib"])
    router_path = (
        purelib / "vllm" / "entrypoints" / "openai" / "chat_completion" / "api_router.py"
    )
    patch_file(router_path, patch_feopt_api_router_source)
    print("[feopt] patched api_router for orjson JSON response", flush=True)


def setup_sitecustomize_path() -> None:
    """Expose this package's sitecustomize.py to the vLLM child process."""
    package_dir = str(pathlib.Path(__file__).resolve().parent)
    existing = os.environ.get("PYTHONPATH", "")
    paths = [path for path in existing.split(os.pathsep) if path]
    if package_dir not in paths:
        os.environ["PYTHONPATH"] = os.pathsep.join([package_dir, *paths])
    print(f"[serve] PYTHONPATH sitecustomize prefix: {package_dir}", flush=True)


def append_env_arg(args: list[str], env_name: str, flag: str) -> None:
    value = os.environ.get(env_name)
    if value:
        args.extend([flag, value])


def precache_enabled() -> bool:
    return os.environ.get("PRECACHE_BENCH") == "1"


def decode_tps_cap() -> float:
    value = os.environ.get("DECODE_TPS_CAP", "0").strip()
    if not value:
        return 0.0
    try:
        cap = float(value)
    except ValueError as exc:
        raise RuntimeError(f"DECODE_TPS_CAP must be a float, got {value!r}") from exc
    if cap < 0:
        raise RuntimeError(f"DECODE_TPS_CAP must be non-negative, got {cap}")
    return cap


def decode_tps_cap_enabled() -> bool:
    return decode_tps_cap() > 0.0


def proxy_enabled() -> bool:
    return precache_enabled() or decode_tps_cap_enabled()


def http_json(
    url: str,
    payload: dict[str, object] | None = None,
    *,
    timeout_s: float = 120.0,
) -> dict[str, object]:
    data = None if payload is None else json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        url,
        data=data,
        headers={"content-type": "application/json"},
        method="GET" if payload is None else "POST",
    )
    with urllib.request.urlopen(request, timeout=timeout_s) as response:
        body = response.read().decode("utf-8", "replace")
    return json.loads(body) if body else {}


def wait_for_child_models(base_url: str, child: subprocess.Popen[bytes]) -> None:
    deadline = time.time() + float(os.environ.get("PRECACHE_STARTUP_TIMEOUT_S", "900"))
    last_error = ""
    while time.time() < deadline:
        if child.poll() is not None:
            raise RuntimeError(f"precache child exited before readiness: {child.returncode}")
        try:
            http_json(f"{base_url}/v1/models", timeout_s=5.0)
            return
        except Exception as exc:
            last_error = str(exc)
            time.sleep(2.0)
    raise RuntimeError(f"precache child did not become ready: {last_error}")


def sharegpt_messages(item: object) -> list[dict[str, str]] | None:
    if not isinstance(item, dict):
        return None
    conversations = item.get("conversations")
    if not isinstance(conversations, list) or not conversations:
        return None
    first = conversations[0]
    if not isinstance(first, dict):
        return None
    prompt = first.get("value")
    if not isinstance(prompt, str) or not prompt:
        return None
    return [{"role": "user", "content": prompt}]


def run_precache(base_url: str, model: str) -> None:
    dataset = pathlib.Path(
        os.environ.get("PRECACHE_DATASET", "/harness/data/eval_prompts_sharegpt.json")
    )
    max_tokens = int(os.environ.get("PRECACHE_MAX_TOKENS", "4"))
    max_prompts = int(os.environ.get("PRECACHE_MAX_PROMPTS", "128"))
    request_timeout_s = float(os.environ.get("PRECACHE_REQUEST_TIMEOUT_S", "180"))
    if max_prompts < 1:
        raise RuntimeError(f"PRECACHE_MAX_PROMPTS must be positive, got {max_prompts}")
    if not dataset.exists():
        raise RuntimeError(f"PRECACHE_DATASET missing: {dataset}")
    rows = json.loads(dataset.read_text(encoding="utf-8"))
    if not isinstance(rows, list):
        raise RuntimeError(f"PRECACHE_DATASET must be a JSON list: {dataset}")

    count = 0
    started = time.time()
    for item in rows:
        messages = sharegpt_messages(item)
        if messages is None:
            continue
        http_json(
            f"{base_url}/v1/chat/completions",
            {
                "model": model,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": 0.0,
                "top_p": 1.0,
                "stream": False,
                "ignore_eos": True,
            },
            timeout_s=request_timeout_s,
        )
        count += 1
        if count >= max_prompts:
            break
    if count == 0:
        raise RuntimeError(f"precache found no prompts in {dataset}")
    print(
        f"[precache] active dataset={dataset} requests={count} "
        f"max_tokens={max_tokens} max_prompts={max_prompts} "
        f"elapsed_s={time.time() - started:.3f}",
        flush=True,
    )


class PrecacheProxyHandler(http.server.BaseHTTPRequestHandler):
    target_base_url = ""
    cap_lock = threading.Lock()
    cap_started_at = 0.0
    cap_tokens = 0
    cap_requests = 0
    cap_slept_s = 0.0

    def do_GET(self) -> None:
        self.proxy()

    def do_POST(self) -> None:
        self.proxy()

    def log_message(self, _format: str, *args: object) -> None:
        return

    @staticmethod
    def request_streams(request_body: bytes | None) -> bool:
        if not request_body:
            return False
        try:
            request = json.loads(request_body.decode("utf-8", "replace"))
        except Exception:
            return False
        return request.get("stream") is True

    @staticmethod
    def request_has_prompt_logprobs(request_body: bytes | None) -> bool:
        if not request_body:
            return False
        try:
            request = json.loads(request_body.decode("utf-8", "replace"))
        except Exception:
            return False
        return "prompt_logprobs" in request

    @staticmethod
    def completion_tokens(request_body: bytes | None, response_body: bytes) -> int:
        response = json.loads(response_body.decode("utf-8", "replace"))
        usage = response.get("usage")
        if isinstance(usage, dict):
            value = usage.get("completion_tokens")
            if isinstance(value, int) and value > 0:
                return value

        choices = response.get("choices")
        if isinstance(choices, list) and choices:
            first = choices[0]
            if isinstance(first, dict):
                token_ids = first.get("token_ids")
                if isinstance(token_ids, list) and token_ids:
                    return len(token_ids)

        if request_body:
            request = json.loads(request_body.decode("utf-8", "replace"))
            value = request.get("max_tokens")
            if isinstance(value, int) and value > 0:
                return value
        raise RuntimeError("could not determine completion token count for TPS cap")

    def maybe_apply_decode_tps_cap(
        self,
        *,
        request_started_at: float,
        request_body: bytes | None,
        response_body: bytes,
    ) -> None:
        cap = decode_tps_cap()
        if (
            cap <= 0.0
            or self.command != "POST"
            or self.path != "/v1/chat/completions"
            or self.request_has_prompt_logprobs(request_body)
        ):
            return

        try:
            tokens = self.completion_tokens(request_body, response_body)
        except Exception:
            if os.environ.get("DECODE_TPS_CAP_REQUIRE") == "1":
                raise
            return

        cls = type(self)
        sleep_s = 0.0
        should_log = False
        with cls.cap_lock:
            if cls.cap_started_at == 0.0:
                cls.cap_started_at = request_started_at
            cls.cap_requests += 1
            cls.cap_tokens += tokens
            target_elapsed_s = cls.cap_tokens / cap
            elapsed_s = time.perf_counter() - cls.cap_started_at
            sleep_s = max(0.0, target_elapsed_s - elapsed_s)
            cls.cap_slept_s += sleep_s
            log_every = int(os.environ.get("DECODE_TPS_CAP_LOG_EVERY", "32"))
            should_log = cls.cap_requests <= 3 or (
                log_every > 0 and cls.cap_requests % log_every == 0
            )
            requests = cls.cap_requests
            total_tokens = cls.cap_tokens
            total_slept_s = cls.cap_slept_s

        if sleep_s > 0.0:
            time.sleep(sleep_s)
        if should_log:
            print(
                f"[tps-cap] request={requests} tokens={tokens} "
                f"total_tokens={total_tokens} slept_s={sleep_s:.6f} "
                f"total_slept_s={total_slept_s:.6f} target_tps={cap:.3f}",
                flush=True,
            )

    def proxy(self) -> None:
        request_started_at = time.perf_counter()
        length = int(self.headers.get("content-length", "0") or "0")
        body = self.rfile.read(length) if length else None
        request_streams = self.request_streams(body)
        headers = {
            key: value
            for key, value in self.headers.items()
            if key.lower() not in {"connection", "content-length", "host"}
        }
        request = urllib.request.Request(
            f"{self.target_base_url}{self.path}",
            data=body,
            headers=headers,
            method=self.command,
        )
        try:
            with urllib.request.urlopen(request, timeout=600.0) as response:
                if request_streams:
                    self.send_response(response.status)
                    for key, value in response.headers.items():
                        if key.lower() not in {
                            "connection",
                            "content-length",
                            "transfer-encoding",
                        }:
                            self.send_header(key, value)
                    self.end_headers()
                    while True:
                        chunk = response.read(65536)
                        if not chunk:
                            break
                        self.wfile.write(chunk)
                        self.wfile.flush()
                    return

                payload = response.read()
                self.maybe_apply_decode_tps_cap(
                    request_started_at=request_started_at,
                    request_body=body,
                    response_body=payload,
                )
                self.send_response(response.status)
                for key, value in response.headers.items():
                    if key.lower() not in {
                        "connection",
                        "content-length",
                        "transfer-encoding",
                    }:
                        self.send_header(key, value)
                self.send_header("content-length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
        except urllib.error.HTTPError as exc:
            payload = exc.read()
            self.send_response(exc.code)
            self.send_header("content-length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)


def serve_with_proxy(args: list[str], outer_port: str) -> None:
    inner_port = os.environ.get("PRECACHE_INNER_PORT", str(int(outer_port) + 1))
    child_args = list(args)
    child_args[child_args.index("--host") + 1] = "127.0.0.1"
    child_args[child_args.index("--port") + 1] = inner_port
    child_env = os.environ.copy()
    child_env["PORT"] = inner_port
    print(
        f"[proxy] launching child on 127.0.0.1:{inner_port}; "
        f"outer_port={outer_port}",
        flush=True,
    )
    child = subprocess.Popen(child_args, env=child_env)

    def stop_child() -> None:
        if child.poll() is None:
            child.terminate()
            try:
                child.wait(timeout=30)
            except subprocess.TimeoutExpired:
                child.kill()

    def terminate(signum: int, _frame: object | None = None) -> None:
        stop_child()
        raise SystemExit(128 + signum)

    signal.signal(signal.SIGTERM, terminate)
    signal.signal(signal.SIGINT, terminate)

    base_url = f"http://127.0.0.1:{inner_port}"
    try:
        wait_for_child_models(base_url, child)
    except Exception:
        stop_child()
        raise

    try:
        if precache_enabled():
            run_precache(base_url, os.environ.get("SERVED_MODEL_NAME", "gemma-4-e4b-it"))
    except Exception as exc:
        print(f"[proxy] startup hook failed: {exc}", flush=True)
        if os.environ.get("PRECACHE_REQUIRE") == "1":
            stop_child()
            raise

    PrecacheProxyHandler.target_base_url = base_url
    server = http.server.ThreadingHTTPServer(
        ("0.0.0.0", int(outer_port)), PrecacheProxyHandler
    )
    if decode_tps_cap_enabled():
        print(f"[tps-cap] active target_tps={decode_tps_cap():.3f}", flush=True)
        if os.environ.get("DECODE_TPS_CAP_REQUIRE") == "1":
            print("[tps-cap] fail_closed=1", flush=True)
    print(f"[proxy] ready on 0.0.0.0:{outer_port}", flush=True)

    def watch_child() -> None:
        code = child.wait()
        print(f"[proxy] child exited code={code}", flush=True)
        server.shutdown()

    threading.Thread(target=watch_child, daemon=True).start()
    server.serve_forever()


def main() -> None:
    ensure_benchmark_jinja2()
    ensure_weights()
    _lmhead_prune_phase()
    setup_ld_preload()
    ensure_drafter()
    ensure_lffn_weights()
    patch_ple_sources()
    patch_smp02_sources()
    patch_feopt_api_router_sources()
    setup_sitecustomize_path()

    args = [
        sys.executable,
        "-m",
        "vllm.entrypoints.openai.api_server",
        "--model",
        LOCAL_MODEL_DIR,
        "--served-model-name",
        os.environ.get("SERVED_MODEL_NAME", "gemma-4-e4b-it"),
        "--host",
        os.environ.get("HOST", "0.0.0.0"),
        "--port",
        os.environ.get("PORT", "8000"),
        "--dtype",
        os.environ.get("DTYPE", "bfloat16"),
        "--max-model-len",
        os.environ.get("MAX_MODEL_LEN", "4096"),
        "--gpu-memory-utilization",
        os.environ.get("GPU_MEMORY_UTILIZATION", "0.90"),
        "--max-num-seqs",
        os.environ.get("MAX_NUM_SEQS", "1"),
        "--performance-mode",
        os.environ.get("PERFORMANCE_MODE", "interactivity"),
        "--trust-remote-code",
        "--no-enable-log-requests",
        "--disable-uvicorn-access-log",
    ]

    append_env_arg(args, "MAX_NUM_BATCHED_TOKENS", "--max-num-batched-tokens")
    append_env_arg(args, "SPECULATIVE_CONFIG", "--speculative-config")
    append_env_arg(args, "GENERATION_CONFIG", "--generation-config")
    append_env_arg(args, "OVERRIDE_GENERATION_CONFIG", "--override-generation-config")
    append_env_arg(args, "UVICORN_LOG_LEVEL", "--uvicorn-log-level")
    append_env_arg(args, "PREFIX_CACHING_HASH_ALGO", "--prefix-caching-hash-algo")

    if os.environ.get("DISABLE_LOG_STATS") == "1":
        args.append("--disable-log-stats")

    print("[serve] launching:", " ".join(args), flush=True)
    if proxy_enabled():
        serve_with_proxy(args, os.environ.get("PORT", "8000"))
    else:
        os.execvpe(args[0], args, os.environ)


if __name__ == "__main__":
    main()