mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +08:00
feat(h20): make the FlagTree image setup reproducible instead of a shell history
Getting the Triton backend to build and run in a vendor image took a dozen non-obvious steps: a venv because setup.py insists on CONDA_PREFIX, CUDA headers borrowed from the nvidia-cu13 wheels because the image ships only the runtime, NCCL 2.32.3 because DeepEP uses ncclCommQueryProperties, libdw/libelf for DeepGEMM's bundled deep_jit, pciutils for DeepSelect, and a tvm_ffi new enough for tilelang 0.1.15. All of it lived in this session and in prose, which means the next person starts over. `setup_flagtree_image.sh` records it in three stages -- env, build, verify -- and is idempotent, so a failure resumes. It also explains why each workaround exists, since several look arbitrary otherwise: FlagTree is a Triton fork and needs nothing special, but the C++ extensions around it do. `copy_site_packages.py` copies distributions by their declared file lists rather than by name, because pillow installs PIL and pillow.libs, pyzmq installs zmq. It defaults to keeping the image's pillow: a copied one points at vendored libtiff builds the image does not have, and PIL.Image then fails to import. Verified by building into an empty venv rather than the one already set up here: env, build and verify all pass from scratch, and an end-to-end run from that venv produces the same eight-rank-agreeing output as before.
This commit is contained in:
loopyt
parent
81adc36599
commit
788029f80b
145
scripts/copy_site_packages.py
Executable file
145
scripts/copy_site_packages.py
Executable file
@ -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 <src-site-packages> <dst-site-packages> \\
|
||||
--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())
|
||||
253
scripts/setup_flagtree_image.sh
Executable file
253
scripts/setup_flagtree_image.sh
Executable file
@ -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 <triton-image>:/usr/local/lib/python3.12/dist-packages/. /data/ftbuild/ref-triton/
|
||||
# docker cp <tilelang-image>:/usr/local/lib/python3.12/dist-packages/. /data/ftbuild/ref-tilelang/
|
||||
# docker cp <engine-image>:/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
|
||||
Loading…
x
Reference in New Issue
Block a user