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(input_args: argparse.Namespace, *,
 538                       base: str,
 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,
 559                                               _distro=hpccm_distro_name(input_args),
 560                                               _as='pyenv')
 561     pyenv_stage += hpccm.building_blocks.packages(ospackages=_python_extra_packages)
 562
 563     for version in [StrictVersion(py_ver) for py_ver in sorted(input_args.venvs)]:
 564         stage_name = 'py' + str(version)
 565         stage = hpccm.Stage()
 566         stage += hpccm.primitives.baseimage(image=base,
 567                                             _distro=hpccm_distro_name(input_args),
 568                                             _as=stage_name)
 569         stage += hpccm.building_blocks.packages(ospackages=_python_extra_packages)
 570
 571         # TODO: Use a non-root user for testing and Python virtual environments.
 572         stage += hpccm.primitives.shell(commands=[
 573             'curl https://pyenv.run | bash',
 574             """echo 'export PYENV_ROOT="$HOME/.pyenv"' >> $HOME/.bashrc""",
 575             """echo 'export PATH="$PYENV_ROOT/bin:$PATH"' >> $HOME/.bashrc""",
 576             """echo 'eval "$(pyenv init -)"' >> $HOME/.bashrc""",
 577             """echo 'eval "$(pyenv virtualenv-init -)"' >> $HOME/.bashrc"""])
 578         pyenv = '$HOME/.pyenv/bin/pyenv'
 579         commands = [f'PYTHON_CONFIGURE_OPTS="--enable-shared" {pyenv} install -s {version}']
 580         stage += hpccm.primitives.shell(commands=commands)
 581
 582         commands = prepare_venv(version)
 583         stage += hpccm.primitives.shell(commands=commands)
 584
 585         # TODO: Update user home directory.
 586         pyenv_stage += hpccm.primitives.copy(_from=stage_name, _mkdir=True, src=['/root/'],
 587                                              dest='/root')
 588
 589         # Add the intermediate build stage to the sequence
 590         output_stages[stage_name] = stage
 591
 592     # TODO: If we activate pyenv for login shells, the `global` "version" should be full-featured.
 593     # # `version` should be a system installation or pyenv environment (or pyenv-virtualenv)
 594     # # with the dependencies for all of the Python aspects of CMake-driven builds.
 595     # commands = '{pyenv} global {version}'.format(
 596     #             pyenv=pyenv,
 597     #             version=...)
 598     # pyenv_stage += hpccm.primitives.shell(commands=commands)
 599
 600     # Add the aggregating build stage to the sequence. This allows the main stage to copy
 601     # the files in a single stage, potentially reducing the overall output image size.
 602     output_stages['pyenv'] = pyenv_stage
 603
 604
 605 def add_documentation_dependencies(input_args,
 606                                    output_stages: typing.MutableMapping[str, hpccm.Stage]):
 607     """Add appropriate layers according to doxygen input arguments."""
 608     if input_args.doxygen is None:
 609         return
 610     # Always clone the same version of linkchecker (latest release at June 1, 2021)
 611     output_stages['main'] += hpccm.building_blocks.pip(pip='pip3', packages=[
 612         'git+https://github.com/linkchecker/linkchecker.git@v10.0.1'])
 613     output_stages['main'] += hpccm.primitives.shell(
 614         commands=[
 615             'sed -i \'/\"XPS\"/d;/\"PDF\"/d;/\"PS\"/d;/\"EPS\"/d;/disable ghostscript format types/d\' /etc/ImageMagick-6/policy.xml'])
 616     if input_args.doxygen == '1.8.5':
 617         doxygen_commit = 'ed4ed873ab0e7f15116e2052119a6729d4589f7a'
 618         output_stages['main'] += hpccm.building_blocks.generic_autotools(
 619             repository='https://github.com/westes/flex.git',
 620             commit='f7788a9a0ecccdc953ed12043ccb59ca25714018',
 621             prefix='/tmp/install-of-flex',
 622             configure_opts=['--disable-shared'],
 623             preconfigure=['./autogen.sh'])
 624         output_stages['main'] += hpccm.building_blocks.generic_autotools(
 625             repository='https://github.com/doxygen/doxygen.git',
 626             commit=doxygen_commit,
 627             prefix='',
 628             configure_opts=[
 629                 '--flex /tmp/install-of-flex/bin/flex',
 630                 '--static'])
 631     else:
 632         version = input_args.doxygen
 633         archive_name = f'doxygen-{version}.linux.bin.tar.gz'
 634         archive_url = f'https://sourceforge.net/projects/doxygen/files/rel-{version}/{archive_name}'
 635         binary_path = f'doxygen-{version}/bin/doxygen'
 636         commands = [
 637             'mkdir doxygen && cd doxygen',
 638             f'wget {archive_url}',
 639             f'tar xf {archive_name} {binary_path}',
 640             f'cp {binary_path} /usr/local/bin/',
 641             'cd .. && rm -rf doxygen'
 642         ]
 643         output_stages['main'] += hpccm.primitives.shell(commands=commands)
 644
 645
 646 def add_base_stage(name: str,
 647                    input_args,
 648                    output_stages: typing.MutableMapping[str, hpccm.Stage]):
 649     """Establish dependencies that are shared by multiple parallel stages."""
 650     # Building blocks are chunks of container-builder instructions that can be
 651     # copied to any build stage with the addition operator.
 652     building_blocks = collections.OrderedDict()
 653     building_blocks['base_packages'] = hpccm.building_blocks.packages(
 654         ospackages=_common_packages)
 655
 656     # These are the most expensive and most reusable layers, so we put them first.
 657     building_blocks['compiler'] = get_compiler(input_args, compiler_build_stage=output_stages.get('compiler_build'))
 658     building_blocks['gdrcopy'] = get_gdrcopy(input_args, building_blocks['compiler'])
 659     building_blocks['ucx'] = get_ucx(input_args, building_blocks['compiler'], building_blocks['gdrcopy'])
 660     building_blocks['mpi'] = get_mpi(input_args, building_blocks['compiler'], building_blocks['ucx'])
 661
 662     # Create the stage from which the targeted image will be tagged.
 663     output_stages[name] = hpccm.Stage()
 664
 665     output_stages[name] += hpccm.primitives.baseimage(image=base_image_tag(input_args),
 666                                                       _distro=hpccm_distro_name(input_args),
 667                                                       _as=name)
 668     for bb in building_blocks.values():
 669         if bb is not None:
 670             output_stages[name] += bb
 671
 672
 673 def build_stages(args) -> typing.Iterable[hpccm.Stage]:
 674     """Define and sequence the stages for the recipe corresponding to *args*."""
 675
 676     # A Dockerfile or Singularity recipe can have multiple build stages.
 677     # The main build stage can copy files from previous stages, though only
 678     # the last stage is included in the tagged output image. This means that
 679     # large or expensive sets of build instructions can be isolated in
 680     # local/temporary images, but all of the stages need to be output by this
 681     # script, and need to occur in the correct order, so we create a sequence
 682     # object early in this function.
 683     stages = collections.OrderedDict()
 684
 685     # If we need TSAN or oneAPI support the early build is more complex,
 686     # so that our compiler images don't have all the cruft needed to get those things
 687     # installed.
 688     if args.llvm is not None and args.tsan is not None:
 689         add_tsan_compiler_build_stage(input_args=args, output_stages=stages)
 690     if args.oneapi is not None:
 691         add_oneapi_compiler_build_stage(input_args=args, output_stages=stages)
 692
 693     add_base_stage(name='build_base', input_args=args, output_stages=stages)
 694
 695     # Add Python environments to MPI images, only, so we don't have to worry
 696     # about whether to install mpi4py.
 697     if args.mpi is not None and len(args.venvs) > 0:
 698         add_python_stages(base='build_base', input_args=args, output_stages=stages)
 699
 700     # Building blocks are chunks of container-builder instructions that can be
 701     # copied to any build stage with the addition operator.
 702     building_blocks = collections.OrderedDict()
 703
 704     for i, cmake in enumerate(args.cmake):
 705         building_blocks['cmake' + str(i)] = hpccm.building_blocks.cmake(
 706             eula=True,
 707             prefix=f'/usr/local/cmake-{cmake}',
 708             version=cmake)
 709
 710     # Install additional packages early in the build to optimize Docker build layer cache.
 711     os_packages = list(get_llvm_packages(args)) + get_opencl_packages(args) + get_rocm_packages(args)
 712     if args.doxygen is not None:
 713         os_packages += _docs_extra_packages
 714     if args.oneapi is not None:
 715         os_packages += ['lsb-release']
 716     if args.hipsycl is not None:
 717         os_packages += ['libboost-fiber-dev']
 718     building_blocks['extra_packages'] = []
 719     if args.intel_compute_runtime:
 720         building_blocks['extra_packages'] += hpccm.building_blocks.packages(
 721             apt_keys=['https://repositories.intel.com/graphics/intel-graphics.key'],
 722             apt_repositories=[f'deb [arch=amd64] https://repositories.intel.com/graphics/ubuntu focal main']
 723         )
 724         os_packages += _intel_compute_runtime_extra_packages
 725     if args.rocm is not None:
 726         building_blocks['extra_packages'] += hpccm.building_blocks.packages(
 727             apt_keys=['http://repo.radeon.com/rocm/rocm.gpg.key'],
 728             apt_repositories=[f'deb [arch=amd64] http://repo.radeon.com/rocm/apt/{args.rocm}/ xenial main']
 729         )
 730     building_blocks['extra_packages'] += hpccm.building_blocks.packages(
 731         ospackages=os_packages,
 732         apt_ppas=['ppa:intel-opencl/intel-opencl'])
 733
 734     if args.cuda is not None and args.llvm is not None:
 735         # Hack to tell clang what version of CUDA we're using
 736         # based on https://github.com/llvm/llvm-project/blob/1fdec59bffc11ae37eb51a1b9869f0696bfd5312/clang/lib/Driver/ToolChains/Cuda.cpp#L43
 737         cuda_version_split = args.cuda.split('.')
 738         # LLVM requires having the version in x.y.z format, while args.cuda be be either x.y or x.y.z
 739         cuda_version_str = '{}.{}.{}'.format(
 740             cuda_version_split[0],
 741             cuda_version_split[1],
 742             cuda_version_split[2] if len(cuda_version_split) > 2 else 0
 743         )
 744         building_blocks['cuda-clang-workaround'] = hpccm.primitives.shell(commands=[
 745             f'echo "CUDA Version {cuda_version_str}" > /usr/local/cuda/version.txt'
 746         ])
 747
 748     building_blocks['clfft'] = get_clfft(args)
 749
 750     building_blocks['heffte'] = get_heffte(args)
 751
 752     building_blocks['hipSYCL'] = get_hipsycl(args)
 753
 754     # Add Python environments to MPI images, only, so we don't have to worry
 755     # about whether to install mpi4py.
 756     if args.mpi is not None and len(args.venvs) > 0:
 757         add_python_stages(building_blocks=building_blocks, input_args=args, output_stages=stages)
 758
 759     # Create the stage from which the targeted image will be tagged.
 760     stages['main'] = hpccm.Stage()
 761
 762     stages['main'] += hpccm.primitives.baseimage(image='build_base',
 763                                                  _distro=hpccm_distro_name(args),
 764                                                  _as='main')
 765     for bb in building_blocks.values():
 766         if bb is not None:
 767             stages['main'] += bb
 768
 769     # We always add Python3 and Pip
 770     stages['main'] += hpccm.building_blocks.python(python3=True, python2=False)
 771
 772     # Add documentation requirements (doxygen and sphinx + misc).
 773     if args.doxygen is not None:
 774         add_documentation_dependencies(args, stages)
 775
 776     if 'pyenv' in stages and stages['pyenv'] is not None:
 777         stages['main'] += hpccm.primitives.copy(_from='pyenv', _mkdir=True, src=['/root/.pyenv/'],
 778                                                 dest='/root/.pyenv')
 779         stages['main'] += hpccm.primitives.copy(_from='pyenv', _mkdir=True, src=['/root/venv/'],
 780                                                 dest='/root/venv')
 781         # TODO: Update user home directory.
 782         # TODO: If we activate pyenv for login shells, the `global` "version" should be full-featured.
 783         # stages['main'] += hpccm.primitives.copy(_from='pyenv', src=['/root/.bashrc'],
 784         #                                         dest='/root/')
 785
 786     # Make sure that `python` resolves to something.
 787     stages['main'] += hpccm.primitives.shell(commands=['test -x /usr/bin/python || '
 788                                                        'update-alternatives --install /usr/bin/python python /usr/bin/python3 1 && '
 789                                                        '/usr/bin/python --version'])
 790
 791     # Note that the list of stages should be sorted in dependency order.
 792     for build_stage in stages.values():
 793         if build_stage is not None:
 794             yield build_stage
 795
 796
 797 if __name__ == '__main__':
 798     args = parser.parse_args()
 799
 800     # Set container specification output format
 801     hpccm.config.set_container_format(args.format)
 802
 803     container_recipe = build_stages(args)
 804
 805     # Output container specification
 806     for stage in container_recipe:
 807         print(stage)