Build and release sgl-deep-ep wheels (#33498)
This commit is contained in:
Baizhou Zhang
Executable
+197
@@ -0,0 +1,197 @@
|
||||
#!/usr/bin/env bash
|
||||
# Build sgl-deep-ep wheels in a CUDA-versioned manylinux container.
|
||||
#
|
||||
# Usage:
|
||||
# build_sgl_deepep.sh <python-version> <cuda-version> <deepep-source> <packaging-overlay> [architecture]
|
||||
#
|
||||
# Writes CUDA-tagged wheels to <deepep-source>/dist. CUDA 13 builds also write
|
||||
# PyPI-ready wheels without the local CUDA version to <deepep-source>/dist-pypi.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage: build_sgl_deepep.sh <python-version> <cuda-version> <deepep-source> <packaging-overlay> [architecture]
|
||||
|
||||
python-version: 3.10, 3.11, 3.12, or 3.13
|
||||
cuda-version: 12.9 or 13.0
|
||||
deepep-source: checkout of the selected DeepEP implementation branch
|
||||
packaging-overlay: path to the shared DeepEP sgl_deep_ep directory
|
||||
architecture: x86_64 or aarch64 (defaults to the current machine)
|
||||
EOF
|
||||
}
|
||||
|
||||
if [[ $# -lt 4 || $# -gt 5 ]]; then
|
||||
usage >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
PYTHON_VERSION="$1"
|
||||
CUDA_VERSION="$2"
|
||||
DEEPEP_SOURCE="$(cd "$3" && pwd)"
|
||||
PACKAGING_OVERLAY="$(cd "$4" && pwd)"
|
||||
ARCHITECTURE="${5:-$(uname -m)}"
|
||||
|
||||
if [[ "${ARCHITECTURE}" == arm64 ]]; then
|
||||
ARCHITECTURE=aarch64
|
||||
fi
|
||||
|
||||
case "${PYTHON_VERSION}" in
|
||||
3.10|3.11|3.12|3.13) ;;
|
||||
*)
|
||||
echo "Unsupported Python version: ${PYTHON_VERSION}" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
|
||||
case "${CUDA_VERSION}" in
|
||||
12.9)
|
||||
CUDA_TAG=cu129
|
||||
;;
|
||||
13.0)
|
||||
CUDA_TAG=cu130
|
||||
;;
|
||||
*)
|
||||
echo "Unsupported CUDA version: ${CUDA_VERSION}" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
|
||||
case "${ARCHITECTURE}" in
|
||||
x86_64)
|
||||
BASE_IMAGE=pytorch/manylinux2_28-builder
|
||||
;;
|
||||
aarch64)
|
||||
BASE_IMAGE=pytorch/manylinuxaarch64-builder
|
||||
;;
|
||||
*)
|
||||
echo "Unsupported architecture: ${ARCHITECTURE}" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
|
||||
for required_file in setup.py deep_ep/__init__.py; do
|
||||
if [[ ! -f "${DEEPEP_SOURCE}/${required_file}" ]]; then
|
||||
echo "DeepEP source is missing ${required_file}: ${DEEPEP_SOURCE}" >&2
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
for required_file in build_sgl_deep_ep.sh setup.py VERSION; do
|
||||
if [[ ! -f "${PACKAGING_OVERLAY}/${required_file}" ]]; then
|
||||
echo "Packaging overlay is missing ${required_file}: ${PACKAGING_OVERLAY}" >&2
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
if ! command -v docker >/dev/null; then
|
||||
echo "docker is required to build sgl-deep-ep" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
PYTHON_TAG="cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}"
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPOSITORY_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
DOCKERFILE="${REPOSITORY_ROOT}/docker/sgl-deep-ep.Dockerfile"
|
||||
IMAGE_TAG="sgl-deep-ep-builder:cuda${CUDA_VERSION}-${PYTHON_TAG}-${ARCHITECTURE}"
|
||||
DIST_DIR="${DEEPEP_SOURCE}/dist"
|
||||
PYPI_DIST_DIR="${DEEPEP_SOURCE}/dist-pypi"
|
||||
|
||||
mkdir -p "${DIST_DIR}" "${PYPI_DIST_DIR}"
|
||||
|
||||
echo "----------------------------------------"
|
||||
echo "Python: ${PYTHON_VERSION} (${PYTHON_TAG})"
|
||||
echo "CUDA: ${CUDA_VERSION} (${CUDA_TAG})"
|
||||
echo "Architecture: ${ARCHITECTURE}"
|
||||
echo "Base image: ${BASE_IMAGE}:cuda${CUDA_VERSION}"
|
||||
echo "DeepEP source: ${DEEPEP_SOURCE}"
|
||||
echo "Packaging overlay: ${PACKAGING_OVERLAY}"
|
||||
echo "Builder image: ${IMAGE_TAG}"
|
||||
echo "----------------------------------------"
|
||||
|
||||
docker build \
|
||||
--file "${DOCKERFILE}" \
|
||||
--build-arg BASE_IMAGE="${BASE_IMAGE}" \
|
||||
--build-arg CUDA_VERSION="${CUDA_VERSION}" \
|
||||
--build-arg CUDA_TAG="${CUDA_TAG}" \
|
||||
--build-arg PYTHON_TAG="${PYTHON_TAG}" \
|
||||
--build-arg ARCHITECTURE="${ARCHITECTURE}" \
|
||||
--build-arg TORCH_VERSION="${TORCH_VERSION:-2.11.0}" \
|
||||
--tag "${IMAGE_TAG}" \
|
||||
--network=host \
|
||||
"${REPOSITORY_ROOT}/docker"
|
||||
|
||||
docker run --rm \
|
||||
--network=host \
|
||||
--env ARCHITECTURE="${ARCHITECTURE}" \
|
||||
--env CUDA_TAG="${CUDA_TAG}" \
|
||||
--env CUDA_VERSION="${CUDA_VERSION}" \
|
||||
--env MAX_JOBS="${MAX_JOBS:-8}" \
|
||||
--volume "${DEEPEP_SOURCE}:/deepep:ro" \
|
||||
--volume "${PACKAGING_OVERLAY}:/packaging:ro" \
|
||||
--volume "${DIST_DIR}:/output/dist" \
|
||||
--volume "${PYPI_DIST_DIR}:/output/dist-pypi" \
|
||||
"${IMAGE_TAG}" \
|
||||
bash -euo pipefail -c '
|
||||
find /output/dist -maxdepth 1 -type f -name "sgl_deep_ep-*.whl" -delete
|
||||
find /output/dist-pypi -maxdepth 1 -type f -name "sgl_deep_ep-*.whl" -delete
|
||||
raw_dir="$(mktemp -d -t sgl-deep-ep-raw.XXXXXX)"
|
||||
trap '\''rm -rf -- "${raw_dir}"'\'' EXIT
|
||||
|
||||
bash /packaging/build_sgl_deep_ep.sh \
|
||||
/deepep /packaging "${raw_dir}" "${CUDA_VERSION}" "${ARCHITECTURE}"
|
||||
|
||||
shopt -s nullglob
|
||||
raw_wheels=("${raw_dir}"/*.whl)
|
||||
if [[ ${#raw_wheels[@]} -ne 1 ]]; then
|
||||
echo "Expected exactly one raw wheel, found ${#raw_wheels[@]}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
auditwheel repair \
|
||||
--plat "manylinux_2_28_${ARCHITECTURE}" \
|
||||
--wheel-dir /output/dist \
|
||||
--exclude libcuda.so.1 \
|
||||
--exclude libcudart.so.12 \
|
||||
--exclude libcudart.so.13 \
|
||||
--exclude libc10.so \
|
||||
--exclude libc10_cuda.so \
|
||||
--exclude libtorch.so \
|
||||
--exclude libtorch_cpu.so \
|
||||
--exclude libtorch_cuda.so \
|
||||
--exclude libtorch_python.so \
|
||||
--exclude libnvshmem_host.so.1 \
|
||||
--exclude libnvshmem_host.so.2 \
|
||||
--exclude libnvshmem_host.so.3 \
|
||||
--exclude libnccl.so.2 \
|
||||
--exclude libgdrapi.so.2 \
|
||||
--exclude libnvToolsExt.so.1 \
|
||||
"${raw_wheels[0]}"
|
||||
|
||||
if [[ "${CUDA_TAG}" == cu130 ]]; then
|
||||
tagged_wheels=(/output/dist/sgl_deep_ep-*+cu130-*.whl)
|
||||
if [[ ${#tagged_wheels[@]} -ne 1 ]]; then
|
||||
echo "Expected exactly one CUDA 13 wheel, found ${#tagged_wheels[@]}" >&2
|
||||
exit 1
|
||||
fi
|
||||
unpack_root="$(mktemp -d -t sgl-deep-ep-pypi.XXXXXX)"
|
||||
python -m wheel unpack "${tagged_wheels[0]}" --dest "${unpack_root}"
|
||||
unpacked="$(find "${unpack_root}" -mindepth 1 -maxdepth 1 -type d | head -1)"
|
||||
dist_info="$(find "${unpacked}" -maxdepth 1 -type d -name "*.dist-info" | head -1)"
|
||||
metadata="${dist_info}/METADATA"
|
||||
original_version="$(sed -n "s/^Version:[[:space:]]*//p" "${metadata}" | head -1)"
|
||||
public_version="${original_version%+cu130}"
|
||||
if [[ "${original_version}" == "${public_version}" ]]; then
|
||||
echo "CUDA 13 wheel metadata lacks the +cu130 local version" >&2
|
||||
exit 1
|
||||
fi
|
||||
sed -i "s/^Version:.*/Version: ${public_version}/" "${metadata}"
|
||||
old_dist_info="$(basename "${dist_info}")"
|
||||
new_dist_info="${old_dist_info/${original_version}/${public_version}}"
|
||||
mv "${dist_info}" "$(dirname "${dist_info}")/${new_dist_info}"
|
||||
python -m wheel pack "${unpacked}" --dest-dir /output/dist-pypi
|
||||
fi
|
||||
|
||||
ls -lh /output/dist/*.whl
|
||||
if [[ "${CUDA_TAG}" == cu130 ]]; then
|
||||
ls -lh /output/dist-pypi/*.whl
|
||||
fi
|
||||
'
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Update the PEP 503 indexes for sgl-deep-ep release wheels."""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import pathlib
|
||||
import re
|
||||
|
||||
SUPPORTED_CUDA_VERSIONS = ("129", "130")
|
||||
WHEEL_PATTERN = re.compile(
|
||||
r"^sgl_deep_ep-(?P<version>[0-9][^-]*)-[^-]+-[^-]+-[^-]+\.whl$"
|
||||
)
|
||||
ANCHOR_PATTERN = re.compile(r'^<a href="[^"]+">(?P<filename>[^<]+)</a><br>$')
|
||||
|
||||
|
||||
def _sha256(path: pathlib.Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as wheel_file:
|
||||
for chunk in iter(lambda: wheel_file.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _ensure_root_link(cuda_root: pathlib.Path) -> None:
|
||||
root_index = cuda_root / "index.html"
|
||||
lines = root_index.read_text().splitlines() if root_index.exists() else []
|
||||
doctype = [line for line in lines if line.startswith("<!DOCTYPE")]
|
||||
anchors = [line for line in lines if line.startswith("<a ")]
|
||||
link = '<a href="sgl-deep-ep/">sgl-deep-ep</a>'
|
||||
if link not in anchors:
|
||||
anchors.append(link)
|
||||
content = [*(doctype or ["<!DOCTYPE html>"]), *sorted(set(anchors))]
|
||||
root_index.write_text("\n".join(content) + "\n")
|
||||
|
||||
|
||||
def update_wheel_index(
|
||||
cuda_version: str,
|
||||
wheel_dir: pathlib.Path,
|
||||
repository: pathlib.Path,
|
||||
) -> None:
|
||||
if cuda_version not in SUPPORTED_CUDA_VERSIONS:
|
||||
raise ValueError(f"Unsupported CUDA version: {cuda_version}")
|
||||
|
||||
cuda_root = repository / f"cu{cuda_version}"
|
||||
index_dir = cuda_root / "sgl-deep-ep"
|
||||
index_dir.mkdir(exist_ok=True, parents=True)
|
||||
_ensure_root_link(cuda_root)
|
||||
|
||||
index_path = index_dir / "index.html"
|
||||
entries = {}
|
||||
if index_path.exists():
|
||||
for line in index_path.read_text().splitlines():
|
||||
match = ANCHOR_PATTERN.match(line)
|
||||
if match:
|
||||
entries[match.group("filename")] = line
|
||||
|
||||
suffix = f"+cu{cuda_version}"
|
||||
release_base = "https://github.com/sgl-project/whl/releases/download"
|
||||
for path in sorted(wheel_dir.glob("*.whl")):
|
||||
match = WHEEL_PATTERN.match(path.name)
|
||||
if not match or suffix not in match.group("version"):
|
||||
continue
|
||||
public_version = match.group("version").split("+", 1)[0]
|
||||
url = f"{release_base}/v{public_version}/{path.name}#sha256={_sha256(path)}"
|
||||
entries[path.name] = f'<a href="{url}">{path.name}</a><br>'
|
||||
|
||||
content = ["<!DOCTYPE html>", *(entries[name] for name in sorted(entries))]
|
||||
index_path.write_text("\n".join(content) + "\n")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--cuda", required=True, choices=SUPPORTED_CUDA_VERSIONS)
|
||||
parser.add_argument("--wheel-dir", type=pathlib.Path, default=pathlib.Path("dist"))
|
||||
parser.add_argument(
|
||||
"--repository", type=pathlib.Path, default=pathlib.Path("sgl-whl")
|
||||
)
|
||||
args = parser.parse_args()
|
||||
update_wheel_index(args.cuda, args.wheel_dir, args.repository)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user