From 57f6a9dbf8bc4477e705895c0a32e024f3cb236d Mon Sep 17 00:00:00 2001
From: august-knox <knox10@llnl.gov>
Date: Wed, 13 Nov 2024 11:29:44 -0800
Subject: [PATCH] progress on updating osu-microbenchmarks

---
 .../osu-micro-benchmarks/experiment.py        | 100 +++++++++++++-----
 1 file changed, 72 insertions(+), 28 deletions(-)

diff --git a/experiments/osu-micro-benchmarks/experiment.py b/experiments/osu-micro-benchmarks/experiment.py
index 60b3e753..f469f98b 100644
--- a/experiments/osu-micro-benchmarks/experiment.py
+++ b/experiments/osu-micro-benchmarks/experiment.py
@@ -1,9 +1,18 @@
 from benchpark.directives import variant
+from benchpark.error import BenchparkError
 from benchpark.experiment import Experiment
+from benchpark.scaling import StrongScaling
 import re
 
 
 class OsuMicroBenchmarks(Experiment):
+
+    variant(
+        "version",
+        default="develop",
+        description="app version",
+    )
+
     variant(
         "workload",
         default="osu_latency",
@@ -83,8 +92,19 @@ class OsuMicroBenchmarks(Experiment):
     )
 
     def compute_applications_section(self):
-        variables = {}
+	scaling_modes = {
+            "strong": self.spec.satisfies("strong=oui"),
+            "single_node": self.spec.satisfies("single_node=oui"),
+        }
+
+        scaling_mode_enabled = [key for key, value in scaling_modes.items() if value]
+        if len(scaling_mode_enabled) != 1:
+            raise BenchparkError(
+                f"Only one type of scaling per experiment is allowed for application package {self.name}"
+            )
 
+        variables = {}
+	num_nodes = {"n_nodes", 1}
         variables["scaling_nodes"] = "2"
         variables["n_nodes"] = "{scaling_nodes}"
         variables["n_ranks_per_node"] = "36"
@@ -93,37 +113,61 @@ def compute_applications_section(self):
         workload = re.search(pattern, str(self.spec.variants))
         workload_string = workload.group(1)
 
-        return {
-            "osu-micro-benchmarks": {
-                # TODO replace with a hash once we have one?
-                "workloads": {
-                    workload_string: {
-                        "experiments": {
-                            "scaling_{n_nodes}nodes_medium": {
-			        "variants": {
-                                    "package_manager": "spack",
-                                },
-                                "variables": variables,
-                            }
-                        }
-                    }
-                }
-            }
+	if self.spec.satisfies("single_node=oui"):
+            for pk, pv in num_nodes.items():
+                self.add_experiment_variable(pk, pv, True)
+
+	elif self.spec.satisfies("strong=oui"):
+            scaled_variables = self.generate_strong_scaling_params(
+                {tuple(num_nodes.keys()): list(num_nodes.values())},
+                int(self.spec.variants["scaling-factor"][0]),
+                int(self.spec.variants["scaling-iterations"][0]),
+            )
+            for pk, pv in scaled_variables.items():
+                self.add_experiment_variable(pk, pv, True)
+       # return {
+        #    "osu-micro-benchmarks": {
+         #       # TODO replace with a hash once we have one?
+          #      "workloads": {
+           #         workload_string: {
+            #            "experiments": {
+             #               "scaling_{n_nodes}nodes_medium": {
+	#		        "variants": {
+         #                           "package_manager": "spack",
+          #                      },
+           #                     "variables": variables,
+            #                }
+             #           }
+              #      }
+               # }
+           # }
         }
 
     def compute_spack_section(self):
+	app_version = self.spec.variants["version"][0]
+        # get system config options
+        # TODO: Get compiler/mpi/package handles directly from system.py
+        system_specs = {}
+        system_specs["compiler"] = "default-compiler"
+        system_specs["mpi"] = "default-mpi"
+        # set package spack specs
+        # empty package_specs value implies external package
+        self.add_spack_spec(system_specs["mpi"])
+        self.add_spack_spec(
+            self.name, [f"osu-micro-benchmarks@{app_version}", system_specs["compiler"]]
+        )
         # TODO: express that we need certain variables from system
         # Does not need to happen before merge, separate task
-        osu_microbenchmarks_spack_spec = "osu-micro-benchmarks"
-        packages = ["default-mpi", self.spec.name]
+        #osu_microbenchmarks_spack_spec = "osu-micro-benchmarks"
+       # packages = ["default-mpi", self.spec.name]
 
-        return {
-            "packages": {
-                "osu-micro-benchmarks": {
-                    "pkg_spec": osu_microbenchmarks_spack_spec,
-                    "compiler": "default-compiler",
-                }
-            },
-            "environments": {"osu-micro-benchmarks": {"packages": packages}},
+#        return {
+ #           "packages": {
+  #              "osu-micro-benchmarks": {
+   #                 "pkg_spec": osu_microbenchmarks_spack_spec,
+    #                "compiler": "default-compiler",
+     #           }
+      #      },
+       #     "environments": {"osu-micro-benchmarks": {"packages": packages}},
             
-        }
+       # }