admin/containers/scripted_gmx_docker_builds.py

   1 #!/usr/bin/env python
   2 #
   3 # This file is part of the GROMACS molecular simulation package.
   4 #
   5 # Copyright (c) 2020,2021, by the GROMACS development team, led by
   6 # Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   7 # and including many others, as listed in the AUTHORS file in the
   8 # top-level source directory and at http://www.gromacs.org.
   9 #
  10 # GROMACS is free software; you can redistribute it and/or
  11 # modify it under the terms of the GNU Lesser General Public License
  12 # as published by the Free Software Foundation; either version 2.1
  13 # of the License, or (at your option) any later version.
  14 #
  15 # GROMACS is distributed in the hope that it will be useful,
  16 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  17 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18 # Lesser General Public License for more details.
  19 #
  20 # You should have received a copy of the GNU Lesser General Public
  21 # License along with GROMACS; if not, see
  22 # http://www.gnu.org/licenses, or write to the Free Software Foundation,
  23 # Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  24 #
  25 # If you want to redistribute modifications to GROMACS, please
  26 # consider that scientific software is very special. Version
  27 # control is crucial - bugs must be traceable. We will be happy to
  28 # consider code for inclusion in the official distribution, but
  29 # derived work must not be called official GROMACS. Details are found
  30 # in the README & COPYING files - if they are missing, get the
  31 # official version at http://www.gromacs.org.
  32 #
  33 # To help us fund GROMACS development, we humbly ask that you cite
  34 # the research papers on the package. Check out http://www.gromacs.org.
  35
  36 """Building block based Dockerfile generation for CI testing images.
  37
  38 Generates a set of docker images used for running GROMACS CI on Gitlab.
  39 The images are prepared according to a selection of build configuration targets
  40 that hope to cover a broad enough scope of different possible systems,
  41 allowing us to check compiler types and versions, as well as libraries used
  42 for accelerators and parallel communication systems. Each combinations is
  43 described as an entry in the build_configs dictionary, with the script
  44 analysing the logic and adding build stages as needed.
  45
  46 Based on the example script provided by the NVidia HPCCM repository.
  47
  48 Reference:
  49     `NVidia HPC Container Maker <https://github.com/NVIDIA/hpc-container-maker>`__
  50
  51 Authors:
  52     * Paul Bauer <paul.bauer.q@gmail.com>
  53     * Eric Irrgang <ericirrgang@gmail.com>
  54     * Joe Jordan <e.jjordan12@gmail.com>
  55     * Mark Abraham <mark.j.abraham@gmail.com>
  56     * Gaurav Garg <gaugarg@nvidia.com>
  57
  58 Usage::
  59
  60     $ python3 scripted_gmx_docker_builds.py --help
  61     $ python3 scripted_gmx_docker_builds.py --format docker > Dockerfile && docker build .
  62     $ python3 scripted_gmx_docker_builds.py | docker build -
  63
  64 See Also:
  65     :file:`buildall.sh`
  66
  67 """
  68
  69 import argparse
  70 import collections
  71 import collections.abc
  72 import typing
  73 from distutils.version import StrictVersion
  74
  75 import hpccm
  76 import hpccm.config
  77 from hpccm.building_blocks.base import bb_base
  78
  79 try:
  80     import utility
  81 except ImportError:
  82     raise RuntimeError(
  83         'This module assumes availability of supporting modules in the same directory. Add the directory to '
  84         'PYTHONPATH or invoke Python from within the module directory so module location can be resolved.')
  85
  86 # Basic packages for all final images.
  87 _common_packages = ['build-essential',
  88                     'ca-certificates',
  89                     'ccache',
  90                     'git',
  91                     'gnupg',
  92                     'gpg-agent',
  93                     'libfftw3-dev',
  94                     'libhwloc-dev',
  95                     'liblapack-dev',
  96                     'libx11-dev',
  97                     'moreutils',
  98                     'ninja-build',
  99                     'rsync',
 100                     'valgrind',
 101                     'vim',
 102                     'wget',
 103                     'xsltproc']
 104
 105 _opencl_extra_packages = [
 106     'nvidia-opencl-dev',
 107     # The following require apt_ppas=['ppa:intel-opencl/intel-opencl']
 108     'intel-opencl-icd',
 109     'ocl-icd-libopencl1',
 110     'ocl-icd-opencl-dev',
 111     'opencl-headers',
 112 ]
 113
 114 _rocm_extra_packages = [
 115     # The following require
 116     #             apt_keys=['http://repo.radeon.com/rocm/rocm.gpg.key'],
 117     #             apt_repositories=['deb [arch=amd64] http://repo.radeon.com/rocm/apt/4.0.1/ xenial main']
 118     'clinfo',
 119     'hipfft',
 120     'libelf1',
 121     'rocfft',
 122     'rocm-opencl',
 123     'rocm-dev',
 124 ]
 125
 126 # Extra packages needed to build Intel Compute Runtime
 127 _intel_compute_runtime_extra_packages = ['intel-opencl-icd',
 128                                          'intel-level-zero-gpu',
 129                                          'level-zero',
 130                                          'libmfx1']
 131
 132 # Extra packages needed to build Python installations from source.
 133 _python_extra_packages = ['build-essential',
 134                           'ca-certificates',
 135                           'ccache',
 136                           'curl',
 137                           'git',
 138                           'libbz2-dev',
 139                           'libffi-dev',
 140                           'liblzma-dev',
 141                           'libncurses5-dev',
 142                           'libncursesw5-dev',
 143                           'libreadline-dev',
 144                           'libsqlite3-dev',
 145                           'libssl-dev',
 146                           'llvm',
 147                           'python-openssl',
 148                           'vim',
 149                           'wget',
 150                           'zlib1g-dev']
 151
 152 # Extra packages needed for images for building documentation.
 153 _docs_extra_packages = ['autoconf',
 154                         'automake',
 155                         'autopoint',
 156                         'autotools-dev',
 157                         'bison',
 158                         'flex',
 159                         'ghostscript',
 160                         'graphviz',
 161                         'help2man',
 162                         'imagemagick',
 163                         'libtool',
 164                         'mscgen',
 165                         'm4',
 166                         'openssh-client',
 167                         'texinfo',
 168                         'texlive-latex-base',
 169                         'texlive-latex-extra',
 170                         'texlive-fonts-recommended',
 171                         'texlive-fonts-extra',
 172                         'tex-gyre']
 173
 174 # Parse command line arguments
 175 parser = argparse.ArgumentParser(description='GROMACS CI image creation script',
 176                                  parents=[utility.parser])
 177
 178 parser.add_argument('--format', type=str, default='docker',
 179                     choices=['docker', 'singularity'],
 180                     help='Container specification format (default: docker)')
 181
 182
 183 def base_image_tag(args) -> str:
 184     """Generate *image* for hpccm.baseimage()."""
 185     # Check if we use CUDA images or plain linux images
 186     if args.cuda is not None:
 187         cuda_version_tag = 'nvidia/cuda:' + args.cuda + '-devel'
 188         if args.centos is not None:
 189             cuda_version_tag += '-centos' + args.centos
 190         elif args.ubuntu is not None:
 191             cuda_version_tag += '-ubuntu' + args.ubuntu
 192         else:
 193             raise RuntimeError('Logic error: no Linux distribution selected.')
 194
 195         base_image_tag = cuda_version_tag
 196     else:
 197         if args.centos is not None:
 198             base_image_tag = 'centos:centos' + args.centos
 199         elif args.ubuntu is not None:
 200             base_image_tag = 'ubuntu:' + args.ubuntu
 201         else:
 202             raise RuntimeError('Logic error: no Linux distribution selected.')
 203     return base_image_tag
 204
 205
 206 def hpccm_distro_name(args) -> str:
 207     """Generate *_distro* for hpccm.baseimage().
 208
 209     Convert the linux distribution variables into something that hpccm
 210     understands.
 211
 212     The same format is used by the lower level hpccm.config.set_linux_distro().
 213     """
 214     if args.centos is not None:
 215         name_mapping = {'7': 'centos7',
 216                         '8': 'centos8'}
 217         if args.centos in name_mapping:
 218             hpccm_name = name_mapping[args.centos]
 219         else:
 220             raise RuntimeError('Logic error: unsupported CentOS distribution selected.')
 221     elif args.ubuntu is not None:
 222         name_mapping = {'20.04': 'ubuntu20',
 223                         '18.04': 'ubuntu18',
 224                         '16.04': 'ubuntu16'}
 225         if args.ubuntu in name_mapping:
 226             hpccm_name = name_mapping[args.ubuntu]
 227         else:
 228             raise RuntimeError('Logic error: unsupported Ubuntu distribution selected.')
 229     else:
 230         raise RuntimeError('Logic error: no Linux distribution selected.')
 231     return hpccm_name
 232
 233
 234 def get_llvm_packages(args) -> typing.Iterable[str]:
 235     # If we use the package version of LLVM, we need to install extra packages for it.
 236     if (args.llvm is not None) and (args.tsan is None):
 237         packages = [f'libomp-{args.llvm}-dev',
 238                     f'libomp5-{args.llvm}',
 239                     'clang-format-' + str(args.llvm),
 240                     'clang-tidy-' + str(args.llvm)]
 241         if args.hipsycl is not None:
 242             packages += [f'llvm-{args.llvm}-dev',
 243                          f'libclang-{args.llvm}-dev',
 244                          f'lld-{args.llvm}']
 245         return packages
 246     else:
 247         return []
 248
 249
 250 def get_opencl_packages(args) -> typing.List[str]:
 251     if (args.doxygen is None) and (args.oneapi is None):
 252         return _opencl_extra_packages
 253     else:
 254         return []
 255
 256
 257 def get_rocm_packages(args) -> typing.List[str]:
 258     if (args.rocm is None):
 259         return []
 260     else:
 261         return _rocm_extra_packages
 262
 263
 264 def get_compiler(args, compiler_build_stage: hpccm.Stage = None) -> bb_base:
 265     # Compiler
 266     if args.llvm is not None:
 267         # Build our own version instead to get TSAN + OMP
 268         if args.tsan is not None:
 269             if compiler_build_stage is not None:
 270                 compiler = compiler_build_stage.runtime(_from='tsan')
 271             else:
 272                 raise RuntimeError('No TSAN compiler build stage!')
 273         # Build the default compiler if we don't need special support
 274         else:
 275             # Currently the focal apt repositories do not contain
 276             # llvm higher than 11, so we work around that. This will
 277             # need further work when we start supporting ubuntu 22.04
 278             compiler = hpccm.building_blocks.llvm(version=args.llvm, upstream=True if int(args.llvm) > 11 else False)
 279
 280     elif args.oneapi is not None:
 281         if compiler_build_stage is not None:
 282             compiler = compiler_build_stage.runtime(_from='oneapi')
 283             # Prepare the toolchain (needed only for builds done within the Dockerfile, e.g.
 284             # OpenMPI builds, which don't currently work for other reasons)
 285             oneapi_toolchain = hpccm.toolchain(CC=f'/opt/intel/oneapi/compiler/{args.oneapi}/linux/bin/intel64/icx',
 286                                                CXX=f'/opt/intel/oneapi/compiler/{args.oneapi}/linux/bin/intel64/icpx')
 287             setattr(compiler, 'toolchain', oneapi_toolchain)
 288
 289         else:
 290             raise RuntimeError('No oneAPI compiler build stage!')
 291
 292     elif args.gcc is not None:
 293         compiler = hpccm.building_blocks.gnu(extra_repository=True,
 294                                              version=args.gcc,
 295                                              fortran=False)
 296     else:
 297         raise RuntimeError('Logic error: no compiler toolchain selected.')
 298     return compiler
 299
 300
 301 def get_gdrcopy(args, compiler):
 302     if args.cuda is not None:
 303         if hasattr(compiler, 'toolchain'):
 304             # Version last updated June 7, 2021
 305             return hpccm.building_blocks.gdrcopy(toolchain=compiler.toolchain, version="2.2")
 306         else:
 307             raise RuntimeError('compiler is not an HPCCM compiler building block!')
 308     else:
 309         return None
 310
 311
 312 def get_ucx(args, compiler, gdrcopy):
 313     if args.cuda is not None:
 314         if hasattr(compiler, 'toolchain'):
 315             use_gdrcopy = (gdrcopy is not None)
 316             # Version last updated June 7, 2021
 317             return hpccm.building_blocks.ucx(toolchain=compiler.toolchain, gdrcopy=use_gdrcopy, version="1.10.1",
 318                                              cuda=True)
 319         else:
 320             raise RuntimeError('compiler is not an HPCCM compiler building block!')
 321     else:
 322         return None
 323
 324
 325 def get_mpi(args, compiler, ucx):
 326     # If needed, add MPI to the image
 327     if args.mpi is not None:
 328         if args.mpi == 'openmpi':
 329             if hasattr(compiler, 'toolchain'):
 330                 if args.oneapi is not None:
 331                     raise RuntimeError('oneAPI building OpenMPI is not supported')
 332                 use_cuda = (args.cuda is not None)
 333                 use_ucx = (ucx is not None)
 334                 # Version last updated June 7, 2021
 335                 return hpccm.building_blocks.openmpi(toolchain=compiler.toolchain, version="4.1.1", cuda=use_cuda,
 336                                                      ucx=use_ucx, infiniband=False)
 337             else:
 338                 raise RuntimeError('compiler is not an HPCCM compiler building block!')
 339
 340         elif args.mpi == 'impi':
 341             # TODO Intel MPI from the oneAPI repo is not working reliably,
 342             # reasons are unclear. When solved, add packagages called:
 343             # 'intel-oneapi-mpi', 'intel-oneapi-mpi-devel'
 344             # during the compiler stage.
 345             # TODO also consider hpccm's intel_mpi package if that doesn't need
 346             # a license to run.
 347             raise RuntimeError('Intel MPI recipe not implemented yet.')
 348         else:
 349             raise RuntimeError('Requested unknown MPI implementation.')
 350     else:
 351         return None
 352
 353
 354 def get_clfft(args):
 355     if (args.clfft is not None):
 356         return hpccm.building_blocks.generic_cmake(
 357             repository='https://github.com/clMathLibraries/clFFT.git',
 358             prefix='/usr/local', recursive=True, branch=args.clfft, directory='clFFT/src')
 359     else:
 360         return None
 361
 362
 363 def get_heffte(args):
 364     if (args.heffte is not None):
 365         return hpccm.building_blocks.generic_cmake(
 366             cmake_opts=['-D CMAKE_BUILD_TYPE=Release',
 367                         '-D CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda',
 368                         '-D Heffte_ENABLE_CUDA=ON',
 369                         '-D Heffte_ENABLE_FFTW=OFF',
 370                         '-D BUILD_SHARED_LIBS=ON'],
 371             repository='https://bitbucket.org/icl/heffte.git',
 372             prefix='/usr/local', recursive=True, commit=args.heffte, directory='heffte')
 373     else:
 374         return None
 375
 376
 377 def get_hipsycl(args):
 378     if args.hipsycl is None:
 379         return None
 380     if args.llvm is None:
 381         raise RuntimeError('Can not build hipSYCL without llvm')
 382
 383     if args.rocm is None:
 384         raise RuntimeError('hipSYCL requires the rocm packages')
 385
 386     cmake_opts = ['-DLLVM_DIR=/opt/rocm/llvm/lib/cmake/llvm',
 387                   '-DCMAKE_PREFIX_PATH=/opt/rocm/lib/cmake',
 388                   '-DWITH_ROCM_BACKEND=ON']
 389     if args.cuda is not None:
 390         cmake_opts += ['-DCUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda',
 391                        '-DWITH_CUDA_BACKEND=ON']
 392
 393     postinstall = [
 394         # https://github.com/illuhad/hipSYCL/issues/361#issuecomment-718943645
 395         'for f in /opt/rocm/amdgcn/bitcode/*.bc; do ln -s "$f" "/opt/rocm/lib/$(basename $f .bc).amdgcn.bc"; done'
 396     ]
 397     if args.cuda is not None:
 398         postinstall += [
 399             # https://github.com/illuhad/hipSYCL/issues/410#issuecomment-743301929
 400             f'sed s/_OPENMP/__OPENMP_NVPTX__/ -i /usr/lib/llvm-{args.llvm}/lib/clang/*/include/__clang_cuda_complex_builtins.h',
 401             # Not needed unless we're building with CUDA 11.x, but no harm in doing always
 402             'ln -s /usr/local/cuda/compat/* /usr/local/cuda/lib64/'
 403         ]
 404
 405     return hpccm.building_blocks.generic_cmake(
 406         repository='https://github.com/illuhad/hipSYCL.git',
 407         directory='/var/tmp/hipSYCL',
 408         prefix='/usr/local', recursive=True, commit=args.hipsycl,
 409         cmake_opts=['-DCMAKE_BUILD_TYPE=Release', *cmake_opts],
 410         postinstall=postinstall)
 411
 412
 413 def add_tsan_compiler_build_stage(input_args, output_stages: typing.Mapping[str, hpccm.Stage]):
 414     """Isolate the expensive TSAN preparation stage.
 415
 416     This is a very expensive stage, but has few and disjoint dependencies, and
 417     its output is easily compartmentalized (/usr/local) so we can isolate this
 418     build stage to maximize build cache hits and reduce rebuild time, bookkeeping,
 419     and final image size.
 420     """
 421     if not isinstance(output_stages, collections.abc.MutableMapping):
 422         raise RuntimeError('Need output_stages container.')
 423     if 'compiler_build' in output_stages:
 424         raise RuntimeError('"compiler_build" output stage is already present.')
 425     tsan_stage = hpccm.Stage()
 426     tsan_stage += hpccm.primitives.baseimage(image=base_image_tag(input_args),
 427                                              _distro=hpccm_distro_name(input_args),
 428                                              _as='tsan')
 429
 430     tsan_stage += hpccm.building_blocks.packages(ospackages=['git', 'ca-certificates', 'build-essential', 'cmake'])
 431     # CMake will get duplicated later, but this is an expensive image, and it isn't worth optimizing
 432     # out that duplication...
 433     tsan_stage += hpccm.building_blocks.python(python3=True, python2=False, devel=False)
 434
 435     compiler_branch = 'release/' + str(input_args.llvm) + '.x'
 436     tsan_stage += hpccm.building_blocks.generic_cmake(
 437         repository='https://github.com/llvm/llvm-project.git',
 438         directory='/var/tmp/llvm-project/llvm/',
 439         prefix='/usr/local', recursive=True, branch=compiler_branch,
 440         cmake_opts=['-D CMAKE_BUILD_TYPE=Release',
 441                     '-D LLVM_ENABLE_PROJECTS="clang;openmp;clang-tools-extra;compiler-rt;lld"',
 442                     '-D LIBOMP_TSAN_SUPPORT=on'],
 443         postinstall=['ln -s /usr/local/bin/clang++ /usr/local/bin/clang++-' + str(input_args.llvm),
 444                      'ln -s /usr/local/bin/clang-format /usr/local/bin/clang-format-' + str(input_args.llvm),
 445                      'ln -s /usr/local/bin/clang-tidy /usr/local/bin/clang-tidy-' + str(input_args.llvm),
 446                      'ln -s /usr/local/share/clang/run-clang-tidy.py /usr/local/bin/run-clang-tidy-'
 447                      + str(input_args.llvm) + '.py',
 448                      'ln -s /usr/local/bin/run-clang-tidy-'
 449                      + str(input_args.llvm) + '.py /usr/local/bin/run-clang-tidy-' + str(input_args.llvm),
 450                      'ln -s /usr/local/libexec/c++-analyzer /usr/local/bin/c++-analyzer-' + str(input_args.llvm)])
 451     output_stages['compiler_build'] = tsan_stage
 452
 453
 454 def oneapi_runtime(_from='0'):
 455     oneapi_runtime_stage = hpccm.Stage()
 456     oneapi_runtime_stage += hpccm.primitives.copy(_from='oneapi-build',
 457                                                   files={"/opt/intel": "/opt/intel",
 458                                                          "/etc/bash.bashrc": "/etc/bash.bashrc"})
 459     return oneapi_runtime_stage
 460
 461
 462 def add_oneapi_compiler_build_stage(input_args, output_stages: typing.Mapping[str, hpccm.Stage]):
 463     """Isolate the oneAPI preparation stage.
 464
 465     This stage is isolated so that its installed components are minimized in the
 466     final image (chiefly /opt/intel) and its environment setup script can be
 467     sourced. This also helps with rebuild time and final image size.
 468     """
 469     if not isinstance(output_stages, collections.abc.MutableMapping):
 470         raise RuntimeError('Need output_stages container.')
 471     if 'compiler_build' in output_stages:
 472         raise RuntimeError('"compiler_build" output stage is already present.')
 473     oneapi_stage = hpccm.Stage()
 474     oneapi_stage += hpccm.primitives.baseimage(image=base_image_tag(input_args),
 475                                                _distro=hpccm_distro_name(input_args),
 476                                                _as='oneapi-build')
 477
 478     version = str(input_args.oneapi)
 479
 480     # Add required components for the next stage (both for hpccm and Intel's setvars.sh script)
 481     oneapi_stage += hpccm.building_blocks.packages(ospackages=['wget', 'gnupg2', 'ca-certificates', 'lsb-release'])
 482     oneapi_stage += hpccm.building_blocks.packages(
 483         apt_keys=['https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2023.PUB'],
 484         apt_repositories=['deb https://apt.repos.intel.com/oneapi all main'],
 485         # Add minimal packages (not the whole HPC toolkit!)
 486         ospackages=[f'intel-oneapi-dpcpp-cpp-{version}',
 487                     f'intel-oneapi-openmp-{version}',
 488                     f'intel-oneapi-mkl-{version}',
 489                     f'intel-oneapi-mkl-devel-{version}']
 490     )
 491     # Ensure that all bash shells on the final container will have access to oneAPI
 492     oneapi_stage += hpccm.primitives.shell(
 493         commands=['echo "source /opt/intel/oneapi/setvars.sh" >> /etc/bash.bashrc',
 494                   'unlink /opt/intel/oneapi/compiler/latest',
 495                   f'ln -sf /opt/intel/oneapi/compiler/{version} /opt/intel/oneapi/compiler/latest']
 496     )
 497     setattr(oneapi_stage, 'runtime', oneapi_runtime)
 498
 499     output_stages['compiler_build'] = oneapi_stage
 500
 501
 502 def prepare_venv(version: StrictVersion) -> typing.Sequence[str]:
 503     """Get shell commands to set up the venv for the requested Python version."""
 504     major = version.version[0]
 505     minor = version.version[1]  # type: int
 506
 507     pyenv = '$HOME/.pyenv/bin/pyenv'
 508
 509     py_ver = f'{major}.{minor}'
 510     venv_path = f'$HOME/venv/py{py_ver}'
 511     commands = [f'$({pyenv} prefix `{pyenv} whence python{py_ver}`)/bin/python -m venv {venv_path}']
 512
 513     commands.append(f'{venv_path}/bin/python -m pip install --upgrade pip setuptools')
 514     # Install dependencies for building and testing gmxapi Python package.
 515     # WARNING: Please keep this list synchronized with python_packaging/requirements-test.txt
 516     # TODO: Get requirements.txt from an input argument.
 517     commands.append(f"""{venv_path}/bin/python -m pip install --upgrade \
 518             'breathe' \
 519             'cmake>=3.16.3' \
 520             'flake8>=3.7.7' \
 521             'gcovr>=4.2' \
 522             'mpi4py>=3.0.3' \
 523             'networkx>=2.0' \
 524             'numpy>=1' \
 525             'pip>=10.1' \
 526             'pybind11>2.6' \
 527             'Pygments>=2.2.0' \
 528             'pytest>=3.9' \
 529             'setuptools>=42' \
 530             'scikit-build>=0.10' \
 531             'Sphinx>=1.6.3' \
 532             'sphinxcontrib-plantuml>=0.14' \
 533             'wheel'""")
 534     return commands
 535
 536
 537 def add_python_stages(building_blocks: typing.Mapping[str, bb_base],
 538                       input_args,
 539                       output_stages: typing.MutableMapping[str, hpccm.Stage]):
 540     """Add the stage(s) necessary for the requested venvs.
 541
 542     One intermediate build stage is created for each venv (see --venv option).
 543
 544     Each stage partially populates Python installations and venvs in the home
 545     directory. The home directory is collected by the 'pyenv' stage for use by
 546     the main build stage.
 547     """
 548     if len(input_args.venvs) < 1:
 549         raise RuntimeError('No venvs to build...')
 550     if output_stages is None or not isinstance(output_stages, collections.abc.Mapping):
 551         raise RuntimeError('Need a container for output stages.')
 552
 553     # Main Python stage that collects the environments from individual stages.
 554     # We collect the stages individually, rather than chaining them, because the
 555     # copy is a bit slow and wastes local Docker image space for each filesystem
 556     # layer.
 557     pyenv_stage = hpccm.Stage()
 558     pyenv_stage += hpccm.primitives.baseimage(image=base_image_tag(input_args),
 559                                               _distro=hpccm_distro_name(input_args),
 560                                               _as='pyenv')
 561     pyenv_stage += building_blocks['compiler']
 562     if building_blocks['gdrcopy'] is not None:
 563         pyenv_stage += building_blocks['gdrcopy']
 564     if building_blocks['ucx'] is not None:
 565         pyenv_stage += building_blocks['ucx']
 566     pyenv_stage += building_blocks['mpi']
 567     pyenv_stage += hpccm.building_blocks.packages(ospackages=_python_extra_packages)
 568
 569     for version in [StrictVersion(py_ver) for py_ver in sorted(input_args.venvs)]:
 570         stage_name = 'py' + str(version)
 571         stage = hpccm.Stage()
 572         stage += hpccm.primitives.baseimage(image=base_image_tag(input_args),
 573                                             _distro=hpccm_distro_name(input_args),
 574                                             _as=stage_name)
 575         stage += building_blocks['compiler']
 576         if building_blocks['gdrcopy'] is not None:
 577             stage += building_blocks['gdrcopy']
 578         if building_blocks['ucx'] is not None:
 579             stage += building_blocks['ucx']
 580         stage += building_blocks['mpi']
 581         stage += hpccm.building_blocks.packages(ospackages=_python_extra_packages)
 582
 583         # TODO: Use a non-root user for testing and Python virtual environments.
 584         stage += hpccm.primitives.shell(commands=[
 585             'curl https://pyenv.run | bash',
 586             """echo 'export PYENV_ROOT="$HOME/.pyenv"' >> $HOME/.bashrc""",
 587             """echo 'export PATH="$PYENV_ROOT/bin:$PATH"' >> $HOME/.bashrc""",
 588             """echo 'eval "$(pyenv init -)"' >> $HOME/.bashrc""",
 589             """echo 'eval "$(pyenv virtualenv-init -)"' >> $HOME/.bashrc"""])
 590         pyenv = '$HOME/.pyenv/bin/pyenv'
 591         commands = [f'PYTHON_CONFIGURE_OPTS="--enable-shared" {pyenv} install -s {version}']
 592         stage += hpccm.primitives.shell(commands=commands)
 593
 594         commands = prepare_venv(version)
 595         stage += hpccm.primitives.shell(commands=commands)
 596
 597         # TODO: Update user home directory.
 598         pyenv_stage += hpccm.primitives.copy(_from=stage_name, _mkdir=True, src=['/root/'],
 599                                              dest='/root')
 600
 601         # Add the intermediate build stage to the sequence
 602         output_stages[stage_name] = stage
 603
 604     # TODO: If we activate pyenv for login shells, the `global` "version" should be full-featured.
 605     # # `version` should be a system installation or pyenv environment (or pyenv-virtualenv)
 606     # # with the dependencies for all of the Python aspects of CMake-driven builds.
 607     # commands = '{pyenv} global {version}'.format(
 608     #             pyenv=pyenv,
 609     #             version=...)
 610     # pyenv_stage += hpccm.primitives.shell(commands=commands)
 611
 612     # Add the aggregating build stage to the sequence. This allows the main stage to copy
 613     # the files in a single stage, potentially reducing the overall output image size.
 614     output_stages['pyenv'] = pyenv_stage
 615
 616
 617 def add_documentation_dependencies(input_args,
 618                                    output_stages: typing.MutableMapping[str, hpccm.Stage]):
 619     """Add appropriate layers according to doxygen input arguments."""
 620     if input_args.doxygen is None:
 621         return
 622     # Always clone the same version of linkchecker (latest release at June 1, 2021)
 623     output_stages['main'] += hpccm.building_blocks.pip(pip='pip3', packages=[
 624         'git+https://github.com/linkchecker/linkchecker.git@v10.0.1'])
 625     output_stages['main'] += hpccm.primitives.shell(
 626         commands=[
 627             'sed -i \'/\"XPS\"/d;/\"PDF\"/d;/\"PS\"/d;/\"EPS\"/d;/disable ghostscript format types/d\' /etc/ImageMagick-6/policy.xml'])
 628     if input_args.doxygen == '1.8.5':
 629         doxygen_commit = 'ed4ed873ab0e7f15116e2052119a6729d4589f7a'
 630         output_stages['main'] += hpccm.building_blocks.generic_autotools(
 631             repository='https://github.com/westes/flex.git',
 632             commit='f7788a9a0ecccdc953ed12043ccb59ca25714018',
 633             prefix='/tmp/install-of-flex',
 634             configure_opts=['--disable-shared'],
 635             preconfigure=['./autogen.sh'])
 636         output_stages['main'] += hpccm.building_blocks.generic_autotools(
 637             repository='https://github.com/doxygen/doxygen.git',
 638             commit=doxygen_commit,
 639             prefix='',
 640             configure_opts=[
 641                 '--flex /tmp/install-of-flex/bin/flex',
 642                 '--static'])
 643     else:
 644         version = input_args.doxygen
 645         archive_name = f'doxygen-{version}.linux.bin.tar.gz'
 646         archive_url = f'https://sourceforge.net/projects/doxygen/files/rel-{version}/{archive_name}'
 647         binary_path = f'doxygen-{version}/bin/doxygen'
 648         commands = [
 649             'mkdir doxygen && cd doxygen',
 650             f'wget {archive_url}',
 651             f'tar xf {archive_name} {binary_path}',
 652             f'cp {binary_path} /usr/local/bin/',
 653             'cd .. && rm -rf doxygen'
 654         ]
 655         output_stages['main'] += hpccm.primitives.shell(commands=commands)
 656
 657
 658 def build_stages(args) -> typing.Iterable[hpccm.Stage]:
 659     """Define and sequence the stages for the recipe corresponding to *args*."""
 660
 661     # A Dockerfile or Singularity recipe can have multiple build stages.
 662     # The main build stage can copy files from previous stages, though only
 663     # the last stage is included in the tagged output image. This means that
 664     # large or expensive sets of build instructions can be isolated in
 665     # local/temporary images, but all of the stages need to be output by this
 666     # script, and need to occur in the correct order, so we create a sequence
 667     # object early in this function.
 668     stages = collections.OrderedDict()
 669
 670     # If we need TSAN or oneAPI support the early build is more complex,
 671     # so that our compiler images don't have all the cruft needed to get those things
 672     # installed.
 673     if args.llvm is not None and args.tsan is not None:
 674         add_tsan_compiler_build_stage(input_args=args, output_stages=stages)
 675     if args.oneapi is not None:
 676         add_oneapi_compiler_build_stage(input_args=args, output_stages=stages)
 677
 678     # Building blocks are chunks of container-builder instructions that can be
 679     # copied to any build stage with the addition operator.
 680     building_blocks = collections.OrderedDict()
 681     building_blocks['base_packages'] = hpccm.building_blocks.packages(
 682         ospackages=_common_packages)
 683
 684     # These are the most expensive and most reusable layers, so we put them first.
 685     building_blocks['compiler'] = get_compiler(args, compiler_build_stage=stages.get('compiler_build'))
 686     building_blocks['gdrcopy'] = get_gdrcopy(args, building_blocks['compiler'])
 687     building_blocks['ucx'] = get_ucx(args, building_blocks['compiler'], building_blocks['gdrcopy'])
 688     building_blocks['mpi'] = get_mpi(args, building_blocks['compiler'], building_blocks['ucx'])
 689     for i, cmake in enumerate(args.cmake):
 690         building_blocks['cmake' + str(i)] = hpccm.building_blocks.cmake(
 691             eula=True,
 692             prefix=f'/usr/local/cmake-{cmake}',
 693             version=cmake)
 694
 695     # Install additional packages early in the build to optimize Docker build layer cache.
 696     os_packages = list(get_llvm_packages(args)) + get_opencl_packages(args) + get_rocm_packages(args)
 697     if args.doxygen is not None:
 698         os_packages += _docs_extra_packages
 699     if args.oneapi is not None:
 700         os_packages += ['lsb-release']
 701     if args.hipsycl is not None:
 702         os_packages += ['libboost-fiber-dev']
 703     building_blocks['extra_packages'] = []
 704     if args.intel_compute_runtime:
 705         building_blocks['extra_packages'] += hpccm.building_blocks.packages(
 706             apt_keys=['https://repositories.intel.com/graphics/intel-graphics.key'],
 707             apt_repositories=[f'deb [arch=amd64] https://repositories.intel.com/graphics/ubuntu focal main']
 708         )
 709         os_packages += _intel_compute_runtime_extra_packages
 710     if args.rocm is not None:
 711         building_blocks['extra_packages'] += hpccm.building_blocks.packages(
 712             apt_keys=['http://repo.radeon.com/rocm/rocm.gpg.key'],
 713             apt_repositories=[f'deb [arch=amd64] http://repo.radeon.com/rocm/apt/{args.rocm}/ xenial main']
 714         )
 715     building_blocks['extra_packages'] += hpccm.building_blocks.packages(
 716         ospackages=os_packages,
 717         apt_ppas=['ppa:intel-opencl/intel-opencl'])
 718
 719     if args.cuda is not None and args.llvm is not None:
 720         # Hack to tell clang what version of CUDA we're using
 721         # based on https://github.com/llvm/llvm-project/blob/1fdec59bffc11ae37eb51a1b9869f0696bfd5312/clang/lib/Driver/ToolChains/Cuda.cpp#L43
 722         cuda_version_split = args.cuda.split('.')
 723         # LLVM requires having the version in x.y.z format, while args.cuda be be either x.y or x.y.z
 724         cuda_version_str = '{}.{}.{}'.format(
 725             cuda_version_split[0],
 726             cuda_version_split[1],
 727             cuda_version_split[2] if len(cuda_version_split) > 2 else 0
 728         )
 729         building_blocks['cuda-clang-workaround'] = hpccm.primitives.shell(commands=[
 730             f'echo "CUDA Version {cuda_version_str}" > /usr/local/cuda/version.txt'
 731         ])
 732
 733     building_blocks['clfft'] = get_clfft(args)
 734
 735     building_blocks['heffte'] = get_heffte(args)
 736
 737     building_blocks['hipSYCL'] = get_hipsycl(args)
 738
 739     # Add Python environments to MPI images, only, so we don't have to worry
 740     # about whether to install mpi4py.
 741     if args.mpi is not None and len(args.venvs) > 0:
 742         add_python_stages(building_blocks=building_blocks, input_args=args, output_stages=stages)
 743
 744     # Create the stage from which the targeted image will be tagged.
 745     stages['main'] = hpccm.Stage()
 746
 747     stages['main'] += hpccm.primitives.baseimage(image=base_image_tag(args))
 748     for bb in building_blocks.values():
 749         if bb is not None:
 750             stages['main'] += bb
 751
 752     # We always add Python3 and Pip
 753     stages['main'] += hpccm.building_blocks.python(python3=True, python2=False)
 754
 755     # Add documentation requirements (doxygen and sphinx + misc).
 756     if args.doxygen is not None:
 757         add_documentation_dependencies(args, stages)
 758
 759     if 'pyenv' in stages and stages['pyenv'] is not None:
 760         stages['main'] += hpccm.primitives.copy(_from='pyenv', _mkdir=True, src=['/root/.pyenv/'],
 761                                                 dest='/root/.pyenv')
 762         stages['main'] += hpccm.primitives.copy(_from='pyenv', _mkdir=True, src=['/root/venv/'],
 763                                                 dest='/root/venv')
 764         # TODO: Update user home directory.
 765         # TODO: If we activate pyenv for login shells, the `global` "version" should be full-featured.
 766         # stages['main'] += hpccm.primitives.copy(_from='pyenv', src=['/root/.bashrc'],
 767         #                                         dest='/root/')
 768
 769     # Make sure that `python` resolves to something.
 770     stages['main'] += hpccm.primitives.shell(commands=['test -x /usr/bin/python || '
 771                                                        'update-alternatives --install /usr/bin/python python /usr/bin/python3 1 && '
 772                                                        '/usr/bin/python --version'])
 773
 774     # Note that the list of stages should be sorted in dependency order.
 775     for build_stage in stages.values():
 776         if build_stage is not None:
 777             yield build_stage
 778
 779
 780 if __name__ == '__main__':
 781     args = parser.parse_args()
 782
 783     # Set container specification output format
 784     hpccm.config.set_container_format(args.format)
 785
 786     container_recipe = build_stages(args)
 787
 788     # Output container specification
 789     for stage in container_recipe:
 790         print(stage)