Improve Dockerfile recipe generation (#35187)

- Update default image to Ubuntu 22.04 (previously was still Ubuntu 18.04)
- Optionally use depfiles to install the environment within the container
- Allow extending Dockerfile Jinja2 template
- Allow extending Singularity definition file Jinja2 template
- Deprecate previous options to add extra instructions

6 files changed, 163 insertions, 15 deletions

diff --git a/lib/spack/docs/containers.rst b/lib/spack/docs/containers.rst
index a919db0642..2c097d366c 100644
--- a/lib/spack/docs/containers.rst
+++ b/lib/spack/docs/containers.rst
@@ -444,6 +444,120 @@ attribute:
 The minimum version of Singularity required to build a SIF (Singularity Image Format)
 image from the recipes generated by Spack is ``3.5.3``.
 
+------------------------------
+Extending the Jinja2 Templates
+------------------------------
+
+The Dockerfile and the Singularity definition file that Spack can generate are based on
+a few Jinja2 templates that are rendered according to the environment being containerized.
+Even though Spack allows a great deal of customization by just setting appropriate values for
+the configuration options, sometimes that is not enough.
+
+In those cases, a user can directly extend the template that Spack uses to render the image
+to e.g. set additional environment variables or perform specific operations either before or
+after a given stage of the build. Let's consider as an example the following structure:
+
+.. code-block:: console
+
+   $ tree /opt/environment
+   /opt/environment
+   ├── data
+   │     └── data.csv
+   ├── spack.yaml
+   ├── data
+   └── templates
+       └── container
+           └── CustomDockerfile
+
+containing both the custom template extension and the environment manifest file. To use a custom
+template, the environment must register the directory containing it, and declare its use under the
+``container`` configuration:
+
+.. code-block:: yaml
+   :emphasize-lines: 7-8,12
+
+   spack:
+     specs:
+     - hdf5~mpi
+     concretizer:
+       unify: true
+     config:
+       template_dirs:
+       - /opt/environment/templates
+     container:
+       format: docker
+       depfile: true
+       template: container/CustomDockerfile
+
+The template extension can override two blocks, named ``build_stage`` and ``final_stage``, similarly to
+the example below:
+
+.. code-block::
+   :emphasize-lines: 3,8
+
+   {% extends "container/Dockerfile" %}
+   {% block build_stage %}
+   RUN echo "Start building"
+   {{ super() }}
+   {% endblock %}
+   {% block final_stage %}
+   {{ super() }}
+   COPY data /share/myapp/data
+   {% endblock %}
+
+The recipe that gets generated contains the two extra instruction that we added in our template extension:
+
+.. code-block:: Dockerfile
+   :emphasize-lines: 4,43
+
+   # Build stage with Spack pre-installed and ready to be used
+   FROM spack/ubuntu-jammy:latest as builder
+
+   RUN echo "Start building"
+
+   # 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 "  - hdf5~mpi" \
+   &&   echo "  concretizer:" \
+   &&   echo "    unify: true" \
+   &&   echo "  config:" \
+   &&   echo "    template_dirs:" \
+   &&   echo "    - /tmp/environment/templates" \
+   &&   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 concretize && spack env depfile -o Makefile && make -j $(nproc) && 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 ubuntu:22.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 /opt/view /opt/view
+   COPY --from=builder /etc/profile.d/z10_spack_environment.sh /etc/profile.d/z10_spack_environment.sh
+
+   COPY data /share/myapp/data
+
+   ENTRYPOINT ["/bin/bash", "--rcfile", "/etc/profile", "-l", "-c", "$*", "--" ]
+   CMD [ "/bin/bash" ]
+
+
 .. _container_config_options:
 
 -----------------------
@@ -464,6 +578,10 @@ to customize the generation of container recipes:
      - The format of the recipe
      - ``docker`` or ``singularity``
      - Yes
+   * - ``depfile``
+     - Whether to use a depfile for installation, or not
+     - True or False (default)
+     - No
    * - ``images:os``
      - Operating system used as a base for the image
      - See :ref:`containers-supported-os`
@@ -512,14 +630,6 @@ to customize the generation of container recipes:
      - 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
-     - Anything understood by the current ``format``
-     - No
-   * - ``extra_instructions:final``
-     - Extra instructions (e.g. `RUN`, `COPY`, etc.) at the end of the ``final`` stage
-     - Anything understood by the current ``format``
-     - No
    * - ``labels``
      - Labels to tag the image
      - Pairs of key-value strings
diff --git a/lib/spack/spack/container/__init__.py b/lib/spack/spack/container/__init__.py
index aa27109c46..e7d487cbb4 100644
--- a/lib/spack/spack/container/__init__.py
+++ b/lib/spack/spack/container/__init__.py
@@ -39,10 +39,10 @@ def validate(configuration_file):
     # Ensure we have a "container" attribute with sensible defaults set
     env_dict = ev.config_dict(config)
     env_dict.setdefault(
-        "container", {"format": "docker", "images": {"os": "ubuntu:18.04", "spack": "develop"}}
+        "container", {"format": "docker", "images": {"os": "ubuntu:22.04", "spack": "develop"}}
     )
     env_dict["container"].setdefault("format", "docker")
-    env_dict["container"].setdefault("images", {"os": "ubuntu:18.04", "spack": "develop"})
+    env_dict["container"].setdefault("images", {"os": "ubuntu:22.04", "spack": "develop"})
 
     # Remove attributes that are not needed / allowed in the
     # container recipe
diff --git a/lib/spack/spack/container/writers/__init__.py b/lib/spack/spack/container/writers/__init__.py
index 306a6ebdd2..7508a67fd7 100644
--- a/lib/spack/spack/container/writers/__init__.py
+++ b/lib/spack/spack/container/writers/__init__.py
@@ -7,6 +7,7 @@ convenience functions.
 """
 import collections
 import copy
+from typing import Optional
 
 import spack.environment as ev
 import spack.schema.env
@@ -131,6 +132,9 @@ class PathContext(tengine.Context):
     directly via PATH.
     """
 
+    # Must be set by derived classes
+    template_name: Optional[str] = None
+
     def __init__(self, config, last_phase):
         self.config = ev.config_dict(config)
         self.container_config = self.config["container"]
@@ -147,6 +151,10 @@ class PathContext(tengine.Context):
         self.last_phase = last_phase
 
     @tengine.context_property
+    def depfile(self):
+        return self.container_config.get("depfile", False)
+
+    @tengine.context_property
     def run(self):
         """Information related to the run image."""
         Run = collections.namedtuple("Run", ["image"])
@@ -280,7 +288,8 @@ class PathContext(tengine.Context):
     def __call__(self):
         """Returns the recipe as a string"""
         env = tengine.make_environment()
-        t = env.get_template(self.template_name)
+        template_name = self.container_config.get("template", self.template_name)
+        t = env.get_template(template_name)
         return t.render(**self.to_dict())
 
 
diff --git a/lib/spack/spack/schema/container.py b/lib/spack/spack/schema/container.py
index 21d8bba8ef..37912d4c55 100644
--- a/lib/spack/spack/schema/container.py
+++ b/lib/spack/spack/schema/container.py
@@ -63,6 +63,8 @@ container_schema = {
         },
         # Add labels to the image
         "labels": {"type": "object"},
+        # Use a custom template to render the recipe
+        "template": {"type": "string", "default": None},
         # Add a custom extra section at the bottom of a stage
         "extra_instructions": {
             "type": "object",
@@ -82,6 +84,16 @@ container_schema = {
             },
         },
         "docker": {"type": "object", "additionalProperties": False, "default": {}},
+        "depfile": {"type": "boolean", "default": False},
+    },
+    "deprecatedProperties": {
+        "properties": ["extra_instructions"],
+        "message": (
+            "container:extra_instructions has been deprecated and will be removed "
+            "in Spack v0.21. Set container:template appropriately to use custom Jinja2 "
+            "templates instead."
+        ),
+        "error": False,
     },
 }
 
diff --git a/share/spack/templates/container/Dockerfile b/share/spack/templates/container/Dockerfile
index b10dcf084a..9116590480 100644
--- a/share/spack/templates/container/Dockerfile
+++ b/share/spack/templates/container/Dockerfile
@@ -6,6 +6,7 @@
 # Build stage with Spack pre-installed and ready to be used
 FROM {{ build.image }} as builder
 
+{% block build_stage %}
 {% if os_packages_build %}
 # Install OS packages needed to build the software
 RUN {% if os_package_update %}{{ os_packages_build.update }} \
@@ -19,7 +20,11 @@ RUN mkdir {{ paths.environment }} \
 {{ manifest }} > {{ paths.environment }}/spack.yaml
 
 # Install the software, remove unnecessary deps
+{% if depfile %}
+RUN cd {{ paths.environment }} && spack env activate . && spack concretize && spack env depfile -o Makefile && make -j $(nproc) && spack gc -y
+{% else %}
 RUN cd {{ paths.environment }} && spack env activate . && spack install --fail-fast && spack gc -y
+{% endif %}
 {% if strip %}
 
 # Strip all the binaries
@@ -37,7 +42,9 @@ RUN cd {{ paths.environment }} && \
 {% if extra_instructions.build %}
 {{ extra_instructions.build }}
 {% endif %}
+{% endblock build_stage %}
 {% endif %}
+
 {% if render_phase.final %}
 # Bare OS image to run the installed executables
 FROM {{ run.image }}
@@ -48,6 +55,8 @@ COPY --from=builder {{ paths.hidden_view }} {{ paths.hidden_view }}
 COPY --from=builder {{ paths.view }} {{ paths.view }}
 COPY --from=builder /etc/profile.d/z10_spack_environment.sh /etc/profile.d/z10_spack_environment.sh
 
+{% block final_stage %}
+
 {% if os_packages_final %}
 RUN {% if os_package_update %}{{ os_packages_final.update }} \
  && {% endif %}{{ os_packages_final.install }} {{ os_packages_final.list | join | replace('\n', ' ') }} \
@@ -57,6 +66,7 @@ RUN {% if os_package_update %}{{ os_packages_final.update }} \
 
 {{ extra_instructions.final }}
 {% endif %}
+{% endblock final_stage %}
 {% for label, value in labels.items() %}
 LABEL "{{ label }}"="{{ value }}"
 {% endfor %}
diff --git a/share/spack/templates/container/singularity.def b/share/spack/templates/container/singularity.def
index 684fe9f988..d5778f0d41 100644
--- a/share/spack/templates/container/singularity.def
+++ b/share/spack/templates/container/singularity.def
@@ -3,6 +3,7 @@ From: {{ build.image }}
 Stage: build
 
 %post
+{% block build_stage %}
 {% if os_packages_build.list %}
   # Update, install and cleanup of system packages needed at build-time
   {% if os_package_update %}
@@ -20,10 +21,14 @@ EOF
 
   # Install all the required software
   . /opt/spack/share/spack/setup-env.sh
-  spack env activate .
-  spack install --fail-fast
+  spack -e . concretize
+{% if depfile %}
+  spack -e . env depfile -o Makefile
+  make -j $(nproc)
+{% else %}
+  spack -e . install
+{% endif %}
   spack gc -y
-  spack env deactivate
   spack env activate --sh -d . >> {{ paths.environment }}/environment_modifications.sh
 {% if strip %}
 
@@ -37,7 +42,7 @@ EOF
 {% if extra_instructions.build %}
 {{ extra_instructions.build }}
 {% endif %}
-
+{% endblock build_stage %}
 {% if apps %}
 {% for application, help_text in apps.items() %}
 
@@ -61,6 +66,7 @@ Stage: final
   {{ paths.environment }}/environment_modifications.sh {{ paths.environment }}/environment_modifications.sh
 
 %post
+{% block final_stage %}
 {% if os_packages_final.list %}
   # Update, install and cleanup of system packages needed at run-time
   {% if os_package_update %}
@@ -74,6 +80,7 @@ Stage: final
 {% if extra_instructions.final %}
 {{ extra_instructions.final }}
 {% endif %}
+{% endblock final_stage %}
 {% if runscript %}
 
 %runscript