runtime_bootstrap.py

"""Install the audited, repository-local Qwen runtime into an isolated /tmp path."""

from __future__ import annotations

import hashlib
import importlib
import os
from pathlib import Path
import subprocess
import sys


REQUIRED_VERSIONS = {
    "accelerate": "1.14.0",
    "bitsandbytes": "0.49.2",
    "huggingface-hub": "1.19.0",
    "peft": "0.19.1",
    "safetensors": "0.8.0",
    "tokenizers": "0.22.2",
    "transformers": "5.3.0",
    "typer": "0.25.1",
}
PYTHON310_REQUIRED_VERSIONS = {"regex": "2026.7.19"}


def _read_manifest(root: Path) -> tuple[str, list[tuple[str, Path]]]:
    manifest = root / "wheel_manifest.sha256"
    raw = manifest.read_bytes()
    entries: list[tuple[str, Path]] = []
    for line in raw.decode("utf-8").splitlines():
        if not line.strip():
            continue
        digest, relative = line.split(None, 1)
        wheel = root / relative.strip().lstrip("*")
        if wheel.parent != root / "wheels":
            raise RuntimeError(f"invalid wheel manifest path: {relative}")
        entries.append((digest.lower(), wheel))
    if not entries:
        raise RuntimeError("wheel manifest is empty")
    return hashlib.sha256(raw).hexdigest(), entries


def _verify_wheels(entries: list[tuple[str, Path]]) -> list[Path]:
    wheels: list[Path] = []
    for expected, wheel in entries:
        if not wheel.is_file():
            raise RuntimeError(f"bundled wheel is missing: {wheel.name}")
        actual = hashlib.sha256(wheel.read_bytes()).hexdigest()
        if actual != expected:
            raise RuntimeError(f"bundled wheel hash mismatch: {wheel.name}")
        wheels.append(wheel)
    return wheels


def _compatible_wheels(wheels: list[Path]) -> tuple[list[Path], list[Path]]:
    """Select wheels accepted by the active interpreter and platform.

    The competition uses CPython 3.10 and therefore installs every bundled
    wheel. Colab currently uses a newer Python, so its smoke test skips only
    interpreter-specific wheels (currently the CPython-3.10 regex build) and
    uses the already installed equivalent dependency.
    """

    from packaging.tags import sys_tags
    from packaging.utils import parse_wheel_filename

    supported = set(sys_tags())
    compatible: list[Path] = []
    skipped: list[Path] = []
    for wheel in wheels:
        _, _, _, tags = parse_wheel_filename(wheel.name)
        if supported.intersection(tags):
            compatible.append(wheel)
        else:
            skipped.append(wheel)
    return compatible, skipped


def bootstrap_local_runtime(root: Path) -> Path:
    """Verify and install wheels before importing Transformers or PEFT."""

    manifest_digest, entries = _read_manifest(root)
    verified = _verify_wheels(entries)
    wheels, skipped = _compatible_wheels(verified)
    if not wheels:
        raise RuntimeError("none of the bundled runtime wheels match this interpreter")
    interpreter = sys.implementation.cache_tag or f"py{sys.version_info.major}{sys.version_info.minor}"
    vendor = Path("/tmp") / f"iol_ai_q35_runtime_{manifest_digest[:12]}_{interpreter}"
    marker = vendor / ".complete"
    if not marker.is_file() or marker.read_text(encoding="utf-8").strip() != manifest_digest:
        vendor.mkdir(parents=True, exist_ok=True)
        command = [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--quiet",
            "--disable-pip-version-check",
            "--no-index",
            "--no-deps",
            "--target",
            str(vendor),
            *[str(wheel) for wheel in wheels],
        ]
        subprocess.run(command, check=True, timeout=180)
        marker.write_text(manifest_digest + "\n", encoding="utf-8")
    sys.path.insert(0, str(vendor))
    if skipped:
        print(
            "runtime bootstrap skipped incompatible smoke-test wheel(s): "
            + ", ".join(wheel.name for wheel in skipped),
            flush=True,
        )
    return vendor


def patch_torch24_for_single_gpu() -> list[str]:
    """Bridge APIs added after torch 2.4 but used during local model loading.

    These bridges affect module replacement and import-time type references
    only. They do not alter tensor kernels or distributed execution.
    """

    import torch

    applied: list[str] = []
    if not hasattr(torch.nn.Module, "set_submodule"):

        def set_submodule(module_self, target: str, module) -> None:
            if not target or target.startswith(".") or target.endswith("."):
                raise ValueError(f"invalid submodule target: {target!r}")
            parts = target.split(".")
            parent = module_self
            for part in parts[:-1]:
                parent = getattr(parent, part)
                if not isinstance(parent, torch.nn.Module):
                    raise AttributeError(f"{part!r} does not resolve to a module")
            setattr(parent, parts[-1], module)

        torch.nn.Module.set_submodule = set_submodule
        applied.append("nn.Module.set_submodule")

    try:
        importlib.import_module("torch.distributed.tensor")
    except ImportError:
        try:
            legacy = importlib.import_module("torch.distributed._tensor")
        except ImportError:
            legacy = None
        if legacy is not None:
            sys.modules["torch.distributed.tensor"] = legacy
            for child in ("_utils", "placement_types"):
                try:
                    sys.modules[f"torch.distributed.tensor.{child}"] = importlib.import_module(
                        f"torch.distributed._tensor.{child}"
                    )
                except ImportError:
                    pass
            applied.append("torch.distributed.tensor")
    return applied


def assert_runtime_versions(vendor: Path) -> dict[str, str]:
    from importlib.metadata import PackageNotFoundError, distributions, version
    from packaging.utils import canonicalize_name

    discovered: dict[str, str] = {}
    for distribution in distributions(path=[str(vendor)]):
        name = canonicalize_name(distribution.metadata["Name"])
        discovered[name] = distribution.version
    required = {
        canonicalize_name(package): expected
        for package, expected in REQUIRED_VERSIONS.items()
    }
    if sys.version_info[:2] == (3, 10):
        required.update(
            {
                canonicalize_name(package): expected
                for package, expected in PYTHON310_REQUIRED_VERSIONS.items()
            }
        )
    for package, expected in required.items():
        actual = discovered.get(package)
        if actual is None:
            raise PackageNotFoundError(package)
        if actual != expected:
            raise RuntimeError(f"{package}=={actual}; expected {expected}")
    if sys.version_info[:2] != (3, 10):
        # Colab's newer interpreter cannot load the bundled CPython-3.10 regex
        # binary. Confirm that its base environment provides the dependency.
        discovered["regex"] = version("regex")
    return {package: discovered[package] for package in sorted(discovered)}


def configure_offline_environment() -> None:
    os.environ["HF_HUB_OFFLINE"] = "1"
    os.environ["TRANSFORMERS_OFFLINE"] = "1"
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"