| |
| |
| |
| """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: |
| |
| |
| 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()) |
|
|
