spack · scheibelp · Nov 17, 2020 · Feb 17, 2020 · Feb 17, 2020 · Feb 17, 2020
diff --git a/lib/spack/docs/containers.rst b/lib/spack/docs/containers.rst
diff --git 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
@@ -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
@@ -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
@@ -13,7 +13,8 @@
 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 @@ def __init__(self, config):
     @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 @@ def manifest(self):
         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
@@ -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,34 +52,22 @@
         },
         # 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',
             'default': True
         },
         # 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
@@ -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
@@ -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
@@ -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
@@ -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