import argparse
import os
import pickle
import shutil
import sys
import warnings
from typing import Dict, List
from polaris import Task, provenance
from polaris.component_graph import (
get_components_in_use,
get_steps_by_component,
)
from polaris.config import PolarisConfigParser
from polaris.io import symlink
from polaris.job import write_job_script
from polaris.machines import discover_machine
from polaris.parallel import set_parallel_systems
from polaris.tasks import get_components
[docs]
def setup_tasks(
work_dir,
task_list=None,
numbers=None,
config_file=None,
machine=None,
baseline_dir=None,
component_path=None,
suite_name='custom',
suite_component=None,
cached=None,
free_running=None,
copy_executable=False,
clean_tasks=False,
model=None,
build=None,
branch=None,
clean_build=None,
quiet_build=None,
cmake_flags=None,
debug=None,
component_args=None,
):
"""
Set up one or more tasks
Parameters
----------
work_dir : str
A directory that will serve as the base for creating task
directories
task_list : list of str, optional
Relative paths for a tasks to set up
numbers : list of str, optional
Task numbers to setup, as listed from ``polaris list``, optionally with
a suffix ``c`` to indicate that all steps in that task should be
cached
config_file : {str, None}, optional
Configuration file with custom options for setting up and running tasks
machine : str, optional
The name of one of the machines with defined config options, which can
be listed with ``polaris list --machines``
baseline_dir : str, optional
Location of baselines that can be compared to
component_path : str, optional
The relative or absolute path to the location where the model and
default namelists have been (or will be) built
suite_name : str, optional
The name of the suite if tasks are being set up through a suite or
``'custom'`` if not
suite_component : str, optional
The name of the component the suite belongs to, which is where its
config file comes from. A suite may contain tasks from other
components, so this is not necessarily the component of any of the
tasks. Defaults to the component of the first task.
cached : list of list of str, optional
For each task in ``tasks``, which steps (if any) should read their
outputs from the cache, or a list with "_all" as the first entry if
all steps in the task should use cached outputs. These override
the task's ``free_running_steps``.
free_running : list of list of str, optional
For the single task in ``tasks``, which steps (if any) should be
forced free-running (not cached), overriding ``default_cached``,
or a list with ``"_all"`` as the first entry if all steps should
be free-running. A step may not appear in both ``cached`` and
``free_running``.
copy_executable : bool, optional
Whether to copy the model executable to the work directory
clean_tasks : bool, optional
Whether to delete the contents of the task work directories before
setting up tasks
model : str, optional
The model to run
build : bool, optional
Whether to build the model
branch : str, optional
The relative or absolute path to the base of the component branch to
build. If not provided, ``component_path`` must point to an existing
build.
clean_build : bool, optional
Whether to clean the build directory before building the model.
``clean_build = True`` implies ``build = True``.
quiet_build : bool, optional
Whether to build the model without output. Implies ``build = True``.
cmake_flags : str, optional
Additional flags to pass to make or CMake when building the model
debug : bool, optional
Whether to build the model in debug mode
component_args : dict of dict, optional
The model, component path and branch for components other than the one
that owns the tasks, with the component names as keys. These come from
the ``--<component>_model``, ``--<component>_path`` and
``--<component>_branch`` flags.
Returns
-------
tasks : dict of polaris.Task
A dictionary of tasks, with the relative path in the work
directory as keys
"""
machine = __get_machine_and_check_params(
machine, config_file, task_list, numbers, cached
)
if work_dir is None:
print(
'Warning: no base work directory was provided so setting up in '
'the current directory.'
)
work_dir = os.getcwd()
work_dir = os.path.abspath(work_dir)
all_tasks = dict()
for component in get_components():
for task in component.tasks.values():
all_tasks[task.path] = task
tasks: Dict[str, Task] = dict()
cached_steps: Dict[str, List[str]] = dict()
_add_tasks_by_number(numbers, all_tasks, tasks, cached_steps)
_add_tasks_by_name(task_list, all_tasks, cached, tasks, cached_steps)
# the unqualified command-line options apply to the component of the
# first task
first_path = next(iter(tasks))
component = tasks[first_path].component
if suite_component is None:
suite_component = component.name
basic_config = _get_basic_config(
config_file=config_file,
machine=machine,
build=build,
cmake_flags=cmake_flags,
debug=debug,
clean_build=clean_build,
quiet_build=quiet_build,
work_dir=work_dir,
)
components_in_use = get_components_in_use(tasks)
component_args = _get_component_args(
component=component,
components_in_use=components_in_use,
model=model,
component_path=component_path,
branch=branch,
component_args=component_args or dict(),
)
component_configs = _get_component_configs(
basic_config=basic_config,
components_in_use=components_in_use,
component_args=component_args,
)
component_config = component_configs[component.name]
steps_by_component = get_steps_by_component(tasks)
for component_in_use in components_in_use:
name = component_in_use.name
component_in_use.configure(
component_configs[name], steps_by_component[name]
)
set_parallel_systems(tasks, basic_config)
provenance.write(
work_dir,
tasks,
config=component_config,
machine=machine,
baseline_dir=baseline_dir,
)
_build_models(
components_in_use=components_in_use,
component_configs=component_configs,
machine=machine,
)
if clean_tasks:
print('')
print('Cleaning task and step work directories:')
_clean_tasks_and_steps(tasks, work_dir)
print('')
_setup_configs(
component_configs=component_configs,
tasks=tasks,
work_dir=work_dir,
copy_executable=copy_executable,
)
# Apply CLI --free_running before caching resolution so that
# _expand_and_mark_cached_steps picks them up.
if free_running is not None:
_apply_free_running(tasks, free_running, cached_steps)
# do this after _setup_configs() in case tasks mark additional steps
# as cached in their configure() methods
_expand_and_mark_cached_steps(tasks, cached_steps)
print('Setting up tasks:')
for path, task in tasks.items():
setup_task(
path,
task,
machine,
work_dir,
baseline_dir,
cached_steps=cached_steps[path],
)
_check_dependencies(tasks)
suite = {'name': suite_name, 'tasks': tasks, 'work_dir': work_dir}
# pickle the task or step dictionary for use at runtime
pickle_file = os.path.join(suite['work_dir'], f'{suite_name}.pickle')
with open(pickle_file, 'wb') as handle:
pickle.dump(suite, handle, protocol=pickle.HIGHEST_PROTOCOL)
_symlink_load_script(work_dir)
max_cores, max_of_min_cores, max_gpus, max_of_min_gpus = (
_get_required_resources(tasks)
)
print(f'target cores: {max_cores}')
print(f'minimum cores: {max_of_min_cores}')
if max_gpus > 0 or max_of_min_gpus > 0:
print(f'target gpus: {max_gpus}')
print(f'minimum gpus: {max_of_min_gpus}')
job_options = None
if machine is not None:
suite_config = _get_suite_config(
basic_config, suite_component, suite_name
)
job_options = write_job_script(
config=suite_config,
machine=machine,
target_cores=max_cores,
min_cores=max_of_min_cores,
target_gpus=max_gpus,
min_gpus=max_of_min_gpus,
work_dir=work_dir,
suite=suite_name,
)
# Rewrite provenance (it was written earlier so that it exists even if
# setup fails) now that the model has been built, recording the version
# it was built from, and we know which scheduler options the suite's job
# script ended up using.
provenance.write(
work_dir,
tasks,
config=component_config,
machine=machine,
baseline_dir=baseline_dir,
job_options=job_options,
)
return tasks
[docs]
def setup_task(path, task, machine, work_dir, baseline_dir, cached_steps):
"""
Set up one or more tasks
Parameters
----------
path : str
Relative path for a tasks to set up
task : polaris.Task
A task to set up
machine : str
The name of one of the machines with defined config options, which can
be listed with ``polaris list --machines``
work_dir : str
A directory that will serve as the base for creating task directories
baseline_dir : str
Location of baselines that can be compared to
cached_steps : list of str
Which steps (if any) should read their outputs from the cache,
identified by a list of subdirectories in the component
"""
print(f' {path}')
task_dir = os.path.join(work_dir, path)
try:
os.makedirs(task_dir)
except FileExistsError:
pass
task.work_dir = task_dir
task.base_work_dir = work_dir
# add the baseline directory for this task
if baseline_dir is not None:
task.baseline_dir = os.path.join(baseline_dir, path)
if len(cached_steps) > 0:
print_steps = ' '.join(cached_steps)
print(f' steps with cached outputs: {print_steps}')
# iterate over steps
for step in task.steps.values():
_setup_step(task, step, work_dir, baseline_dir, task_dir)
# wait until we've set up all the steps before pickling because steps may
# need other steps to be set up
for step in task.steps.values():
if step.setup_complete:
# this is a shared step that has already been set up
continue
# pickle the task and step for use at runtime
pickle_filename = os.path.join(step.work_dir, 'step.pickle')
with open(pickle_filename, 'wb') as handle:
pickle.dump(step, handle, protocol=pickle.HIGHEST_PROTOCOL)
_symlink_load_script(step.work_dir)
if machine is not None:
cores = step.cores
min_cores = step.min_cores
gpus = step.gpus
min_gpus = step.min_gpus
write_job_script(
config=step.config,
machine=machine,
target_cores=cores,
min_cores=min_cores,
target_gpus=gpus,
min_gpus=min_gpus,
work_dir=step.work_dir,
)
step.setup_complete = True
# pickle the task and step for use at runtime
pickle_filename = os.path.join(task.work_dir, 'task.pickle')
with open(pickle_filename, 'wb') as handle:
suite = {
'name': 'task',
'tasks': {task.path: task},
'work_dir': task.work_dir,
}
pickle.dump(suite, handle, protocol=pickle.HIGHEST_PROTOCOL)
_symlink_load_script(task_dir)
if machine is not None:
max_cores, max_of_min_cores, max_gpus, max_of_min_gpus = (
_get_required_resources({path: task})
)
write_job_script(
config=task.config,
machine=machine,
target_cores=max_cores,
min_cores=max_of_min_cores,
target_gpus=max_gpus,
min_gpus=max_of_min_gpus,
work_dir=task_dir,
)
def main():
parser = argparse.ArgumentParser(
description='Set up one or more tasks', prog='polaris setup'
)
parser.add_argument(
'-t',
'--tasks',
nargs='+',
dest='tasks',
help='Relative path for a task(s) to set up.',
metavar='PATH',
)
parser.add_argument(
'-n',
'--task_number',
nargs='+',
dest='task_num',
type=str,
help='Task number(s) to setup, as listed from '
"'polaris list'. Can be a space-separated "
"list of task numbers. A suffix 'c' indicates "
'that all steps in the task should use cached '
'outputs.',
metavar='NUM',
)
parser.add_argument(
'-f',
'--config_file',
dest='config_file',
help='Configuration file for task setup.',
metavar='FILE',
)
parser.add_argument(
'-m',
'--machine',
dest='machine',
help='The name of the machine for loading machine-'
'related config options.',
metavar='MACH',
)
parser.add_argument(
'-w',
'--work_dir',
dest='work_dir',
required=True,
help='A base directory for setting up tasks.',
metavar='PATH',
)
parser.add_argument(
'-b',
'--baseline_dir',
dest='baseline_dir',
help='Location of baselines that can be compared to.',
metavar='PATH',
)
parser.add_argument(
'-p',
'--component_path',
dest='component_path',
help='The path where the component executable and '
'default namelists have been (or will be) built.',
metavar='PATH',
)
parser.add_argument(
'--suite_name',
dest='suite_name',
default='custom',
help="The name to use for the 'custom' suite "
'containing all setup tasks.',
metavar='SUITE',
)
parser.add_argument(
'--cached',
dest='cached',
nargs='+',
help='A list of steps in a single task supplied with '
'--tasks or --task_number that should use cached '
'outputs, even if the task makes them free-running, '
"or '_all' if all steps should use cached outputs.",
metavar='STEP',
)
parser.add_argument(
'--free_running',
dest='free_running',
nargs='+',
help='A list of steps in a single task supplied with '
'--tasks or --task_number that should be free-running '
"(not cached), overriding default_cached, or '_all' "
'if all steps should be free-running.',
metavar='STEP',
)
parser.add_argument(
'--copy_executable',
dest='copy_executable',
action='store_true',
help='If the model executable should be copied to the work directory.',
)
parser.add_argument(
'--clean_tasks',
dest='clean_tasks',
action='store_true',
help='If the task work directories should be deleted '
'before setting up the tasks.',
)
parser.add_argument(
'--model',
dest='model',
help="The model to run (one of 'mpas-ocean', 'omega', "
"or 'mpas-seaice')",
)
parser.add_argument(
'--build',
dest='build',
action='store_true',
help='If the model should be built.',
)
parser.add_argument(
'--branch',
dest='branch',
help='The branch of the model to build. The default is the submodule '
'associated with the model',
)
parser.add_argument(
'--clean_build',
dest='clean_build',
action='store_true',
help='If the model should be cleaned before building. Implies '
'--build.',
)
parser.add_argument(
'--quiet_build',
dest='quiet_build',
action='store_true',
help='If the model should be built without output. Implies --build.',
)
parser.add_argument(
'--cmake_flags',
dest='cmake_flags',
help='Additional flags to pass to make or CMake when building the '
'model.',
)
parser.add_argument(
'--debug',
dest='debug',
action='store_true',
help='If the model should be built in debug mode.',
)
add_component_model_args(parser)
args = parser.parse_args(sys.argv[2:])
cached = None
if args.cached is not None:
if args.tasks is not None and len(args.tasks) != 1:
raise ValueError(
'You can only cache steps for one task at at time.'
)
if args.task_num is not None and len(args.task_num) != 1:
raise ValueError(
'You can only cache steps for one task at at time.'
)
# cached is a list of lists
cached = [args.cached]
free_running = None
if args.free_running is not None:
if args.tasks is not None and len(args.tasks) != 1:
raise ValueError(
'You can only set free-running steps for one task at a time.'
)
if args.task_num is not None and len(args.task_num) != 1:
raise ValueError(
'You can only set free-running steps for one task at a time.'
)
# free_running is a list of lists
free_running = [args.free_running]
setup_tasks(
task_list=args.tasks,
numbers=args.task_num,
config_file=args.config_file,
machine=args.machine,
work_dir=args.work_dir,
baseline_dir=args.baseline_dir,
component_path=args.component_path,
suite_name=args.suite_name,
cached=cached,
free_running=free_running,
copy_executable=args.copy_executable,
clean_tasks=args.clean_tasks,
model=args.model,
build=args.build,
branch=args.branch,
clean_build=args.clean_build,
quiet_build=args.quiet_build,
cmake_flags=args.cmake_flags,
debug=args.debug,
component_args=get_component_model_args(args),
)
def _apply_free_running(tasks, free_running, cached_steps):
"""
Add steps requested with ``--free_running`` to the task's
``free_running_steps``, raising an error for any step that was also
requested with ``--cached``
"""
# Exactly one task is guaranteed by earlier validation.
path = next(iter(tasks.keys()))
task = tasks[path]
fr_steps = free_running[0]
fr_step_list = (
list(task.steps.keys()) if fr_steps[0] == '_all' else fr_steps
)
for step_name in fr_step_list:
if step_name not in task.steps:
raise ValueError(
f'Step {step_name!r} is not in the task. '
f'Available steps: {list(task.steps.keys())}'
)
task.free_running_steps.add(task.steps[step_name].subdir)
# Raise an error if any step is listed in both --cached and
# --free_running (after expanding _all in cached_steps).
cached_list = (
list(task.steps.keys())
if cached_steps.get(path, []) == ['_all']
else cached_steps.get(path, [])
)
conflicts = sorted(set(cached_list) & set(fr_step_list))
if conflicts:
raise ValueError(
f'Steps {conflicts} are listed in both --cached and '
'--free_running. Each step must be unambiguously cached, '
'free-running, or left at its default.'
)
def _expand_and_mark_cached_steps(tasks, cached_steps):
"""
Mark any cached steps that will get their outputs from the cache database,
rather than being run.
Resolution order (highest to lowest priority):
1. Explicit requests: steps named with ``--cached`` or in a suite's
``cached`` lines are cached, and steps named with ``--free_running``
run. A step may not be requested both ways.
2. Task preference: if any selected task added a step subdir to
``free_running_steps``, that step runs (never reads from cache).
3. Factory default: steps whose ``default_cached`` flag is True.
On return, ``cached_steps`` lists the steps of each task that are cached.
"""
# Expand _all shorthand
for path, task in tasks.items():
cached_names = cached_steps[path]
if len(cached_names) > 0 and cached_names[0] == '_all':
cached_steps[path] = list(task.steps.keys())
# Phase 1: apply explicitly requested caching
explicit_cached_subdirs: set[str] = set()
for path, task in tasks.items():
for step_name in cached_steps[path]:
step = task.steps[step_name]
step.cached = True
explicit_cached_subdirs.add(step.subdir)
# Phase 2: collect free-running preferences from all selected tasks
# (includes additions made by configure() and --free_running, which are
# applied before this)
free_running_subdirs: set[str] = set()
for task in tasks.values():
free_running_subdirs.update(getattr(task, 'free_running_steps', set()))
free_running_subdirs -= explicit_cached_subdirs
# Phase 3: apply default_cached to steps no task wants free-running
for task in tasks.values():
for step in task.steps.values():
if step.subdir in free_running_subdirs:
step.cached = False
elif getattr(step, 'default_cached', False):
step.cached = True
# Report the steps that actually ended up cached
for path, task in tasks.items():
cached_steps[path] = [
step_name for step_name, step in task.steps.items() if step.cached
]
def _setup_configs(
component_configs,
tasks,
work_dir,
copy_executable,
):
"""Set up config parsers for each component in use"""
common_configs = dict()
for name, component_config in component_configs.items():
common_config = component_config.copy()
if copy_executable:
common_config.set('setup', 'copy_executable', 'True')
if 'POLARIS_BRANCH' in os.environ:
polaris_branch = os.environ['POLARIS_BRANCH']
common_config.set('paths', 'polaris_branch', polaris_branch)
else:
common_config.set('paths', 'polaris_branch', os.getcwd())
common_configs[name] = common_config
initial_configs = _add_task_configs(tasks, common_configs)
# okay, we're finally ready to configure all the tasks and add configs
# to the "owned" steps
configs = _configure_tasks_and_add_step_configs(
tasks, initial_configs, common_configs
)
_write_configs(common_configs, configs, work_dir)
_symlink_configs(tasks, work_dir)
def _add_task_configs(tasks, common_configs):
"""
Add config parsers for tasks and steps that don't already have shared ones
"""
# get a list of shared steps and add config files for tasks to the
# component that owns them
configs = dict()
owners = dict()
for task in tasks.values():
component = task.component
if task.config.filepath is None:
task.config_filename = f'{task.name}.cfg'
task.config.filepath = os.path.join(
component.name, task.subdir, task.config_filename
)
component.add_config(task.config)
configs[task.config.filepath] = task.config
owners[task.config.filepath] = component
# now go through all the configs and prepend the config options of the
# component that owns each one, then run the setup() method for each in
# case there is some customization
for filepath, config in configs.items():
config.prepend(common_configs[owners[filepath].name])
config.setup()
return configs
def _configure_tasks_and_add_step_configs(
tasks, initial_configs, common_configs
):
"""
Call the configure() method for each task and add configs to "owned" steps
"""
for config in initial_configs.values():
for task in config.tasks:
task.configure()
config.set(
section=f'{task.name}',
option='steps_to_run',
value=' '.join(task.steps_to_run),
comment=f'A list of steps to include when running the '
f'{task.name} task',
)
# add configs to steps after calling task.configure() on all tasks in case
# new steps were added
configs = dict()
new_configs = dict()
owners = dict()
for task in tasks.values():
configs[task.config.filepath] = task.config
owners[task.config.filepath] = task.component
for step in task.steps.values():
if step.has_shared_config:
if step.config.filepath is None:
step.config_filename = f'{step.name}.cfg'
step.config.filepath = os.path.join(
step.component.name, step.subdir, step.config_filename
)
_check_config_owner(step, owners)
configs[step.config.filepath] = step.config
if step.config.filepath not in initial_configs:
new_configs[step.config.filepath] = step.config
step.component.add_config(step.config)
else:
_check_step_component(step, task)
step._set_config(task.config, link=task.config_filename)
for filepath, config in new_configs.items():
config.prepend(common_configs[owners[filepath].name])
config.setup()
return configs
def _check_config_owner(step, owners):
"""
Make sure a shared config is used by steps from only one component, and
keep track of which component owns it
"""
filepath = step.config.filepath
owner = owners.get(filepath)
if owner is None:
owners[filepath] = step.component
elif owner is not step.component:
raise ValueError(
f'The shared config file {filepath} is used by steps from more '
f'than one component: {owner.name} and {step.component.name}. A '
f'shared config file belongs to a single component, which '
f'provides its config options.'
)
def _check_step_component(step, task):
"""
Make sure a step that will get its config options from the task belongs to
the same component as the task
"""
if step.component is not task.component:
raise ValueError(
f'Step {step.name} in {task.path} belongs to component '
f'{step.component.name} but the task belongs to component '
f'{task.component.name}. A step from another component must be '
f'a shared step with a shared config file, so that it gets the '
f'config options of its own component.'
)
def _write_configs(common_configs, configs, work_dir):
"""Write out all the config files"""
# add the common config for each component at the component level
for name, common_config in common_configs.items():
common_config.filepath = f'{name}.cfg'
configs[common_config.filepath] = common_config
# finally, write out the config files
for config in configs.values():
config_filepath = os.path.join(work_dir, config.filepath)
config_dir = os.path.dirname(config_filepath)
try:
os.makedirs(config_dir)
except FileExistsError:
pass
with open(config_filepath, 'w') as f:
config.write(f)
def _symlink_configs(tasks, work_dir):
"""Symlink config files for requested tasks and steps"""
symlinks = dict()
for task in tasks.values():
config = task.config
config_filepath = os.path.join(work_dir, config.filepath)
link_path = os.path.join(
work_dir, task.component.name, task.subdir, task.config_filename
)
if not os.path.exists(link_path) and link_path not in symlinks:
symlinks[link_path] = config_filepath
for step in task.steps.values():
config = step.config
config_filepath = os.path.join(work_dir, config.filepath)
link_path = os.path.join(
work_dir,
step.component.name,
step.subdir,
step.config_filename,
)
if not os.path.exists(link_path) and link_path not in symlinks:
symlinks[link_path] = config_filepath
for link_path, config_filepath in symlinks.items():
link_dir = os.path.dirname(link_path)
try:
os.makedirs(link_dir)
except FileExistsError:
pass
symlink(config_filepath, link_path)
def _clean_tasks_and_steps(tasks, base_work_dir):
"""
Remove contents of task and step work directories to start fresh
"""
print(f'{base_work_dir}:')
for path, task in tasks.items():
task_work_dir = os.path.join(base_work_dir, path)
try:
shutil.rmtree(task_work_dir)
print(f' {path}')
except FileNotFoundError:
pass
for step in task.steps.values():
step_work_dir = os.path.join(base_work_dir, step.path)
try:
shutil.rmtree(step_work_dir)
print(f' {step.path}')
except FileNotFoundError:
pass
def _get_required_resources(tasks):
"""
Get max target and minimum CPU and GPU resource counts across task steps
"""
max_cores = 0
max_of_min_cores = 0
max_gpus = 0
max_of_min_gpus = 0
for task in tasks.values():
for step_name in task.steps_to_run:
step = task.steps[step_name]
if step.cached:
continue
if step.ntasks is None:
raise ValueError(
f'The number of tasks (ntasks) was never set for '
f'{task.path} step {step_name}'
)
if step.cpus_per_task is None:
raise ValueError(
f'The number of CPUs per task (cpus_per_task) was never '
f'set for {task.path} step {step_name}'
)
cores = step.cores
min_cores = step.min_cores
gpus = step.gpus
min_gpus = step.min_gpus
max_cores = max(max_cores, cores)
max_of_min_cores = max(max_of_min_cores, min_cores)
max_gpus = max(max_gpus, gpus)
max_of_min_gpus = max(max_of_min_gpus, min_gpus)
return max_cores, max_of_min_cores, max_gpus, max_of_min_gpus
def __get_machine_and_check_params(
machine, config_file, tasks, numbers, cached
):
if machine is None and 'POLARIS_MACHINE' in os.environ:
machine = os.environ['POLARIS_MACHINE']
if machine is None:
machine = discover_machine()
if config_file is None and machine is None:
raise ValueError('At least one of config_file and machine is needed.')
if config_file is not None and not os.path.exists(config_file):
raise FileNotFoundError(
f"The user config file wasn't found: {config_file}"
)
if tasks is None and numbers is None:
raise ValueError('At least one of tasks or numbers is needed.')
if cached is not None:
if tasks is None:
warnings.warn(
'Ignoring "cached" argument because "tasks" was not provided',
stacklevel=2,
)
elif len(cached) != len(tasks):
raise ValueError(
'A list of cached steps must be provided for each task in'
+ '"tasks"'
)
return machine
def _get_basic_config(
config_file,
machine,
build,
cmake_flags,
debug,
clean_build,
quiet_build,
work_dir,
):
"""
Get a base config parser for the machine and the command line but not for
any component, task or step
"""
config = PolarisConfigParser()
if config_file is not None:
config.add_user_config(config_file)
# start with default polaris config options
config.add_from_package('polaris', 'default.cfg')
# add the E3SM config options from mache
if machine is not None:
config.add_from_package(
'mache.machines', f'{machine}.cfg', exception=False
)
# add the polaris machine config file
if machine is None:
machine = 'default'
config.add_from_package('polaris.machines', f'{machine}.cfg')
if 'POLARIS_BRANCH' in os.environ:
polaris_branch = os.environ['POLARIS_BRANCH']
config.set('paths', 'polaris_branch', polaris_branch)
else:
config.set('paths', 'polaris_branch', os.getcwd())
if build is not None:
config.set('build', 'build', str(build), user=True)
compiler = os.environ['POLARIS_COMPILER']
mpi = os.environ['POLARIS_MPI']
config.set('build', 'machine', machine, user=True)
config.set('build', 'compiler', compiler, user=True)
config.set('build', 'mpi', mpi, user=True)
if cmake_flags is not None:
config.set('build', 'cmake_flags', cmake_flags, user=True)
if debug is not None:
config.set('build', 'debug', str(debug), user=True)
if clean_build is not None:
config.set('build', 'clean', str(clean_build), user=True)
if quiet_build is not None:
config.set('build', 'quiet', str(quiet_build), user=True)
build = config.getboolean('build', 'build')
clean_build = config.getboolean('build', 'clean')
if clean_build and not build:
# the user presumably expects to clean the build, notice that it's
# absent, and build again
config.set('build', 'build', 'True', user=True)
if quiet_build and not build:
# the user presumably expects to build, since they want to build
# quietly
config.set('build', 'build', 'True', user=True)
config.set('paths', 'base_work_dir', work_dir, user=True)
return config
def add_component_model_args(parser):
"""
Add ``--<component>_model``, ``--<component>_path`` and
``--<component>_branch`` flags for each component that has a model.
The unqualified ``--model``, ``-p``/``--component_path`` and ``--branch``
flags apply to the component that owns the tasks being set up. These flags
are how the model, path or branch of any other component is set.
Parameters
----------
parser : argparse.ArgumentParser
The parser to add the flags to
"""
for component in get_components(add_tasks=False):
if not component.has_model():
continue
name = component.name
prefix = name.replace('/', '_')
parser.add_argument(
f'--{prefix}_model',
dest=f'{prefix}_model',
help=f'The model to run for the {name} component.',
)
parser.add_argument(
f'--{prefix}_path',
dest=f'{prefix}_path',
help=f'The path where the {name} component executable and '
f'default namelists have been (or will be) built.',
metavar='PATH',
)
parser.add_argument(
f'--{prefix}_branch',
dest=f'{prefix}_branch',
help=f'The branch of the {name} component model to build.',
metavar='PATH',
)
def get_component_model_args(args):
"""
Get the model, component path and branch for each component they were
provided for on the command line
Parameters
----------
args : argparse.Namespace
The parsed command-line arguments
Returns
-------
component_args : dict of dict
The options for each component they were provided for, with the
component names as keys
"""
options = dict(model='model', path='component_path', branch='branch')
component_args: Dict[str, Dict[str, str]] = dict()
for component in get_components(add_tasks=False):
if not component.has_model():
continue
prefix = component.name.replace('/', '_')
for flag, option in options.items():
value = getattr(args, f'{prefix}_{flag}', None)
if value is not None:
component_args.setdefault(component.name, dict())[option] = (
value
)
return component_args
def _get_component_configs(basic_config, components_in_use, component_args):
"""
Get the config options for each component in use: the component that owns
the tasks being set up, the components that own steps in those tasks and
the components of any dependencies of those steps.
Parameters
----------
basic_config : polaris.config.PolarisConfigParser
The config options for the machine and the command line
components_in_use : list of polaris.Component
The components in use for the tasks being set up
component_args : dict of dict
The model, component path and branch from the command line for each
component they were provided for
Returns
-------
component_configs : dict of polaris.config.PolarisConfigParser
The config options for each component in use, with the component names
as keys
"""
component_configs = dict()
for component in components_in_use:
component_configs[component.name] = _get_component_config(
basic_config=basic_config,
component=component,
**component_args.get(component.name, dict()),
)
return component_configs
def _get_component_config(
basic_config, component, model=None, component_path=None, branch=None
):
"""
Get the config options for a component: the basic config options plus the
component's own config file.
This is the config that the component's ``configure()`` method modifies
and that gets prepended to the config of each task and shared step the
component owns.
Parameters
----------
basic_config : polaris.config.PolarisConfigParser
The config options for the machine and the command line
component : polaris.Component
The component to get config options for
model : str, optional
The model to run, if provided on the command line
component_path : str, optional
The path to the component's build, if provided on the command line
branch : str, optional
The branch to build the component's model from, if provided on the
command line
Returns
-------
config : polaris.config.PolarisConfigParser
The config options for the component
"""
config = PolarisConfigParser()
config.prepend(basic_config)
# set the model from the command line if provided
if model is not None:
config.set(component.name, 'model', model, user=True)
# add the config options for the component
config.add_from_package(
f'polaris.{component.name.replace("/", ".")}',
f'{component.name}.cfg',
exception=False,
)
# set the component_path from the command line if provided
if component_path is not None:
config.set(
'paths',
'component_path',
os.path.abspath(component_path),
user=True,
)
if branch is not None:
config.set('build', 'branch', os.path.abspath(branch), user=True)
return config
def _get_component_args(
component, components_in_use, model, component_path, branch, component_args
):
"""
Work out which component each of the model, component path and branch
options from the command line applies to.
The unqualified options apply to the component that owns the tasks being
set up. The ``--<component>_*`` options apply to the component they name.
Returns
-------
component_args : dict of dict
The model, component path and branch for each component they were
provided for, with the component names as keys
"""
resolved: Dict[str, Dict[str, str]] = dict()
unqualified = dict(
model=model, component_path=component_path, branch=branch
)
unqualified = {
option: value
for option, value in unqualified.items()
if value is not None
}
if len(unqualified) > 0:
if not component.has_model():
flags = _get_model_flags(components_in_use)
raise ValueError(
f'The {component.name} component has no model, so --model, '
f'-p/--component_path and --branch do not apply to the tasks '
f'being set up. Use the flags for the component whose model '
f'you mean: {flags}.'
)
resolved[component.name] = unqualified
names_in_use = [other.name for other in components_in_use]
for name, options in component_args.items():
if name not in names_in_use:
prefix = name.replace('/', '_')
raise ValueError(
f'--{prefix}_* was provided but the {name} component is not '
f'used by the tasks being set up.'
)
resolved.setdefault(name, dict()).update(options)
return resolved
def _check_model_executables(components_in_use, component_configs):
"""
Make sure no two components in use expect their model executable in the
same place, which would mean one of them building over the other
"""
executables: Dict[str, str] = dict()
for component in components_in_use:
if not component.has_model():
continue
config = component_configs[component.name]
if not config.has_option('executables', 'component'):
continue
executable = config.get('executables', 'component')
other = executables.get(executable)
if other is not None:
other_prefix = other.replace('/', '_')
prefix = component.name.replace('/', '_')
raise ValueError(
f'The {other} and {component.name} components both expect '
f'their model executable at {executable}. Use '
f'--{other_prefix}_path and --{prefix}_path to give them '
f'different build directories.'
)
executables[executable] = component.name
def _build_models(components_in_use, component_configs, machine):
"""
Build the model of each component that has one and whose config options
ask for a build, and check that each model is compatible with polaris
"""
_check_model_executables(components_in_use, component_configs)
for component in components_in_use:
if not component.has_model():
continue
config = component_configs[component.name]
if config.getboolean('build', 'build'):
component.build_model(config=config, machine=machine)
component.check_model_version(config=config)
def _get_model_flags(components_in_use):
"""
Get a printable list of the per-component command-line flags for the
components in use that have a model
"""
flags = list()
for component in components_in_use:
if not component.has_model():
continue
prefix = component.name.replace('/', '_')
flags.extend(
[f'--{prefix}_model', f'--{prefix}_path', f'--{prefix}_branch']
)
if len(flags) == 0:
return 'none of the components in use has a model'
return ', '.join(flags)
def _add_tasks_by_number(numbers, all_tasks, tasks, cached_steps):
if numbers is not None:
keys = list(all_tasks)
for number in numbers:
cache_all = False
if number.endswith('c'):
cache_all = True
number = int(number[:-1])
else:
number = int(number)
if number >= len(keys):
raise ValueError(
f'task number {number} is out of range. '
f'There are only {len(keys)} tasks.'
)
path = keys[number]
if cache_all:
cached_steps[path] = ['_all']
else:
cached_steps[path] = list()
tasks[path] = all_tasks[path]
def _add_tasks_by_name(task_list, all_tasks, cached, tasks, cached_steps):
if task_list is not None:
for index, path in enumerate(task_list):
if path not in all_tasks:
raise ValueError(f'Task with path {path} is not in tasks')
if cached is not None:
cached_steps[path] = cached[index]
else:
cached_steps[path] = list()
tasks[path] = all_tasks[path]
def _setup_step(task, step, work_dir, baseline_dir, task_dir):
"""Set up a step in a task"""
# make the step directory if it doesn't exist
step_dir = os.path.join(work_dir, step.path)
if step.name in task.step_symlinks:
symlink(
step_dir, os.path.join(task_dir, task.step_symlinks[step.name])
)
if step.setup_complete:
# this is a shared step that has already been set up
return
try:
os.makedirs(step_dir)
except FileExistsError:
pass
step.work_dir = step_dir
step.base_work_dir = work_dir
# set up the step
step.setup()
# add the baseline directory for this step
if baseline_dir is not None:
step.baseline_dir = os.path.join(baseline_dir, step.path)
# process input, output, namelist and streams files
step.process_inputs_and_outputs()
def _symlink_load_script(work_dir):
"""make a symlink to the script for loading the polaris env."""
if 'POLARIS_LOAD_SCRIPT' in os.environ:
script_filename = os.environ['POLARIS_LOAD_SCRIPT']
symlink(script_filename, os.path.join(work_dir, 'load_polaris_env.sh'))
def _check_dependencies(tasks):
for task in tasks.values():
for step in task.steps.values():
for name, dependency in step.dependencies.items():
if dependency.work_dir == '':
raise ValueError(
f'The dependency {name} of '
f'{task.path} step {step.name} was '
f'not set up.'
)
def _get_suite_config(basic_config, component_name, suite_name):
"""
Get a copy of ``basic_config`` with the suite's config options added.
The suite's config options are used only for the suite's job script. They
must not be added to ``basic_config`` itself, because that would propagate
them to the config options of every task and step in the suite.
Parameters
----------
basic_config : polaris.config.PolarisConfigParser
The config options for the machine and the command line
component_name : str
The name of the component that the suite belongs to
suite_name : str
The name of the suite, or ``'custom'`` if tasks are not being set
up as part of a suite
Returns
-------
suite_config : polaris.config.PolarisConfigParser
The config options to use for the suite's job script
"""
suite_config = basic_config.copy()
_add_suite_config(suite_config, component_name, suite_name)
return suite_config
def _add_suite_config(config, component_name, suite_name):
"""
Add config options from the suite's optional ``<suite_name>.cfg``
file, if it exists, to ``config``
Parameters
----------
config : polaris.config.PolarisConfigParser
The config options to add the suite's config options to
component_name : str
The name of the component that the suite belongs to
suite_name : str
The name of the suite, or ``'custom'`` if tasks are not being set
up as part of a suite
"""
if not suite_name or suite_name == 'custom':
return
package = f'polaris.suites.{component_name.replace("/", ".")}'
config.add_from_package(package, f'{suite_name}.cfg', exception=False)