spack containerize: allow users to customize the base image (#15028)

This PR reworks a few attributes in the container subsection of
spack.yaml to permit the injection of custom base images when
generating containers with Spack. In more detail, users can still
specify the base operating system and Spack version they want to use:

  spack:
    container:
      images:
        os: ubuntu:18.04
        spack: develop

in which case the generated recipe will use one of the Spack images
built on Docker Hub for the build stage and the base OS image in the
final stage. Alternatively, they can specify explicitly the two
base images:

  spack:
    container:
      images:
        build: spack/ubuntu-bionic:latest
        final: ubuntu:18.04

and it will be up to them to ensure their consistency.

Additional changes:

* This commit adds documentation on the two approaches.
* Users can now specify OS packages to install (e.g. with apt or yum)
  prior to the build (previously this was only available for the
  finalized image).
* Handles to avoid an update of the available system packages have been
  added to the configuration to facilitate the generation of recipes
  permitting deterministic builds.

11 files changed, 479 insertions, 219 deletions

diff --git a/lib/spack/docs/containers.rst b/lib/spack/docs/containers.rst
index 2590fec8b6..f5bd5602bc 100644
--- a/lib/spack/docs/containers.rst
+++ b/lib/spack/docs/containers.rst
@@ -9,28 +9,48 @@
 Container Images
 ================
 
-Spack can be an ideal tool to setup images for containers since all the
-features discussed in :ref:`environments` can greatly help to manage
-the installation of software during the image build process. Nonetheless,
-building a production image from scratch still requires a lot of
-boilerplate to:
+Spack :ref:`environments` are a great tool to create container images, but
+preparing one that is suitable for production requires some more boilerplate
+than just:
 
-- Get Spack working within the image, possibly running as root
-- Minimize the physical size of the software installed
-- Properly update the system software in the base image
+.. code-block:: docker
+
+   COPY spack.yaml /environment
+   RUN spack -e /environment install
+
+Additional actions may be needed to minimize the size of the
+container, or to update the system software that is installed in the base
+image, or to set up a proper entrypoint to run the image. These tasks are
+usually both necessary and repetitive, so Spack comes with a command
+to generate recipes for container images starting from a ``spack.yaml``.
+
+--------------------
+A Quick Introduction
+--------------------
+
+Consider having a Spack environment like the following:
+
+.. code-block:: yaml
+
+   spack:
+     specs:
+     - gromacs+mpi
+     - mpich
 
-To facilitate users with these tedious tasks, Spack provides a command
-to automatically generate recipes for container images based on
-Environments:
+Producing a ``Dockerfile`` from it is as simple as moving to the directory
+where the ``spack.yaml`` file is stored and giving the following command:
 
 .. code-block:: console
 
-   $ ls
-   spack.yaml
+   $ spack containerize > Dockerfile
+
+The ``Dockerfile`` that gets created uses multi-stage builds and
+other techniques to minimize the size of the final image:
+
+.. code-block:: docker
 
-   $ spack containerize
    # Build stage with Spack pre-installed and ready to be used
-   FROM spack/centos7:latest as builder
+   FROM spack/ubuntu-bionic:latest as builder
 
    # What we want to install and how we want to install it
    # is specified in a manifest file (spack.yaml)
@@ -45,7 +65,7 @@ Environments:
    &&   echo "  view: /opt/view") > /opt/spack-environment/spack.yaml
 
    # Install the software, remove unnecessary deps
-   RUN cd /opt/spack-environment && spack env activate . && spack install && spack gc -y
+   RUN cd /opt/spack-environment && spack env activate . && spack install --fail-fast && spack gc -y
 
    # Strip all the binaries
    RUN find -L /opt/view/* -type f -exec readlink -f '{}' \; | \
@@ -58,45 +78,34 @@ Environments:
    RUN cd /opt/spack-environment && \
        spack env activate --sh -d . >> /etc/profile.d/z10_spack_environment.sh
 
-
    # Bare OS image to run the installed executables
-   FROM centos:7
+   FROM ubuntu:18.04
 
    COPY --from=builder /opt/spack-environment /opt/spack-environment
    COPY --from=builder /opt/software /opt/software
    COPY --from=builder /opt/view /opt/view
    COPY --from=builder /etc/profile.d/z10_spack_environment.sh /etc/profile.d/z10_spack_environment.sh
 
-   RUN yum update -y && yum install -y epel-release && yum update -y                                   \
-    && yum install -y libgomp \
-    && rm -rf /var/cache/yum  && yum clean all
-
-   RUN echo 'export PS1="\[$(tput bold)\]\[$(tput setaf 1)\][gromacs]\[$(tput setaf 2)\]\u\[$(tput sgr0)\]:\w $ "' >> ~/.bashrc
-
-
-   LABEL "app"="gromacs"
-   LABEL "mpi"="mpich"
-
    ENTRYPOINT ["/bin/bash", "--rcfile", "/etc/profile", "-l"]
 
-In order to build and run the image, execute:
+The image itself can then be built and run in the usual way, with any of the
+tools suitable for the task. For instance, if we decided to use ``docker``:
 
 .. code-block:: bash
 
    $ spack containerize > Dockerfile
    $ docker build -t myimage .
+   [ ... ]
    $ docker run -it myimage
 
-The bits that make this automation possible are discussed in details
-below. All the images generated in this way will be based on
-multi-stage builds with:
+The various components involved in the generation of the recipe and their
+configuration are discussed in details in the sections below.
 
-- A fat ``build`` stage containing common build tools and Spack itself
-- A minimal ``final`` stage containing only the software requested by the user
+.. _container_spack_images:
 
------------------
-Spack Base Images
------------------
+--------------------------
+Spack Images on Docker Hub
+--------------------------
 
 Docker images with Spack preinstalled and ready to be used are
 built on `Docker Hub <https://hub.docker.com/u/spack>`_
@@ -131,19 +140,20 @@ All the images are tagged with the corresponding release of Spack:
 with the exception of the ``latest`` tag that points to the HEAD
 of the ``develop`` branch. These images are available for anyone
 to use and take care of all the repetitive tasks that are necessary
-to setup Spack within a container. All the container recipes generated
-automatically by Spack use them as base images for their ``build`` stage.
-
+to setup Spack within a container. The container recipes generated
+by Spack use them as default base images for their ``build`` stage,
+even though handles to use custom base images provided by users are
+available to accommodate complex use cases.
 
--------------------------
-Environment Configuration
--------------------------
+---------------------------------
+Creating Images From Environments
+---------------------------------
 
 Any Spack Environment can be used for the automatic generation of container
 recipes. Sensible defaults are provided for things like the base image or the
-version of Spack used in the image. If a finer tuning is needed it can be
-obtained by adding the relevant metadata under the ``container`` attribute
-of environments:
+version of Spack used in the image.
+If a finer tuning is needed it can be obtained by adding the relevant metadata
+under the ``container`` attribute of environments:
 
 .. code-block:: yaml
 
@@ -157,9 +167,10 @@ of environments:
        # singularity or anything else that is currently supported
        format: docker
 
-       # Select from a valid list of images
-       base:
-         image: "centos:7"
+       # Sets the base images for the stages where Spack builds the
+       # software or where the software gets installed after being built..
+       images:
+         os: "centos:7"
          spack: develop
 
        # Whether or not to strip binaries
@@ -167,7 +178,8 @@ of environments:
 
        # Additional system packages that are needed at runtime
        os_packages:
-       - libgomp
+         final:
+         - libgomp
 
        # Extra instructions
        extra_instructions:
@@ -179,7 +191,210 @@ of environments:
          app: "gromacs"
          mpi: "mpich"
 
-The tables below describe the configuration options that are currently supported:
+A detailed description of the options available can be found in the
+:ref:`container_config_options` section.
+
+-------------------
+Setting Base Images
+-------------------
+
+The ``images`` subsection is used to select both the image where
+Spack builds the software and the image where the built software
+is installed. This attribute can be set in two different ways and
+which one to use depends on the use case at hand.
+
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Use Official Spack Images From Dockerhub
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To generate a recipe that uses an official Docker image from the
+Spack organization to build the software and the corresponding official OS image
+to install the built software, all the user has to do is specify:
+
+1. An operating system under ``images:os``
+2. A Spack version under ``images:spack``
+
+Any combination of these two values that can be mapped to one of the images
+discussed in :ref:`container_spack_images` is allowed. For instance, the
+following ``spack.yaml``:
+
+.. code-block:: yaml
+
+   spack:
+     specs:
+     - gromacs+mpi
+     - mpich
+
+     container:
+       images:
+         os: centos/7
+         spack: 0.15.4
+
+uses ``spack/centos7:0.15.4``  and ``centos:7`` for the stages where the
+software is respectively built and installed:
+
+.. code-block:: docker
+
+   # Build stage with Spack pre-installed and ready to be used
+   FROM spack/centos7:0.15.4 as builder
+
+   # What we want to install and how we want to install it
+   # is specified in a manifest file (spack.yaml)
+   RUN mkdir /opt/spack-environment \
+   &&  (echo "spack:" \
+   &&   echo "  specs:" \
+   &&   echo "  - gromacs+mpi" \
+   &&   echo "  - mpich" \
+   &&   echo "  concretization: together" \
+   &&   echo "  config:" \
+   &&   echo "    install_tree: /opt/software" \
+   &&   echo "  view: /opt/view") > /opt/spack-environment/spack.yaml
+   [ ... ]
+   # Bare OS image to run the installed executables
+   FROM centos:7
+
+   COPY --from=builder /opt/spack-environment /opt/spack-environment
+   COPY --from=builder /opt/software /opt/software
+   COPY --from=builder /opt/view /opt/view
+   COPY --from=builder /etc/profile.d/z10_spack_environment.sh /etc/profile.d/z10_spack_environment.sh
+
+   ENTRYPOINT ["/bin/bash", "--rcfile", "/etc/profile", "-l"]
+
+This method of selecting base images is the simplest of the two, and we advise
+to use it whenever possible. There are cases though where using Spack official
+images is not enough to fit production needs. In these situations users can manually
+select which base image to start from in the recipe, as we'll see next.
+
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Use Custom Images Provided by Users
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Consider, as an example, building a production grade image for a CUDA
+application. The best strategy would probably be to build on top of
+images provided by the vendor and regard CUDA as an external package.
+
+Spack doesn't currently provide an official image with CUDA configured
+this way, but users can build it on their own and then configure the
+environment to explicitly pull it. This requires users to:
+
+1. Specify the image used to build the software under ``images:build``
+2. Specify the image used to install the built software under ``images:final``
+
+A ``spack.yaml`` like the following:
+
+.. code-block:: yaml
+
+   spack:
+     specs:
+     - gromacs@2019.4+cuda build_type=Release
+     - mpich
+     - fftw precision=float
+     packages:
+       cuda:
+         buildable: False
+         externals:
+         - spec: cuda%gcc
+           prefix: /usr/local/cuda
+
+     container:
+       images:
+         build: custom/cuda-10.1-ubuntu18.04:latest
+         final: nvidia/cuda:10.1-base-ubuntu18.04
+
+produces, for instance, the following ``Dockerfile``:
+
+.. code-block:: docker
+
+   # Build stage with Spack pre-installed and ready to be used
+   FROM custom/cuda-10.1-ubuntu18.04:latest as builder
+
+   # What we want to install and how we want to install it
+   # is specified in a manifest file (spack.yaml)
+   RUN mkdir /opt/spack-environment \
+   &&  (echo "spack:" \
+   &&   echo "  specs:" \
+   &&   echo "  - gromacs@2019.4+cuda build_type=Release" \
+   &&   echo "  - mpich" \
+   &&   echo "  - fftw precision=float" \
+   &&   echo "  packages:" \
+   &&   echo "    cuda:" \
+   &&   echo "      buildable: false" \
+   &&   echo "      externals:" \
+   &&   echo "      - spec: cuda%gcc" \
+   &&   echo "        prefix: /usr/local/cuda" \
+   &&   echo "  concretization: together" \
+   &&   echo "  config:" \
+   &&   echo "    install_tree: /opt/software" \
+   &&   echo "  view: /opt/view") > /opt/spack-environment/spack.yaml
+
+   # Install the software, remove unnecessary deps
+   RUN cd /opt/spack-environment && spack env activate . && spack install --fail-fast && spack gc -y
+
+   # Strip all the binaries
+   RUN find -L /opt/view/* -type f -exec readlink -f '{}' \; | \
+       xargs file -i | \
+       grep 'charset=binary' | \
+       grep 'x-executable\|x-archive\|x-sharedlib' | \
+       awk -F: '{print $1}' | xargs strip -s
+
+   # Modifications to the environment that are necessary to run
+   RUN cd /opt/spack-environment && \
+       spack env activate --sh -d . >> /etc/profile.d/z10_spack_environment.sh
+
+   # Bare OS image to run the installed executables
+   FROM nvidia/cuda:10.1-base-ubuntu18.04
+
+   COPY --from=builder /opt/spack-environment /opt/spack-environment
+   COPY --from=builder /opt/software /opt/software
+   COPY --from=builder /opt/view /opt/view
+   COPY --from=builder /etc/profile.d/z10_spack_environment.sh /etc/profile.d/z10_spack_environment.sh
+
+   ENTRYPOINT ["/bin/bash", "--rcfile", "/etc/profile", "-l"]
+
+where the base images for both stages are completely custom.
+
+This second mode of selection for base images is more flexible than just
+choosing an operating system and a Spack version, but is also more demanding.
+Users may need to generate by themselves their base images and it's also their
+responsibility to ensure that:
+
+1. Spack is available in the ``build`` stage and set up correctly to install the required software
+2. The artifacts produced in the ``build`` stage can be executed in the ``final`` stage
+
+Therefore we don't recommend its use in cases that can be otherwise
+covered by the simplified mode shown first.
+
+----------------------------
+Singularity Definition Files
+----------------------------
+
+In addition to producing recipes in ``Dockerfile`` format Spack can produce
+Singularity Definition Files by just changing the value of the ``format``
+attribute:
+
+.. code-block:: console
+
+   $ cat spack.yaml
+   spack:
+     specs:
+     - hdf5~mpi
+     container:
+       format: singularity
+
+   $ spack containerize > hdf5.def
+   $ sudo singularity build hdf5.sif hdf5.def
+
+The minimum version of Singularity required to build a SIF (Singularity Image Format)
+image from the recipes generated by Spack is ``3.5.3``.
+
+.. _container_config_options:
+
+-----------------------
+Configuration Reference
+-----------------------
+
+The tables below describe all the configuration options that are currently supported
+to customize the generation of container recipes:
 
 .. list-table:: General configuration options for the ``container`` section of ``spack.yaml``
    :header-rows: 1
@@ -192,21 +407,41 @@ The tables below describe the configuration options that are currently supported
      - The format of the recipe
      - ``docker`` or ``singularity``
      - Yes
-   * - ``base:image``
-     - Base image for ``final`` stage
+   * - ``images:os``
+     - Operating system used as a base for the image
      - See :ref:`containers-supported-os`
-     - Yes
-   * - ``base:spack``
-     - Version of Spack
+     - Yes, if using constrained selection of base images
+   * - ``images:spack``
+     - Version of Spack use in the ``build`` stage
      - Valid tags for ``base:image``
-     - Yes
+     - Yes, if using constrained selection of base images
+   * - ``images:build``
+     - Image to be used in the ``build`` stage
+     - Any valid container image
+     - Yes, if using custom selection of base images
+   * - ``images:final``
+     - Image to be used in the ``build`` stage
+     - Any valid container image
+     - Yes, if using custom selection of base images
    * - ``strip``
      - Whether to strip binaries
      - ``true`` (default) or ``false``
      - No
-   * - ``os_packages``
-     - System packages to be installed
-     - Valid packages for the ``final`` OS
+   * - ``os_packages:command``
+     - Tool used to manage system packages
+     - ``apt``, ``yum``
+     - Only with custom base images
+   * - ``os_packages:update``
+     - Whether or not to update the list of available packages
+     - True or False (default: True)
+     - No
+   * - ``os_packages:build``
+     - System packages needed at build-time
+     - Valid packages for the current OS
+     - No
+   * - ``os_packages:final``
+     - System packages needed at run-time
+     - Valid packages for the current OS
      - No
    * - ``extra_instructions:build``
      - Extra instructions (e.g. `RUN`, `COPY`, etc.) at the end of the ``build`` stage
@@ -245,74 +480,6 @@ The tables below describe the configuration options that are currently supported
      - Description string
      - No
 
-Once the Environment is properly configured a recipe for a container
-image can be printed to standard output by issuing the following
-command from the directory where the ``spack.yaml`` resides:
-
-.. code-block:: console
-
-   $ spack containerize
-
-The example ``spack.yaml`` above would produce for instance the
-following ``Dockerfile``:
-
-.. code-block:: docker
-
-   # Build stage with Spack pre-installed and ready to be used
-   FROM spack/centos7:latest as builder
-
-   # What we want to install and how we want to install it
-   # is specified in a manifest file (spack.yaml)
-   RUN mkdir /opt/spack-environment \
-   &&  (echo "spack:" \
-   &&   echo "  specs:" \
-   &&   echo "  - gromacs+mpi" \
-   &&   echo "  - mpich" \
-   &&   echo "  concretization: together" \
-   &&   echo "  config:" \
-   &&   echo "    install_tree: /opt/software" \
-   &&   echo "  view: /opt/view") > /opt/spack-environment/spack.yaml
-
-   # Install the software, remove unnecessary deps
-   RUN cd /opt/spack-environment && spack env activate . && spack install && spack gc -y
-
-   # Strip all the binaries
-   RUN find -L /opt/view/* -type f -exec readlink -f '{}' \; | \
-       xargs file -i | \
-       grep 'charset=binary' | \
-       grep 'x-executable\|x-archive\|x-sharedlib' | \
-       awk -F: '{print $1}' | xargs strip -s
-
-   # Modifications to the environment that are necessary to run
-   RUN cd /opt/spack-environment && \
-       spack env activate --sh -d . >> /etc/profile.d/z10_spack_environment.sh
-
-
-   # Bare OS image to run the installed executables
-   FROM centos:7
-
-   COPY --from=builder /opt/spack-environment /opt/spack-environment
-   COPY --from=builder /opt/software /opt/software
-   COPY --from=builder /opt/view /opt/view
-   COPY --from=builder /etc/profile.d/z10_spack_environment.sh /etc/profile.d/z10_spack_environment.sh
-
-   RUN yum update -y && yum install -y epel-release && yum update -y                                   \
-    && yum install -y libgomp \
-    && rm -rf /var/cache/yum  && yum clean all
-
-   RUN echo 'export PS1="\[$(tput bold)\]\[$(tput setaf 1)\][gromacs]\[$(tput setaf 2)\]\u\[$(tput sgr0)\]:\w $ "' >> ~/.bashrc
-
-
-   LABEL "app"="gromacs"
-   LABEL "mpi"="mpich"
-
-   ENTRYPOINT ["/bin/bash", "--rcfile", "/etc/profile", "-l"]
-
-.. note::
-   Spack can also produce Singularity definition files to build the image. The
-   minimum version of Singularity required to build a SIF (Singularity Image Format)
-   from them is ``3.5.3``.
-
 --------------
 Best Practices
 --------------
diff --git a/lib/spack/spack/container/__init__.py b/lib/spack/spack/container/__init__.py
index fc3750355a..8206efef01 100644
--- a/lib/spack/spack/container/__init__.py
+++ b/lib/spack/spack/container/__init__.py
@@ -38,11 +38,11 @@ def validate(configuration_file):
     env_dict = spack.environment.config_dict(config)
     env_dict.setdefault('container', {
         'format': 'docker',
-        'base': {'image': 'ubuntu:18.04', 'spack': 'develop'}
+        'images': {'os': 'ubuntu:18.04', 'spack': 'develop'}
     })
     env_dict['container'].setdefault('format', 'docker')
     env_dict['container'].setdefault(
-        'base', {'image': 'ubuntu:18.04', 'spack': 'develop'}
+        'images', {'os': 'ubuntu:18.04', 'spack': 'develop'}
     )
 
     # Remove attributes that are not needed / allowed in the
diff --git a/lib/spack/spack/container/images.json b/lib/spack/spack/container/images.json
index efaedc68b8..cb495908c9 100644
--- a/lib/spack/spack/container/images.json
+++ b/lib/spack/spack/container/images.json
@@ -1,42 +1,45 @@
 {
-  "ubuntu:18.04": {
-    "update": "apt-get -yqq update && apt-get -yqq upgrade",
-    "install": "apt-get -yqq install",
-    "clean": "rm -rf /var/lib/apt/lists/*",
-    "environment": [],
-    "build": "spack/ubuntu-bionic",
-    "build_tags": {
-      "develop": "latest"
+  "images": {
+    "ubuntu:18.04": {
+      "os_package_manager": "apt",
+      "build": "spack/ubuntu-bionic",
+      "build_tags": {
+        "develop": "latest"
+      }
+    },
+    "ubuntu:16.04": {
+      "os_package_manager": "apt",
+      "build": "spack/ubuntu-xenial",
+      "build_tags": {
+        "develop": "latest"
+      }
+    },
+    "centos:7": {
+      "os_package_manager": "yum",
+      "environment": [],
+      "build": "spack/centos7",
+      "build_tags": {
+        "develop": "latest"
+      }
+    },
+    "centos:6": {
+      "os_package_manager": "yum",
+      "build": "spack/centos6",
+      "build_tags": {
+        "develop": "latest"
+      }
     }
   },
-  "ubuntu:16.04": {
-    "update": "apt-get -yqq update && apt-get -yqq upgrade",
-    "install": "apt-get -yqq install",
-    "clean": "rm -rf /var/lib/apt/lists/*",
-    "environment": [],
-    "build": "spack/ubuntu-xenial",
-    "build_tags": {
-      "develop": "latest"
-    }
-  },
-  "centos:7": {
-    "update": "yum update -y && yum install -y epel-release && yum update -y",
-    "install": "yum install -y",
-    "clean": "rm -rf /var/cache/yum  && yum clean all",
-    "environment": [],
-    "build": "spack/centos7",
-    "build_tags": {
-      "develop": "latest"
-    }
-  },
-  "centos:6": {
-    "update": "yum update -y && yum install -y epel-release && yum update -y",
-    "install": "yum install -y",
-    "clean": "rm -rf /var/cache/yum  && yum clean all",
-    "environment": [],
-    "build": "spack/centos6",
-    "build_tags": {
-      "develop": "latest"
+  "os_package_managers": {
+    "apt": {
+      "update": "apt-get -yqq update && apt-get -yqq upgrade",
+      "install": "apt-get -yqq install",
+      "clean": "rm -rf /var/lib/apt/lists/*"
+    },
+    "yum": {
+      "update": "yum update -y && yum install -y epel-release && yum update -y",
+      "install": "yum install -y",
+      "clean": "rm -rf /var/cache/yum  && yum clean all"
     }
   }
 }
diff --git a/lib/spack/spack/container/images.py b/lib/spack/spack/container/images.py
index f702781d32..32fcff2770 100644
--- a/lib/spack/spack/container/images.py
+++ b/lib/spack/spack/container/images.py
@@ -38,7 +38,7 @@ def build_info(image, spack_version):
     """
     # Don't handle error here, as a wrong image should have been
     # caught by the JSON schema
-    image_data = data()[image]
+    image_data = data()["images"][image]
     build_image = image_data['build']
 
     # Try to check if we have a tag for this Spack version
@@ -55,19 +55,30 @@ def build_info(image, spack_version):
     return build_image, build_tag
 
 
-def package_info(image):
-    """Returns the commands used to update system repositories, install
-    system packages and clean afterwards.
+def os_package_manager_for(image):
+    """Returns the name of the OS package manager for the image
+    passed as argument.
 
     Args:
         image (str): image to be used at run-time. Should be of the form
             <image_name>:<image_tag> e.g. "ubuntu:18.04"
 
     Returns:
+        Name of the package manager, e.g. "apt" or "yum"
+    """
+    name = data()["images"][image]["os_package_manager"]
+    return name
+
+
+def commands_for(package_manager):
+    """Returns the commands used to update system repositories, install
+    system packages and clean afterwards.
+
+    Args:
+        package_manager (str): package manager to be used
+
+    Returns:
         A tuple of (update, install, clean) commands.
     """
-    image_data = data()[image]
-    update = image_data['update']
-    install = image_data['install']
-    clean = image_data['clean']
-    return update, install, clean
+    info = data()["os_package_managers"][package_manager]
+    return info['update'], info['install'], info['clean']
diff --git a/lib/spack/spack/container/writers/__init__.py b/lib/spack/spack/container/writers/__init__.py
index a1d2fa3102..176dd7a50c 100644
--- a/lib/spack/spack/container/writers/__init__.py
+++ b/lib/spack/spack/container/writers/__init__.py
@@ -13,7 +13,8 @@ import spack.schema.env
 import spack.tengine as tengine
 import spack.util.spack_yaml as syaml
 
-from spack.container.images import build_info, package_info
+from spack.container.images import build_info, commands_for
+from spack.container.images import os_package_manager_for
 
 #: Caches all the writers that are currently supported
 _writer_factory = {}
@@ -63,21 +64,34 @@ class PathContext(tengine.Context):
     @tengine.context_property
     def run(self):
         """Information related to the run image."""
-        image = self.container_config['base']['image']
+        images_config = self.container_config['images']
+
+        # Check if we have custom images
+        image = images_config.get('final', None)
+        # If not use the base OS image
+        if image is None:
+            image = images_config['os']
+
         Run = collections.namedtuple('Run', ['image'])
         return Run(image=image)
 
     @tengine.context_property
     def build(self):
         """Information related to the build image."""
+        images_config = self.container_config['images']
+
+        # Check if we have custom images
+        image = images_config.get('build', None)
 
-        # Map the final image to the correct build image
-        run_image = self.container_config['base']['image']
-        spack_version = self.container_config['base']['spack']
-        image, tag = build_info(run_image, spack_version)
+        # If not select the correct build image based on OS and Spack version
+        if image is None:
+            operating_system = images_config['os']
+            spack_version = images_config['spack']
+            image_name, tag = build_info(operating_system, spack_version)
+            image = ':'.join([image_name, tag])
 
-        Build = collections.namedtuple('Build', ['image', 'tag'])
-        return Build(image=image, tag=tag)
+        Build = collections.namedtuple('Build', ['image'])
+        return Build(image=image)
 
     @tengine.context_property
     def strip(self):
@@ -116,14 +130,50 @@ class PathContext(tengine.Context):
         return syaml.dump(manifest, default_flow_style=False).strip()
 
     @tengine.context_property
-    def os_packages(self):
+    def os_packages_final(self):
         """Additional system packages that are needed at run-time."""
-        package_list = self.container_config.get('os_packages', None)
+        return self._os_packages_for_stage('final')
+
+    @tengine.context_property
+    def os_packages_build(self):
+        """Additional system packages that are needed at build-time."""
+        return self._os_packages_for_stage('build')
+
+    @tengine.context_property
+    def os_package_update(self):
+        """Whether or not to update the OS package manager cache."""
+        os_packages = self.container_config.get('os_packages', {})
+        return os_packages.get('update', True)
+
+    def _os_packages_for_stage(self, stage):
+        os_packages = self.container_config.get('os_packages', {})
+        package_list = os_packages.get(stage, None)
+        return self._package_info_from(package_list)
+
+    def _package_info_from(self, package_list):
+        """Helper method to pack a list of packages with the additional
+        information required by the template.
+
+        Args:
+            package_list: list of packages
+
+        Returns:
+            Enough information to know how to update the cache, install
+            a list opf packages, and clean in the end.
+        """
         if not package_list:
             return package_list
 
-        image = self.container_config['base']['image']
-        update, install, clean = package_info(image)
+        image_config = self.container_config['images']
+        image = image_config.get('build', None)
+
+        if image is None:
+            os_pkg_manager = os_package_manager_for(image_config['os'])
+        else:
+            os_pkg_manager = self.container_config['os_packages']['command']
+
+        update, install, clean = commands_for(os_pkg_manager)
+
         Packages = collections.namedtuple(
             'Packages', ['update', 'install', 'list', 'clean']
         )
diff --git a/lib/spack/spack/schema/container.py b/lib/spack/spack/schema/container.py
index d534afee08..7e14efd75c 100644
--- a/lib/spack/spack/schema/container.py
+++ b/lib/spack/spack/schema/container.py
@@ -4,6 +4,42 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 """Schema for the 'container' subsection of Spack environments."""
 
+_stages_from_dockerhub = {
+    'type': 'object',
+    'additionalProperties': False,
+    'properties': {
+        'os': {
+            'type': 'string',
+            'enum': ['ubuntu:18.04',
+                     'ubuntu:16.04',
+                     'centos:7',
+                     'centos:6']
+        },
+        'spack': {
+            'type': 'string',
+        },
+    },
+    'required': ['os', 'spack']
+}
+
+_custom_stages = {
+    'type': 'object',
+    'additionalProperties': False,
+    'properties': {
+        'build': {'type': 'string'},
+        'final': {'type': 'string'}
+    },
+    'required': ['build', 'final']
+}
+
+#: List of packages for the schema below
+_list_of_packages = {
+    'type': 'array',
+    'items': {
+        'type': 'string'
+    }
+}
+
 #: Schema for the container attribute included in Spack environments
 container_schema = {
     'type': 'object',
@@ -16,23 +52,7 @@ container_schema = {
         },
         # Describes the base image to start from and the version
         # of Spack to be used
-        'base': {
-            'type': 'object',
-            'additionalProperties': False,
-            'properties': {
-                'image': {
-                    'type': 'string',
-                    'enum': ['ubuntu:18.04',
-                             'ubuntu:16.04',
-                             'centos:7',
-                             'centos:6']
-                },
-                'spack': {
-                    'type': 'string',
-                },
-            },
-            'required': ['image', 'spack']
-        },
+        'images': {'anyOf': [_stages_from_dockerhub, _custom_stages]},
         # Whether or not to strip installed binaries
         'strip': {
             'type': 'boolean',
@@ -40,10 +60,14 @@ container_schema = {
         },
         # Additional system packages that are needed at runtime
         'os_packages': {
-            'type': 'array',
-            'items': {
-                'type': 'string'
-            }
+            'type': 'object',
+            'properties': {
+                'command': {'type': 'string', 'enum': ['apt', 'yum']},
+                'update': {'type': 'boolean'},
+                'build': _list_of_packages,
+                'final': _list_of_packages
+            },
+            'additionalProperties': False
         },
         # Add labels to the image
         'labels': {
diff --git a/lib/spack/spack/tengine.py b/lib/spack/spack/tengine.py
index e1e059c6a6..8581c8b41e 100644
--- a/lib/spack/spack/tengine.py
+++ b/lib/spack/spack/tengine.py
@@ -80,7 +80,9 @@ def make_environment(dirs=None):
     # Loader for the templates
     loader = jinja2.FileSystemLoader(dirs)
     # Environment of the template engine
-    env = jinja2.Environment(loader=loader, trim_blocks=True)
+    env = jinja2.Environment(
+        loader=loader, trim_blocks=True, lstrip_blocks=True
+    )
     # Custom filters
     _set_filters(env)
     return env
diff --git a/lib/spack/spack/test/container/conftest.py b/lib/spack/spack/test/container/conftest.py
index 65a89b0748..70a1ae2547 100644
--- a/lib/spack/spack/test/container/conftest.py
+++ b/lib/spack/spack/test/container/conftest.py
@@ -18,8 +18,8 @@ def minimal_configuration():
             ],
             'container': {
                 'format': 'docker',
-                'base': {
-                    'image': 'ubuntu:18.04',
+                'images': {
+                    'os': 'ubuntu:18.04',
                     'spack': 'develop'
                 }
             }
diff --git a/lib/spack/spack/test/container/docker.py b/lib/spack/spack/test/container/docker.py
index f75ae7de2a..6392f97db2 100644
--- a/lib/spack/spack/test/container/docker.py
+++ b/lib/spack/spack/test/container/docker.py
@@ -22,20 +22,22 @@ def test_build_and_run_images(minimal_configuration):
 
     # Test the output of the build property
     build = writer.build
-    assert build.image == 'spack/ubuntu-bionic'
-    assert build.tag == 'latest'
+    assert build.image == 'spack/ubuntu-bionic:latest'
 
 
 def test_packages(minimal_configuration):
     # In this minimal configuration we don't have packages
     writer = writers.create(minimal_configuration)
-    assert writer.os_packages is None
+    assert writer.os_packages_build is None
+    assert writer.os_packages_final is None
 
     # If we add them a list should be returned
     pkgs = ['libgomp1']
-    minimal_configuration['spack']['container']['os_packages'] = pkgs
+    minimal_configuration['spack']['container']['os_packages'] = {
+        'final': pkgs
+    }
     writer = writers.create(minimal_configuration)
-    p = writer.os_packages
+    p = writer.os_packages_final
     assert p.update
     assert p.install
     assert p.clean
diff --git a/lib/spack/spack/test/container/images.py b/lib/spack/spack/test/container/images.py
index 6af69e0c83..6cec1cc592 100644
--- a/lib/spack/spack/test/container/images.py
+++ b/lib/spack/spack/test/container/images.py
@@ -22,7 +22,8 @@ def test_build_info(image, spack_version, expected):
     'ubuntu:18.04'
 ])
 def test_package_info(image):
-    update, install, clean = spack.container.images.package_info(image)
+    pkg_manager = spack.container.images.os_package_manager_for(image)
+    update, install, clean = spack.container.images.commands_for(pkg_manager)
     assert update
     assert install
     assert clean
diff --git a/lib/spack/spack/test/container/schema.py b/lib/spack/spack/test/container/schema.py
index 3f33a3f9f7..e86f9d3fe8 100644
--- a/lib/spack/spack/test/container/schema.py
+++ b/lib/spack/spack/test/container/schema.py
@@ -10,7 +10,7 @@ import spack.schema.container
 def test_images_in_schema():
     properties = spack.schema.container.container_schema['properties']
     allowed_images = set(
-        properties['base']['properties']['image']['enum']
+        properties['images']['anyOf'][0]['properties']['os']['enum']
     )
-    images_in_json = set(x for x in spack.container.images.data())
+    images_in_json = set(x for x in spack.container.images.data()['images'])
     assert images_in_json == allowed_images