Skip to content

[WIP] Use Enums to Enforce Syntax #945

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

[WIP] Use Enums to Enforce Syntax #945

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 70 additions & 29 deletions lib/benchpark/experiment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,31 +85,59 @@ def get_helper_name_prefix(self):
return "single_node" if self.spec.satisfies("+single_node") else ""


class AffinityVariantValues(str, Enum):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this could be AffinityEnum and we could make a base enum class for the tuple_values methods etc. that are shared across enums

AFFINITY = "affinity"
NAME = "name"
MODE = "mode"
PKG_SPEC = "pkg_spec"
COMPILER = "compiler"

NONE = "none"
ON = "on"

CUDA = "cuda"
ROCM = "rocm"
MPI = "mpi"
MODE = "mode"


@classmethod
def tuple_values(cls):
"""Return only the original values for tuple creation"""
return tuple(v.value for v in [cls.NONE, cls.ON])

@classmethod
def variant(cls, name):
"""Return variant name with + prefix"""
return f"+{name}"


class Affinity:
variant(
"affinity",
default="none",
values=(
"none",
"on",
),
AffinityVariantValues.AFFINITY.value,
default=AffinityVariantValues.NONE.value,
values=AffinityVariantValues.tuple_values(),
multi=False,
description="Build and run the affinity package",
)

class Helper(ExperimentHelper):
def compute_modifiers_section(self):
modifier_list = []
if not self.spec.satisfies("affinity=none"):
if not self.spec.satisfies(f"{AffinityVariantValues.AFFINITY.value}={AffinityVariantValues.NONE.value}"):
affinity_modifier_modes = {}
affinity_modifier_modes["name"] = "affinity"
if self.spec.satisfies("+cuda"):
affinity_modifier_modes["mode"] = "cuda"
elif self.spec.satisfies("+rocm"):
affinity_modifier_modes["mode"] = "rocm"
affinity_modifier_modes[AffinityVariantValues.NAME.value] = AffinityVariantValues.AFFINITY.value

if self.spec.satisfies(AffinityVariantValues.variant(AffinityVariantValues.CUDA.value)):
affinity_modifier_modes[AffinityVariantValues.MODE.val] = AffinityVariantValues.CUDA.value

elif self.spec.satisfies(AffinityVariantValues.variant(AffinityVariantValues.ROCM.value)):
affinity_modifier_modes[AffinityVariantValues.MODE.val] = AffinityVariantValues.ROCM.value

else:
affinity_modifier_modes["mode"] = "mpi"
affinity_modifier_modes[AffinityVariantValues.MODE.val] = AffinityVariantValues.MPI.value
modifier_list.append(affinity_modifier_modes)

return modifier_list

def compute_package_section(self):
Expand All @@ -119,39 +147,52 @@ def compute_package_section(self):
# get system config options
# TODO: Get compiler/mpi/package handles directly from system.py
system_specs = {}
system_specs["compiler"] = "default-compiler"
if self.spec.satisfies("+cuda"):
system_specs[AffinityVariantValues.COMPILER.value] = "default-compiler"

if self.spec.satisfies(AffinityVariantValues.variant(AffinityVariantValues.CUDA.value)):
system_specs["cuda_arch"] = "{cuda_arch}"
if self.spec.satisfies("+rocm"):

if self.spec.satisfies(AffinityVariantValues.variant(AffinityVariantValues.ROCM.value)):
system_specs["rocm_arch"] = "{rocm_arch}"

# set package spack specs
package_specs = {}

if not self.spec.satisfies("affinity=none"):
package_specs["affinity"] = {
"pkg_spec": f"affinity@{affinity_version}+mpi",
"compiler": system_specs["compiler"],
if not self.spec.satisfies(f"{AffinityVariantValues.AFFINITY.value}={AffinityVariantValues.NONE.value}"):
package_specs[AffinityVariantValues.AFFINITY.value] = {
AffinityVariantValues.PKG_SPEC.value: f"{AffinityVariantValues.AFFINITY.value}@{affinity_version}{AffinityVariantValues.variant(AffinityVariantValues.MPI.value)}",
AffinityVariantValues.COMPILER.value: system_specs[AffinityVariantValues.COMPILER.value],
}
if self.spec.satisfies("+cuda"):
package_specs["affinity"]["pkg_spec"] += "+cuda"
elif self.spec.satisfies("+rocm"):
package_specs["affinity"]["pkg_spec"] += "+rocm"

if self.spec.satisfies(AffinityVariantValues.variant(AffinityVariantValues.CUDA.value)):
package_specs[AffinityVariantValues.AFFINITY.value][AffinityVariantValues.PKG_SPEC.value] += AffinityVariantValues.variant(AffinityVariantValues.CUDA.value)

elif self.spec.satisfies(AffinityVariantValues.variant(AffinityVariantValues.ROCM.value)):
package_specs[AffinityVariantValues.AFFINITY.value][AffinityVariantValues.PKG_SPEC.value] += AffinityVariantValues.variant(AffinityVariantValues.ROCM.value)

return {
"packages": {k: v for k, v in package_specs.items() if v},
"environments": {"affinity": {"packages": list(package_specs.keys())}},
"environments": {AffinityVariantValues.AFFINITY.value: {"packages": list(package_specs.keys())}},
}


class HwlocVariantValues(str, Enum):
HWLOC = "hwloc"
NAME = "name"
MODE = "mode"

NONE = "none"
ON = "on"

@classmethod
def tuple_values(cls):
"""Return only the original values for tuple creation"""
return tuple(v.value for v in [cls.NONE, cls.ON])


class Hwloc:
variant(
"hwloc",
HwlocVariantValues.HWLOC.value,
default=HwlocVariantValues.NONE.value,
values=tuple(v.value for v in HwlocVariantValues),
multi=False,
Expand All @@ -162,10 +203,10 @@ class Helper(ExperimentHelper):
def compute_modifiers_section(self):
modifier_list = []

if not self.spec.satisfies(f"hwloc={HwlocVariantValues.NONE.value}"):
if not self.spec.satisfies(f"{HwlocVariantValues.HWLOC.value}={HwlocVariantValues.NONE.value}"):
affinity_modifier_modes = {}
affinity_modifier_modes["name"] = "hwloc"
affinity_modifier_modes["mode"] = self.spec.variants["hwloc"][0]
affinity_modifier_modes[HwlocVariantValues.NAME.value] = HwlocVariantValues.HWLOC.value
affinity_modifier_modes[HwlocVariantValues.MODE.value] = self.spec.variants[HwlocVariantValues.HWLOC.value][0]
modifier_list.append(affinity_modifier_modes)

return modifier_list
Expand Down
Loading