allow conditional variant definition (#379)

* allow conditional variant definition

* style

---------

Co-authored-by: pearce8 <[email protected]>

lib/benchpark/directives.py

@@ -53,6 +53,53 @@ def __init__(self, *args, **kwargs):
 benchpark_directive = DirectiveMeta.directive
 
 
+def _make_when_spec(
+    value: Optional[Union["benchpark.spec.Spec", str, bool]]
+) -> Optional["benchpark.spec.Spec"]:
+    """Create a ``Spec`` that indicates when a directive should be applied.
+
+    Directives with ``when`` specs, e.g.:
+
+        patch('foo.patch', when='@4.5.1:')
+        depends_on('mpi', when='+mpi')
+        depends_on('readline', when=sys.platform() != 'darwin')
+
+    are applied conditionally depending on the value of the ``when``
+    keyword argument.  Specifically:
+
+      1. If the ``when`` argument is ``True``, the directive is always applied
+      2. If it is ``False``, the directive is never applied
+      3. If it is a ``Spec`` string, it is applied when the package's
+         concrete spec satisfies the ``when`` spec.
+
+    The first two conditions are useful for the third example case above.
+    It allows package authors to include directives that are conditional
+    at package definition time, in additional to ones that are evaluated
+    as part of concretization.
+
+    Arguments:
+        value: a conditional Spec, constant ``bool``, or None if not supplied
+           value indicating when a directive should be applied.
+
+    """
+    if isinstance(value, benchpark.spec.Spec):
+        return value
+
+    # Unsatisfiable conditions are discarded by the caller, and never
+    # added to the package class
+    if value is False:
+        return None
+
+    # If there is no constraint, the directive should always apply;
+    # represent this by returning the unconstrained `Spec()`, which is
+    # always satisfied.
+    if value is None or value is True:
+        return benchpark.spec.Spec()
+
+    # This is conditional on the spec
+    return benchpark.spec.Spec(value)
+
+
 @benchpark_directive("variants")
 def variant(
     name: str,
@@ -82,6 +129,8 @@ def variant(
     Raises:
         DirectiveError: If arguments passed to the directive are invalid
     """
+    if sticky:
+        raise NotImplementedError("Sticky variants are not yet implemented in Ramble")
 
     def format_error(msg, pkg):
         msg += " @*r{{[{0}, variant '{1}']}}"
@@ -140,12 +189,15 @@ def _raise_default_not_set(pkg):
     description = str(description).strip()
 
     def _execute_variant(pkg):
+        when_spec = _make_when_spec(when)
+
         if not re.match(benchpark.spec.IDENTIFIER, name):
             directive = "variant"
             msg = "Invalid variant name in {0}: '{1}'"
             raise DirectiveError(directive, msg.format(pkg.name, name))
 
-        pkg.variants[name] = benchpark.variant.Variant(
+        variants_by_name = pkg.variants.setdefault(when_spec, {})
+        variants_by_name[name] = benchpark.variant.Variant(
             name, default, description, values, multi, validator, sticky
         )
 

lib/benchpark/experiment.py

@@ -3,7 +3,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Dict, Tuple
+from typing import Dict
 import yaml  # TODO: some way to ensure yaml available
 
 from benchpark.directives import ExperimentSystemBase
@@ -44,8 +44,8 @@ class Experiment(ExperimentSystemBase):
 
     # This allows analysis tools to correctly interpret the class attributes.
     variants: Dict[
-        str,
-        Tuple["benchpark.variant.Variant", "benchpark.spec.ConcreteExperimentSpec"],
+        "benchpark.spec.Spec",
+        Dict[str, benchpark.variant.Variant],
     ]
 
     def __init__(self, spec):

lib/benchpark/spec.py

@@ -46,6 +46,12 @@ def satisfies(self, other: "VariantMap") -> bool:
             name in self and set(self[name]) >= set(other[name]) for name in other
         )
 
+    def constrain(self, other: "VariantMap") -> None:
+        for name in other:
+            self_values = list(self.dict.get(name, []))
+            values = llnl.util.lang.dedupe(self_values + list(other[name]))
+            self[name] = values
+
     @staticmethod
     def stringify(name: str, values: tuple) -> str:
         if len(values) == 1:
@@ -124,6 +130,10 @@ def __eq__(self, other: "Spec"):
             and self.variants == other.variants
         )
 
+    def __hash__(self):
+        # If hashing specs, client code is responsible for ensuring they do not mutate
+        return hash((self.name, self.namespace, self.variants))
+
     def __str__(self):
         string = ""
         if self.namespace is not None:
@@ -176,6 +186,21 @@ def satisfies(self, other: Union[str, "Spec"]) -> bool:
             and self.variants.satisfies(other.variants)
         )
 
+    def constrain(self, other: Union[str, "Spec"]) -> None:
+        if not isinstance(other, Spec):
+            other = Spec(other)
+        if other.name:
+            if self.name and self.name != other.name:
+                raise Exception
+            self.name == other.name
+
+        if other.namespace:
+            if self.namespace and self.namespace != other.namespace:
+                raise Exception
+            self.namespace == other.namespace
+
+        self.variants.constrain(other.variants)
+
     def concretize(self):
         raise NotImplementedError("Spec.concretize must be implemented by subclass")
 
@@ -205,9 +230,6 @@ def __init__(self, str_or_spec: Union[str, Spec]):
         super().__init__(str_or_spec)
         self._concretize()
 
-    def __hash__(self):
-        return hash((self.name, self.namespace, self.variants))
-
     @property
     def name(self):
         return self._name
@@ -237,19 +259,62 @@ def _concretize(self):
             raise AnonymousSpecError(f"Cannot concretize anonymous {type(self)} {self}")
 
         if not self.namespace:
-            # TODO interface combination ##
             self._namespace = self.object_class.namespace
 
-        # TODO interface combination ##
-        for name, variant in self.object_class.variants.items():
-            if name not in self.variants:
-                self._variants[name] = variant.default
+        # For variants that are set, set whatever they imply
+        variants_to_check = set(
+            (name, values) for name, values in self.variants.items()
+        )
+        checked = set()
+        while variants_to_check:
+            name, values = variants_to_check.pop()
+            checked.add((name, values))
+
+            conditions = [
+                w
+                for w, v_by_n in self.object_class.variants.items()
+                for n, v in v_by_n.items()
+                if n == name and v.validate_values_bool(values)
+            ]
+
+            if not conditions:
+                raise Exception(f"{name} is not a valid variant of {self.name}")
+
+            # This variant is already valid on self
+            if any(self.satisfies(c) for c in conditions):
+                continue
 
+            # If there is not a condition that allows this variant already, add one
+            cond = conditions[0]
+            for n, v in cond.variants.items():
+                if (n, v) not in checked:
+                    variants_to_check.add((n, v))
+            self.constrain(cond)
+
+        # Concretize variants that aren't set
+        changed = True
+        while changed:
+            changed = False
+
+            for when, variants_by_name in self.object_class.variants.items():
+                for name, variant in variants_by_name.items():
+                    if self.satisfies(when):
+                        if name not in self.variants:
+                            changed = True
+                            self._variants[name] = variant.default
+
+        # Validate all set variant values
         for name, values in self.variants.items():
-            if name not in self.object_class.variants:
+            try:
+                variant = next(
+                    v
+                    for w, v_by_n in self.object_class.variants.items()
+                    for n, v in v_by_n.items()
+                    if self.satisfies(w) and n == name
+                )
+            except StopIteration:
                 raise Exception(f"{name} is not a valid variant of {self.name}")
 
-            variant = self.object_class.variants[name]
             variant.validate_values(self.variants[name])
 
         # Convert to immutable type

lib/benchpark/variant.py

@@ -86,6 +86,14 @@ def validate_values(self, variant_values, pkg_cls=None):
         if not_allowed_values:
             raise ValueError(f"{not_allowed_values} are not valid values for {pkg_cls}")
 
+    def validate_values_bool(self, *args, **kwargs):
+        """Wrapper around ``validate_values`` that returns boolean instead of raising."""
+        try:
+            self.validate_values(*args, **kwargs)
+            return True
+        except Exception:
+            return False
+
     @property
     def allowed_values(self):
         """Returns a string representation of the allowed values for