serving/receipt_bootstrap.py

#!/usr/bin/env python3
# Modification notice: DevelopingDad, 2026-10-08. Experimental isolated FP8 draft-head integration and/or byte-preserving per-boot diagnostic archival.
# SPDX-License-Identifier: Apache-2.0. Upstream copyright and SPDX notices retained below.
"""Archive a prior draft-head startup receipt, then exec the image's vLLM serve.

This small entrypoint is mounted read-only into the persistent container.  The
receipt itself lives in the writable cache, so an automatic Docker restart can
start from the same cache without making the overlay's write-once receipt
writer fail.  Archive files are written with exclusive creation and the
original bytes are preserved before the current path is removed.
"""

from __future__ import annotations

import datetime as _dt
import errno
import hashlib
import os
from pathlib import Path
import stat
import sys


DEFAULT_RECEIPT = "/cache/draft-head-fp8-receipt.json"
ARCHIVE_PREFIX = "draft-head-fp8-receipt.boot-"


class BootstrapError(RuntimeError):
    """The boot receipt could not be archived safely."""


def _receipt_path() -> Path:
    value = os.environ.get("VLLM_DRAFT_HEAD_FP8_RECEIPT", DEFAULT_RECEIPT)
    path = Path(value)
    if not path.is_absolute() or path.name != Path(DEFAULT_RECEIPT).name:
        raise BootstrapError("VLLM_DRAFT_HEAD_FP8_RECEIPT must be the absolute draft-head receipt path")
    return path


def _read_regular(path: Path) -> tuple[bytes, os.stat_result]:
    flags = os.O_RDONLY
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    try:
        fd = os.open(path, flags)
    except OSError as exc:
        if exc.errno == errno.ENOENT:
            raise FileNotFoundError(path) from exc
        raise BootstrapError(f"cannot open receipt safely: {path}: {exc}") from exc
    try:
        first = os.fstat(fd)
        if not stat.S_ISREG(first.st_mode):
            raise BootstrapError(f"receipt is not a regular file: {path}")
        chunks: list[bytes] = []
        while True:
            chunk = os.read(fd, 1024 * 1024)
            if not chunk:
                break
            chunks.append(chunk)
        return b"".join(chunks), first
    finally:
        os.close(fd)


def _exclusive_write(path: Path, data: bytes, mode: int = 0o600) -> None:
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
    fd = os.open(path, flags, mode)
    try:
        view = memoryview(data)
        while view:
            count = os.write(fd, view)
            if count <= 0:
                raise BootstrapError(f"short write while creating archive: {path}")
            view = view[count:]
        os.fsync(fd)
    finally:
        os.close(fd)


def _archive_name(parent: Path, digest: str) -> tuple[Path, Path]:
    stamp = _dt.datetime.now(_dt.timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")
    base = f"{ARCHIVE_PREFIX}{stamp}-{os.getpid()}-{digest[:16]}"
    for suffix in range(1000):
        extra = "" if suffix == 0 else f"-{suffix}"
        archive = parent / f"{base}{extra}.json"
        sidecar = parent / f"{base}{extra}.json.sha256"
        if not archive.exists() and not sidecar.exists():
            return archive, sidecar
    raise BootstrapError("could not allocate a unique receipt archive name")


def archive_existing_receipt(receipt: Path | None = None) -> dict[str, object]:
    """Archive and remove one stale receipt, or report an absent receipt.

    The archive is an exact byte copy and its sidecar records the SHA-256.  A
    stale receipt is removed only after both files are durably written.  Any
    error leaves the current receipt in place and raises, so the vLLM process
    cannot start with an unreviewed receipt collision.
    """
    path = receipt or _receipt_path()
    try:
        payload, before = _read_regular(path)
    except FileNotFoundError:
        return {"present": False, "receipt": str(path), "archived": False}
    digest = hashlib.sha256(payload).hexdigest()
    try:
        after = path.stat()
    except OSError as exc:
        raise BootstrapError(f"receipt disappeared before archive: {path}") from exc
    if (before.st_ino, before.st_size, before.st_mtime_ns) != (after.st_ino, after.st_size, after.st_mtime_ns):
        raise BootstrapError("receipt changed while it was being archived")
    archive, sidecar = _archive_name(path.parent, digest)
    try:
        _exclusive_write(archive, payload)
        _exclusive_write(sidecar, f"{digest}  {archive.name}\n".encode("ascii"))
        path.unlink()
        try:
            fd = os.open(path.parent, os.O_RDONLY)
            try:
                os.fsync(fd)
            finally:
                os.close(fd)
        except OSError:
            # Directory fsync is useful but not available on every mounted
            # filesystem.  The archive and receipt removal already succeeded.
            pass
    except OSError as exc:
        raise BootstrapError(f"could not archive stale receipt: {exc}") from exc
    return {"present": True, "receipt": str(path), "archived": True,
            "archive": str(archive), "sidecar": str(sidecar), "sha256": digest,
            "bytes": len(payload)}


def vllm_argv(original_args: list[str]) -> list[str]:
    """Return the exact image entrypoint plus its original command arguments."""
    if not original_args:
        raise BootstrapError("original vLLM serve arguments are missing")
    return ["vllm", "serve", *original_args]


def main(argv: list[str] | None = None) -> int:
    original_args = list(sys.argv[1:] if argv is None else argv)
    try:
        archive_existing_receipt()
        os.execvp(vllm_argv(original_args)[0], vllm_argv(original_args))
    except BootstrapError as exc:
        print(f"receipt bootstrap refused to start vLLM: {exc}", file=sys.stderr)
        return 78
    except OSError as exc:
        print(f"receipt bootstrap could not exec vLLM: {exc}", file=sys.stderr)
        return 127
    return 0


if __name__ == "__main__":
    raise SystemExit(main())