import importlib.resources as imp_res
import json
import os
import warnings
from configparser import RawConfigParser
from mache.parallel import ParallelSystem, get_parallel_system
from mpas_tools.io import open_dataset, write_netcdf
from mpas_tools.logging import check_call
from polaris.config import PolarisConfigParser
# attributes that describe how a component was set up rather than what a step
# needs to run, so they are left out of pickles
_SETUP_ONLY_ATTRIBUTES = ('tasks', 'steps', 'configs')
[docs]
class Component:
"""
The base class for housing all the tasks for a given component, such as
ocean, landice, or seaice
Attributes
----------
name : str
the name of the component
tasks : dict
A dictionary of tasks in the component with the subdirectories of the
tasks in the component as keys. Setup only: this is empty in a
component that has been unpickled (see ``__getstate__()``).
steps : dict
A dictionary of steps in the component with the subdirectories of the
steps in the component as keys. Setup only, like ``tasks``.
configs : dict
A dictionary of the shared configs in the component with the
filepaths the configs are written to as keys. Setup only, like
``tasks``.
cached_files : dict
A dictionary that maps from output file names in steps within tasks to
cached files in the ``polaris_cache`` database for the component. These
file mappings are read in from ``cached_files.json`` in the component.
"""
[docs]
def __init__(self, name):
"""
Create a new container for the tasks for a given component
Parameters
----------
name : str
the name of the component
"""
self.name = name
# tasks are added with add_task()
self.tasks = dict()
# steps are added with add_step()
self.steps = dict()
# configs are added with add_config()
self.configs = dict()
self.cached_files = dict()
self.parallel_system: ParallelSystem | None = None
self._has_model: bool | None = None
self._read_cached_files()
def __getstate__(self):
"""
Get the state to pickle, leaving out ``tasks``, ``steps`` and
``configs``
Those three are built up by ``add_task()``, ``add_step()`` and
``add_config()`` as a component is set up, and nothing reads them
once setup is over. Every task and step in them points back at this
component, so pickling them makes each step's pickle a copy of the
whole component rather than a description of the step.
Returns
-------
state : dict
The attributes to pickle
"""
state = dict(self.__dict__)
for attribute in _SETUP_ONLY_ATTRIBUTES:
state.pop(attribute, None)
return state
def __setstate__(self, state):
"""
Restore a pickled component, giving the attributes that were left out
of the pickle empty dictionaries so that a component is always usable
Parameters
----------
state : dict
The unpickled attributes
"""
self.__dict__.update(state)
for attribute in _SETUP_ONLY_ATTRIBUTES:
if attribute not in state:
setattr(self, attribute, dict())
[docs]
def has_model(self):
"""
Whether this component has a model executable, and therefore whether
options like the path to that executable apply to it.
This is determined by whether any of the component's config files
defines ``component`` in the ``[executables]`` section, so a component
that gains a model does not also have to be listed somewhere.
Returns
-------
has_model : bool
Whether this component has a model executable
"""
if self._has_model is None:
self._has_model = self._detect_model_configs()
return self._has_model
[docs]
def set_parallel_system(self, config: PolarisConfigParser) -> None:
"""
Construct and store the active parallel system for this component
Parameters
----------
config : polaris.config.PolarisConfigParser
The config to use in constructing the parallel system
"""
if config.combined is None:
config.combine()
assert config.combined is not None
self.parallel_system = get_parallel_system(config.combined)
[docs]
def get_available_resources(self, placement=None):
"""
Get available resources from the active parallel system
Parameters
----------
placement : mache.parallel.ResourcePlacement, optional
The part of the allocation a step is confined to. When it is
given, the resources described are that subset's rather than the
whole allocation's, so that a confined step is told about what it
can actually use. Resources withheld from a step have to be
genuinely withheld, not merely subtracted from a number.
Returns
-------
available_resources : dict
Available CPU and GPU resources and machine capabilities
"""
if self.parallel_system is None:
raise ValueError(
f'Parallel system has not been set for component {self.name}'
)
if placement is not None:
return _placement_resources(placement, self.parallel_system)
memory_per_node = _get_memory_per_node(self.parallel_system)
nodes = self.parallel_system.nodes
memory = None
if memory_per_node is not None and nodes is not None:
memory = memory_per_node * nodes
return dict(
cores=self.parallel_system.cores,
nodes=nodes,
cores_per_node=self.parallel_system.cores_per_node,
gpus=self.parallel_system.gpus,
gpus_per_node=self.parallel_system.gpus_per_node,
memory=memory,
memory_per_node=memory_per_node,
mpi_allowed=self.parallel_system.mpi_allowed,
)
[docs]
def run_parallel_command(
self,
args,
cpus_per_task,
ntasks,
openmp_threads,
logger,
gpus=0,
placement=None,
gpus_per_task=None,
):
"""
Run a command using the active parallel system
Parameters
----------
args : list of str
Command line arguments for the executable
cpus_per_task : int
Number of CPUs per task
ntasks : int
Number of parallel tasks
openmp_threads : int
Number of OpenMP threads
logger : logging.Logger
Logger to output command-line execution info
gpus : int, optional
Number of GPUs this launch needs, as a total rather than a count
per task
placement : mache.parallel.ResourcePlacement, optional
The part of the allocation to confine this launch to. Passing
none gives exactly the command Polaris has always built.
gpus_per_task : int, optional
Number of GPUs per task
.. deprecated:: 1.1.0
Use ``gpus`` instead
"""
if self.parallel_system is None:
raise ValueError(
f'Parallel system has not been set for component {self.name}'
)
if gpus_per_task:
warnings.warn(
'gpus_per_task is deprecated. Use gpus, which says how many '
'GPUs the launch needs in total.',
DeprecationWarning,
stacklevel=2,
)
gpus = max(gpus, gpus_per_task * ntasks)
if placement is not None and placement.gpus != gpus:
# a placement carries the GPUs itself, so the two saying
# different things means whatever built the placement and the
# step have drifted apart. Much cheaper to catch here than as
# silent oversubscription once the work is running.
raise ValueError(
f'This launch asks for {gpus} GPUs but its placement asks '
f'for {placement.gpus}. They must agree.'
)
env = dict(os.environ)
env['OMP_NUM_THREADS'] = f'{openmp_threads}'
if openmp_threads > 1:
logger.info(f'Running with {openmp_threads} OpenMP threads')
command_line_args = self.parallel_system.get_parallel_command(
args=args,
ntasks=ntasks,
cpus_per_task=cpus_per_task,
gpus_per_task=_gpus_per_task(gpus, ntasks),
placement=placement,
)
check_call(command_line_args, logger, env=env)
[docs]
def add_task(self, task):
"""
Add a task to the component
Parameters
----------
task : polaris.Task
The task to add
"""
if task.subdir in self.tasks:
raise ValueError(
f'A task has already been added with path {task.subdir}'
)
self.tasks[task.subdir] = task
[docs]
def add_step(self, step):
"""
Add a step to the component
Parameters
----------
step : polaris.Step
The step to add
"""
if step.subdir in self.steps and self.steps[step.subdir] != step:
raise ValueError(
f'A different step has already been added with '
f'path {step.subdir}'
)
self.steps[step.subdir] = step
[docs]
def remove_step(self, step):
"""
Remove the given step from this component
Parameters
----------
step : polaris.Step
The step to add if adding by Step object, not subdirectory
"""
if step.subdir not in self.steps:
raise ValueError(
f'step {step.name} at {step.subdir} not in the '
f'{self.name} component'
)
self.steps.pop(step.subdir)
[docs]
def add_config(self, config):
"""
Add a shared config to the component
Parameters
----------
config : polaris.config.PolarisConfigParser
The config to add
"""
if config.filepath is None:
raise ValueError(
'The filepath attribute of a config must be set '
'before it can be added to a component'
)
if (
config.filepath in self.configs
and self.configs[config.filepath] != config
):
raise ValueError(
f'A different shared config has already been '
f'added at {config.filepath}'
)
self.configs[config.filepath] = config
[docs]
def get_or_create_shared_config(self, filepath, create=None):
"""
Get a shared config from the component if it exists, otherwise build
and add it.
The companion to :py:meth:`get_or_create_shared_step`, and needed for
the same reason. A ``get_*_steps()`` helper is called once per
consumer, so building its config unconditionally hands the second
caller a config that the shared steps -- created on the first call --
are not using. Options set on it reach nothing, and passing it to
:py:meth:`polaris.Task.set_shared_config` raises, because a different
config is already registered at that path.
Parameters
----------
filepath : str
The path of the config file relative to the base work directory,
which is what identifies a shared config
create : callable, optional
Called with no arguments to build the config, and only when it does
not already exist -- so it must not be relied on for anything a
later caller needs. It must return a
:py:class:`polaris.config.PolarisConfigParser` whose ``filepath``
is ``filepath``. A config builder that returns a new parser, as
several in Polaris do, can be passed directly. The default builds
an empty parser.
Returns
-------
config : polaris.config.PolarisConfigParser
The shared config parser
"""
if filepath in self.configs:
return self.configs[filepath]
if create is None:
config = PolarisConfigParser(filepath=filepath)
else:
config = create()
if config.filepath != filepath:
raise ValueError(
f'The config built for {filepath} has filepath '
f'{config.filepath}; it would be registered under a path '
f'it does not know about.'
)
self.add_config(config)
return config
[docs]
def build_model(self, config, machine):
"""
Build the model for this component. Components with a model that
polaris knows how to build should override this method.
Parameters
----------
config : polaris.config.PolarisConfigParser
the config options for this component
machine : str
The name of the machine to build the model on
"""
raise ValueError(
f'An automated build is not implemented for the {self.name} '
f'component'
)
[docs]
def check_model_version(self, config):
"""
Check that the model this component will run is compatible with this
version of polaris. Components that have such a check should override
this method.
Parameters
----------
config : polaris.config.PolarisConfigParser
the config options for this component
"""
pass
[docs]
def get_or_create_shared_step(
self, step_cls, subdir, config, config_filename=None, **kwargs
):
"""
Get a shared step from the component if it exists, otherwise create and
add it.
Parameters
----------
step_cls : type
The class of the step to create if it doesn't exist
subdir : str
The subdirectory for the step
config : polaris.config.PolarisConfigParser
The shared config options for this step
config_filename : str
the name of the symlink to the shared configuration file in this
step
kwargs : dict
Arguments to pass to the step constructor
Returns
-------
step : polaris.Step
The shared step instance
"""
if subdir in self.steps:
return self.steps[subdir]
step = step_cls(component=self, subdir=subdir, **kwargs)
step.set_shared_config(config, link=config_filename)
self.add_step(step)
return step
def open_model_dataset(self, filename, config, **kwargs):
"""
Open the given dataset, mapping variable and dimension names from Omega
to MPAS-Ocean names if appropriate
Parameters
----------
filename : str
The path for the NetCDF file to open
config : polaris.config.PolarisConfigParser
Configuration for the task; passed to model-specific subclasses
that need it (e.g. to compute layer thickness from pseudo-thickness
for Omega). The base implementation ignores it.
kwargs
keyword arguments passed to `xarray.open_dataset()`
Returns
-------
ds : xarray.Dataset
The dataset with variables named as expected in MPAS-Ocean
"""
ds = open_dataset(filename, **kwargs)
return ds
def write_model_dataset(self, ds, filename, config):
"""
Write out the given dataset, mapping dimension and variable names from
MPAS-Ocean to Omega names if appropriate
Parameters
----------
ds : xarray.Dataset
A dataset containing MPAS-Ocean variable names
filename : str
The path for the NetCDF file to write
config : polaris.config.PolarisConfigParser
Configuration for the task; passed to model-specific subclasses
that need it (e.g. to convert layer thickness to pseudo-thickness
for Omega). The base implementation ignores it.
"""
write_netcdf(ds=ds, fileName=filename)
def _detect_model_configs(self):
"""
Determine whether any of the component's config files describes a
model executable
"""
package = f'polaris.{self.name.replace("/", ".")}'
try:
pkg_files = list(imp_res.files(package).iterdir())
except (ModuleNotFoundError, FileNotFoundError, NotADirectoryError):
return False
for pkg_file in pkg_files:
if not pkg_file.name.endswith('.cfg'):
continue
# the config files use extended interpolation, which we don't
# want to resolve here
config = RawConfigParser()
config.read_string(pkg_file.read_text())
if config.has_option('executables', 'component'):
return True
return False
def _read_cached_files(self):
"""Read in the dictionary of cached files from cached_files.json"""
package = f'polaris.{self.name.replace("/", ".")}'
filename = 'cached_files.json'
try:
pkg_file = imp_res.files(package).joinpath(filename)
with pkg_file.open('r') as data_file:
self.cached_files = json.load(data_file)
except (FileNotFoundError, ModuleNotFoundError, TypeError):
# the component has no package of its own or no cached files in
# it, so there is nothing to read
pass
def _gpus_per_task(gpus, ntasks):
"""
Convert a launch's total GPUs into the per-task count mache asks for.
Without a placement to carry a total, mache expresses GPUs per task, so
a step's total has to be divided back out. Rounding up rather than down
keeps a launch from being handed fewer GPUs than it asked for.
"""
if gpus <= 0 or ntasks <= 0:
return 0
return -(-gpus // ntasks)
def _placement_resources(placement, parallel_system):
"""Describe what a placement gives a step, in the usual resource terms."""
nodes = max(len(placement.nodes), 1)
cores_per_node = len(placement.cores)
# a placement carries no memory, because no launcher acts on one. What
# a placement does imply is a share of the nodes it names, in the same
# proportion as the cores it took from them.
memory_per_node = _get_memory_per_node(parallel_system)
memory = None
machine_cores_per_node = parallel_system.cores_per_node
if memory_per_node is not None and machine_cores_per_node:
memory = (
cores_per_node * nodes * memory_per_node // machine_cores_per_node
)
return dict(
cores=cores_per_node * nodes,
nodes=nodes,
cores_per_node=cores_per_node,
gpus=placement.gpus,
gpus_per_node=placement.gpus // nodes,
memory=memory,
memory_per_node=memory_per_node,
mpi_allowed=parallel_system.mpi_allowed,
)
def _get_memory_per_node(parallel_system):
"""
Get the memory a node has, in MB, or ``None`` if this machine has not
said.
Read from the attribute where a newer mache offers one and from the
config option otherwise, because the option is arriving in mache while
this is being written and Polaris should work either side of it. A
machine that says nothing leaves memory undeclared rather than guessed
at: a wrong figure here would propagate into every step's default.
"""
memory_per_node = getattr(parallel_system, 'memory_per_node', None)
if memory_per_node:
return int(memory_per_node)
memory_per_node = parallel_system.get_config_int('memory_per_node')
if memory_per_node:
return int(memory_per_node)
return None