def __publishing_context(self):
"""Get a context manager for resolving the data dependencies of this node.
The returned object is a Python context manager (used to open a `with` block)
to define the scope in which the operation's output can be published.
'output' type resources can be published exactly once, and only while the
publishing context is active. (See operation.function_wrapper())
Used internally to implement ResourceManager.publishing_resources()
Responsibilities of the context manager are to:
* (TODO) Make sure dependencies are resolved.
* Make sure outputs are marked 'done' when leaving the context.
"""
# TODO:
# if self._data.done():
# raise exceptions.ProtocolError('Resources have already been published.')
resource = PublishingDataProxy(weakref.proxy(self))
# ref: https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager
try:
yield resource
except Exception as e:
message = 'Uncaught exception while providing output-publishing resources for {}.'.format(self._runner)
raise exceptions.ApiError(message) from e
finally:
self.done = True
def check_element(element):
if element.workspec is None:
self.workspec.add_element(element)
assert element.workspec is self.workspec
assert element.name in self.workspec.elements
elif element.workspec is not self.workspec:
raise exceptions.ApiError(
"Element will need to be moved to the same workspec.")
return True
def __init__(self, tprfile: typing.Union[str, TprFile]):
if not isinstance(tprfile, TprFile):
try:
tprfile = TprFile(tprfile)
except Exception as e:
# This class is an implementation detail of TPR file I/O...
raise exceptions.ApiError("Must be initialized from a TprFile.") from e
assert isinstance(tprfile, TprFile)
self.__tprfile = tprfile
self.__parameters = None
def __init__(self, tprfile: str = None):
if not isinstance(tprfile, TprFile):
try:
tprfile = TprFile(tprfile)
except Exception:
# This class is an implementation detail of TPR file I/O...
raise exceptions.ApiError(
"Must be initialized from a gmx.fileio.TprFile object.")
self.__tprfile = tprfile
self.__parameters = None
def future(self, name: str = None, dtype=None):
"""Retrieve a Future for a named output.
TODO: (FR5+) Normalize this part of the interface between operation definitions and
resource managers.
"""
if not isinstance(name, str) or name not in self._data:
raise exceptions.ValueError('"name" argument must name an output.')
assert dtype is not None
if dtype != self._data[name].dtype:
message = 'Requested Future of type {} is not compatible with available type {}.'
message = message.format(dtype, self._data[name].dtype)
raise exceptions.ApiError(message)
return Future(self, name, dtype)
def __init__(self, resource_manager: _op.ResourceManager):
from .context import Context as LegacyContext
import gmxapi._gmxapi as _gmxapi
self._gmxapi = _gmxapi
assert isinstance(resource_manager, _op.ResourceManager)
# We use several details of the gmxapi.operation.Context.ResourceManager.
# These dependencies can evolve into the subscription protocol between Contexts.
# Configure and run a gmxapi 0.0.7 session.
# 0. Determine ensemble width.
# 1. Choose, create/check working directories.
# 2. Create source TPR.
# 3. Create workspec.
# 3.5 Add plugin potentials, if any.
# 4. Run.
# 5. Collect outputs from context (including plugin outputs) and be ready to publish them.
# Determine ensemble width
ensemble_width = resource_manager.ensemble_width
# Choose working directories
# TODO: operation working directory naming scheme should be centrally well-defined.
# Note that workflow.WorkSpec.uid is currently dependent on the input file parameter,
# so we cannot create the input file path in the working directory based on WorkSpec.uid.
workdir_list = [
'{node}_{member}'.format(node=resource_manager.operation_id,
member=member)
for member in range(ensemble_width)
]
parameters_dict_list = [{}] * ensemble_width
# This is a reasonable place to start using MPI ensemble implementation details.
# We will want better abstraction in the future, but it is best if related filesystem
# accesses occur from the same processes, consistently. Note that we already
# handle some optimization and dependency reduction when the ensemble size is 1.
# TODO: multithread and multiprocess alternatives to MPI ensemble management.
# TODO: Allow user to provide communicator instead of implicitly getting COMM_WORLD
with scoped_communicator(None) as context_comm:
context_rank = context_comm.Get_rank()
with scoped_communicator(context_comm,
ensemble_width) as ensemble_comm:
# Note that in the current implementation, extra ranks have nothing to do,
# but they may have a dummy communicator, so be sure to skip those members
# of the context_comm.
if context_rank < ensemble_width:
assert ensemble_comm.Get_size() == ensemble_width
ensemble_rank = ensemble_comm.Get_rank()
# TODO: This should be a richer object that includes at least host information
# and preferably the full gmxapi Future interface.
self.workdir = os.path.abspath(workdir_list[ensemble_rank])
with resource_manager.local_input(
member=ensemble_rank) as input_pack:
source_file = input_pack.kwargs['_simulation_input']
parameters = input_pack.kwargs['parameters']
# If there are any other key word arguments to process from the gmxapi.mdrun
# factory call, do it here.
# TODO: We should really name this file with a useful input-dependent tag.
tprfile = os.path.join(self.workdir, 'topol.tpr')
expected_working_files = [tprfile]
if os.path.exists(self.workdir):
if os.path.isdir(self.workdir):
# Confirm that this is a restarted simulation.
# It is unspecified by the API, but at least through gmxapi 0.1,
# all simulations are initialized with a checkpoint file named state.cpt
# (see src/api/cpp/context.cpp)
checkpoint_file = os.path.join(
self.workdir, 'state.cpt')
expected_working_files.append(checkpoint_file)
for file in expected_working_files:
if not os.path.exists(file):
raise exceptions.ApiError(
'Cannot determine working directory state: {}'
.format(self.workdir))
else:
raise exceptions.ApiError(
'Chosen working directory path exists but is not a directory: {}'
.format(self.workdir))
else:
# Build the working directory and input files.
os.mkdir(self.workdir)
sim_input = fileio.read_tpr(source_file)
for key, value in parameters.items():
try:
sim_input.parameters.set(key=key, value=value)
except _gmxapi.Exception as e:
raise exceptions.ApiError(
'Bug encountered. Unknown error when trying to set simulation '
'parameter {} to {}'.format(key,
value)) from e
fileio.write_tpr_file(output=tprfile, input=sim_input)
logger.info('Created {} on rank {}'.format(
tprfile, context_rank))
# Gather the actual outputs from the ensemble members.
if hasattr(ensemble_comm, 'allgather'):
# We should not assume that abspath expands the same on different MPI ranks.
workdir_list = ensemble_comm.allgather(self.workdir)
tpr_filenames = ensemble_comm.allgather(tprfile)
parameters = fileio.read_tpr(
tprfile).parameters.extract()
parameters_dict_list = ensemble_comm.allgather(
parameters)
else:
workdir_list = [
os.path.abspath(workdir)
for workdir in workdir_list
]
# TODO: If we use better input file names, they need to be updated in multiple places.
tpr_filenames = [
os.path.join(workdir, 'topol.tpr')
for workdir in workdir_list
]
parameters_dict_list = [
fileio.read_tpr(tprfile).parameters.extract()
for tprfile in tpr_filenames
]
logger.debug(
'Context rank {} acknowledges working directories {}'.
format(context_rank, workdir_list))
logger.debug('Operation {}:{} will use {}'.format(
resource_manager.operation_id, ensemble_rank,
self.workdir))
# TODO: (#3718) Normalize the way we pass run time parameters to mdrun.
kwargs = getattr(resource_manager, 'mdrun_kwargs', {})
for key, value in kwargs.items():
logger.debug(
'Adding mdrun run time argument: {}'.format(
key + '=' + str(value)))
work = workflow.from_tpr(tpr_filenames, **kwargs)
self.workspec = work.workspec
context = LegacyContext(work=self.workspec,
workdir_list=workdir_list,
communicator=ensemble_comm)
self.simulation_module_context = context
# Go ahead and execute immediately. No need for lazy initialization in this basic case.
with self.simulation_module_context as session:
session.run()
# TODO: There may be some additional results that we need to extract...
# end: if context_rank < ensemble_width
# end scoped_communicator: ensemble_comm
if context_comm.Get_size() > 1:
context_comm.bcast(workdir_list, root=0)
# end scoped_communicator: context_comm
self.workdir = workdir_list
self.parameters = parameters_dict_list
def state(self):
raise exceptions.ApiError("property not implemented.")
def topology(self):
raise exceptions.ApiError("property not implemented.")
"""Provide helpers to access package data.
Some sample files are provided with the module for testing purposes.
"""
# At some point, these will probably be accessed in the newer get_data()
# idiom, but an unzippable package and direct file access is okay for now.
import os
from gmxapi import exceptions
tpr_filename = None
_tpr_filename = 'data/topol.tpr'
try:
from pkg_resources import resource_filename
_tpr_filename = resource_filename(__name__, _tpr_filename)
except Exception as e:
raise exceptions.ApiError(
"Need pkg_resources from setuptools package to access gmx "
"package data.") from e
if os.path.exists(_tpr_filename) and os.path.isfile(_tpr_filename):
tpr_filename = os.path.abspath(_tpr_filename)
else:
raise exceptions.ApiError(
'Package data file data/topol.tpr not accessible at {}'.format(
_tpr_filename))
def keys(self):
for key in dict.keys(self):
if not isinstance(key, str):
raise exceptions.ApiError(
'Invalid key type found: {} {}'.format(key, type(key)))
yield key
def local_input(self):
"""In an API session, get a handle to fully resolved locally available input data.
Execution dependencies are resolved on creation of the context manager. Input data
becomes available in the ``as`` object when entering the context manager, which
becomes invalid after exiting the context manager. Resources allocated to hold the
input data may be released when exiting the context manager.
It is left as an implementation detail whether the context manager is reusable and
under what circumstances one may be obtained.
"""
# Localize data
for dependency in self._dependencies:
dependency()
# TODO: (FR3+) be more rigorous.
# This should probably also use a sort of Context-based observer pattern rather than
# the result() method, which is explicitly for moving data across the API boundary.
args = []
try:
for arg in self._input_fingerprint.args:
if hasattr(arg, 'result'):
args.append(arg.result())
else:
args.append(arg)
except Exception as E:
raise exceptions.ApiError('input_fingerprint not iterating on "args" attr as expected.') from E
kwargs = {}
try:
for key, value in self._input_fingerprint.kwargs.items():
if hasattr(value, 'result'):
kwargs[key] = value.result()
else:
kwargs[key] = value
if isinstance(kwargs[key], list):
new_list = []
for item in kwargs[key]:
if hasattr(item, 'result'):
new_list.append(item.result())
else:
new_list.append(item)
kwargs[key] = new_list
try:
for item in kwargs[key]:
# TODO: This should not happen. Need proper tools for NDArray Futures.
# assert not hasattr(item, 'result')
if hasattr(item, 'result'):
kwargs[key][item] = item.result()
except TypeError:
# This is only a test for iterables
pass
except Exception as E:
raise exceptions.ApiError('input_fingerprint not iterating on "kwargs" attr as expected.') from E
assert 'input' not in kwargs
for key, value in kwargs.items():
if key == 'command':
if type(value) == list:
for item in value:
assert not hasattr(item, 'result')
input_pack = collections.namedtuple('InputPack', ('args', 'kwargs'))(args, kwargs)
# Prepare input data structure
yield input_pack
