Merge pull request #38 from douglasjacobsen/mirror_spack

Update mirror logic, and mirror software

Author: rfbgo

lib/ramble/ramble/application.py

@@ -56,6 +56,7 @@ def __init__(self, file_path):
         self._setup_phases = []
         self._analyze_phases = []
         self._archive_phases = ['archive_experiments']
+        self._mirror_phases = ['mirror_inputs']
 
         self._file_path = file_path
 
@@ -345,9 +346,19 @@ def _add_expand_vars(self, expander):
 
         expander.set_var('spack_setup', '', 'experiment')
 
-    def all_inputs_and_fetchers(self):
-        expected_files = {}
-        for workload_name, workload in self.workloads.items():
+    def _inputs_and_fetchers(self, workload=None):
+        """Extract all inputs for a given workload
+
+        Take a workload name and extract all inputs for the workload.
+        If the workload is set to None, extract all inputs for all workloads.
+        """
+
+        workload_names = [workload] if workload else self.workloads.keys()
+
+        inputs = {}
+        for workload_name in workload_names:
+            workload = self.workloads[workload_name]
+
             for input_file in workload['inputs']:
                 input_conf = self.inputs[input_file]
 
@@ -361,49 +372,40 @@ def all_inputs_and_fetchers(self):
 
                 namespace = f'{self.name}.{workload_name}'
 
-                expected_files[file_name] = {'fetcher': fetcher,
-                                             'namespace': namespace,
-                                             'extension': fetcher.extension,
-                                             'input_name': input_file}
-
-        return expected_files
+                inputs[file_name] = {'fetcher': fetcher,
+                                     'namespace': namespace,
+                                     'target_dir': input_conf['target_dir'],
+                                     'extension': fetcher.extension,
+                                     'input_name': input_file}
+        return inputs
 
-    def mirror_inputs(self, mirror_root, mirror, mirror_stats):
-        required_inputs = self.all_inputs_and_fetchers()
-        for input_name, conf in required_inputs.items():
+    def _mirror_inputs(self, workspace, expander):
+        for input_file, input_conf in self._inputs_and_fetchers(expander.workload_name).items():
             mirror_paths = ramble.mirror.mirror_archive_paths(
-                conf['fetcher'], os.path.join(self.name, input_name))
-            stage = ramble.stage.InputStage(conf['fetcher'], name=conf['namespace'],
-                                            path=mirror_root,
-                                            mirror_paths=mirror_paths, lock=False)
+                input_conf['fetcher'], os.path.join(self.name, input_file))
+            fetch_dir = os.path.join(workspace._input_mirror_path, self.name)
+            fs.mkdirp(fetch_dir)
+            stage = ramble.stage.InputStage(input_conf['fetcher'], name=input_conf['namespace'],
+                                            path=fetch_dir, mirror_paths=mirror_paths, lock=False)
 
-            stage.cache_mirror(mirror, mirror_stats)
+            stage.cache_mirror(workspace._input_mirror_cache, workspace._input_mirror_stats)
 
     def _get_inputs(self, workspace, expander):
         workload_namespace = '%s.%s' % (expander.application_name,
                                         expander.workload_name)
 
-        workload = self.workloads[expander.workload_name]
-
-        for input_file in workload['inputs']:
-            input_conf = self.inputs[input_file]
-            input_url = input_conf['url']
-
+        for input_file, input_conf in self._inputs_and_fetchers(expander.workload_name).items():
             if not workspace.dry_run:
-                fetcher = ramble.fetch_strategy.URLFetchStrategy(**input_conf)
-
-                file_name = '.'.join(os.path.basename(input_url).split('.')[0:-1])
-
                 mirror_paths = ramble.mirror.mirror_archive_paths(
-                    fetcher, os.path.join(self.name, file_name))
+                    input_conf['fetcher'], os.path.join(self.name, input_file))
 
-                with ramble.stage.InputStage(fetcher, name=workload_namespace,
+                with ramble.stage.InputStage(input_conf['fetcher'], name=workload_namespace,
                                              path=expander.application_input_dir,
                                              mirror_paths=mirror_paths) \
                         as stage:
                     stage.set_subdir(expander.expand_var(input_conf['target_dir']))
                     stage.fetch()
-                    if input_conf['sha256']:
+                    if input_conf['fetcher'].digest:
                         stage.check()
                     stage.cache_local()
                     if input_conf['expand']:
@@ -412,7 +414,7 @@ def _get_inputs(self, workspace, expander):
                         except spack.util.executable.ProcessError:
                             pass
             else:
-                tty.msg('DRY-RUN: Would download %s' % input_url)
+                tty.msg('DRY-RUN: Would download %s' % input_conf['fetcher'].url)
 
     def _make_experiments(self, workspace, expander):
         experiment_run_dir = expander.experiment_run_dir

lib/ramble/ramble/application_types/spack.py

@@ -6,6 +6,8 @@
 # option. This file may not be copied, modified, or distributed
 # except according to those terms.
 
+import os
+
 import llnl.util.tty as tty
 
 from ramble.application import ApplicationBase, ApplicationError
@@ -44,13 +46,19 @@ class SpackApplication(ApplicationBase):
     def __init__(self, file_path):
         super().__init__(file_path)
         self._setup_phases = [
+            'create_spack_env',
             'install_software',
             'get_inputs',
             'make_experiments'
         ]
 
         self._analyze_phases = ['analyze_experiments']
         self._archive_phases = ['archive_experiments']
+        self._mirror_phases = [
+            'mirror_inputs',
+            'create_spack_env',
+            'mirror_software'
+        ]
 
         self.application_class = 'SpackApplication'
 
@@ -84,39 +92,43 @@ def _add_expand_vars(self, expander):
         except ramble.spack_runner.RunnerError as e:
             tty.die(e)
 
-    def _install_software(self, workspace, expander):
-        import os
+    def _extract_specs(self, workspace, expander, spec_name, app_name):
+        """Build a list of all specs which the named spec requires
+
+        Traverse a spec and all of its dependencies to extract a list
+        of specs
+        """
+        spec_list = []
+        spec = workspace.get_named_spec(spec_name, app_name)
+        if 'dependencies' in spec:
+            for dep in spec['dependencies']:
+                spec_list.extend(
+                    self._extract_specs(workspace,
+                                        expander,
+                                        dep, app_name))
+        spec['application_name'] = app_name
+        spec_list.append((spec_name, spec))
+        return spec_list
+
+    def _create_spack_env(self, workspace, expander):
+        """Create the spack environment for this experiment
+
+        Extract all specs this experiment uses, and write the spack environment
+        file for it.
+        """
 
         # See if we cached this already, and if so return
         namespace = expander.spec_namespace
         if not namespace:
             raise ApplicationError('Ramble spec_namespace is set to None.')
 
-        cache_tupl = ('spack', namespace)
+        cache_tupl = ('spack-env', namespace)
         if workspace.check_cache(cache_tupl):
             tty.debug('{} already in cache.'.format(cache_tupl))
             return
         else:
             workspace.add_to_cache(cache_tupl)
 
-        def extract_specs(workspace, expander, spec_name, app_name):
-            """Build a list of all specs the named spec requires
-
-            Traverse a spec and all of its dependencies to extract a list
-            of specs
-            """
-            spec_list = []
-            spec = workspace.get_named_spec(spec_name, app_name)
-            if 'dependencies' in spec:
-                for dep in spec['dependencies']:
-                    spec_list.extend(
-                        extract_specs(workspace,
-                                      expander,
-                                      dep, app_name))
-            spec['application_name'] = app_name
-            spec_list.append((spec_name, spec))
-            return spec_list
-
         try:
             runner = ramble.spack_runner.SpackRunner(dry_run=workspace.dry_run)
 
@@ -145,8 +157,8 @@ def extract_specs(workspace, expander, spec_name, app_name):
                     runner.add_spec(workspace.spec_string(mpi_spec,
                                                           use_custom_specifier=True))
 
-                pkg_specs = extract_specs(workspace, expander, name,
-                                          app_context)
+                pkg_specs = self._extract_specs(workspace, expander, name,
+                                                app_context)
                 for pkg_name, pkg_info in pkg_specs:
                     if pkg_name not in added_specs:
                         added_specs[pkg_name] = True
@@ -170,20 +182,44 @@ def extract_specs(workspace, expander, spec_name, app_name):
                                 'defined in ramble.yaml')
 
             runner.concretize()
+
+        except ramble.spack_runner.RunnerError as e:
+            tty.die(e)
+
+    def _install_software(self, workspace, expander):
+        """Install application's software using spack"""
+
+        # See if we cached this already, and if so return
+        namespace = expander.spec_namespace
+        if not namespace:
+            raise ApplicationError('Ramble spec_namespace is set to None.')
+
+        cache_tupl = ('spack-install', namespace)
+        if workspace.check_cache(cache_tupl):
+            tty.debug('{} already in cache.'.format(cache_tupl))
+            return
+        else:
+            workspace.add_to_cache(cache_tupl)
+
+        try:
+            runner = ramble.spack_runner.SpackRunner(dry_run=workspace.dry_run)
+            runner.set_env(expander.expand_var('{spack_env}'))
+
+            runner.activate()
             runner.install()
 
+            app_context = expander.expand_var('{spec_name}')
             for name, spec_info in \
                     workspace.all_application_specs(app_context):
-                if 'mpi' in spec_info and \
-                        spec_info['mpi'] not in mpi_added:
+                if 'mpi' in spec_info:
                     mpi_spec = workspace.get_named_spec(spec_info['mpi'],
                                                         'mpi_library')
                     spec_str = workspace.spec_string(mpi_spec)
                     package_path = runner.get_package_path(spec_str)
                     expander.set_package_path(name, package_path)
 
-                pkg_specs = extract_specs(workspace, expander, name,
-                                          app_context)
+                pkg_specs = self._extract_specs(workspace, expander, name,
+                                                app_context)
                 for pkg_name, pkg_info in pkg_specs:
                     spec = workspace._build_spec_dict(pkg_info,
                                                       app_name=app_context)
@@ -194,3 +230,61 @@ def extract_specs(workspace, expander, spec_name, app_name):
 
         except ramble.spack_runner.RunnerError as e:
             tty.die(e)
+
+    def _mirror_software(self, workspace, expander):
+        """Mirror software source for this experiment using spack"""
+        import re
+
+        # See if we cached this already, and if so return
+        namespace = expander.spec_namespace
+        if not namespace:
+            raise ApplicationError('Ramble spec_namespace is set to None.')
+
+        cache_tupl = ('spack-mirror', namespace)
+        if workspace.check_cache(cache_tupl):
+            tty.debug('{} already in cache.'.format(cache_tupl))
+            return
+        else:
+            workspace.add_to_cache(cache_tupl)
+
+        try:
+            runner = ramble.spack_runner.SpackRunner(dry_run=workspace.dry_run)
+            runner.set_env(expander.expand_var('{spack_env}'))
+
+            runner.activate()
+
+            mirror_output = runner.mirror_environment(workspace._software_mirror_path)
+
+            present = 0
+            added = 0
+            failed = 0
+
+            present_regex = re.compile(r'\s+(?P<num>[0-9]+)\s+already present')
+            present_match = present_regex.search(mirror_output)
+            if present_match:
+                present = int(present_match.group('num'))
+
+            added_regex = re.compile(r'\s+(?P<num>[0-9]+)\s+added')
+            added_match = added_regex.search(mirror_output)
+            if added_match:
+                added = int(added_match.group('num'))
+
+            failed_regex = re.compile(r'\s+(?P<num>[0-9]+)\s+failed to fetch.')
+            failed_match = failed_regex.search(mirror_output)
+            if failed_match:
+                failed = int(failed_match.group('num'))
+
+            added_start = len(workspace._software_mirror_stats.new)
+            for i in range(added_start, added_start + added):
+                workspace._software_mirror_stats.new[i] = i
+
+            present_start = len(workspace._software_mirror_stats.present)
+            for i in range(present_start, present_start + present):
+                workspace._software_mirror_stats.present[i] = i
+
+            error_start = len(workspace._software_mirror_stats.errors)
+            for i in range(error_start, error_start + failed):
+                workspace._software_mirror_stats.errors.add(i)
+
+        except ramble.spack_runner.RunnerError as e:
+            tty.die(e)

lib/ramble/ramble/cmd/mirror.py

@@ -5,8 +5,6 @@
 # <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
 # option. This file may not be copied, modified, or distributed
 # except according to those terms.
-import sys
-
 import llnl.util.tty as tty
 
 import ramble.cmd.common.arguments as arguments
@@ -32,11 +30,6 @@ def setup_parser(subparser):
     sp = subparser.add_subparsers(
         metavar='SUBCOMMAND', dest='mirror_command')
 
-    # Create
-    create_parser = sp.add_parser('create', help=mirror_create.__doc__)
-    create_parser.add_argument('-d', '--directory', default=None,
-                               help="directory in which to create mirror")
-
     # Destroy
     destroy_parser = sp.add_parser('destroy', help=mirror_destroy.__doc__)
 
@@ -182,36 +175,6 @@ def _read_specs_from_file(filename):
     return specs
 
 
-def mirror_create(args):
-    """Create a directory to be used as a ramble mirror, and fill it with
-       input archives."""
-    ws = ramble.cmd.require_active_workspace(cmd_name='mirror create')
-
-    mirror = ramble.mirror.Mirror(
-        args.directory or ramble.config.get('config:input_cache'))
-
-    directory = url_util.format(mirror.push_url)
-
-    existed = web_util.url_exists(directory)
-
-    # Actually do the work to create the mirror
-    present, mirrored, error = ramble.mirror.create(
-        directory, ws)
-    p, m, e = len(present), len(mirrored), len(error)
-
-    verb = "updated" if existed else "created"
-    tty.msg(
-        "Successfully %s mirror in %s" % (verb, directory),
-        "Archive stats:",
-        "  %-4d already present"  % p,
-        "  %-4d added"            % m,
-        "  %-4d failed to fetch." % e)
-    if error:
-        tty.error("Failed downloads:")
-        tty.colify(s.cformat("{name}") for s in error)
-        sys.exit(1)
-
-
 def mirror_destroy(args):
     """Given a url, recursively delete everything under it."""
     mirror_url = None
@@ -226,8 +189,7 @@ def mirror_destroy(args):
 
 
 def mirror(parser, args):
-    action = {'create': mirror_create,
-              'destroy': mirror_destroy,
+    action = {'destroy': mirror_destroy,
               'add': mirror_add,
               'remove': mirror_remove,
               'rm': mirror_remove,