@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 127 lines 5.31 KB
Raw
1 #!/usr/bin/env bash
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 #
5 # Shared helper for rewriting a wheel's platform tag from the generic
6 # ``linux_<arch>`` to the correct ``manylinux_<major>_<minor>_<arch>``.
7 # After sourcing, call ``apply_manylinux_tag <wheel>`` on each wheel
8 # that still carries the generic tag; the renamed path is printed on
9 # stdout (logs go to stderr).
10 #
11 # Why a pinned Docker container instead of using whatever Python
12 # happens to be on the agent:
13 # - vLLM's release agents are heterogeneous -- they don't agree on
14 # a Python minor version, and we can't rely on a particular
15 # ``auditwheel`` being installed.
16 # - ``detect-manylinux-tag.py`` reads ``auditwheel.wheel_abi`` and
17 # ``Policy.sym_policy``, which are *internal* APIs without a
18 # stability promise. Pinning both Python and auditwheel makes the
19 # detected tag a function of the inputs alone, and shifts version
20 # bumps from "implicit drift" to "deliberate, retested change".
21 # - Other release scripts (``generate-and-upload-nightly-index.sh``,
22 # ``upload-rocm-wheels.sh``) already use the python:3-slim image
23 # when the agent's interpreter is too old; this is the same idea
24 # made stricter.
25 #
26 # To keep the per-wheel cost down (the ROCm upload retags ~10 wheels
27 # each run), we install auditwheel into a long-lived helper container
28 # once on source, then ``docker exec`` into it for each call.
29 #
30 # Trap behaviour:
31 # - Sourcing installs an EXIT trap that calls ``manylinux_cleanup`` to
32 # tear down the helper container. Any EXIT trap that was already in
33 # place when this file was sourced is captured and run AFTER our
34 # cleanup, so we don't silently clobber it.
35 # - If a caller sets a new EXIT trap *after* sourcing, that trap will
36 # replace ours; in that case the caller should call
37 # ``manylinux_cleanup`` from their own handler.
38
39 if [[ -n "${_MANYLINUX_LIB_SOURCED:-}" ]]; then
40 return 0
41 fi
42 _MANYLINUX_LIB_SOURCED=1
43
44 # Pin both sides. Bump these deliberately and re-run a representative
45 # wheel from each build target through the detection.
46 _MANYLINUX_PYTHON_IMAGE="python:3.12-slim"
47 _MANYLINUX_AUDITWHEEL_VERSION="6.6.0"
48
49 # Resolve our own directory (and the sibling detect script) using the
50 # canonical, symlink-resolved path. The container mounts cwd at the
51 # same absolute path on both sides, so all paths we hand to it -- the
52 # script, the wheel -- must canonicalise to a location under cwd.
53 _MANYLINUX_LIB_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
54 _MANYLINUX_DETECT_SCRIPT="$(cd "${_MANYLINUX_LIB_DIR}/.." && pwd -P)/detect-manylinux-tag.py"
55 _MANYLINUX_CWD="$(pwd -P)"
56
57 docker pull --quiet "$_MANYLINUX_PYTHON_IMAGE" >/dev/null
58
59 # Spin up a long-lived helper container so we install auditwheel once
60 # and then ``docker exec`` into it for each wheel.
61 #
62 # The container runs as root so ``pip install`` can write into the
63 # system site-packages; individual ``docker exec`` calls below pin
64 # themselves to the host UID so any file rename happens with host
65 # ownership, not root.
66 _MANYLINUX_CONTAINER="$(docker run -d --rm \
67 -v "$_MANYLINUX_CWD:$_MANYLINUX_CWD" \
68 -w "$_MANYLINUX_CWD" \
69 "$_MANYLINUX_PYTHON_IMAGE" \
70 sleep infinity)"
71 docker exec "$_MANYLINUX_CONTAINER" \
72 pip install --quiet --disable-pip-version-check \
73 --root-user-action=ignore \
74 "auditwheel==${_MANYLINUX_AUDITWHEEL_VERSION}"
75
76 # Public cleanup -- safe to call multiple times.
77 manylinux_cleanup() {
78 if [[ -n "${_MANYLINUX_CONTAINER:-}" ]]; then
79 docker rm -f "$_MANYLINUX_CONTAINER" >/dev/null 2>&1 || true
80 _MANYLINUX_CONTAINER=""
81 fi
82 }
83
84 # Capture any EXIT trap that was already in place so we can chain to
85 # it rather than overwrite it. ``trap -p EXIT`` prints the handler in
86 # eval-able form (``trap -- 'CMD' EXIT``) or nothing if unset; we
87 # strip the wrapper to recover ``CMD``. Handles the common case --
88 # CMDs without embedded single quotes -- and degrades gracefully (we
89 # still run our own cleanup) for the pathological case.
90 _manylinux_prev_exit_trap_cmd=""
91 _manylinux_existing_exit_trap="$(trap -p EXIT)"
92 if [[ -n "$_manylinux_existing_exit_trap" ]]; then
93 _tmp="${_manylinux_existing_exit_trap#trap -- \'}"
94 _manylinux_prev_exit_trap_cmd="${_tmp%\' EXIT}"
95 unset _tmp
96 fi
97 unset _manylinux_existing_exit_trap
98
99 _manylinux_run_exit_chain() {
100 manylinux_cleanup
101 if [[ -n "$_manylinux_prev_exit_trap_cmd" ]]; then
102 eval "$_manylinux_prev_exit_trap_cmd"
103 fi
104 }
105 trap _manylinux_run_exit_chain EXIT
106
107 # Detect the manylinux platform tag for a single wheel and rename it
108 # in place, printing the renamed wheel path on stdout. Returns
109 # non-zero on failure (which under ``set -e`` propagates to caller).
110 #
111 # The wheel must be reachable via a path under the host cwd so it's
112 # visible inside the helper container; in CI the wheels always live
113 # under ``artifacts/`` so this is fine.
114 apply_manylinux_tag() {
115 local wheel="$1"
116 local abs_wheel
117 abs_wheel="$(realpath "$wheel")"
118 local new_wheel
119 new_wheel="$(docker exec -u "$(id -u):$(id -g)" \
120 "$_MANYLINUX_CONTAINER" \
121 python "$_MANYLINUX_DETECT_SCRIPT" "$abs_wheel")"
122 if [[ -z "$new_wheel" || ! -f "$new_wheel" ]]; then
123 echo "apply_manylinux_tag: detect-manylinux-tag.py did not produce a valid wheel path for $wheel" >&2
124 return 1
125 fi
126 printf '%s\n' "$new_wheel"
127 }