pkg_helpers.bash 14 KB
Newer Older
1
2
3
4
# A set of useful bash functions for common functionality we need to do in
# many build scripts


Edward Z. Yang's avatar
Edward Z. Yang committed
5
6
7
8
9
# Setup CUDA environment variables, based on CU_VERSION
#
# Inputs:
#   CU_VERSION (cpu, cu92, cu100)
#   NO_CUDA_PACKAGE (bool)
10
#   BUILD_TYPE (conda, wheel)
Edward Z. Yang's avatar
Edward Z. Yang committed
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
#
# Outputs:
#   VERSION_SUFFIX (e.g., "")
#   PYTORCH_VERSION_SUFFIX (e.g., +cpu)
#   WHEEL_DIR (e.g., cu100/)
#   CUDA_HOME (e.g., /usr/local/cuda-9.2, respected by torch.utils.cpp_extension)
#   FORCE_CUDA (respected by torchvision setup.py)
#   NVCC_FLAGS (respected by torchvision setup.py)
#
# Precondition: CUDA versions are installed in their conventional locations in
# /usr/local/cuda-*
#
# NOTE: Why VERSION_SUFFIX versus PYTORCH_VERSION_SUFFIX?  If you're building
# a package with CUDA on a platform we support CUDA on, VERSION_SUFFIX ==
# PYTORCH_VERSION_SUFFIX and everyone is happy.  However, if you are building a
# package with only CPU bits (e.g., torchaudio), then VERSION_SUFFIX is always
# empty, but PYTORCH_VERSION_SUFFIX is +cpu (because that's how you get a CPU
# version of a Python package.  But that doesn't apply if you're on OS X,
# since the default CU_VERSION on OS X is cpu.
setup_cuda() {
31
32
33
34
35
36
37
38

  # First, compute version suffixes.  By default, assume no version suffixes
  export VERSION_SUFFIX=""
  export PYTORCH_VERSION_SUFFIX=""
  export WHEEL_DIR=""
  # Wheel builds need suffixes (but not if they're on OS X, which never has suffix)
  if [[ "$BUILD_TYPE" == "wheel" ]] && [[ "$(uname)" != Darwin ]]; then
    # The default CUDA has no suffix
39
    if [[ "$CU_VERSION" != "cu102" ]]; then
40
      export PYTORCH_VERSION_SUFFIX="+$CU_VERSION"
41
    fi
42
43
44
45
    # Match the suffix scheme of pytorch, unless this package does not have
    # CUDA builds (in which case, use default)
    if [[ -z "$NO_CUDA_PACKAGE" ]]; then
      export VERSION_SUFFIX="$PYTORCH_VERSION_SUFFIX"
46
      export WHEEL_DIR="$CU_VERSION/"
Edward Z. Yang's avatar
Edward Z. Yang committed
47
    fi
48
  fi
49
50
51

  # Now work out the CUDA settings
  case "$CU_VERSION" in
peterjc123's avatar
peterjc123 committed
52
53
54
55
56
57
58
59
60
61
62
    cu110)
      if [[ "$OSTYPE" == "msys" ]]; then
        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.0"
      else
        export CUDA_HOME=/usr/local/cuda-11.0/
      fi
      export FORCE_CUDA=1
      # Hard-coding gencode flags is temporary situation until
      # https://github.com/pytorch/pytorch/pull/23408 lands
      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_50,code=compute_50"
      ;;
63
64
65
66
67
68
69
70
71
72
73
    cu102)
      if [[ "$OSTYPE" == "msys" ]]; then
        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v10.2"
      else
        export CUDA_HOME=/usr/local/cuda-10.2/
      fi
      export FORCE_CUDA=1
      # Hard-coding gencode flags is temporary situation until
      # https://github.com/pytorch/pytorch/pull/23408 lands
      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
      ;;
Francisco Massa's avatar
Francisco Massa committed
74
    cu101)
75
76
77
78
79
      if [[ "$OSTYPE" == "msys" ]]; then
        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v10.1"
      else
        export CUDA_HOME=/usr/local/cuda-10.1/
      fi
Francisco Massa's avatar
Francisco Massa committed
80
81
82
83
84
      export FORCE_CUDA=1
      # Hard-coding gencode flags is temporary situation until
      # https://github.com/pytorch/pytorch/pull/23408 lands
      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
      ;;
85
    cu100)
86
87
88
89
90
      if [[ "$OSTYPE" == "msys" ]]; then
        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v10.0"
      else
        export CUDA_HOME=/usr/local/cuda-10.0/
      fi
91
92
93
94
95
96
      export FORCE_CUDA=1
      # Hard-coding gencode flags is temporary situation until
      # https://github.com/pytorch/pytorch/pull/23408 lands
      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
      ;;
    cu92)
97
98
99
100
101
      if [[ "$OSTYPE" == "msys" ]]; then
        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v9.2"
      else
        export CUDA_HOME=/usr/local/cuda-9.2/
      fi
102
103
104
105
106
107
108
109
110
111
      export FORCE_CUDA=1
      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_50,code=compute_50"
      ;;
    cpu)
      ;;
    *)
      echo "Unrecognized CU_VERSION=$CU_VERSION"
      exit 1
      ;;
  esac
112
113
}

Edward Z. Yang's avatar
Edward Z. Yang committed
114
115
116
117
# Populate build version if necessary, and add version suffix
#
# Inputs:
#   BUILD_VERSION (e.g., 0.2.0 or empty)
118
#   VERSION_SUFFIX (e.g., +cpu)
Edward Z. Yang's avatar
Edward Z. Yang committed
119
120
121
122
123
124
#
# Outputs:
#   BUILD_VERSION (e.g., 0.2.0.dev20190807+cpu)
#
# Fill BUILD_VERSION if it doesn't exist already with a nightly string
# Usage: setup_build_version 0.2.0
125
126
setup_build_version() {
  if [[ -z "$BUILD_VERSION" ]]; then
Edward Z. Yang's avatar
Edward Z. Yang committed
127
128
129
    export BUILD_VERSION="$1.dev$(date "+%Y%m%d")$VERSION_SUFFIX"
  else
    export BUILD_VERSION="$BUILD_VERSION$VERSION_SUFFIX"
130
  fi
131
132
133
134
135
136

  # Set build version based on tag if on tag
  if [[ -n "${CIRCLE_TAG}" ]]; then
    # Strip tag
    export BUILD_VERSION="$(echo "${CIRCLE_TAG}" | sed -e 's/^v//' -e 's/-.*$//')"
  fi
137
138
139
140
141
142
143
144
145
}

# Set some useful variables for OS X, if applicable
setup_macos() {
  if [[ "$(uname)" == Darwin ]]; then
    export MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++
  fi
}

146
147
148
149
150
151
152
153
154
# set variable to determine whether the typing library needs to be built in
setup_typing() {
  if [[ "$PYTHON_VERSION" == 3.5 ]]; then
    export CONDA_TYPING_CONSTRAINT="- typing"
  else
    export CONDA_TYPING_CONSTRAINT=""
  fi
}

Edward Z. Yang's avatar
Edward Z. Yang committed
155
156
157
158
159
160
161
# Top-level entry point for things every package will need to do
#
# Usage: setup_env 0.2.0
setup_env() {
  setup_cuda
  setup_build_version "$1"
  setup_macos
162
  setup_typing
Edward Z. Yang's avatar
Edward Z. Yang committed
163
164
}

165
166
167
168
169
# Function to retry functions that sometimes timeout or have flaky failures
retry () {
    $*  || (sleep 1 && $*) || (sleep 2 && $*) || (sleep 4 && $*) || (sleep 8 && $*)
}

Edward Z. Yang's avatar
Edward Z. Yang committed
170
171
172
173
174
175
176
177
178
# Inputs:
#   PYTHON_VERSION (2.7, 3.5, 3.6, 3.7)
#   UNICODE_ABI (bool)
#
# Outputs:
#   PATH modified to put correct Python version in PATH
#
# Precondition: If Linux, you are in a soumith/manylinux-cuda* Docker image
setup_wheel_python() {
179
  if [[ "$(uname)" == Darwin || "$OSTYPE" == "msys" ]]; then
Edward Z. Yang's avatar
Edward Z. Yang committed
180
181
182
183
    eval "$(conda shell.bash hook)"
    conda env remove -n "env$PYTHON_VERSION" || true
    conda create -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
    conda activate "env$PYTHON_VERSION"
184
    # Install libpng from Anaconda (defaults)
185
    conda install libpng jpeg -y
Edward Z. Yang's avatar
Edward Z. Yang committed
186
  else
187
188
    # Install native CentOS libJPEG, LAME, freetype and GnuTLS
    yum install -y libjpeg-turbo-devel lame freetype gnutls
Edward Z. Yang's avatar
Edward Z. Yang committed
189
190
191
192
193
194
195
196
197
198
199
    case "$PYTHON_VERSION" in
      2.7)
        if [[ -n "$UNICODE_ABI" ]]; then
          python_abi=cp27-cp27mu
        else
          python_abi=cp27-cp27m
        fi
        ;;
      3.5) python_abi=cp35-cp35m ;;
      3.6) python_abi=cp36-cp36m ;;
      3.7) python_abi=cp37-cp37m ;;
200
      3.8) python_abi=cp38-cp38 ;;
Edward Z. Yang's avatar
Edward Z. Yang committed
201
202
203
204
205
      *)
        echo "Unrecognized PYTHON_VERSION=$PYTHON_VERSION"
        exit 1
        ;;
    esac
206
207
208
209
210
211
212
    # Download all the dependencies required to compile image and video_reader
    # extensions

    mkdir -p ext_libraries
    pushd ext_libraries
    popd
    export PATH="/opt/python/$python_abi/bin:$(pwd)/ext_libraries/bin:$PATH"
Edward Z. Yang's avatar
Edward Z. Yang committed
213
214
215
  fi
}

216
217
218
219
220
221
222
223
224
225
226
# Install with pip a bit more robustly than the default
pip_install() {
  retry pip install --progress-bar off "$@"
}

# Install torch with pip, respecting PYTORCH_VERSION, and record the installed
# version into PYTORCH_VERSION, if applicable
setup_pip_pytorch_version() {
  if [[ -z "$PYTORCH_VERSION" ]]; then
    # Install latest prerelease version of torch, per our nightlies, consistent
    # with the requested cuda version
Edward Z. Yang's avatar
Edward Z. Yang committed
227
    pip_install --pre torch -f "https://download.pytorch.org/whl/nightly/${WHEEL_DIR}torch_nightly.html"
228
229
230
231
232
233
234
235
    if [[ "$CUDA_VERSION" == "cpu" ]]; then
      # CUDA and CPU are ABI compatible on the CPU-only parts, so strip
      # in this case
      export PYTORCH_VERSION="$(pip show torch | grep ^Version: | sed 's/Version:  *//' | sed 's/+.\+//')"
    else
      export PYTORCH_VERSION="$(pip show torch | grep ^Version: | sed 's/Version:  *//')"
    fi
  else
236
    pip_install "torch==$PYTORCH_VERSION$PYTORCH_VERSION_SUFFIX" \
237
238
239
      -f "https://download.pytorch.org/whl/${CU_VERSION}/torch_stable.html" \
      -f "https://download.pytorch.org/whl/test/${CU_VERSION}/torch_test.html" \
      -f "https://download.pytorch.org/whl/nightly/${CU_VERSION}/torch_nightly.html"
240
241
242
243
244
245
  fi
}

# Fill PYTORCH_VERSION with the latest conda nightly version, and
# CONDA_CHANNEL_FLAGS with appropriate flags to retrieve these versions
#
246
# You MUST have populated PYTORCH_VERSION_SUFFIX before hand.
247
248
setup_conda_pytorch_constraint() {
  if [[ -z "$PYTORCH_VERSION" ]]; then
249
    export CONDA_CHANNEL_FLAGS="-c pytorch-nightly -c pytorch"
Francisco Massa's avatar
Francisco Massa committed
250
251
    export PYTORCH_VERSION="$(conda search --json 'pytorch[channel=pytorch-nightly]' | \
                              python -c "import os, sys, json, re; cuver = os.environ.get('CU_VERSION'); \
252
253
                               cuver_1 = cuver.replace('cu', 'cuda') if cuver != 'cpu' else cuver; \
                               cuver_2 = (cuver[:-1] + '.' + cuver[-1]).replace('cu', 'cuda') if cuver != 'cpu' else cuver; \
Francisco Massa's avatar
Francisco Massa committed
254
255
                               print(re.sub(r'\\+.*$', '', \
                                [x['version'] for x in json.load(sys.stdin)['pytorch'] \
256
                                  if (x['platform'] == 'darwin' or cuver_1 in x['fn'] or cuver_2 in x['fn']) \
Francisco Massa's avatar
Francisco Massa committed
257
258
259
260
261
262
                                    and 'py' + os.environ['PYTHON_VERSION'] in x['fn']][-1]))")"
    if [[ -z "$PYTORCH_VERSION" ]]; then
      echo "PyTorch version auto detection failed"
      echo "No package found for CU_VERSION=$CU_VERSION and PYTHON_VERSION=$PYTHON_VERSION"
      exit 1
    fi
263
  else
264
    export CONDA_CHANNEL_FLAGS="-c pytorch -c pytorch-nightly -c pytorch-test"
265
  fi
Edward Z. Yang's avatar
Edward Z. Yang committed
266
267
  if [[ "$CU_VERSION" == cpu ]]; then
    export CONDA_PYTORCH_BUILD_CONSTRAINT="- pytorch==$PYTORCH_VERSION${PYTORCH_VERSION_SUFFIX}"
268
269
    export CONDA_PYTORCH_CONSTRAINT="- pytorch==$PYTORCH_VERSION"
  else
Edward Z. Yang's avatar
Edward Z. Yang committed
270
271
    export CONDA_PYTORCH_BUILD_CONSTRAINT="- pytorch==${PYTORCH_VERSION}${PYTORCH_VERSION_SUFFIX}"
    export CONDA_PYTORCH_CONSTRAINT="- pytorch==${PYTORCH_VERSION}${PYTORCH_VERSION_SUFFIX}"
272
  fi
273
274
275
  if [[ "$OSTYPE" == msys && "$CU_VERSION" == cu92 ]]; then
    export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c defaults -c numba/label/dev"
  fi
276
277
278
279
280
281
282
283
}

# Translate CUDA_VERSION into CUDA_CUDATOOLKIT_CONSTRAINT
setup_conda_cudatoolkit_constraint() {
  export CONDA_CPUONLY_FEATURE=""
  if [[ "$(uname)" == Darwin ]]; then
    export CONDA_CUDATOOLKIT_CONSTRAINT=""
  else
Edward Z. Yang's avatar
Edward Z. Yang committed
284
    case "$CU_VERSION" in
peterjc123's avatar
peterjc123 committed
285
286
287
      cu110)
        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.0,<11.1 # [not osx]"
        ;;
288
289
290
      cu102)
        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.2,<10.3 # [not osx]"
        ;;
Francisco Massa's avatar
Francisco Massa committed
291
292
293
      cu101)
        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.1,<10.2 # [not osx]"
        ;;
Edward Z. Yang's avatar
Edward Z. Yang committed
294
      cu100)
295
296
        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.0,<10.1 # [not osx]"
        ;;
Edward Z. Yang's avatar
Edward Z. Yang committed
297
      cu92)
298
299
300
301
302
303
        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=9.2,<9.3 # [not osx]"
        ;;
      cpu)
        export CONDA_CUDATOOLKIT_CONSTRAINT=""
        export CONDA_CPUONLY_FEATURE="- cpuonly"
        ;;
Edward Z. Yang's avatar
Edward Z. Yang committed
304
305
306
307
      *)
        echo "Unrecognized CU_VERSION=$CU_VERSION"
        exit 1
        ;;
308
309
310
    esac
  fi
}
311

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
setup_conda_cudatoolkit_plain_constraint() {
  export CONDA_CPUONLY_FEATURE=""
  export CMAKE_USE_CUDA=1
  if [[ "$(uname)" == Darwin ]]; then
    export CONDA_CUDATOOLKIT_CONSTRAINT=""
    export CMAKE_USE_CUDA=0
  else
    case "$CU_VERSION" in
      cu102)
        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=10.2"
        ;;
      cu101)
        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=10.1"
        ;;
      cu100)
        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=10.0"
        ;;
      cu92)
        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=9.2"
        ;;
      cpu)
        export CONDA_CUDATOOLKIT_CONSTRAINT=""
        export CONDA_CPUONLY_FEATURE="cpuonly"
        export CMAKE_USE_CUDA=0
        ;;
      *)
        echo "Unrecognized CU_VERSION=$CU_VERSION"
        exit 1
        ;;
    esac
  fi
}

345
346
347
# Build the proper compiler package before building the final package
setup_visual_studio_constraint() {
  if [[ "$OSTYPE" == "msys" ]]; then
348
      export VSTOOLCHAIN_PACKAGE=vs$VC_YEAR
349
350
351
352
      conda build $CONDA_CHANNEL_FLAGS --no-anaconda-upload packaging/$VSTOOLCHAIN_PACKAGE
      cp packaging/$VSTOOLCHAIN_PACKAGE/conda_build_config.yaml packaging/torchvision/conda_build_config.yaml
  fi
}
353
354
355
356
357
358

setup_junit_results_folder() {
  if [[ "$CI" == "true" ]]; then
    export CONDA_PYTORCH_BUILD_RESULTS_DIRECTORY="${SOURCE_ROOT_DIR}/build_results/results.xml"
  fi
}
359
360
361
362


download_copy_ffmpeg() {
  if [[ "$OSTYPE" == "msys" ]]; then
363
    # conda install -yq ffmpeg=4.2 -c pytorch
364
365
366
367
368
369
370
371
372
    # curl -L -q https://anaconda.org/pytorch/ffmpeg/4.3/download/win-64/ffmpeg-4.3-ha925a31_0.tar.bz2 --output ffmpeg-4.3-ha925a31_0.tar.bz2
    # bzip2 --decompress --stdout ffmpeg-4.3-ha925a31_0.tar.bz2 | tar -x --file=-
    # cp Library/bin/*.dll ../torchvision
    echo "FFmpeg is disabled currently on Windows"
  else
    if [[ "$(uname)" == Darwin ]]; then
      conda install -yq ffmpeg=4.2 -c pytorch
      conda install -yq wget
    else
373
374
375
376
377
378
379
380
      # pushd ext_libraries
      # wget -q https://anaconda.org/pytorch/ffmpeg/4.2/download/linux-64/ffmpeg-4.2-hf484d3e_0.tar.bz2
      # tar -xjvf ffmpeg-4.2-hf484d3e_0.tar.bz2
      # rm -rf ffmpeg-4.2-hf484d3e_0.tar.bz2
      # ldconfig
      # which ffmpeg
      # popd
      echo "FFmpeg is disabled currently on Linux"
381
382
383
    fi
  fi
}