Augmentations: Allow configuring of augmentations per workload

Add support for enabling and disabling of augmentations (instruments and
output processors) on a per workload basis.

wa/framework/execution.py

@@ -273,7 +273,7 @@ class Executor(object):
 
         Params::
 
-            :state: a ``ConfigManager`` containing processed configuraiton
+            :state: a ``ConfigManager`` containing processed configuration
             :output: an initialized ``RunOutput`` that will be used to
                      store the results.
 
@@ -402,6 +402,7 @@ class Runner(object):
     def finalize_run(self):
         self.logger.info('Finalizing run')
         self.context.end_run()
+        self.pm.enable_all()
         self.pm.process_run_output(self.context)
         self.pm.export_run_output(self.context)
         self.pm.finalize()
@@ -446,6 +447,9 @@ class Runner(object):
         job.set_status(Status.RUNNING)
         self.send(signal.JOB_STARTED)
 
+        self.logger.info('Configuring augmentations')
+        job.configure_augmentations(context, self.pm)
+
         with signal.wrap('JOB_TARGET_CONFIG', self, context):
             job.configure_target(context)
 

wa/framework/job.py

@@ -1,7 +1,7 @@
 import logging
 from datetime import datetime
 
-from wa.framework import pluginloader, signal
+from wa.framework import pluginloader, signal, instruments
 from wa.framework.configuration.core import Status
 
 # Because of use of Enum (dynamic attrs)
@@ -64,6 +64,33 @@ class Job(object):
         self.set_status(Status.PENDING)
         context.update_job_state(self)
 
+    def configure_augmentations(self, context, pm):
+        instruments_to_enable = set()
+        output_processors_to_enable = set()
+        enabled_instruments = set(i.name for i in instruments.get_enabled())
+        enabled_output_processors = set(p.name for p in pm.get_enabled())
+
+        for augmentation in self.spec.augmentations.values():
+            augmentation_cls = context.cm.plugin_cache.get_plugin_class(augmentation)
+            if augmentation_cls.kind == 'instrument':
+                instruments_to_enable.add(augmentation)
+            elif augmentation_cls.kind == 'output_processor':
+                output_processors_to_enable.add(augmentation)
+
+        # Disable unrequired instruments
+        for instrument in enabled_instruments.difference(instruments_to_enable):
+            instruments.disable(instrument)
+        # Enable additional instruments
+        for instrument in instruments_to_enable.difference(enabled_instruments):
+            instruments.enable(instrument)
+
+        # Disable unrequired output_processors
+        for processor in enabled_output_processors.difference(output_processors_to_enable):
+            pm.disable(processor)
+        # Enable additional output_processors
+        for processor in output_processors_to_enable.difference(enabled_output_processors):
+            pm.enable(processor)
+
     def configure_target(self, context):
         self.logger.info('Configuring target for job {} [{}]'.format(self.id, self.iteration))
         context.tm.commit_runtime_parameters(self.spec.runtime_parameters)

wa/framework/output_processor.py

@@ -5,6 +5,7 @@ from wa.framework.exception import ConfigError
 from wa.framework.instruments import is_installed
 from wa.framework.plugin import Plugin
 from wa.utils.log import log_error, indent, dedent
+from wa.utils.misc import isiterable
 
 
 class OutputProcessor(Plugin):
@@ -12,6 +13,10 @@ class OutputProcessor(Plugin):
     kind = 'output_processor'
     requires = []
 
+    def __init__(self, **kwargs):
+        super(OutputProcessor, self).__init__(**kwargs)
+        self.is_enabled = True
+
     def validate(self):
         super(OutputProcessor, self).validate()
         for instrument in self.requires:
@@ -40,6 +45,43 @@ class ProcessorManager(object):
         processor.logger.context = context
         self.processors.append(processor)
 
+    def disable_all(self):
+        for output_processor in self.processors:
+            self._disable_output_processor(output_processor)
+
+    def enable_all(self):
+        for output_processor in self.processors:
+            self._enable_output_processor(output_processor)
+
+    def enable(self, to_enable):
+        if isiterable(to_enable):
+            for inst in to_enable:
+                self._enable_output_processor(inst)
+        else:
+            self._enable_output_processor(to_enable)
+
+    def disable(self, to_disable):
+        if isiterable(to_disable):
+            for inst in to_disable:
+                self._disable_output_processor(inst)
+        else:
+            self._disable_output_processor(to_disable)
+
+    def get_output_processor(self, processor):
+        if isinstance(processor, OutputProcessor):
+            return processor
+
+        for p in self.processors:
+            if processor == p.name:
+                return p
+        raise ValueError('Output processor {} is not installed'.format(processor))
+
+    def get_enabled(self):
+        return [p for p in self.processors if p.is_enabled]
+
+    def get_disabled(self):
+        return [p for p in self.processors if not p.is_enabled]
+
     def validate(self):
         for proc in self.processors:
             proc.validate()
@@ -74,15 +116,29 @@ class ProcessorManager(object):
         try:
             indent()
             for proc in self.processors:
-                proc_func = getattr(proc, method_name, None)
-                if proc_func is None:
-                    continue
-                try:
-                    self.logger.info(message.format(proc.name))
-                    proc_func(*args)
-                except Exception as e:
-                    if isinstance(e, KeyboardInterrupt):
-                        raise
-                    log_error(e, self.logger)
+                if proc.is_enabled:
+                    proc_func = getattr(proc, method_name, None)
+                    if proc_func is None:
+                        continue
+                    try:
+                        self.logger.info(message.format(proc.name))
+                        proc_func(*args)
+                    except Exception as e:
+                        if isinstance(e, KeyboardInterrupt):
+                            raise
+                        log_error(e, self.logger)
         finally:
             dedent()
+
+    def _enable_output_processor(self, inst):
+        inst = self.get_output_processor(inst)
+        self.logger.debug('Enabling output processor {}'.format(inst.name))
+        if not inst.is_enabled:
+            inst.is_enabled = True
+
+    def _disable_output_processor(self, inst):
+        inst = self.get_output_processor(inst)
+        self.logger.debug('Disabling output processor {}'.format(inst.name))
+        if inst.is_enabled:
+            inst.is_enabled = False
+