-
Notifications
You must be signed in to change notification settings - Fork 28
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Abstract container modifiers to use a container-base modifier
This commit extracts some of the shared logic for container based modifiers into a common container-base base_modifier. This allows more explicit standardization between the container modifiers.
- Loading branch information
1 parent
11d96f1
commit 6ac4bca
Showing
4 changed files
with
209 additions
and
410 deletions.
There are no files selected for viewing
172 changes: 172 additions & 0 deletions
172
var/ramble/repos/builtin/base_modifiers/container-base/base_modifier.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,172 @@ | ||
# Copyright 2022-2025 The Ramble Authors | ||
# | ||
# Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or | ||
# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license | ||
# <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your | ||
# option. This file may not be copied, modified, or distributed | ||
# except according to those terms. | ||
|
||
import os | ||
|
||
from ramble.modkit import * | ||
|
||
|
||
class ContainerBase(BasicModifier): | ||
"""This base modifier contains many of the variable and method definitions | ||
that containerized runtimes need to implement individual modifiers. It is | ||
used as a layer to standarize the interface for the containerized runtimes, | ||
to give them a consistent behavior within Ramble.""" | ||
|
||
name = "container-base" | ||
|
||
tags("container") | ||
|
||
maintainers("douglasjacobsen") | ||
|
||
mode("standard", description="Standard execution mode") | ||
default_mode("standard") | ||
|
||
required_variable( | ||
"container_uri", | ||
description="The variable controls the URI the container is pulled from. " | ||
"This should be of the format accepted by this container runtime.", | ||
) | ||
|
||
modifier_variable( | ||
"container_mounts", | ||
default="", | ||
description="Comma delimited list of mount points for the container. Filled in by modifier", | ||
modes=["standard"], | ||
) | ||
|
||
modifier_variable( | ||
"container_env_vars", | ||
default="", | ||
description="Comma delimited list of environments to import into container. Filled in by modifier", | ||
modes=["standard"], | ||
) | ||
|
||
modifier_variable( | ||
"container_extract_dir", | ||
default="{workload_input_dir}", | ||
description="Directory where the extracted paths will be stored", | ||
modes=["standard"], | ||
) | ||
|
||
modifier_variable( | ||
"container_extract_paths", | ||
default="[]", | ||
description="List of paths to extract from the sqsh file into the {workload_input_dir}. " | ||
+ "Will have paths of {workload_input_dir}/enroot_extractions/{path_basename}", | ||
modes=["standard"], | ||
track_used=False, | ||
) | ||
|
||
def _build_runner( | ||
self, runtime, check_software_env=False, app_inst=None, dry_run=False | ||
): | ||
"""Construct command runner for container runtime""" | ||
|
||
runner_name = f"{runtime}_runner" | ||
|
||
if ( | ||
not hasattr(self, runner_name) | ||
or getattr(self, runner_name) is None | ||
): | ||
path = None | ||
# If using spack, load spack environment before getting container runtime exec path | ||
if check_software_env and app_inst.package_manager is not None: | ||
if app_inst.package_manager.spec_prefix() == "spack": | ||
app_inst.package_manager.runner.activate() | ||
_, base = app_inst.package_manager.runner.get_package_path( | ||
runtime | ||
) | ||
app_inst.package_manager.runner.deactivate() | ||
|
||
if base and os.path.exists(base): | ||
test_path = os.path.join(base, "bin") | ||
if os.path.isdir(test_path): | ||
path = test_path | ||
|
||
exec_runner = CommandRunner( | ||
name=runtime, | ||
command=runtime, | ||
dry_run=dry_run, | ||
path=path, | ||
) | ||
|
||
setattr(self, runner_name, exec_runner) | ||
|
||
register_phase( | ||
"define_container_variables", | ||
pipeline="setup", | ||
run_before=["get_inputs"], | ||
) | ||
|
||
def _define_container_variables(self, workspace, app_inst=None): | ||
"""Define helper variables for working with containerized experiments | ||
To ensure it is defined properly, construct a comma delimited list of | ||
environment variable names that will be added into the | ||
container_env_vars variable. | ||
""" | ||
|
||
def extract_names(itr, name_set=set()): | ||
"""Extract names of environment variables from the environment variable action sets | ||
Given an iterator over environment variable action sets, extract | ||
the names of the environment variables. | ||
Modifies the name_set argument inplace. | ||
""" | ||
for action, conf in itr: | ||
if action in ["set", "unset"]: | ||
for name in conf: | ||
name_set.add(name) | ||
elif action == "prepend": | ||
for group in conf: | ||
for name in group["paths"]: | ||
name_set.add(name) | ||
elif action == "append": | ||
for group in conf: | ||
for name in group["vars"]: | ||
name_set.add(name) | ||
|
||
# Only define variables if mode is standard | ||
if self._usage_mode == "standard": | ||
# Define container_env-vars | ||
set_names = set() | ||
|
||
for env_var_set in app_inst._env_variable_sets: | ||
extract_names(env_var_set.items(), set_names) | ||
|
||
for mod_inst in app_inst._modifier_instances: | ||
extract_names(mod_inst.all_env_var_modifications(), set_names) | ||
|
||
env_var_list = ",".join(set_names) | ||
app_inst.define_variable("container_env_vars", env_var_list) | ||
|
||
# Define container_mounts | ||
input_mounts = app_inst.expander.expand_var("{container_mounts}") | ||
|
||
exp_mount = "{experiment_run_dir}:{experiment_run_dir}" | ||
expanded_exp_mount = app_inst.expander.expand_var(exp_mount) | ||
|
||
if ( | ||
exp_mount not in input_mounts | ||
and expanded_exp_mount not in input_mounts | ||
): | ||
add_mod = self._usage_mode not in self.variable_modifications | ||
add_mod = add_mod or ( | ||
self._usage_mode in self.variable_modifications | ||
and "container_mounts" | ||
not in self.variable_modifications[self._usage_mode] | ||
) | ||
if add_mod: | ||
self.variable_modification( | ||
"container_mounts", | ||
modification=exp_mount, | ||
separator=",", | ||
method="append", | ||
mode=self._usage_mode, | ||
) |
Oops, something went wrong.