Add Missing Helper from Refactor and Cleanup Logic (#1265)

* Add get_spack_variants back

* Simplify class and checks. Only add spack variant if defined. Only define variant if used by experiment

* Simplify checks since error will be caught by spec logic

Author: michaelmckinsey1

lib/benchpark/experiment.py

@@ -255,34 +255,22 @@ def __init__(self, spec):
 
         self.package_specs = {}
 
-        # Set available programming models for checks
-        models = set()
-        for cls in self.__class__.mro():
-            models.update(getattr(cls, "_available_programming_models", ()))
-        self.programming_models = list(models)
+        # Get available programming models for checks
+        self.programming_models = list(self._available_programming_models)
 
         # Explicitly ordered list. "mpi" first
         models = ["mpi"] + ["openmp", "cuda", "rocm"]
-        valid_models = []
-        invalid_models = []
-        for model in models:
-            if self.spec.satisfies("+" + model):
-                valid_models.append(model)
-                # Experiment specifying model in add_package_spec that it doesn't implement
-                if model not in self.programming_models:
-                    invalid_models.append(model)
-        # MPI is always valid if with another programming model, even if no mpionly
-        if "mpi" in invalid_models and len(valid_models) > 1:
-            invalid_models.remove("mpi")
         # Case where there are no experiments specified in experiment.py
         if len(self.programming_models) == 0:
             raise BenchparkError(
                 f"Please specify a programming model in your {self.name}/experiment.py (e.g. ProgrammingModelType.Mpionly, ProgrammingModelType.Openmp, ProgrammingModelType.Cuda, ProgrammingModelType.Rocm). See other experiments for examples."
             )
-        elif len(invalid_models) > 0:
-            print(self.spec)
+        # Check if experiment is trying to run in MpiOnly mode without being an ProgrammingModelType.Mpionly experiment
+        elif "mpi" not in str(self.spec) and not any(
+            self.spec.satisfies("+" + model) for model in models[1:]
+        ):
             raise BenchparkError(
-                f'{invalid_models} are not valid programming models for "{self.name}". Choose from {self.programming_models}.'
+                f'"{self.name}" cannot run with MPI only without inheriting from ProgrammingModelType.Mpionly. Choose from {self.programming_models}'
             )
 
         if (

lib/benchpark/programming_model.py

@@ -25,16 +25,23 @@ def ProgrammingModel(*types):
     # Normalize once so we can reuse
     _available = tuple(t.value for t in types)
 
-    class BaseModel:
-        requires("mpi", when="+mpi")
-        requires("rocm", when="+rocm")
-        requires("cuda", when="+cuda")
-        requires("openmp", when="+openmp")
+    @staticmethod
+    def is_available(mod):
+        return mod in _available
 
-        variant("mpi", default=True, description="Run with MPI")
-        variant("rocm", default=False, description="Build and run with ROCm")
-        variant("cuda", default=False, description="Build and run with CUDA")
-        variant("openmp", default=False, description="Build and run with OpenMP")
+    class BaseModel:
+        if "mpi" in _available:
+            requires("mpi", when="+mpi")
+            variant("mpi", default=True, description="Run with MPI")
+        if "rocm" in _available:
+            requires("rocm", when="+rocm")
+            variant("rocm", default=False, description="Build and run with ROCm")
+        if "cuda" in _available:
+            requires("cuda", when="+cuda")
+            variant("cuda", default=False, description="Build and run with CUDA")
+        if "openmp" in _available:
+            requires("openmp", when="+openmp")
+            variant("openmp", default=False, description="Build and run with OpenMP")
 
         # Class-level list of supported models for any class that includes this mixin
         _available_programming_models = _available
@@ -48,12 +55,7 @@ def available_programming_models(self):
                 models.update(getattr(cls, "_available_programming_models", ()))
             return tuple(sorted(models))
 
-        # Quick check helper
-        @staticmethod
-        def supports_model(name: str) -> bool:
-            return name in _available
-
-    # Helper class (unchanged except for optional new method)
+    # Helper class
     class Helper(ExperimentHelper):
         def get_helper_name_prefix(self):
             models = []
@@ -69,12 +71,31 @@ def get_helper_name_prefix(self):
                 return models
             return "no_model"
 
-        # Optional: expose *available* (not selected) models via helper, too
-        def get_available_models(self):
-            models = set()
-            for cls in type(self).__mro__:
-                models.update(getattr(cls, "_available_programming_models", ()))
-            return tuple(sorted(models))
+        def get_spack_variants(self):
+            models = []
+            model_dict = {
+                ProgrammingModelType.Openmp.value: ["+openmp", "~openmp"],
+                ProgrammingModelType.Cuda.value: [
+                    "+cuda cuda_arch={cuda_arch}",
+                    "~cuda",
+                ],
+                ProgrammingModelType.Rocm.value: [
+                    "+rocm amdgpu_target={rocm_arch}",
+                    "~rocm",
+                ],
+            }
+            for s in [
+                ProgrammingModelType.Openmp.value,
+                ProgrammingModelType.Cuda.value,
+                ProgrammingModelType.Rocm.value,
+            ]:
+                if is_available(s):
+                    if self.spec.satisfies("+" + s):
+                        models.append(model_dict[s][0])
+                    else:
+                        models.append(model_dict[s][1])
+
+            return " ".join(models)
 
     return type(
         "ProgrammingModelType",

lib/benchpark/test/experiment_errors.py

@@ -12,15 +12,14 @@
 def test_programming_model_checks():
     # babelstream mpi-only not valid
     with pytest.raises(
-        BenchparkError, match=r"mpi.*are not valid programming models for.*babelstream"
+        BenchparkError,
+        match=r"babelstream.*cannot run with MPI only without inheriting from",
     ):
         spec = benchpark.spec.ExperimentSpec("babelstream").concretize()
         experiment = spec.experiment  # noqa: F841
 
     # stream+openmp not valid
-    with pytest.raises(
-        Exception, match=r"openmp.*are not valid programming models for.*stream.*mpi"
-    ):
+    with pytest.raises(Exception, match=r"openmp is not a valid variant of stream"):
         spec = benchpark.spec.ExperimentSpec(
             "stream+openmp workload=stream"
         ).concretize()