diff --git a/scripts/copy_site_packages.py b/scripts/copy_site_packages.py new file mode 100755 index 0000000..a98e14d --- /dev/null +++ b/scripts/copy_site_packages.py @@ -0,0 +1,145 @@ +"""Copy installed distributions between Python environments, file for file. + +Why this exists rather than `pip install`: pip on the H20 host stalls indefinitely on PyPI +(the wheel for `posix_ipc` never finishes downloading), and the packages we need are already +installed and pinned in a reference image. Copying is exact and takes seconds. + +Two things it handles that a naive `cp -r` does not: + +* A distribution's files are not always under its own name. `pillow` installs `PIL` and + `pillow.libs`, `pyzmq` installs `zmq` and `pyzmq.libs`. The file list comes from the + distribution's metadata, so those sibling directories come along. +* Some packages must *not* be copied even when present, because the target image ships a + version linked against libs the source does not have. `pillow` is the case in hand: copying + 11.3 into an image whose `pillow.libs` holds a different libtiff set leaves `PIL.Image` + unable to import. Pass it in `--keep-system` and the target's own copy is reused. + + python scripts/copy_site_packages.py \\ + --packages einops numpy pillow \\ + --keep-system pillow +""" + +import argparse +import importlib.metadata as md +import shutil +import sys +from pathlib import Path + +# Everything `requirements.txt` asks for that the engine imports at load time. torch, +# tilelang, tile-kernels and tvm-ffi are deliberately absent: they are either provided by +# the image or installed by setup_flagtree_image.sh from a reference container, because +# their builds are ABI-sensitive in ways a file copy only hides. +DEFAULT_PACKAGES = [ + "einops", + "huggingface_hub", + "ninja", + "numpy", + "orjson", + "pillow", + "posix_ipc", + "psutil", + "pybind11", + "pyzmq", + "requests", + "safetensors", + "scipy", + "sympy", + "tabulate", + "tokenizers", + "tqdm", + "transformers", + "xgrammar", +] + + +def owned_paths(dist: md.Distribution) -> list[Path]: + """The top-level entries a distribution installs, plus its metadata directory.""" + root = Path(dist.locate_file("")) + found: dict[str, Path] = {} + + # top_level.txt names the import roots, which is what we want when they differ from the + # distribution name. + try: + declared = dist.read_text("top_level.txt") or "" + except (FileNotFoundError, UnicodeDecodeError): + declared = "" + for line in declared.splitlines(): + name = line.strip() + if name and (root / name).exists(): + found[name] = root / name + + # Fall back to the file list; single-module distributions have no top_level entry. + if not found: + for f in dist.files or []: + first = f.parts[0] + if first.endswith((".dist-info", ".egg-info", ".data")) or first.startswith(".."): + continue + if (root / first).exists(): + found[first] = root / first + + found[dist._path.name] = dist._path.resolve() + return list(found.values()) + + +def install(src: Path, dst: Path, name: str, dist: md.Distribution) -> list[str]: + copied = [] + for path in owned_paths(dist): + target = dst / path.name + if target.is_symlink() or target.is_file(): + target.unlink() + elif target.is_dir(): + shutil.rmtree(target) + if path.is_dir() and not path.is_symlink(): + shutil.copytree(path, target, symlinks=True) + else: + # `.so` files and metadata are files; copy_file keeps the mode bits, which the + # extension modules need. + shutil.copy2(path, target, follow_symlinks=False) + copied.append(path.name) + return copied + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("src", type=Path, help="source site-packages directory") + parser.add_argument("dst", type=Path, help="destination site-packages directory") + parser.add_argument("--packages", nargs="+", default=DEFAULT_PACKAGES) + parser.add_argument( + "--keep-system", + nargs="*", + default=["pillow"], + help="packages to leave to the destination image, because their vendored libs differ", + ) + args = parser.parse_args() + + for p in (args.src, args.dst): + if not p.is_dir(): + print(f"not a directory: {p}", file=sys.stderr) + return 2 + + # Reading metadata from `src` without installing it: the finder consults sys.path, so the + # source tree goes first for the life of this process. + sys.path.insert(0, str(args.src)) + + missing = [] + for name in args.packages: + if name in args.keep_system: + print(f"skip {name} (keeping the destination image's copy)") + continue + try: + dist = md.distribution(name) + except md.PackageNotFoundError: + missing.append(name) + print(f"missing {name}") + continue + copied = install(args.src, args.dst, name, dist) + print(f"{dist.version:<10}{name}: {', '.join(copied)}") + + if missing: + print(f"\nnot found in {args.src}: {', '.join(missing)}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/setup_flagtree_image.sh b/scripts/setup_flagtree_image.sh new file mode 100755 index 0000000..cf00e6d --- /dev/null +++ b/scripts/setup_flagtree_image.sh @@ -0,0 +1,253 @@ +#!/bin/bash +# Prepare a FlagTree/FlagOS image to build and run DeepInfer's Triton backend. +# +# The kernels themselves need nothing but Triton, and FlagTree *is* a Triton fork -- that part +# works out of the box. What does not is everything around it: `deepinfer._C` is a torch C++ +# extension, and the vendored DeepGEMM/DeepEP/DeepSelect are CUDA extensions. None of them +# involve Triton, and none of them ship in a vendor image, so they have to be built here. +# This script assembles a build environment where that succeeds. +# +# Run it *inside* the target container, from the repo root: +# +# bash scripts/setup_flagtree_image.sh env # system packages, venv, headers, deps +# bash scripts/setup_flagtree_image.sh build # _C and the three vendor wheels +# bash scripts/setup_flagtree_image.sh all +# +# Before `env`, the caller stages reference packages on the host, because `docker cp` only +# exists outside the container: +# +# mkdir -p /data/ftbuild/ref-triton /data/ftbuild/ref-tilelang /data/ftbuild/ref-site +# docker cp :/usr/local/lib/python3.12/dist-packages/. /data/ftbuild/ref-triton/ +# docker cp :/usr/local/lib/python3.12/dist-packages/. /data/ftbuild/ref-tilelang/ +# docker cp :/usr/local/lib/python3.12/dist-packages/. /data/ftbuild/ref-site/ +# +# `ref-triton` needs packages matching this image's torch build; `ref-site` supplies the +# Python packages from requirements.txt that the engine imports at load time. +# +# Every step is idempotent, so re-running after a failure resumes rather than starting over. + +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +SCRATCH="${SCRATCH:-/data/ftbuild}" +VENV_DIR="${VENV_DIR:-$SCRATCH/venv}" +WHEEL_DIR="${WHEEL_DIR:-$SCRATCH/wheels}" + +REF_TRITON_DIR="${REF_TRITON_DIR:-$SCRATCH/ref-triton}" +REF_TILELANG_DIR="${REF_TILELANG_DIR:-$SCRATCH/ref-tilelang}" +REF_SITE_PACKAGES="${REF_SITE_PACKAGES:-$SCRATCH/ref-site}" + +# DeepEP calls ncclCommQueryProperties and reads props.ginMinStride, both of which arrived in +# NCCL 2.32.3. Vendor images ship 2.28.9, so this one has to be replaced. +NCCL_VERSION="${NCCL_VERSION:-2.32.3}" +NCCL_URL="${NCCL_URL:-https://mirrors.aliyun.com/pypi/packages/5b/29/6b277e63c92d91f9cb4d1a3a554e148983de39d54baa652bb52c798af78e/nvidia_nccl_cu13-${NCCL_VERSION}-py3-none-manylinux_2_27_x86_64.whl}" + +MAX_JOBS="${MAX_JOBS:-$(nproc)}" +PY_TAG="python$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')" +SITE="$VENV_DIR/lib/$PY_TAG/site-packages" + +say() { printf '\n\033[1m== %s\033[0m\n' "$*"; } + +# Copy one or more site-packages entries, expanded from shell globs, replacing what is there. +# Globbing the dist-info directory rather than hardcoding its name keeps this working when the +# reference image bumps a version. +stage() { + local label="$1" src="$2" dst="$3" + shift 3 + local count=0 path + for pattern in "$@"; do + for path in $src/$pattern; do + [ -e "$path" ] || continue + rm -rf "${dst:?}/$(basename "$path")" + cp -a "$path" "$dst/" + count=$((count + 1)) + done + done + if [ "$count" -eq 0 ]; then + echo " !! $label: nothing matched under $src (looked for: $*)" >&2 + return 1 + fi + echo " $label: $count entr$([ "$count" -eq 1 ] && echo y || echo ies)" +} + +step_env() { + say "preflight" + command -v nvcc >/dev/null || { echo "nvcc missing: this image has no CUDA toolkit" >&2; exit 1; } + python3 -c 'import torch; assert torch.version.cuda' || { echo "torch is not a CUDA build" >&2; exit 1; } + echo " torch $(python3 -c 'import torch; print(torch.__version__)')" + echo " nvcc $(nvcc --version | tail -1 | sed 's/.*, //')" + echo " gcc $(gcc -dumpversion)" + echo " repo $REPO_ROOT" + echo " venv $VENV_DIR" + + say "system packages" + # libdw/libelf: DeepGEMM's bundled deep_jit resolves stack frames through libdwfl. + # libzstd: DeepEP links it. pciutils: DeepSelect's setup.py shells out to lspci to pick + # the build target. git: vendor submodule tooling. + if command -v apt-get >/dev/null; then + export DEBIAN_FRONTEND=noninteractive + apt-get update -qq || echo " (apt update had warnings; continuing)" + apt-get install -y -qq libdw-dev libelf-dev libzstd-dev pciutils git >/dev/null + echo " installed" + else + echo " no apt-get here; make sure libdw, libelf, libzstd and lspci are present" + fi + + say "venv" + # setup.py refuses to build unless CONDA_PREFIX or VIRTUAL_ENV is set, and only uses the + # value as a guard -- it never reads it. A venv is the lightest way to satisfy it while + # inheriting the image's torch. + if [ ! -x "$VENV_DIR/bin/python" ]; then + python3 -m venv --system-site-packages "$VENV_DIR" + fi + # shellcheck disable=SC1091 + source "$VENV_DIR/bin/activate" + python -c 'import torch; print(" torch visible in venv:", torch.__version__)' + + say "CUDA headers" + # The image ships the CUDA runtime, not the toolkit headers; torch's own headers include + # cusparse.h, so building any extension fails without it. The nvidia-*-cu13 wheels carry + # the headers, so borrow them rather than installing a full toolkit. + if [ ! -f /usr/local/cuda/include/cusparse.h ]; then + for hdr in /usr/local/lib/"$PY_TAG"/dist-packages/nvidia/cu13/include \ + /usr/local/lib/"$PY_TAG"/dist-packages/nvidia/cuda_runtime/include; do + [ -d "$hdr" ] && cp -n "$hdr"/*.h /usr/local/cuda/include/ 2>/dev/null + done + fi + if [ -f /usr/local/cuda/include/cusparse.h ]; then + echo " cusparse.h present" + else + echo " !! cusparse.h still missing; _C will not compile" >&2 + return 1 + fi + + say "NCCL $NCCL_VERSION" + if python -c "import nvidia.nccl" 2>/dev/null && + [ "$(python -c 'import importlib.metadata as m; print(m.version("nvidia-nccl-cu13"))' 2>/dev/null)" = "$NCCL_VERSION" ]; then + echo " already $NCCL_VERSION" + else + mkdir -p "$SCRATCH/pipdl" + local whl="$SCRATCH/pipdl/nvidia_nccl_cu13-${NCCL_VERSION}.whl" + # pip's own download stalls on this host for a 300MB wheel; curl with resume finishes. + [ -s "$whl" ] || curl -sSL -C - --retry 5 --retry-delay 3 --max-time 900 -o "$whl" "$NCCL_URL" + # pip validates the filename, so it has to look like one. + mv "$whl" "$SCRATCH/pipdl/nvidia_nccl_cu13-${NCCL_VERSION}-py3-none-manylinux_2_27_x86_64.whl" + pip install -q --no-deps --force-reinstall \ + "$SCRATCH/pipdl/nvidia_nccl_cu13-${NCCL_VERSION}-py3-none-manylinux_2_27_x86_64.whl" + echo " installed" + fi + + say "Triton (= FlagTree)" + # FlagTree is a fork of Triton: the import name stays `triton`, and `flagtree` only holds + # FlagPrism, which triton imports at load time. Both halves plus their metadata are needed, + # because the driver is registered through an entry point that only dist-info declares. + stage triton "$REF_TRITON_DIR" "$SITE" \ + triton flagtree 'tokenspeed*triton*.dist-info' 'flagtree*.dist-info' + python -c 'import triton, flagtree; print(" triton", triton.__version__, "= flagtree", flagtree.__version__)' + + say "TileLang runtime" + # Only the fallback path needs these -- the kernels Triton does not implement yet are still + # served by TileLang. tile_kernels is the quantiser used by the reference implementations. + # tvm_ffi must come along: the image's own 0.1.9 lacks tvm_ffi.dataclasses.Enum, which + # tilelang 0.1.15 imports during startup. + if [ -d "$REF_TILELANG_DIR/tilelang" ]; then + stage tilelang "$REF_TILELANG_DIR" "$SITE" \ + tilelang tile_kernels tvm_ffi \ + 'tilelang-*.dist-info' 'tile_kernels-*.dist-info' 'apache_tvm_ffi-*.dist-info' + python -c 'import tilelang, tile_kernels; print(" tilelang", tilelang.__version__)' + else + echo " skipped: no $REF_TILELANG_DIR (end-to-end will fail on the MTP fallback path)" + fi + + say "Python dependencies" + if [ -d "$REF_SITE_PACKAGES" ]; then + # pillow stays as the image ships it: its wheels carry vendored libs (libtiff and + # friends) under pillow.libs, and a copy from another image points at different ones. + python "$REPO_ROOT/scripts/copy_site_packages.py" "$REF_SITE_PACKAGES" "$SITE" --keep-system pillow + else + echo " no $REF_SITE_PACKAGES; falling back to pip (may stall on this host)" + pip install -q -r "$REPO_ROOT/requirements.txt" || echo " !! pip failed; stage ref-site instead" + fi + + say "env ready" + echo " source $VENV_DIR/bin/activate" +} + +step_build() { + # shellcheck disable=SC1091 + source "$VENV_DIR/bin/activate" + cd "$REPO_ROOT" + mkdir -p "$SCRATCH/tmp" "$WHEEL_DIR" + export TMPDIR="$SCRATCH/tmp" + export DEEPINFER_ACCELERATOR=nvidia + export MAX_JOBS + + say "_C extension" + bash scripts/build.sh build + + say "vendor wheels" + for pkg in DeepGEMM DeepEP DeepSelect; do + [ -f "vendor/$pkg/setup.py" ] || { + echo " !! vendor/$pkg is empty -- run 'git submodule update --init' on a git checkout" >&2 + return 1 + } + echo " building $pkg" + rm -rf "vendor/$pkg/build" "vendor/$pkg"/*.egg-info + local envs=() + # Matches scripts/build_vendor.sh: DeepGEMM only rebuilds when forced, DeepSelect needs + # its platform target spelled out. + [ "$pkg" = DeepGEMM ] && envs+=(DG_FORCE_BUILD=1) + [ "$pkg" = DeepSelect ] && envs+=(DEEP_SELECT_BUILD_TARGET_PLATFORM=CUDA) + ( + cd "vendor/$pkg" + env "${envs[@]}" python setup.py bdist_wheel --dist-dir "$WHEEL_DIR" >/dev/null + ) + done + ls -lh "$WHEEL_DIR" + + say "install wheels" + # --no-deps and a forced reinstall: the image has its own deep_ep, and pip would otherwise + # keep it because the version number happens to look newer. + pip install -q --no-deps --force-reinstall \ + "$WHEEL_DIR"/deep_gemm-*.whl "$WHEEL_DIR"/deep_ep-*.whl "$WHEEL_DIR"/deep_select-*.whl + pip list 2>/dev/null | grep -iE "^(deep_gemm|deep_ep|deep_select|tilelang|triton) " || true +} + +step_verify() { + # shellcheck disable=SC1091 + source "$VENV_DIR/bin/activate" + cd "$REPO_ROOT" + export PYTHONPATH="$REPO_ROOT:$REPO_ROOT/examples:$REPO_ROOT/scripts" + export DEEPINFER_OPS_BACKEND=triton DEEPINFER_OPS_STRICT=1 + + say "imports" + python - <<'PY' +import torch, triton, flagtree +print(" torch ", torch.__version__, "cap", torch.cuda.get_device_capability(0)) +print(" triton ", triton.__version__, "(= FlagTree", flagtree.__version__ + ")") +import deep_gemm, deep_ep, deep_select +print(" vendor deep_gemm", deep_gemm.__version__, "| deep_ep", deep_ep.__version__) +from deepinfer._C import engine +print(" _C ", len([a for a in dir(engine) if not a.startswith("_")]), "symbols") +PY + + say "ops backend" + python -c ' +from deepinfer.model_executor.layers.tilelang_ops import ops +print(" backend", ops._BACKEND, "strict", ops._STRICT, "kernels", len(ops._SELECTED))' +} + +case "${1:-all}" in + env) step_env ;; + build) step_build ;; + verify) step_verify ;; + all) + step_env + step_build + step_verify + ;; + *) + echo "usage: $0 {env|build|verify|all}" >&2 + exit 2 + ;; +esac