1#!/bin/bash 2 3set -ex 4 5image="$1" 6shift 7 8if [ -z "${image}" ]; then 9 echo "Usage: $0 IMAGE" 10 exit 1 11fi 12 13function extract_version_from_image_name() { 14 eval export $2=$(echo "${image}" | perl -n -e"/$1(\d+(\.\d+)?(\.\d+)?)/ && print \$1") 15 if [ "x${!2}" = x ]; then 16 echo "variable '$2' not correctly parsed from image='$image'" 17 exit 1 18 fi 19} 20 21function extract_all_from_image_name() { 22 # parts $image into array, splitting on '-' 23 keep_IFS="$IFS" 24 IFS="-" 25 declare -a parts=($image) 26 IFS="$keep_IFS" 27 unset keep_IFS 28 29 for part in "${parts[@]}"; do 30 name=$(echo "${part}" | perl -n -e"/([a-zA-Z]+)\d+(\.\d+)?(\.\d+)?/ && print \$1") 31 vername="${name^^}_VERSION" 32 # "py" is the odd one out, needs this special case 33 if [ "x${name}" = xpy ]; then 34 vername=ANACONDA_PYTHON_VERSION 35 fi 36 # skip non-conforming fields such as "pytorch", "linux" or "bionic" without version string 37 if [ -n "${name}" ]; then 38 extract_version_from_image_name "${name}" "${vername}" 39 fi 40 done 41} 42 43# Use the same pre-built XLA test image from PyTorch/XLA 44if [[ "$image" == *xla* ]]; then 45 echo "Using pre-built XLA test image..." 46 exit 0 47fi 48 49if [[ "$image" == *-focal* ]]; then 50 UBUNTU_VERSION=20.04 51elif [[ "$image" == *-jammy* ]]; then 52 UBUNTU_VERSION=22.04 53elif [[ "$image" == *ubuntu* ]]; then 54 extract_version_from_image_name ubuntu UBUNTU_VERSION 55elif [[ "$image" == *centos* ]]; then 56 extract_version_from_image_name centos CENTOS_VERSION 57fi 58 59if [ -n "${UBUNTU_VERSION}" ]; then 60 OS="ubuntu" 61elif [ -n "${CENTOS_VERSION}" ]; then 62 OS="centos" 63else 64 echo "Unable to derive operating system base..." 65 exit 1 66fi 67 68DOCKERFILE="${OS}/Dockerfile" 69# When using ubuntu - 22.04, start from Ubuntu docker image, instead of nvidia/cuda docker image. 70if [[ "$image" == *cuda* && "$UBUNTU_VERSION" != "22.04" ]]; then 71 DOCKERFILE="${OS}-cuda/Dockerfile" 72elif [[ "$image" == *rocm* ]]; then 73 DOCKERFILE="${OS}-rocm/Dockerfile" 74elif [[ "$image" == *xpu* ]]; then 75 DOCKERFILE="${OS}-xpu/Dockerfile" 76elif [[ "$image" == *cuda*linter* ]]; then 77 # Use a separate Dockerfile for linter to keep a small image size 78 DOCKERFILE="linter-cuda/Dockerfile" 79elif [[ "$image" == *linter* ]]; then 80 # Use a separate Dockerfile for linter to keep a small image size 81 DOCKERFILE="linter/Dockerfile" 82fi 83 84# CMake 3.18 is needed to support CUDA17 language variant 85CMAKE_VERSION=3.18.5 86 87_UCX_COMMIT=7bb2722ff2187a0cad557ae4a6afa090569f83fb 88_UCC_COMMIT=20eae37090a4ce1b32bcce6144ccad0b49943e0b 89 90# It's annoying to rename jobs every time you want to rewrite a 91# configuration, so we hardcode everything here rather than do it 92# from scratch 93case "$image" in 94 pytorch-linux-focal-cuda12.4-cudnn9-py3-gcc9) 95 CUDA_VERSION=12.4.1 96 CUDNN_VERSION=9 97 ANACONDA_PYTHON_VERSION=3.10 98 GCC_VERSION=9 99 PROTOBUF=yes 100 DB=yes 101 VISION=yes 102 KATEX=yes 103 UCX_COMMIT=${_UCX_COMMIT} 104 UCC_COMMIT=${_UCC_COMMIT} 105 CONDA_CMAKE=yes 106 TRITON=yes 107 ;; 108 pytorch-linux-focal-cuda12.1-cudnn9-py3-gcc9) 109 CUDA_VERSION=12.1.1 110 CUDNN_VERSION=9 111 ANACONDA_PYTHON_VERSION=3.10 112 GCC_VERSION=9 113 PROTOBUF=yes 114 DB=yes 115 VISION=yes 116 KATEX=yes 117 UCX_COMMIT=${_UCX_COMMIT} 118 UCC_COMMIT=${_UCC_COMMIT} 119 CONDA_CMAKE=yes 120 TRITON=yes 121 ;; 122 pytorch-linux-focal-cuda12.4-cudnn9-py3-gcc9-inductor-benchmarks) 123 CUDA_VERSION=12.4.1 124 CUDNN_VERSION=9 125 ANACONDA_PYTHON_VERSION=3.10 126 GCC_VERSION=9 127 PROTOBUF=yes 128 DB=yes 129 VISION=yes 130 KATEX=yes 131 UCX_COMMIT=${_UCX_COMMIT} 132 UCC_COMMIT=${_UCC_COMMIT} 133 CONDA_CMAKE=yes 134 TRITON=yes 135 INDUCTOR_BENCHMARKS=yes 136 ;; 137 pytorch-linux-focal-cuda12.1-cudnn9-py3-gcc9-inductor-benchmarks) 138 CUDA_VERSION=12.1.1 139 CUDNN_VERSION=9 140 ANACONDA_PYTHON_VERSION=3.10 141 GCC_VERSION=9 142 PROTOBUF=yes 143 DB=yes 144 VISION=yes 145 KATEX=yes 146 UCX_COMMIT=${_UCX_COMMIT} 147 UCC_COMMIT=${_UCC_COMMIT} 148 CONDA_CMAKE=yes 149 TRITON=yes 150 INDUCTOR_BENCHMARKS=yes 151 ;; 152 pytorch-linux-focal-cuda12.1-cudnn9-py3.12-gcc9-inductor-benchmarks) 153 CUDA_VERSION=12.1.1 154 CUDNN_VERSION=9 155 ANACONDA_PYTHON_VERSION=3.12 156 GCC_VERSION=9 157 PROTOBUF=yes 158 DB=yes 159 VISION=yes 160 KATEX=yes 161 UCX_COMMIT=${_UCX_COMMIT} 162 UCC_COMMIT=${_UCC_COMMIT} 163 CONDA_CMAKE=yes 164 TRITON=yes 165 INDUCTOR_BENCHMARKS=yes 166 ;; 167 pytorch-linux-focal-cuda12.4-cudnn9-py3.12-gcc9-inductor-benchmarks) 168 CUDA_VERSION=12.4.1 169 CUDNN_VERSION=9 170 ANACONDA_PYTHON_VERSION=3.12 171 GCC_VERSION=9 172 PROTOBUF=yes 173 DB=yes 174 VISION=yes 175 KATEX=yes 176 UCX_COMMIT=${_UCX_COMMIT} 177 UCC_COMMIT=${_UCC_COMMIT} 178 CONDA_CMAKE=yes 179 TRITON=yes 180 INDUCTOR_BENCHMARKS=yes 181 ;; 182 pytorch-linux-focal-cuda11.8-cudnn9-py3-gcc9) 183 CUDA_VERSION=11.8.0 184 CUDNN_VERSION=9 185 ANACONDA_PYTHON_VERSION=3.10 186 GCC_VERSION=9 187 PROTOBUF=yes 188 DB=yes 189 VISION=yes 190 KATEX=yes 191 UCX_COMMIT=${_UCX_COMMIT} 192 UCC_COMMIT=${_UCC_COMMIT} 193 CONDA_CMAKE=yes 194 TRITON=yes 195 ;; 196 pytorch-linux-focal-cuda12.4-cudnn9-py3-gcc9) 197 CUDA_VERSION=12.4.1 198 CUDNN_VERSION=9 199 ANACONDA_PYTHON_VERSION=3.10 200 GCC_VERSION=9 201 PROTOBUF=yes 202 DB=yes 203 VISION=yes 204 KATEX=yes 205 UCX_COMMIT=${_UCX_COMMIT} 206 UCC_COMMIT=${_UCC_COMMIT} 207 CONDA_CMAKE=yes 208 TRITON=yes 209 ;; 210 pytorch-linux-focal-cuda12.1-cudnn9-py3-gcc9) 211 CUDA_VERSION=12.1.1 212 CUDNN_VERSION=9 213 ANACONDA_PYTHON_VERSION=3.10 214 GCC_VERSION=9 215 PROTOBUF=yes 216 DB=yes 217 VISION=yes 218 KATEX=yes 219 UCX_COMMIT=${_UCX_COMMIT} 220 UCC_COMMIT=${_UCC_COMMIT} 221 CONDA_CMAKE=yes 222 TRITON=yes 223 ;; 224 pytorch-linux-focal-cuda12.4-cudnn9-py3-gcc9) 225 CUDA_VERSION=12.4.1 226 CUDNN_VERSION=9 227 ANACONDA_PYTHON_VERSION=3.10 228 GCC_VERSION=9 229 PROTOBUF=yes 230 DB=yes 231 VISION=yes 232 KATEX=yes 233 UCX_COMMIT=${_UCX_COMMIT} 234 UCC_COMMIT=${_UCC_COMMIT} 235 CONDA_CMAKE=yes 236 TRITON=yes 237 ;; 238 pytorch-linux-focal-py3-clang10-onnx) 239 ANACONDA_PYTHON_VERSION=3.9 240 CLANG_VERSION=10 241 PROTOBUF=yes 242 DB=yes 243 VISION=yes 244 CONDA_CMAKE=yes 245 ONNX=yes 246 ;; 247 pytorch-linux-focal-py3-clang9-android-ndk-r21e) 248 ANACONDA_PYTHON_VERSION=3.9 249 CLANG_VERSION=9 250 LLVMDEV=yes 251 PROTOBUF=yes 252 ANDROID=yes 253 ANDROID_NDK_VERSION=r21e 254 GRADLE_VERSION=6.8.3 255 NINJA_VERSION=1.9.0 256 ;; 257 pytorch-linux-focal-py3.9-clang10) 258 ANACONDA_PYTHON_VERSION=3.9 259 CLANG_VERSION=10 260 PROTOBUF=yes 261 DB=yes 262 VISION=yes 263 VULKAN_SDK_VERSION=1.2.162.1 264 SWIFTSHADER=yes 265 CONDA_CMAKE=yes 266 TRITON=yes 267 ;; 268 pytorch-linux-focal-py3.11-clang10) 269 ANACONDA_PYTHON_VERSION=3.11 270 CLANG_VERSION=10 271 PROTOBUF=yes 272 DB=yes 273 VISION=yes 274 VULKAN_SDK_VERSION=1.2.162.1 275 SWIFTSHADER=yes 276 CONDA_CMAKE=yes 277 TRITON=yes 278 ;; 279 pytorch-linux-focal-py3.9-gcc9) 280 ANACONDA_PYTHON_VERSION=3.9 281 GCC_VERSION=9 282 PROTOBUF=yes 283 DB=yes 284 VISION=yes 285 CONDA_CMAKE=yes 286 TRITON=yes 287 ;; 288 pytorch-linux-focal-rocm-n-1-py3) 289 ANACONDA_PYTHON_VERSION=3.10 290 GCC_VERSION=9 291 PROTOBUF=yes 292 DB=yes 293 VISION=yes 294 ROCM_VERSION=6.1 295 NINJA_VERSION=1.9.0 296 CONDA_CMAKE=yes 297 TRITON=yes 298 ;; 299 pytorch-linux-focal-rocm-n-py3) 300 ANACONDA_PYTHON_VERSION=3.10 301 GCC_VERSION=9 302 PROTOBUF=yes 303 DB=yes 304 VISION=yes 305 ROCM_VERSION=6.2 306 NINJA_VERSION=1.9.0 307 CONDA_CMAKE=yes 308 TRITON=yes 309 ;; 310 pytorch-linux-jammy-xpu-2024.0-py3) 311 ANACONDA_PYTHON_VERSION=3.9 312 GCC_VERSION=11 313 PROTOBUF=yes 314 DB=yes 315 VISION=yes 316 XPU_VERSION=0.5 317 NINJA_VERSION=1.9.0 318 CONDA_CMAKE=yes 319 TRITON=yes 320 ;; 321 pytorch-linux-jammy-py3.9-gcc11-inductor-benchmarks) 322 ANACONDA_PYTHON_VERSION=3.9 323 GCC_VERSION=11 324 PROTOBUF=yes 325 DB=yes 326 VISION=yes 327 KATEX=yes 328 CONDA_CMAKE=yes 329 TRITON=yes 330 DOCS=yes 331 INDUCTOR_BENCHMARKS=yes 332 ;; 333 pytorch-linux-jammy-cuda11.8-cudnn9-py3.9-clang12) 334 ANACONDA_PYTHON_VERSION=3.9 335 CUDA_VERSION=11.8 336 CUDNN_VERSION=9 337 CLANG_VERSION=12 338 PROTOBUF=yes 339 DB=yes 340 VISION=yes 341 TRITON=yes 342 ;; 343 pytorch-linux-jammy-py3-clang12-asan) 344 ANACONDA_PYTHON_VERSION=3.9 345 CLANG_VERSION=12 346 PROTOBUF=yes 347 DB=yes 348 VISION=yes 349 CONDA_CMAKE=yes 350 TRITON=yes 351 ;; 352 pytorch-linux-jammy-py3-clang15-asan) 353 ANACONDA_PYTHON_VERSION=3.10 354 CLANG_VERSION=15 355 CONDA_CMAKE=yes 356 VISION=yes 357 ;; 358 pytorch-linux-jammy-py3.9-gcc11) 359 ANACONDA_PYTHON_VERSION=3.9 360 GCC_VERSION=11 361 PROTOBUF=yes 362 DB=yes 363 VISION=yes 364 KATEX=yes 365 CONDA_CMAKE=yes 366 TRITON=yes 367 DOCS=yes 368 UNINSTALL_DILL=yes 369 ;; 370 pytorch-linux-jammy-py3-clang12-executorch) 371 ANACONDA_PYTHON_VERSION=3.10 372 CLANG_VERSION=12 373 CONDA_CMAKE=yes 374 EXECUTORCH=yes 375 ;; 376 pytorch-linux-jammy-py3.12-halide) 377 CUDA_VERSION=12.4 378 ANACONDA_PYTHON_VERSION=3.12 379 GCC_VERSION=11 380 CONDA_CMAKE=yes 381 HALIDE=yes 382 ;; 383 pytorch-linux-focal-linter) 384 # TODO: Use 3.9 here because of this issue https://github.com/python/mypy/issues/13627. 385 # We will need to update mypy version eventually, but that's for another day. The task 386 # would be to upgrade mypy to 1.0.0 with Python 3.11 387 ANACONDA_PYTHON_VERSION=3.9 388 CONDA_CMAKE=yes 389 ;; 390 pytorch-linux-jammy-cuda11.8-cudnn9-py3.9-linter) 391 ANACONDA_PYTHON_VERSION=3.9 392 CUDA_VERSION=11.8 393 CONDA_CMAKE=yes 394 ;; 395 pytorch-linux-jammy-aarch64-py3.10-gcc11) 396 ANACONDA_PYTHON_VERSION=3.10 397 GCC_VERSION=11 398 ACL=yes 399 PROTOBUF=yes 400 DB=yes 401 VISION=yes 402 CONDA_CMAKE=yes 403 # snadampal: skipping sccache due to the following issue 404 # https://github.com/pytorch/pytorch/issues/121559 405 SKIP_SCCACHE_INSTALL=yes 406 # snadampal: skipping llvm src build install because the current version 407 # from pytorch/llvm:9.0.1 is x86 specific 408 SKIP_LLVM_SRC_BUILD_INSTALL=yes 409 ;; 410 pytorch-linux-jammy-aarch64-py3.10-gcc11-inductor-benchmarks) 411 ANACONDA_PYTHON_VERSION=3.10 412 GCC_VERSION=11 413 ACL=yes 414 PROTOBUF=yes 415 DB=yes 416 VISION=yes 417 CONDA_CMAKE=yes 418 # snadampal: skipping sccache due to the following issue 419 # https://github.com/pytorch/pytorch/issues/121559 420 SKIP_SCCACHE_INSTALL=yes 421 # snadampal: skipping llvm src build install because the current version 422 # from pytorch/llvm:9.0.1 is x86 specific 423 SKIP_LLVM_SRC_BUILD_INSTALL=yes 424 INDUCTOR_BENCHMARKS=yes 425 ;; 426 *) 427 # Catch-all for builds that are not hardcoded. 428 PROTOBUF=yes 429 DB=yes 430 VISION=yes 431 echo "image '$image' did not match an existing build configuration" 432 if [[ "$image" == *py* ]]; then 433 extract_version_from_image_name py ANACONDA_PYTHON_VERSION 434 fi 435 if [[ "$image" == *cuda* ]]; then 436 extract_version_from_image_name cuda CUDA_VERSION 437 extract_version_from_image_name cudnn CUDNN_VERSION 438 fi 439 if [[ "$image" == *rocm* ]]; then 440 extract_version_from_image_name rocm ROCM_VERSION 441 NINJA_VERSION=1.9.0 442 TRITON=yes 443 # To ensure that any ROCm config will build using conda cmake 444 # and thus have LAPACK/MKL enabled 445 CONDA_CMAKE=yes 446 fi 447 if [[ "$image" == *centos7* ]]; then 448 NINJA_VERSION=1.10.2 449 fi 450 if [[ "$image" == *gcc* ]]; then 451 extract_version_from_image_name gcc GCC_VERSION 452 fi 453 if [[ "$image" == *clang* ]]; then 454 extract_version_from_image_name clang CLANG_VERSION 455 fi 456 if [[ "$image" == *devtoolset* ]]; then 457 extract_version_from_image_name devtoolset DEVTOOLSET_VERSION 458 fi 459 if [[ "$image" == *glibc* ]]; then 460 extract_version_from_image_name glibc GLIBC_VERSION 461 fi 462 if [[ "$image" == *cmake* ]]; then 463 extract_version_from_image_name cmake CMAKE_VERSION 464 fi 465 ;; 466esac 467 468tmp_tag=$(basename "$(mktemp -u)" | tr '[:upper:]' '[:lower:]') 469 470#when using cudnn version 8 install it separately from cuda 471if [[ "$image" == *cuda* && ${OS} == "ubuntu" ]]; then 472 IMAGE_NAME="nvidia/cuda:${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel-ubuntu${UBUNTU_VERSION}" 473 if [[ ${CUDNN_VERSION} == 9 ]]; then 474 IMAGE_NAME="nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}" 475 fi 476fi 477 478# Build image 479docker build \ 480 --no-cache \ 481 --progress=plain \ 482 --build-arg "BUILD_ENVIRONMENT=${image}" \ 483 --build-arg "PROTOBUF=${PROTOBUF:-}" \ 484 --build-arg "LLVMDEV=${LLVMDEV:-}" \ 485 --build-arg "DB=${DB:-}" \ 486 --build-arg "VISION=${VISION:-}" \ 487 --build-arg "UBUNTU_VERSION=${UBUNTU_VERSION}" \ 488 --build-arg "CENTOS_VERSION=${CENTOS_VERSION}" \ 489 --build-arg "DEVTOOLSET_VERSION=${DEVTOOLSET_VERSION}" \ 490 --build-arg "GLIBC_VERSION=${GLIBC_VERSION}" \ 491 --build-arg "CLANG_VERSION=${CLANG_VERSION}" \ 492 --build-arg "ANACONDA_PYTHON_VERSION=${ANACONDA_PYTHON_VERSION}" \ 493 --build-arg "GCC_VERSION=${GCC_VERSION}" \ 494 --build-arg "CUDA_VERSION=${CUDA_VERSION}" \ 495 --build-arg "CUDNN_VERSION=${CUDNN_VERSION}" \ 496 --build-arg "TENSORRT_VERSION=${TENSORRT_VERSION}" \ 497 --build-arg "ANDROID=${ANDROID}" \ 498 --build-arg "ANDROID_NDK=${ANDROID_NDK_VERSION}" \ 499 --build-arg "GRADLE_VERSION=${GRADLE_VERSION}" \ 500 --build-arg "VULKAN_SDK_VERSION=${VULKAN_SDK_VERSION}" \ 501 --build-arg "SWIFTSHADER=${SWIFTSHADER}" \ 502 --build-arg "CMAKE_VERSION=${CMAKE_VERSION:-}" \ 503 --build-arg "NINJA_VERSION=${NINJA_VERSION:-}" \ 504 --build-arg "KATEX=${KATEX:-}" \ 505 --build-arg "ROCM_VERSION=${ROCM_VERSION:-}" \ 506 --build-arg "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH:-gfx906;gfx90a}" \ 507 --build-arg "IMAGE_NAME=${IMAGE_NAME}" \ 508 --build-arg "UCX_COMMIT=${UCX_COMMIT}" \ 509 --build-arg "UCC_COMMIT=${UCC_COMMIT}" \ 510 --build-arg "CONDA_CMAKE=${CONDA_CMAKE}" \ 511 --build-arg "TRITON=${TRITON}" \ 512 --build-arg "ONNX=${ONNX}" \ 513 --build-arg "DOCS=${DOCS}" \ 514 --build-arg "INDUCTOR_BENCHMARKS=${INDUCTOR_BENCHMARKS}" \ 515 --build-arg "EXECUTORCH=${EXECUTORCH}" \ 516 --build-arg "HALIDE=${HALIDE}" \ 517 --build-arg "XPU_VERSION=${XPU_VERSION}" \ 518 --build-arg "ACL=${ACL:-}" \ 519 --build-arg "SKIP_SCCACHE_INSTALL=${SKIP_SCCACHE_INSTALL:-}" \ 520 --build-arg "SKIP_LLVM_SRC_BUILD_INSTALL=${SKIP_LLVM_SRC_BUILD_INSTALL:-}" \ 521 -f $(dirname ${DOCKERFILE})/Dockerfile \ 522 -t "$tmp_tag" \ 523 "$@" \ 524 . 525 526# NVIDIA dockers for RC releases use tag names like `11.0-cudnn9-devel-ubuntu18.04-rc`, 527# for this case we will set UBUNTU_VERSION to `18.04-rc` so that the Dockerfile could 528# find the correct image. As a result, here we have to replace the 529# "$UBUNTU_VERSION" == "18.04-rc" 530# with 531# "$UBUNTU_VERSION" == "18.04" 532UBUNTU_VERSION=$(echo ${UBUNTU_VERSION} | sed 's/-rc$//') 533 534function drun() { 535 docker run --rm "$tmp_tag" $* 536} 537 538if [[ "$OS" == "ubuntu" ]]; then 539 540 if !(drun lsb_release -a 2>&1 | grep -qF Ubuntu); then 541 echo "OS=ubuntu, but:" 542 drun lsb_release -a 543 exit 1 544 fi 545 if !(drun lsb_release -a 2>&1 | grep -qF "$UBUNTU_VERSION"); then 546 echo "UBUNTU_VERSION=$UBUNTU_VERSION, but:" 547 drun lsb_release -a 548 exit 1 549 fi 550fi 551 552if [ -n "$ANACONDA_PYTHON_VERSION" ]; then 553 if !(drun python --version 2>&1 | grep -qF "Python $ANACONDA_PYTHON_VERSION"); then 554 echo "ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION, but:" 555 drun python --version 556 exit 1 557 fi 558fi 559 560if [ -n "$GCC_VERSION" ]; then 561 if !(drun gcc --version 2>&1 | grep -q " $GCC_VERSION\\W"); then 562 echo "GCC_VERSION=$GCC_VERSION, but:" 563 drun gcc --version 564 exit 1 565 fi 566fi 567 568if [ -n "$CLANG_VERSION" ]; then 569 if !(drun clang --version 2>&1 | grep -qF "clang version $CLANG_VERSION"); then 570 echo "CLANG_VERSION=$CLANG_VERSION, but:" 571 drun clang --version 572 exit 1 573 fi 574fi 575 576if [ -n "$KATEX" ]; then 577 if !(drun katex --version); then 578 echo "KATEX=$KATEX, but:" 579 drun katex --version 580 exit 1 581 fi 582fi 583