diff --git a/pyproject.toml b/pyproject.toml index 94954b47..32032c3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,10 @@ dependencies = [ "jinja2", "pint", "multiprocessing_logging", - "openmm[cuda12]", + "openmm", + "mdtraj>=1.11.0", + "openmm-cuda-12>=8.3.1", + "rich>=14.1.0", ] [project.optional-dependencies] diff --git a/src/wepy/analysis/contig_tree.py b/src/wepy/analysis/contig_tree.py index b46c8ff3..2289a03b 100644 --- a/src/wepy/analysis/contig_tree.py +++ b/src/wepy/analysis/contig_tree.py @@ -175,23 +175,19 @@ def __init__( if continuations is Ellipsis: # add continuations involving both ends of the continuation - self._continuations.update( - [ - (a, b) - for a, b in wepy_h5.continuations - if a in self._run_idxs and b in self._run_idxs - ] - ) + self._continuations.update([ + (a, b) + for a, b in wepy_h5.continuations + if a in self._run_idxs and b in self._run_idxs + ]) # if a subset of continuations was given use only those elif continuations is not None: - self._continuations.update( - [ - (a, b) - for a, b in continuations - if a in self._run_idxs and b in self._run_idxs - ] - ) + self._continuations.update([ + (a, b) + for a, b in continuations + if a in self._run_idxs and b in self._run_idxs + ]) # using the wepy_h5 create a tree of the cycles self._create_tree(wepy_h5) @@ -1564,9 +1560,9 @@ def __init__(self, wepy_h5, **kwargs): # check that the result is a single contig spanning_contig_traces = self.spanning_contig_traces() - assert ( - len(spanning_contig_traces) == 1 - ), "continuations given do not form a single contig" + assert len(spanning_contig_traces) == 1, ( + "continuations given do not form a single contig" + ) # if so we add some useful attributes valid for only a # standalone contig diff --git a/src/wepy/hdf5.py b/src/wepy/hdf5.py index 940dc5fe..202f8cbc 100644 --- a/src/wepy/hdf5.py +++ b/src/wepy/hdf5.py @@ -892,19 +892,21 @@ def __init__( # read only mode elif self._wepy_mode == "r": # if any data was given, warn the user - if any([ - kwarg is not None - for kwarg in [ - topology, - units, - sparse_fields, - feature_shapes, - feature_dtypes, - n_dims, - alt_reps, - main_rep_idxs, + if any( + [ + kwarg is not None + for kwarg in [ + topology, + units, + sparse_fields, + feature_shapes, + feature_dtypes, + n_dims, + alt_reps, + main_rep_idxs, + ] ] - ]): + ): warn("Data was given but opening in read-only mode", RuntimeWarning) # then run the initialization process @@ -973,9 +975,9 @@ def _create_init(self): and set with the new ones if given. """ - assert self._topology is not None, ( - "Topology must be given for a creation constructor" - ) + assert ( + self._topology is not None + ), "Topology must be given for a creation constructor" # initialize the runs group runs_grp = self._h5.create_group(RUNS) @@ -1617,9 +1619,9 @@ def _extend_contiguous_traj_field(self, run_idx, traj_idx, field_path, field_dat field = traj_grp[field_path] # make sure this is a feature vector - assert len(field_data.shape) > 1, ( - "field_data must be a feature vector with the same number of dimensions as the number" - ) + assert ( + len(field_data.shape) > 1 + ), "field_data must be a feature vector with the same number of dimensions as the number" # of datase new frames n_new_frames = field_data.shape[0] @@ -1628,9 +1630,9 @@ def _extend_contiguous_traj_field(self, run_idx, traj_idx, field_path, field_dat if all([i == 0 for i in field.shape]): # check the feature shape against the maxshape which gives # the feature dimensions for an empty dataset - assert field_data.shape[1:] == field.maxshape[1:], ( - "field feature dimensions must be the same, i.e. all but the first dimension" - ) + assert ( + field_data.shape[1:] == field.maxshape[1:] + ), "field feature dimensions must be the same, i.e. all but the first dimension" # if it is empty resize it to make an array the size of # the new field_data with the maxshape for the feature @@ -1644,9 +1646,9 @@ def _extend_contiguous_traj_field(self, run_idx, traj_idx, field_path, field_dat else: # make sure the new data has the right dimensions against # the shape it already has - assert field_data.shape[1:] == field.shape[1:], ( - "field feature dimensions must be the same, i.e. all but the first dimension" - ) + assert ( + field_data.shape[1:] == field.shape[1:] + ), "field feature dimensions must be the same, i.e. all but the first dimension" # append to the dataset on the first dimension, keeping the # others the same, these must be feature vectors and therefore @@ -1688,10 +1690,10 @@ def _extend_sparse_traj_field( if all([i == 0 for i in field_data.shape]): # check the feature shape against the maxshape which gives # the feature dimensions for an empty dataset - assert values.shape[1:] == field_data.maxshape[1:], ( - "input value features have shape {}, expected {}".format( - values.shape[1:], field_data.maxshape[1:] - ) + assert ( + values.shape[1:] == field_data.maxshape[1:] + ), "input value features have shape {}, expected {}".format( + values.shape[1:], field_data.maxshape[1:] ) # if it is empty resize it to make an array the size of @@ -1705,25 +1707,29 @@ def _extend_sparse_traj_field( else: # make sure the new data has the right dimensions - assert values.shape[1:] == field_data.shape[1:], ( - "field feature dimensions must be the same, i.e. all but the first dimension" - ) + assert ( + values.shape[1:] == field_data.shape[1:] + ), "field feature dimensions must be the same, i.e. all but the first dimension" # append to the dataset on the first dimension, keeping the # others the same, these must be feature vectors and therefore # must exist - field_data.resize(( - field_data.shape[0] + n_new_frames, - *field_data.shape[1:], - )) + field_data.resize( + ( + field_data.shape[0] + n_new_frames, + *field_data.shape[1:], + ) + ) # add the new data field_data[-n_new_frames:, ...] = values # add the sparse idxs in the same way - field_sparse_idxs.resize(( - field_sparse_idxs.shape[0] + n_new_frames, - *field_sparse_idxs.shape[1:], - )) + field_sparse_idxs.resize( + ( + field_sparse_idxs.shape[0] + n_new_frames, + *field_sparse_idxs.shape[1:], + ) + ) # add the new data field_sparse_idxs[-n_new_frames:, ...] = sparse_idxs @@ -1871,9 +1877,9 @@ def _extend_run_record_data_field( field = records_grp[field_name] # make sure this is a feature vector - assert len(field_data.shape) > 1, ( - "field_data must be a feature vector with the same number of dimensions as the number" - ) + assert ( + len(field_data.shape) > 1 + ), "field_data must be a feature vector with the same number of dimensions as the number" # of datase new frames n_new_frames = field_data.shape[0] @@ -1914,9 +1920,9 @@ def _extend_run_record_data_field( if all([i == 0 for i in field.shape]): # check the feature shape against the maxshape which gives # the feature dimensions for an empty dataset - assert field_data.shape[1:] == field.maxshape[1:], ( - "field feature dimensions must be the same, i.e. all but the first dimension" - ) + assert ( + field_data.shape[1:] == field.maxshape[1:] + ), "field feature dimensions must be the same, i.e. all but the first dimension" # if it is empty resize it to make an array the size of # the new field_data with the maxshape for the feature @@ -3940,14 +3946,18 @@ def add_continuation(self, continuation_run, base_run): """ continuations_dset = self.settings_grp[CONTINUATIONS] - continuations_dset.resize(( - continuations_dset.shape[0] + 1, - continuations_dset.shape[1], - )) - continuations_dset[continuations_dset.shape[0] - 1] = np.array([ - continuation_run, - base_run, - ]) + continuations_dset.resize( + ( + continuations_dset.shape[0] + 1, + continuations_dset.shape[1], + ) + ) + continuations_dset[continuations_dset.shape[0] - 1] = np.array( + [ + continuation_run, + base_run, + ] + ) def new_run(self, init_walkers, continue_run=None, **kwargs): """Initialize a new run. @@ -4346,9 +4356,9 @@ def add_traj(self, run_idx, data, weights=None, sparse_idxs=None, metadata=None) weights = np.ones((n_frames, 1), dtype=float) else: assert isinstance(weights, np.ndarray), "weights must be a numpy.ndarray" - assert weights.shape[0] == n_frames, ( - "weights and the number of frames must be the same length" - ) + assert ( + weights.shape[0] == n_frames + ), "weights and the number of frames must be the same length" # current traj_idx traj_idx = self.next_run_traj_idx(run_idx) @@ -4374,15 +4384,15 @@ def add_traj(self, run_idx, data, weights=None, sparse_idxs=None, metadata=None) ) # check to make sure the positions are the right shape - assert traj_data[POSITIONS].shape[1] == self.num_atoms, ( - "positions given have different number of atoms: {}, should be {}".format( - traj_data[POSITIONS].shape[1], self.num_atoms - ) + assert ( + traj_data[POSITIONS].shape[1] == self.num_atoms + ), "positions given have different number of atoms: {}, should be {}".format( + traj_data[POSITIONS].shape[1], self.num_atoms ) - assert traj_data[POSITIONS].shape[2] == self.num_dims, ( - "positions given have different number of dims: {}, should be {}".format( - traj_data[POSITIONS].shape[2], self.num_dims - ) + assert ( + traj_data[POSITIONS].shape[2] == self.num_dims + ), "positions given have different number of dims: {}, should be {}".format( + traj_data[POSITIONS].shape[2], self.num_dims ) # add datasets to the traj group @@ -4453,9 +4463,9 @@ def extend_traj(self, run_idx, traj_idx, data, weights=None): """ if self._wepy_mode == "c-": - assert self._append_flags[dataset_key], ( - "dataset is not available for appending to" - ) + assert self._append_flags[ + dataset_key + ], "dataset is not available for appending to" # convenient alias traj_data = data @@ -4479,9 +4489,9 @@ def extend_traj(self, run_idx, traj_idx, data, weights=None): weights = np.ones((n_new_frames, 1), dtype=float) else: assert isinstance(weights, np.ndarray), "weights must be a numpy.ndarray" - assert weights.shape[0] == n_new_frames, ( - "weights and the number of frames must be the same length" - ) + assert ( + weights.shape[0] == n_new_frames + ), "weights and the number of frames must be the same length" # add the weights weights_ds = traj_grp[WEIGHTS] @@ -4489,10 +4499,12 @@ def extend_traj(self, run_idx, traj_idx, data, weights=None): # append to the dataset on the first dimension, keeping the # others the same, if they exist if len(weights_ds.shape) > 1: - weights_ds.resize(( - weights_ds.shape[0] + n_new_frames, - *weights_ds.shape[1:], - )) + weights_ds.resize( + ( + weights_ds.shape[0] + n_new_frames, + *weights_ds.shape[1:], + ) + ) else: weights_ds.resize((weights_ds.shape[0] + n_new_frames,)) @@ -5517,10 +5529,12 @@ def get_contig_trace_fields(self, contig_trace, fields): # (there must be the same number of trajectories in each run) n_trajs_test = self.num_run_trajs(run_idxs[0]) - assert all([ - True if n_trajs_test == self.num_run_trajs(run_idx) else False - for run_idx in run_idxs - ]) + assert all( + [ + True if n_trajs_test == self.num_run_trajs(run_idx) else False + for run_idx in run_idxs + ] + ) # then using this we go run by run and get all the # trajectories diff --git a/src/wepy/orchestration/configuration.py b/src/wepy/orchestration/configuration.py index f01c04ed..d55b54b8 100644 --- a/src/wepy/orchestration/configuration.py +++ b/src/wepy/orchestration/configuration.py @@ -134,9 +134,9 @@ def __init__( ### Monitor options # get the names of the reporters in the order they are - reporter_order = tuple( - [str(reporter_class.__name__) for reporter_class in self._reporter_classes] - ) + reporter_order = tuple([ + str(reporter_class.__name__) for reporter_class in self._reporter_classes + ]) # init the kwargs for the monitor if monitor_partial_kwargs is None: @@ -242,12 +242,10 @@ def _gen_reporters(self): # the number of filenames all_exts = list( - it.chain( - *[ - [ext for ext in rep.SUGGESTED_EXTENSIONS] - for rep in self.reporter_classes - ] - ) + it.chain(*[ + [ext for ext in rep.SUGGESTED_EXTENSIONS] + for rep in self.reporter_classes + ]) ) n_exts = len(all_exts) diff --git a/src/wepy/orchestration/orchestrator.py b/src/wepy/orchestration/orchestrator.py index 14d057a6..34127de5 100644 --- a/src/wepy/orchestration/orchestrator.py +++ b/src/wepy/orchestration/orchestrator.py @@ -482,9 +482,9 @@ def configuration_hash_registered(self, config_hash): """ - if any( - [True if config_hash == h else False for h in self.configuration_hashes] - ): + if any([ + True if config_hash == h else False for h in self.configuration_hashes + ]): return True else: return False @@ -497,9 +497,9 @@ def add_configuration(self, configuration): config_hash = self.hash_snapshot(serialized_config) # check that the hash is not already in the snapshots - if any( - [True if config_hash == md5 else False for md5 in self.configuration_hashes] - ): + if any([ + True if config_hash == md5 else False for md5 in self.configuration_hashes + ]): # just skip the rest of the function and return the hash return config_hash @@ -513,9 +513,9 @@ def add_serial_configuration(self, serial_configuration): snaphash = self.hash_snapshot(serial_configuration) # check that the hash is not already in the configurations - if any( - [True if snaphash == md5 else False for md5 in self.configuration_hashes] - ): + if any([ + True if snaphash == md5 else False for md5 in self.configuration_hashes + ]): # just skip the rest of the function and return the hash return snaphash @@ -644,9 +644,7 @@ def get_run_records(self): get_run_record_query = """ SELECT * FROM runs - """.format( - fields=", ".join(self.RUN_SELECT_FIELDS) - ) + """.format(fields=", ".join(self.RUN_SELECT_FIELDS)) cursor = self._db.cursor() cursor.execute(get_run_record_query) @@ -659,9 +657,7 @@ def get_run_record(self, start_hash, end_hash): SELECT {fields} FROM runs WHERE start_hash=? AND end_hash=? - """.format( - fields=", ".join(self.RUN_SELECT_FIELDS) - ) + """.format(fields=", ".join(self.RUN_SELECT_FIELDS)) params = (start_hash, end_hash) @@ -1265,9 +1261,9 @@ def reconcile_orchestrators(host_path, *orchestrator_paths): """ if not osp.exists(host_path): - assert ( - len(orchestrator_paths) > 1 - ), "If the host path is a new orchestrator, must give at least 2 orchestrators to merge." + assert len(orchestrator_paths) > 1, ( + "If the host path is a new orchestrator, must give at least 2 orchestrators to merge." + ) # open the host orchestrator at the location which will have all # of the new things put into it from the other orchestrators. If @@ -1310,9 +1306,9 @@ def reconcile_orchestrators(host_path, *orchestrator_paths): # snapshots themselves, we trust they are correct for snaphash in orch.snapshot_hashes: # check that the hash is not already in the snapshots - if any( - [True if snaphash == md5 else False for md5 in new_orch.snapshot_hashes] - ): + if any([ + True if snaphash == md5 else False for md5 in new_orch.snapshot_hashes + ]): # skip it and move on continue @@ -1326,12 +1322,10 @@ def reconcile_orchestrators(host_path, *orchestrator_paths): config_hash = orch.run_configuration_hash(*run_id) # check that the hash is not already in the snapshots - if any( - [ - True if config_hash == md5 else False - for md5 in new_orch.configuration_hashes - ] - ): + if any([ + True if config_hash == md5 else False + for md5 in new_orch.configuration_hashes + ]): # skip it and move on continue @@ -1346,9 +1340,7 @@ def reconcile_orchestrators(host_path, *orchestrator_paths): # query to attach the foreign database attach_query = """ ATTACH '{}' AS {} - """.format( - orch_path, attached_table_name - ) + """.format(orch_path, attached_table_name) # query to update the runs tabel with new unique runs union_query = """ @@ -1358,16 +1350,12 @@ def reconcile_orchestrators(host_path, *orchestrator_paths): EXCEPT SELECT * FROM runs ) - """.format( - attached_table_name - ) + """.format(attached_table_name) # query to detach the table detach_query = """ DETACH {} - """.format( - attached_table_name - ) + """.format(attached_table_name) # then run the queries diff --git a/src/wepy/reporter/reporter.py b/src/wepy/reporter/reporter.py index 26630b82..7b09bd95 100644 --- a/src/wepy/reporter/reporter.py +++ b/src/wepy/reporter/reporter.py @@ -77,9 +77,9 @@ def init(self, **kwargs): """ method_name = "init" - assert not hasattr( - super(), method_name - ), f"Superclass with method {method_name} is masked" + assert not hasattr(super(), method_name), ( + f"Superclass with method {method_name} is masked" + ) def report(self, **kwargs): """Given data concerning the main simulation components state, perform @@ -142,9 +142,9 @@ def report(self, **kwargs): """ method_name = "report" - assert not hasattr( - super(), method_name - ), "Superclass with method {} is masked".format(method_name) + assert not hasattr(super(), method_name), ( + "Superclass with method {} is masked".format(method_name) + ) def cleanup(self, **kwargs): """Teardown routines for the reporter at the end of the simulation. @@ -177,9 +177,9 @@ def cleanup(self, **kwargs): """ method_name = "cleanup" - assert not hasattr( - super(), method_name - ), "Superclass with method {} is masked".format(method_name) + assert not hasattr(super(), method_name), ( + "Superclass with method {} is masked".format(method_name) + ) class FileReporter(Reporter): @@ -289,9 +289,9 @@ def __init__( # file paths - assert not ( - (file_paths is not None) and (file_path is not None) - ), "only file_paths or file_path kwargs can be specified" + assert not ((file_paths is not None) and (file_path is not None)), ( + "only file_paths or file_path kwargs can be specified" + ) # if only one file path is given then we handle it as multiple if file_path is not None: @@ -346,13 +346,13 @@ def __init__( # have been given else: # make sure it is in kwargs and valid - assert ( - file_paths is not None - ), "if no explicit file path is given the 'file_paths' must have a value" + assert file_paths is not None, ( + "if no explicit file path is given the 'file_paths' must have a value" + ) - assert len(file_paths) == len( - self.FILE_ORDER - ), "you must give file_paths {} paths".format(len(self.FILE_ORDER)) + assert len(file_paths) == len(self.FILE_ORDER), ( + "you must give file_paths {} paths".format(len(self.FILE_ORDER)) + ) # using the file_path paths we got above we set them as # attributes in this object @@ -364,9 +364,9 @@ def __init__( # modes - assert not ( - (modes is not None) and (mode is not None) - ), "only modes or mode kwargs can be specified" + assert not ((modes is not None) and (mode is not None)), ( + "only modes or mode kwargs can be specified" + ) # if modes is None we make modes, from defaults if we have to if modes is None: diff --git a/src/wepy/resampling/decisions/decision.py b/src/wepy/resampling/decisions/decision.py index a7263a05..48495ac4 100644 --- a/src/wepy/resampling/decisions/decision.py +++ b/src/wepy/resampling/decisions/decision.py @@ -197,14 +197,14 @@ def record(cls, enum_value, **fields): """ - assert ( - enum_value in cls.enum_dict_by_value() - ), "value is not a valid Enumerated value" + assert enum_value in cls.enum_dict_by_value(), ( + "value is not a valid Enumerated value" + ) for field_key in fields.keys(): - assert ( - field_key in cls.FIELDS - ), "The field {} is not a field for that decision".format(field_key) + assert field_key in cls.FIELDS, ( + "The field {} is not a field for that decision".format(field_key) + ) assert field_key != "decision_id", "'decision_id' cannot be an extra field" rec = {"decision_id": enum_value} diff --git a/src/wepy/resampling/resamplers/wexplore.py b/src/wepy/resampling/resamplers/wexplore.py index c35566c1..6bd488a3 100644 --- a/src/wepy/resampling/resamplers/wexplore.py +++ b/src/wepy/resampling/resamplers/wexplore.py @@ -255,10 +255,10 @@ def __init__( self._max_num_walkers = False self._min_num_walkers = False - assert ( - merge_method in self.MERGE_METHODS - ), "the merge method given, '{}', must be one of the methods available {}".format( - merge_method, self.MERGE_METHODS + assert merge_method in self.MERGE_METHODS, ( + "the merge method given, '{}', must be one of the methods available {}".format( + merge_method, self.MERGE_METHODS + ) ) self._merge_method = merge_method @@ -668,15 +668,13 @@ def place_walkers(self, walkers): assignment = self.branch_tree(parent_id, image) # save it to keep track of new branches as they occur - new_branches.append( - { - "distance": np.array([distance]), - "branching_level": np.array([level]), - "new_leaf_id": np.array(assignment), - # NOTE: Skip image for now. It is producing errors - # "image": image, - } - ) + new_branches.append({ + "distance": np.array([distance]), + "branching_level": np.array([level]), + "new_leaf_id": np.array(assignment), + # NOTE: Skip image for now. It is producing errors + # "image": image, + }) # we have made a new branch so we don't need to # continue this loop @@ -1292,9 +1290,9 @@ def _gen_best_donation( # find the largest difference comparing (a,b) and (b,a), # this will give the donor, acceptor pair permutations = [(a, b), (b, a)] - perm_idx = np.argmax( - [children_shares[i] - children_shares[j] for i, j in permutations] - ) + perm_idx = np.argmax([ + children_shares[i] - children_shares[j] for i, j in permutations + ]) donor_acceptor_pair = permutations[perm_idx] @@ -1457,17 +1455,15 @@ def _find_best_donation_pair( # check that both a donor and acceptor were identified and # that values for there shares were given - assert all( - [ - True if val is not None else False - for val in [ - donor_n_shares, - acceptor_n_shares, - donor_child_node_id, - acceptor_child_node_id, - ] + assert all([ + True if val is not None else False + for val in [ + donor_n_shares, + acceptor_n_shares, + donor_child_node_id, + acceptor_child_node_id, ] - ), "A donor or acceptor was not found" + ]), "A donor or acceptor was not found" # if the acceptor's number of shares is not less then the # donor then there is not possible donation @@ -1961,9 +1957,9 @@ def _check_clone_merge_specs(self, merge_groups, walkers_num_clones): keep_merge_walker_idxs.append(walker_idx) # add up the weights of the squashed walkers - squashed_weight = sum( - [self.walker_weights[i] for i in squash_walker_idxs] - ) + squashed_weight = sum([ + self.walker_weights[i] for i in squash_walker_idxs + ]) # add them to the weight for the keep walker walker_weight = self._walker_weights[walker_idx] + squashed_weight @@ -2009,9 +2005,9 @@ def _check_clone_merge_specs(self, merge_groups, walkers_num_clones): overweight_producer_idxs.append(keep_merge_walker_idx) raise ResamplerError( - "Merge specs produce overweight walkers for merge groups {}".format( - [str(i) for i in overweight_producer_idxs] - ) + "Merge specs produce overweight walkers for merge groups {}".format([ + str(i) for i in overweight_producer_idxs + ]) ) # check that all of the weights are less than or equal to the pmin @@ -2031,9 +2027,9 @@ def _check_clone_merge_specs(self, merge_groups, walkers_num_clones): underweight_producer_idxs.append(clone_parent_walker_idx) raise ResamplerError( - "Clone specs produce underweight walkers for clone walkers {}".format( - [str(i) for i in underweight_producer_idxs] - ) + "Clone specs produce underweight walkers for clone walkers {}".format([ + str(i) for i in underweight_producer_idxs + ]) ) def balance_tree(self, delta_walkers=0): @@ -2551,12 +2547,10 @@ def _check_resampling_data(resampling_data): ) # check that all squashes are going to a merge slot - if not all( - [ - False if squash_slot_idx not in keep_merge_slot_idxs else True - for squash_slot_idx in set(squash_slot_idxs) - ] - ): + if not all([ + False if squash_slot_idx not in keep_merge_slot_idxs else True + for squash_slot_idx in set(squash_slot_idxs) + ]): raise ResamplerError("Not all squashes are assigned to keep_merge slots") def _resample_init(self, walkers=None): diff --git a/src/wepy/runners/openmm.py b/src/wepy/runners/openmm.py index 23dfe9f4..d916e880 100644 --- a/src/wepy/runners/openmm.py +++ b/src/wepy/runners/openmm.py @@ -26,9 +26,6 @@ # Standard Library import logging - -logger = logging.getLogger(__name__) -# Standard Library import time from copy import copy from warnings import warn @@ -54,6 +51,9 @@ from wepy.work_mapper.task_mapper import WalkerTaskProcess from wepy.work_mapper.worker import Worker + +logger = logging.getLogger(__name__) + ## Constants KEYS = ( @@ -198,19 +198,18 @@ def get_state_fields_present(sim_state): # RAND_SEED_RANGE_MAX = 1000000 -# the runner for the simulation which runs the actual dynamics -class OpenMMRunner(Runner): +class OpenMMRunnerShared(Runner): """Runner for OpenMM simulations.""" def __init__( self, - system, - topology, - integrator, - platform=None, - platform_kwargs=None, - enforce_box=False, - get_state_kwargs=None, + system: omm.System, + topology: omma.Topology, + integrator: omm.Integrator, + platform: str | None = None, + platform_kwargs: dict | None = None, + enforce_box: bool = False, + get_state_kwargs: dict | None = None, ): """Constructor for OpenMMRunner. @@ -274,9 +273,9 @@ def __init__( """ if platform is not None: - assert isinstance( - platform, str - ), f"platform should be a string, not {type(platform)}" + assert isinstance(platform, str), ( + f"platform should be a string, not {type(platform)}" + ) # we save the different components. However, if we are to make # this runner picklable we have to convert the SWIG objects to @@ -302,7 +301,422 @@ def __init__( else: self.getState_kwargs = dict(GET_STATE_KWARG_DEFAULTS) + self._cycle_platform = None + self._cycle_platform_kwargs = None + + # for special monitoring purposes to get split times to debug + # performance + self._last_cycle_segments_split_times = [] + + def pre_cycle(self, platform=None, platform_kwargs=None, **kwargs): + # choose to use the platform spec in this function call or to + # use the default one saved in the runner + + # if the platform is given locally use this one + if platform is not None: + logger.info( + f"Setting the platform ({platform}) in the 'pre_cycle' OpenMM Runner call" + f"with platform kwargs: {platform_kwargs}" + ) + # set the platform and kwargs for this cycle + self._cycle_platform = platform + self._cycle_platform_kwargs = platform_kwargs + + # otherwise we just don't set this and let resolution of + # platform happen at run segment. + + super().pre_cycle(**kwargs) + + # each segment split times will get appended to this + self._last_cycle_segments_split_times = [] + + def post_cycle(self, **kwargs): + super().post_cycle(**kwargs) + + # remove the platform and kwargs for this cycle + self._cycle_platform = None + self._cycle_platform_kwargs = None + + def _resolve_platform( + self, + platform, + platform_kwargs, + ): + # resolve which platform to use + + # force usage of environmental one + if platform is Ellipsis: + platform_name = None + platform_kwargs = None + + # use the runtime given one + elif platform is not None: + platform_name = platform + platform_kwargs = platform_kwargs + + # if the pre_cycle configured platform is set use this over + # the default + elif self._cycle_platform is not None: + platform_name = self._cycle_platform + platform_kwargs = self._cycle_platform_kwargs + + # use the default one + elif self.platform_name is not None: + platform_name = self.platform_name + platform_kwargs = self.platform_kwargs + + # if the default is not set fall back to the environmental one + else: + platform_name = None + platform_kwargs = None + + return ( + platform_name, + platform_kwargs, + ) + + def run_segment( + self, + walker: Walker, + segment_length: int | float, + getState_kwargs: dict | None = None, + platform: str | None = None, + platform_kwargs: dict | None = None, + **kwargs, + ): + """Run dynamics for the walker. + + Parameters + ---------- + walker : object implementing the Walker interface + The walker for which dynamics will be propagated. + + segment_length : int or float + The numerical value that specifies how much dynamics are to be run. + + getState_kwargs : dict of str : bool, optional + Specify the key-word arguments to pass to + simulation.context.getState when getting simulation + states. If None defaults object values. + + platform : str or None or Ellipsis + The specification for the computational platform to + use. If None will use the default for the runner and + ignore platform_kwargs. If Ellipsis forces the use of the + OpenMM default or environmentally defined platform. See + OpenMM documentation for all value but typical ones are: + Reference, CUDA, OpenCL. If value is None the automatic + platform determining mechanism in OpenMM will be used. + + platform_kwargs : dict of str : bool, optional + key-values to set for a platform with + platform.setPropertyDefaultValue for this segment only. + + + Returns + ------- + new_walker : object implementing the Walker interface + Walker after dynamics was run, only the state should be modified. + + """ + + run_segment_start = time.time() + + # set the kwargs that will be passed to getState + tmp_getState_kwargs = getState_kwargs + + logger.info(f"Default 'getState_kwargs' in runner: {self.getState_kwargs}") + + logger.info(f"'getState_kwargs' passed to 'run_segment' : {getState_kwargs}") + + # start with the object value + getState_kwargs = copy(self.getState_kwargs) + if tmp_getState_kwargs is not None: + getState_kwargs.update(tmp_getState_kwargs) + + logger.info( + "After resolving 'getState_kwargs' that will be used are: " + f"{getState_kwargs}" + ) + + gen_sim_start = time.time() + + # make a copy of the integrator for this particular segment + new_integrator = copy(self.integrator) + # force setting of random seed to 0, which is a special + # value that forces the integrator to choose another + # random number + # new_integrator.setRandomNumberSeed(0) + + ## Platform + + logger.info(f"Default 'platform' in runner: {self.platform_name}") + + logger.info(f"pre_cycle set 'platform' in runner: {self._cycle_platform}") + + logger.info(f"'platform' passed to 'run_segment' : {platform}") + + logger.info(f"Default 'platform_kwargs' in runner: {self.platform_kwargs}") + + logger.info( + f"pre_cycle set 'platform_kwargs' in runner: {self._cycle_platform_kwargs}" + ) + + logger.info(f"'platform_kwargs' passed to 'run_segment' : {platform_kwargs}") + + platform_name, platform_kwargs = self._resolve_platform( + platform, + platform_kwargs, + ) + + logger.info(f"Resolved 'platform' : {platform_name}") + + logger.info(f"Resolved 'platform_kwargs' : {platform_kwargs}") + + # create simulation object + + ## create the platform and customize + + # if a platform was given we use it to make a Simulation object + if platform_name is not None: + logger.info("Using platform configured in code.") + + # get the platform by its name to use + platform = omm.Platform.getPlatformByName(platform_name) + logger.info(f"Platform object created: {platform}") + + if platform_kwargs is None: + platform_kwargs = {} + + # set properties from the kwargs if they apply to the platform + for key, value in platform_kwargs.items(): + if key in platform.getPropertyNames(): + logger.info(f"Setting platform property: {key} : {value}") + platform.setPropertyDefaultValue(key, value) + + else: + warn( + f"Platform kwargs given ({key} : {value}) " + f"but is not valid for this platform ({platform_name})" + ) + + # make a new simulation object + simulation = omma.Simulation( + self.topology, + self.system, + new_integrator, + omm.Platform.getPlatformByName("Reference"), + ) + + # otherwise just use the default or environmentally defined one + else: + logger.info("Using environmental platform.") + simulation = omma.Simulation(self.topology, self.system, new_integrator) + + # set the state to the context from the walker + # simulation.context.setState(walker.state.sim_state) + simulation.context.setPositions(walker.state["positions"]) + + gen_sim_end = time.time() + gen_sim_time = gen_sim_end - gen_sim_start + + logger.info("Time to generate the system: {}".format(gen_sim_time)) + + # actually run the simulation + steps_start = time.time() + logger.debug(f"Running {segment_length} steps of dynamics. {simulation!r}") + + # Run the simulation segment for the number of time steps + simulation.step(segment_length) + + steps_end = time.time() + steps_time = steps_end - steps_start + + logger.info("Time to run {} sim steps: {}".format(segment_length, steps_time)) + + get_state_start = time.time() + + get_state_end = time.time() + get_state_time = get_state_end - get_state_start + logger.info("Getting context state time: {}".format(get_state_time)) + + # generate the new state/walker + # new_state = self.generate_state( + # simulation, + # segment_length, + # walker, + # getState_kwargs, + # ) + new_sim_state = simulation.context.getState(**getState_kwargs) + new_state = PlainOpenMMState.from_openmm_state(new_sim_state) + + # create a new walker for this + new_walker = OpenMMWalker(new_state, walker.weight) + + run_segment_end = time.time() + run_segment_time = run_segment_end - run_segment_start + logger.info("Total internal run_segment time: {}".format(run_segment_time)) + + segment_split_times = { + "gen_sim_time": gen_sim_time, + "steps_time": steps_time, + "get_state_time": get_state_time, + "run_segment_time": run_segment_time, + } + + self._last_cycle_segments_split_times.append(segment_split_times) + + return new_walker + + def generate_state( + self, simulation, segment_length, starting_walker, getState_kwargs + ): + """Method for generating a wepy compliant state from an OpenMM + simulation object and data about the last segment of dynamics run. + + Parameters + ---------- + + simulation : simtk.openmm.app.Simulation object + A complete simulation object from which the state will be extracted. + + segment_length : int + The number of integration steps run in a segment of simulation. + + starting_walker : wepy.walker.Walker subclass object + The walker that was the beginning of this segment of simyulation. + + getState_kwargs : dict of str : bool + Specify the key-word arguments to pass to + simulation.context.getState when getting simulation + states. + + Returns + ------- + + new_state : wepy.runners.openmm.OpenMMState object + A new state from the simulation state. + + This method is meant to be called from within the + `run_segment` method during a simulation. It can be customized + in subclasses to allow for the addition of custom attributes + for a state, in addition to the base ones implemented in the + interface to the openmm simulation state in OpenMMState. + + The extra arguments to this function are data that would allow + for the calculation of integral values over the duration of + the segment, such as time elapsed and differences from the + starting state. + + """ + + # save the state of the system with all possible values + new_sim_state = simulation.context.getState(**getState_kwargs) + + # make an OpenMMState wrapper with this + new_state = OpenMMState(new_sim_state) + + return new_state + + +# the runner for the simulation which runs the actual dynamics +class OpenMMRunner(Runner): + """Runner for OpenMM simulations.""" + + def __init__( + self, + system: omm.System, + topology: omma.Topology, + integrator: omm.Integrator, + platform: str | None = None, + platform_kwargs: dict | None = None, + enforce_box: bool = False, + get_state_kwargs: dict | None = None, + ): + """Constructor for OpenMMRunner. + + Parameters + ---------- + system : simtk.openmm.System object + The system (forcefields) for the simulation. + + topology : simtk.openmm.app.Topology object + The topology for you system. + + integrator : subclass simtk.openmm.Integrator object + Integrator for propagating dynamics. + + platform : str + The specification for the default computational platform + to use. Platform can also be set when run_segment is + called. If None uses OpenMM default platform, see OpenMM + documentation for all value but typical ones are: + Reference, CUDA, OpenCL. If value is None the automatic + platform determining mechanism in OpenMM will be used. + + platform_kwargs : dict of str : bool, optional + key-values to set for a platform with + platform.setPropertyDefaultValue as the default for this + runner. + + enforce_box : bool + Calls 'context.getState' with 'enforcePeriodicBox' if True. + (Default value = False) + + get_state_kwargs : dict of str : bool, optional + key-values to set for getting the state from the OpenMM context. + keys not included will use the values in GET_STATE_KWARG_DEFAULTS. + Will override the enforce_box flag. + + Warnings + -------- + + Regarding the enforce_box option. + + When retrieving states from an OpenMM simulation Context, you + have the option to enforce periodic boundary conditions in the + resulting atomic positions in a topology aware way that + doesn't break bonds through boundaries. This is convenient for + post-processing as this can be a complex task and is not + readily exposed in the OpenMM API as a standalone function. + + However, in some types of simulations the periodic box vectors + are ignored (such as implicit solvent ones) despite there + being no option to not have periodic boundaries in the context + itself. Likely if you are running one of these kinds of + simulations you will not pay attention to the box vectors at + all and the random defaults that exist will be very wrong but + this incorrectness will not show in a non-wepy simulation with + openmm unless you are handling the context states + yourself. Then when you run in wepy the default of True to + enforce the boxes will be applied and confusingly wrong + answers will result that are difficult to find root cause of. + + """ + + # we save the different components. However, if we are to make + # this runner picklable we have to convert the SWIG objects to + # a picklable form + self.system = system + self.integrator = integrator + + # these are not SWIG objects + self.topology = topology + self.platform_name = platform + self.platform_kwargs = platform_kwargs + + self.enforce_box = enforce_box + + self.getState_kwargs = dict(GET_STATE_KWARG_DEFAULTS) + # update with the user based enforce_box + if get_state_kwargs is not None: + for k in get_state_kwargs: + self.getState_kwargs[k] = get_state_kwargs[k] + + # override enforce_box option if specified in get_state_kwargs + if "enforce_box" in get_state_kwargs: + self.enforce_box = get_state_kwargs["enforce_box"] self._cycle_platform = None self._cycle_platform_kwargs = None @@ -380,11 +794,11 @@ def _resolve_platform( def run_segment( self, - walker, - segment_length, - getState_kwargs=None, - platform=None, - platform_kwargs=None, + walker: Walker, + segment_length: int | float, + getState_kwargs: dict | None = None, + platform: str | None = None, + platform_kwargs: dict | None = None, **kwargs, ): """Run dynamics for the walker. @@ -468,7 +882,8 @@ def run_segment( logger.info(f"'platform_kwargs' passed to 'run_segment' : {platform_kwargs}") platform_name, platform_kwargs = self._resolve_platform( - platform, platform_kwargs + platform, + platform_kwargs, ) logger.info(f"Resolved 'platform' : {platform_name}") @@ -504,7 +919,10 @@ def run_segment( # make a new simulation object simulation = omma.Simulation( - self.topology, self.system, new_integrator, platform + self.topology, + self.system, + new_integrator, + platform, ) # otherwise just use the default or environmentally defined one @@ -513,7 +931,8 @@ def run_segment( simulation = omma.Simulation(self.topology, self.system, new_integrator) # set the state to the context from the walker - simulation.context.setState(walker.state.sim_state) + # simulation.context.setState(walker.state.sim_state) + simulation.context.setPositions(walker.state["positions"]) gen_sim_end = time.time() gen_sim_time = gen_sim_end - gen_sim_start @@ -523,6 +942,7 @@ def run_segment( # actually run the simulation steps_start = time.time() + logger.debug(f"Running {segment_length} steps of dynamics. {simulation!r}") # Run the simulation segment for the number of time steps simulation.step(segment_length) @@ -539,12 +959,14 @@ def run_segment( logger.info("Getting context state time: {}".format(get_state_time)) # generate the new state/walker - new_state = self.generate_state( - simulation, segment_length, walker, getState_kwargs - ) + # new_state = self.generate_state( + # simulation, segment_length, walker, getState_kwargs + # ) + new_sim_state = simulation.context.getState(**getState_kwargs) + new_state = PlainOpenMMState.from_openmm_state(new_sim_state) # create a new walker for this - new_walker = OpenMMWalker(new_state, walker.weight) + new_walker = Walker(new_state, walker.weight) run_segment_end = time.time() run_segment_time = run_segment_end - run_segment_start @@ -612,6 +1034,84 @@ def generate_state( return new_state +class PlainOpenMMState(WalkerState): + def __init__( + self, + fields_present: list[str], + **kwargs, + ): + self._fields_present = fields_present + self._data = kwargs + + def mutate_context(self, context: omm.Context): + # TODO: figure out units + for field in self._fields_present: + if field == "positions": + context.setPositions(self._data["positions"]) + elif field == "velocities": + context.setVelocities(self._data["velocities"]) + else: + logger.warning(f"Mutate context does not know how to set field {field}") + + @classmethod + def from_dict(cls, state_dict: dict): + fields_present = list(state_dict.keys()) + return cls(fields_present, **state_dict) + + @classmethod + def from_openmm_state(cls, state: omm.State): + fields_present = get_state_fields_present(state) + + fields = dict() + + if "positions" in fields_present: + positions = state.getPositions(asNumpy=True) + unit = positions.unit + fields["positions"] = positions.value_in_unit(unit) + + if "velocities" in fields_present: + velocities = state.getVelocities(asNumpy=True) + unit = velocities.unit + fields["velocities"] = velocities.value_in_unit(unit) + + if "forces" in fields_present: + forces = state.getForces(asNumpy=True) + unit = forces.unit + fields["forces"] = forces.value_in_unit(unit) + + if "kinetic_energy" in fields_present: + kinetic_energy = state.getKineticEnergy() + unit = kinetic_energy.unit + fields["kinetic_energy"] = kinetic_energy.value_in_unit(unit) + + fields["kinetic_energy"] = np.array([kinetic_energy]) + + if "potential_energy" in fields_present: + potential_energy = state.getPotentialEnergy() + unit = potential_energy.unit + fields["potential_energy"] = np.array([ + potential_energy.value_in_unit(unit) + ]) + + if "time" in fields_present: + fields["time"] = state.getTime() + + if "box_vectors" in fields_present: + box_vectors = state.getPeriodicBoxVectors() + unit = box_vectors.unit + fields["box_vectors"] = box_vectors.value_in_unit(unit) + + if "box_volume" in fields_present: + box_volume = state.getPeriodicBoxVolume() + unit = box_volume.unit + fields["box_volume"] = np.array([box_volume.value_in_unit(unit)]) + + return cls( + fields_present, + **fields, + ) + + class OpenMMState(WalkerState): """Walker state that wraps an simtk.openmm.State object. @@ -661,9 +1161,7 @@ def __init__(self, sim_state, **kwargs): warn( "Key {} in kwargs is already taken by this class, renaming to {}".format( self.OTHER_KEY_TEMPLATE - ).format( - key - ) + ).format(key) ) # make a new key @@ -874,9 +1372,9 @@ def kinetic_energy_value(self): if kinetic_energy is None: return None else: - return np.array( - [self.kinetic_energy.value_in_unit(self.kinetic_energy_unit)] - ) + return np.array([ + self.kinetic_energy.value_in_unit(self.kinetic_energy_unit) + ]) # Potential Energy @property @@ -907,9 +1405,9 @@ def potential_energy_value(self): if potential_energy is None: return None else: - return np.array( - [self.potential_energy.value_in_unit(self.potential_energy_unit)] - ) + return np.array([ + self.potential_energy.value_in_unit(self.potential_energy_unit) + ]) # Time @property @@ -1322,21 +1820,21 @@ class OpenMMWalker(Walker): def __init__(self, state, weight): # documented in superclass - assert isinstance( - state, OpenMMState - ), "state must be an instance of class OpenMMState not {}".format(type(state)) + assert isinstance(state, OpenMMState), ( + "state must be an instance of class OpenMMState not {}".format(type(state)) + ) super().__init__(state, weight) class OpenMMCPUWorker(Worker): - """Worker for OpenMM GPU simulations (CUDA or OpenCL platforms). + """Worker for OpenMM CPU simulations. This is intended to be used with the wepy.work_mapper.WorkerMapper work mapper class. This class must be used in order to ensure OpenMM runs jobs on the - appropriate GPU device. + appropriate CPU configuration. """ @@ -1360,8 +1858,6 @@ def run_task(self, task): # make the platform kwargs dictionary platform_options = {"Threads": str(self.attributes["num_threads"])} - # run the task and pass in the DeviceIndex for OpenMM to - # assign work to the correct GPU return task(platform_kwargs=platform_options) @@ -1402,7 +1898,6 @@ class OpenMMCPUWalkerTaskProcess(WalkerTaskProcess): NAME_TEMPLATE = "OpenMM_CPU_Walker_Task-{}" def run_task(self, task): - print("CPU Walker Task ---->", self.mapper_attributes, task, task.func) if "num_threads" in self.mapper_attributes: num_threads = self.mapper_attributes["num_threads"] @@ -1425,7 +1920,6 @@ class OpenMMGPUWalkerTaskProcess(WalkerTaskProcess): def run_task(self, task): logger.info(f"Starting to run a task as worker {self._worker_idx}") - print("GPU Walker Task ---->", self.mapper_attributes) # get the platform platform = self.mapper_attributes["platform"] @@ -1441,3 +1935,40 @@ def run_task(self, task): platform=platform, platform_kwargs=platform_options, ) + + +def gen_plain_walker_state(positions, system, integrator, getState_kwargs=None): + """Convenience function for generating a wepy walker State object for + an openmm simulation state. + + Parameters + ---------- + + positions : arraylike of float + The positions for the system you want to set + + system : openmm.app.System object + + integrator : openmm.Integrator object + + Returns + ------- + + walker_state : wepy.runners.openmm.OpenMMState object + + """ + + box_vectors = system.getDefaultPeriodicBoxVectors() + box_vectors_values = [] + for v in range(3): + u = box_vectors[v].unit + v3 = box_vectors[v].value_in_unit(u) + box_vectors_values.append([v3.x, v3.y, v3.z]) + + positions_unit = positions.unit + positions_values = positions.value_in_unit(positions_unit) + + return PlainOpenMMState.from_dict({ + "positions": positions_values, + "box_vectors": np.array(box_vectors_values), + }) diff --git a/src/wepy/sim_manager.py b/src/wepy/sim_manager.py index f9413aa6..0c8cfb72 100644 --- a/src/wepy/sim_manager.py +++ b/src/wepy/sim_manager.py @@ -43,10 +43,9 @@ """ # Standard Library -from typing import Final, Any +from typing import Final import logging -logger = logging.getLogger(__name__) # Standard Library import time from copy import deepcopy @@ -61,6 +60,9 @@ from wepy.boundary_conditions.boundary import BoundaryConditions +logger = logging.getLogger(__name__) + + class Manager: """The class that coordinates wepy simulations. @@ -120,7 +122,7 @@ def __init__( work_mapper=None, resampler: Resampler | None = None, boundary_conditions: BoundaryConditions | None = None, - reporters: Reporter | None = None, + reporters: list[Reporter] | None = None, sim_monitor=None, ): """Constructor for Manager. @@ -339,13 +341,16 @@ def run_cycle( def _run_cycle( self, - walkers, - n_segment_steps, - cycle_idx, - runner_opts=None, + walkers: list[Walker], + n_segment_steps: int, + cycle_idx: int, + runner_opts: dict | None = None, ): """See run_cycle.""" + if self.runner is None: + raise ValueError("No runner provided to simulation manager.") + if runner_opts is None: runner_opts = {} @@ -492,9 +497,9 @@ def _run_cycle( # check that all of the keys that are specified for this sim # manager are present - assert all( - [True if rep_key in report else False for rep_key in self.REPORT_ITEM_KEYS] - ) + assert all([ + True if rep_key in report else False for rep_key in self.REPORT_ITEM_KEYS + ]) logger.info("Starting reporting") # report results to the reporters @@ -638,7 +643,12 @@ def cleanup(self): reporters=self.reporters, ) - def run_simulation_by_time(self, run_time, segments_length, num_workers=None): + def run_simulation_by_time( + self, + run_time: float, + segments_length: int, + num_workers: int | None = None, + ) -> tuple[list[Walker], list]: """Run a simulation for a certain amount of time. This starts timing as soon as this is called. If the time @@ -683,6 +693,7 @@ def run_simulation_by_time(self, run_time, segments_length, num_workers=None): self.init(num_workers=num_workers) cycle_idx = 0 walkers = self.init_walkers + filters = [] while time.time() - start_time < run_time: logger.info( "starting cycle {} at time {}".format( @@ -705,10 +716,10 @@ def run_simulation_by_time(self, run_time, segments_length, num_workers=None): def run_simulation( self, - n_cycles, - segment_lengths, - num_workers=None, - ): + n_cycles: int, + segment_lengths: int | list[int], + num_workers: int | None = None, + ) -> tuple[list[Walker], list]: """Run a simulation for an explicit number of cycles. Parameters @@ -716,7 +727,7 @@ def run_simulation( n_cycles : int Number of cycles to perform. - segment_lengths : int + segment_lengths : int or list[int] The number of steps for each runner segment. num_workers : int @@ -743,15 +754,20 @@ def run_simulation( self.init(num_workers=num_workers) - if type(segment_lengths) == int: - segment_lengths = [segment_lengths for _ in range(n_cycles)] + if isinstance(segment_lengths, int): + _segment_lengths = [segment_lengths for _ in range(n_cycles)] + else: + _segment_lengths = segment_lengths walkers = self.init_walkers + filters = [] # the main cycle loop for cycle_idx in range(n_cycles): walkers, filters = self.run_cycle( - walkers, segment_lengths[cycle_idx], cycle_idx + walkers, + _segment_lengths[cycle_idx], + cycle_idx, ) # run the simulation monitor to get metrics on everything @@ -764,10 +780,10 @@ def run_simulation( def continue_run_simulation( self, - run_idx, - n_cycles, - segment_lengths, - num_workers=None, + run_idx: int, + n_cycles: int, + segment_lengths: list[int], + num_workers: int | None = None, ): """Continue a simulation. All this does is provide a run idx to the reporters, which is the run that is intended to be @@ -809,6 +825,8 @@ def continue_run_simulation( self.init(num_workers=num_workers, continue_run=run_idx) walkers = self.init_walkers + filters = [] + # the main cycle loop for cycle_idx in range(n_cycles): walkers, filters = self.run_cycle( @@ -824,8 +842,12 @@ def continue_run_simulation( return walkers, filters def continue_run_simulation_by_time( - self, run_idx, run_time, segments_length, num_workers=None - ): + self, + run_idx: int, + run_time: float, + segments_length: int, + num_workers: int | None = None, + ) -> tuple[list[Walker], list]: """Continue a simulation with a separate run by time. This starts timing as soon as this is called. If the time @@ -859,6 +881,7 @@ def continue_run_simulation_by_time( cycle_idx = 0 walkers = self.init_walkers + filters = [] while time.time() - start_time < run_time: logger.info( "starting cycle {} at time {}".format( diff --git a/src/wepy/walker.py b/src/wepy/walker.py index 1c8119d8..370d9641 100644 --- a/src/wepy/walker.py +++ b/src/wepy/walker.py @@ -32,7 +32,7 @@ logger = logging.getLogger(__name__) # Standard Library import random as rand -from copy import deepcopy +from copy import deepcopy, copy def split(walker, number=2): @@ -125,7 +125,7 @@ def merge(walkers): return new_walker, keep_idx -class Walker(object): +class Walker: """Reference implementation of the Walker interface. A container for: @@ -174,7 +174,7 @@ def clone(self, number=1): split_prob = self.weight / (number + 1) # make the clones clones = [] - for i in range(number + 1): + for _ in range(number + 1): clones.append(type(self)(self.state, split_prob)) return clones @@ -211,7 +211,7 @@ def merge(self, other_walkers): return merge([self] + other_walkers) -class WalkerState(object): +class WalkerState: """Reference implementation of the WalkerState interface. Access all key-value pairs as a dictionary with the dict() method. diff --git a/src/wepy/work_mapper/mapper.py b/src/wepy/work_mapper/mapper.py index e60f9404..db8cf384 100644 --- a/src/wepy/work_mapper/mapper.py +++ b/src/wepy/work_mapper/mapper.py @@ -6,27 +6,34 @@ # Standard Library import logging - -logger = logging.getLogger(__name__) -# Standard Library import multiprocessing as mp import queue as pyq import signal import sys import time import traceback +from collections.abc import Callable +from typing import Any, Literal from warnings import warn # First Party Library -from wepy.util.util import set_loglevel +from wepy.walker import Walker + +# from wepy.work_mapper.worker import Worker + +logger = logging.getLogger(__name__) + +ProcStartMethod = Literal["fork", "spawn", "forkserver"] + +SegmentFunc = Callable[[Walker, int | float], Walker] PY_MAP = map -class ABCMapper(object): +class ABCMapper: """Abstract base class for a Mapper.""" - def __init__(self, segment_func=None, **kwargs): + def __init__(self, segment_func: SegmentFunc | None = None, **kwargs): """Constructor for the Mapper class. No arguments are required. Parameters @@ -44,7 +51,7 @@ def __init__(self, segment_func=None, **kwargs): def attributes(self): return self._attributes - def init(self, segment_func=None, **kwargs): + def init(self, segment_func: SegmentFunc | None = None, **kwargs): """Runtime initialization and setting of function to map over walkers. Parameters @@ -68,7 +75,7 @@ def init(self, segment_func=None, **kwargs): self._func = segment_func @property - def segment_func(self): + def segment_func(self) -> SegmentFunc | None: """The function that will be called for new data in the `map` method.""" return self._func @@ -94,7 +101,7 @@ def map(self, *args, **kwargs): class Mapper(ABCMapper): """Basic non-parallel reference implementation of a mapper.""" - def __init__(self, segment_func=None, **kwargs): + def __init__(self, segment_func: SegmentFunc | None = None, **kwargs): """Constructor for the Mapper class. No arguments are required. Parameters @@ -123,12 +130,14 @@ def map(self, *args, **kwargs): Examples -------- - >>> Mapper(segment_func=sum).map([(0,1,2), (3,4,5)]) [3, 12] """ + if self._func is None: + raise ValueError("No segment_func has been set.") + # expand the generators for the args and kwargs args = [list(arg) for arg in args] kwargs = {key: list(kwarg) for key, kwarg in kwargs.items()} @@ -204,7 +213,7 @@ def worker_segment_times(self): class Task: """Class that composes a function and arguments.""" - def __init__(self, func, *args, **kwargs): + def __init__(self, func: Callable, *args, **kwargs): """Constructor for Task. Parameters @@ -254,876 +263,908 @@ class TaskException(WrapperException): pass -class ABCWorkerMapper(ABCMapper): +# same for the worker in terms of refactoring +class Worker(mp.Process): + """Worker process. + + This is a subclass of process with an overriden `__init__` + constructor that will automatically generate the Process. + + When this class is constructed a new process will be formed. + + """ + + NAME_TEMPLATE = "Worker-{}" + """A string formatting template to identify worker processes in + logs. The field will be filled with the worker index.""" + def __init__( - self, num_workers=None, segment_func=None, proc_start_method="fork", **kwargs + self, + worker_idx: int, + task_queue: mp.JoinableQueue, + result_queue: mp.Queue, + exception_queue: mp.Queue, + interrupt_connection, + mapper_attributes: dict[str, Any] | None = None, + log_level="INFO", + **kwargs, ): - """Constructor for WorkerMapper. - + """Constructor for the Worker class. Parameters ---------- - num_workers : int - The number of worker processes to spawn. - - segment_func : callable, optional - Set a default segment_func. Typically set at runtime. - - proc_start_method : str or None - A string indicating the type of process start method to - use from python multiprocessing typically 'fork', 'spawn', - or 'forkserver', or the platform default for None. See - documentation. Generates a context with the method - multiprocessing.get_context(proc_start_method) on `init`. + worker_idx : int + The index of the worker. Should be unique. - """ + task_queue : multiprocessing.JoinableQueue + The shared task queue the worker will watch for new tasks to complete. - super().__init__(segment_func=segment_func, **kwargs) + result_queue : multiprocessing.Queue + The shared queue that completed task results will be placed on. - self._proc_start_method = proc_start_method + interrupt_connection : multiprocessing.Connection + One end of a pipe to listen for messages specific to this worker. - self._num_workers = num_workers - self._worker_segment_times = None + mapper_attributes : None or dict + A dictionary of the attributes of the mapper for reference in workers. - if num_workers is not None: - self._worker_segment_times = {i: [] for i in range(self.num_workers)} + kwargs : + The worker specific attributes - def init(self, num_workers=None, segment_func=None, **kwargs): - """Runtime initialization and setting of function to map over walkers. + """ - Parameters - ---------- - num_workers : int - The number of worker processes to spawn + # call the Process constructor + super().__init__(name=self.NAME_TEMPLATE.format(worker_idx)) - segment_func : callable implementing the Runner.run_segment interface + self._exception_queue = exception_queue + self._exception = None + self._traceback = None - """ + # the queue that will trigger a shutdown in the event of failure + self._irq_channel = interrupt_connection - super().init(segment_func=segment_func) + # also register the SIGTERM signal handler for graceful + # shutdown with reporting to mapper + signal.signal(signal.SIGTERM, self._sigterm_shutdown) - # create the multiprocessing context to use for spawning - # processes here - self._mp_ctx = mp.get_context(method=self._proc_start_method) + self._worker_idx = worker_idx - # the number of workers must be given here or set as an object attribute - if num_workers is None and self.num_workers is None: - raise ValueError( - "The number of workers must be given, received {}".format(num_workers) - ) + self._mapper_attributes = mapper_attributes - # if the number of walkers was given for this init() call use - # that, otherwise we use the default that was specified when - # the object was created - elif num_workers is not None and self.num_workers is None: - self._num_workers = num_workers + # set all the kwargs into an attributes dictionary + self._attributes = kwargs - # update the worker segment times - self._worker_segment_times = {i: [] for i in range(self.num_workers)} + # the queues for work to be done and work done + self._task_queue = task_queue + self._result_queue = result_queue - def cleanup(self, **kwargs): - # ALERT: is this all we need to do? I have a hunch there is - # more caveats, but these context objects are not really - # documented + logger.debug("{} process created".format(self.name)) - # make sure the context for this work mapper is destroyed - del self._mp_ctx + @property + def worker_idx(self): + """Dictionary of attributes of the worker.""" + return self._worker_idx @property - def num_workers(self): - """The number of worker processes.""" - return self._num_workers + def attributes(self): + """Dictionary of attributes of the worker.""" + return self._attributes @property - def worker_segment_times(self): - """The run timings for each segment for each walker. + def mapper_attributes(self): + """Dictionary of attributes of the worker.""" + return self._mapper_attributes - Returns - ------- - worker_seg_times : dict of int : list of float - Dictionary mapping worker indices to a list of times in - seconds for each segment run. + def run(self): + logger.debug("{}: starting to run".format(self.name)) - """ - return self._worker_segment_times + # try to run the worker and it's task, except either class of + # error that can come from it either from the worker + # (WorkerException) or the task (TaskException) and communicate it + # back to the main process - def _make_task(self, *args, **kwargs): - """Generate a task from 'segment_func' attribute. + # if we get an exception there is some cleanup logic + run_exception = None - Similar to partial evaluation (or currying). + try: + # run the worker, which will retrieve its task from the + # queue attempt to run the task, and if it succeeds will + # put the results on the result queue, if the task fails + # it will catch it and wrap it as a task exception + self._run_worker() - Args will be eventually used as the arguments to the call of - 'segment_func' by the worker processes when they receive the - task from the queue. + except TaskException as task_exception: + logger.error("{}: TaskException caught".format(self.name)) - Returns - ------- - task : Task object + run_exception = task_exception - """ - return Task(self._func, *args, **kwargs) + # anything else is considered a WorkerException so take the + # original exception and generate a worker exception from that + except Exception as exception: + logger.debug("{}: WorkerError caught".format(self.name)) + # get the traceback + tb = sys.exc_info()[2] -# ---------------------------------- -# everything below this logically belongs in worker.py and should be imported from there + msg = "Exception '{}({})' caught in a worker.".format( + type(exception).__name__, exception + ) + traceback_log_msg = """Traceback: +-------------------------------------------------------------------------------- +{} +-------------------------------------------------------------------------------- + """.format( + "".join(traceback.format_exception(type(exception), exception, tb)), + ) + logger.error("{}:".format(self.name) + msg + "\n" + traceback_log_msg) -class WorkerException(WrapperException): - pass + # raise a TaskError to distinguish it from the worker + # errors with the metadata about the original exception + worker_exception = WorkerException( + "Error occured during worker execution.", + wrapped_exception=exception, + tb=tb, + ) -class WorkerKilledError(ChildProcessError): - pass + run_exception = worker_exception + # raise worker_exception + if run_exception is not None: + logger.debug("{}: Putting exception on exception queue".format(self.name)) -# TODO: move this class to the wepy.work_mapper.worker class where it -# belongs. It shouldn't be in this namespace, but we will leave it -# here. Furthermore I would like to rename it since we now have -# different worker mapper implementations with different concurrency -# models -class WorkerMapper(ABCWorkerMapper): - """Work mapper implementation using multiple worker processes and task - queue. + # then put the exception and the traceback onto the queue + # so we can communicate back to the parent process + try: + self._exception_queue.put((self.name, self.pid, run_exception)) + except BrokenPipeError as exc: + logger.error( + "Pipe is broken indicating the root process has already exited:\n{}".format( + exc + ) + ) - Uses the python multiprocessing module to spawn multiple worker - processes which watch a task queue of walker segments. - """ + # TODO: not sure if this is good or not + # then reraise the exception so it can be caught + # raise run_exception - def __init__( - self, - num_workers=None, - worker_type=None, - worker_attributes=None, - segment_func=None, - **kwargs, - ): - """Constructor for WorkerMapper. + def _sigterm_shutdown(self, signum, frame): + logger.debug("Received external SIGTERM kill command.") + logger.debug("Alerting mapper that this will be honored.") - Parameters - ---------- - num_workers : int - The number of worker processes to spawn. + # send an error to the mapper that the worker has been killed + self._irq_channel.send( + WorkerKilledError( + "{} (pid: {}) killed by external SIGTERM signal".format( + self.name, self.pid + ) + ) + ) - worker_type : callable, optional - Callable that generates an object implementing the Worker - interface, typically a type from a Worker class. + logger.debug("Acknowledgment sent") - worker_attributes : dictionary - A dictionary of values that are passed to the worker - constructor as key-word arguments. + logger.debug("Shutting down process") - segment_func : callable, optional - Set a default segment_func. Typically set at runtime. + def _shutdown(self): + logger.debug("Received SIGTERM kill command from mapper") - """ + logger.debug("Acknowledging kill request will be honored") - super().__init__(num_workers=num_workers, segment_func=segment_func, **kwargs) + # report back that we are shutting down with a True + self._irq_channel.send(True) - # since the workers will be their own process classes we - # handle this data + logger.debug("Acknowledgment sent") - # attributes that will be passed to the worker constructors - if worker_attributes is not None: - self._worker_attributes = worker_attributes - else: - self._worker_attributes = {} + logger.debug("Shutting down process") - # choose the type of the worker - if worker_type is None: - self._worker_type = Worker - warn("worker_type not given using the default base class") - logger.warn("worker_type not given using the default base class") - else: - self._worker_type = worker_type + def _run_worker(self): + # run the logic associated with communication and liveness of + # the worker process itself, this is not necessarily a fatal + # (critical) error and restarting a worker might resolve the + # problem. This calls the _run_task method though which is + # always critical since the logic in the code cannot be + # disputed - @property - def worker_type(self): - """The callable that generates a worker object. + # TODO remove when confirmed that this works + # worker_process = mp.current_process() + logger.info( + "{}: Worker process started as name: {}; PID: {}".format( + self.name, self.name, self.pid + ) + ) - Typically this is just the type from the class definition of - the Worker where the constructor is called. + while True: + # check to see if there is any signals in the interrupt channel + if self._irq_channel.poll(): + # get the message + message = self._irq_channel.recv() - """ - return self._worker_type + logger.debug( + "{}: Received message from mapper on filehandle {}: {}".format( + self.name, self._irq_channel.fileno(), message + ) + ) - def init(self, num_workers=None, segment_func=None, **kwargs): - """Runtime initialization and setting of function to map over walkers. + # handle the message - Parameters - ---------- - num_workers : int - The number of worker processes to spawn + # a SIGTERM is a signal to kill the process + # unconditionally + if message is signal.SIGTERM: + self._shutdown() - segment_func : callable implementing the Runner.run_segment interface + # break from the event (while) loop and shut down + break - """ + # anything is not recognized and we will continue and + # report back that we don't recognize the message with + # a ValueError object + else: + logger.error( + "{}: Message not recognized, continuing operations and" + " sending error to mapper".format(self.name) + ) + self._irq_channel.send( + ValueError( + "Message: {} not recognized continuing operations".format( + message + ) + ) + ) - super().init(num_workers=num_workers, segment_func=segment_func, **kwargs) + # get the next task + try: + task_idx, next_task = self._task_queue.get(block=False, timeout=None) - manager = self._mp_ctx.Manager() + logger.debug("{}: Got task {}".format(self.name, task_idx)) - # Establish communication queues + except pyq.Empty: + task_idx = None + next_task = Ellipsis - # A queue for errors - self._exception_queue = manager.Queue() + # # check for the poison pill which is the signal to stop + if next_task is None: + logger.info( + "{}: received {} {}: FINISHED".format( + self.name, task_idx, next_task + ) + ) - # queue for the tasks we know the batch size so we don't need - # a JoinableQueue - self._task_queue = manager.Queue() + # TODO remove since we aren't using joinble queue anymore + # mark the poison pill task as done + # self.task_queue.task_done() - # results queue - self._result_queue = manager.Queue() + # and exit the loop + break - # use pipes for communication channels between this parent - # process and the children for sending specific interrupts - # such as the signal to kill them. Note that the clean way to - # end the process is to send poison pills on the task queue, - # this is for other stuff. IRQ is a common abbreviation for - # interrupts - self._irq_parent_conns = [] + # only execute this if a task was actually receieved from + # the queue; an Ellipsis indicates continue the loop + elif next_task is not Ellipsis: + logger.info( + "{}; task_idx : {}; args : {} ".format( + self.name, + task_idx, + next_task.args, + ) + ) - # Start workers, giving them all the queues - self._workers = [] - for i in range(self.num_workers): - # make a pipe to communicate with this worker for the int - parent_conn, child_conn = self._mp_ctx.Pipe() - self._irq_parent_conns.append(parent_conn) + # run the task + start = time.time() - # create the worker giving it all of the communication - # channels - worker = self.worker_type( - i, - self._task_queue, - self._result_queue, - self._exception_queue, - child_conn, - mapper_attributes=self._attributes, - **self._worker_attributes, - ) - self._workers.append(worker) + answer = self._run_task(next_task) - # start the worker processes - for worker in self._workers: - worker.start() + end = time.time() + task_time = end - start - logger.info( - "Worker process started as name: {}; PID: {}".format( - worker.name, worker.pid + logger.info( + "{}: task_idx : {}; COMPLETED in {} s".format( + self.name, task_idx, task_time + ) ) - ) - # now that we have started the processes register the handler - # for SIGTERM signals that will clean up our children cleanly - signal.signal(signal.SIGTERM, self._sigterm_shutdown) + # put the results into the results queue with it's task + # index so we can sort them later + self._result_queue.put((task_idx, self.worker_idx, task_time, answer)) - def _sigterm_shutdown(self, signum, frame): - logger.critical("Received external SIGTERM, forcing shutdown.") + def run_task(self, task): + """Actually executes the task. - self.force_shutdown() + This default runner simply executes the task thunk. - def force_shutdown(self, **kwargs): - logger.critical("Forcing shutdown") + This can be customized by subclasses in order to allow for + injection of worker specific data. - # our primary job is to shut down all of the running processes - # without just shutting down the queues and breaking the pipes + Parameters + ---------- + task : Task object + The partially evaluated task; function plus arguments - # to do this we send the kill signals to them on the kill - # channel. + Returns + ------- + task_result + Results of running the task. - for worker_idx, worker in enumerate(self._workers): - logger.critical( - "Sending SIGTERM message on {} to worker {}".format( - self._irq_parent_conns[worker_idx].fileno(), worker_idx - ) - ) + """ - # send a kill message to the worker - self._irq_parent_conns[worker_idx].send(signal.SIGTERM) + return task() - logger.critical("All kill messages sent to workers") + def _run_task(self, task: Task): + """Runs the given task and returns the results. - # check that all have exited - alive_workers = [worker.is_alive() for worker in self._workers] - worker_acks = {} - worker_exitcodes = {} - premature_exit = False - while any(alive_workers) and not premature_exit: - for worker_idx, worker in enumerate(self._workers): - # ignore already known dead workers - if not alive_workers[worker_idx]: - continue + This manages handling exceptions and tracebacks from the + actual `run_task` function which is intended to be specialized + by different workers to inject worker specific arguments to + tasks. Such as node and device identification. - if worker.is_alive(): - # if it is still alive and we have an ack from it - # just terminate. There is a bug in the code and - # is out of our control - if worker_idx in worker_acks: - logger.debug( - "Ack received from {} but has not shut down".format( - worker.name - ) - ) - premature_exit = True + Parameters + ---------- + task : Task object + The partially evaluated task; function plus arguments - # otherwise we need to try and receive the ack - elif self._irq_parent_conns[worker_idx].poll(1): - # receive the acknowledgement - ack = self._irq_parent_conns[worker_idx].recv() + Returns + ------- + task_result + Results of running the task. - logger.debug( - "Received {} acknowledgement from {}".format( - ack, worker.name - ) - ) + """ - # make sure the ack is affirmative - if ack is True: - worker_acks[worker_idx] = ack + logger.info("Running task") + try: + return self.run_task(task) - # if it is an exeption wrap it as a worker - # error and use the os to kill the process - elif issubclass(type(ack), Exception): - # wrap it as a worker exception - exception = WorkerException(wrapped_exception=ack) - worker_acks[worker_idx] = exception + except Exception as task_exception: + # get the traceback for the exception + tb = sys.exc_info()[2] - logger.critical( - "{} not responding, terminating with SIGTERM".format( - worker.name - ) - ) + msg = "Exception '{}({})' caught in a task.".format( + type(task_exception).__name__, task_exception + ) + traceback_log_msg = """Traceback: +-------------------------------------------------------------------------------- +{} +-------------------------------------------------------------------------------- + """.format( + "".join( + traceback.format_exception(type(task_exception), task_exception, tb) + ), + ) - worker.terminate() + logger.critical(msg + "\n" + traceback_log_msg) - else: - alive_workers[worker_idx] = False - worker_exitcodes[worker_idx] = worker.exitcode + # raise a TaskException to distinguish it from the worker + # errors with the metadata about the original exception - if any(alive_workers): - logger.critical( - "Terminating main process with running workers {}".format( - ",".join( - [ - str(worker_idx) - for worker_idx in range(len(self._workers)) - if alive_workers[worker_idx] - ] - ) - ) + raise TaskException( + "Error occured during task execution, recovery not possible.", + wrapped_exception=task_exception, + tb=tb, ) - def cleanup(self, **kwargs): - """Runtime post-simulation tasks. - This is run either at the end of a successful simulation or - upon an error in the main process of the simulation manager - call to `run_cycle`. +class ABCWorkerMapper(ABCMapper): + def __init__( + self, + num_workers: int | None = None, + segment_func: SegmentFunc | None = None, + proc_start_method: ProcStartMethod | None = "fork", + **kwargs, + ): + """Constructor for WorkerMapper. - The Mapper class performs no actions here and all arguments - are ignored. + Parameters + ---------- + num_workers : int + The number of worker processes to spawn. + + segment_func : callable, optional + Set a default segment_func. Typically set at runtime. + + proc_start_method : str or None + A string indicating the type of process start method to + use from python multiprocessing typically 'fork', 'spawn', + or 'forkserver', or the platform default for None. See + documentation. Generates a context with the method + multiprocessing.get_context(proc_start_method) on `init`. """ - super().cleanup(**kwargs) + super().__init__(segment_func=segment_func, **kwargs) - # send poison pills (Stop signals) to the queues to stop them in a nice way - # and let them finish up - for i in range(self.num_workers): - self._task_queue.put((None, None)) + self._proc_start_method = proc_start_method - # delete the queues and workers - self._task_queue = None - self._result_queue = None - self._workers = None + self._num_workers = num_workers + self._worker_segment_times = None - def map(self, *args, **kwargs): - # docstring in superclass + if self.num_workers is not None: + self._worker_segment_times = {i: [] for i in range(self.num_workers)} - map_process = self._mp_ctx.current_process() - logger.info( - "Mapping from process {}; PID {}".format(map_process.name, map_process.pid) - ) + def init( + self, + segment_func: SegmentFunc | None = None, + num_workers: int | None = None, + **kwargs, + ): + """Runtime initialization and setting of function to map over walkers. - # make tuples for the arguments to each function call - task_args = zip(*args) - kwargs = {key: list(kwarg) for key, kwarg in kwargs.items()} + Parameters + ---------- + num_workers : int + The number of worker processes to spawn - num_tasks = len(args[0]) - # Enqueue the jobs - for task_idx, task_arg in enumerate(task_args): - task_kwargs = {key: value[task_idx] for key, value in kwargs.items()} + segment_func : callable implementing the Runner.run_segment interface - # a task will be the actual task and its task idx so we can - # sort them later - self._task_queue.put((task_idx, self._make_task(*task_arg, **task_kwargs))) + """ - logger.info("Waiting for tasks to be run") + super().init(segment_func=segment_func) - # poll the exception and result queues for results - n_results_left = num_tasks - results = [] - while n_results_left > 0: - # first check if any errors came back but don't wait, - # since the methods for querying whether it is empty or - # not are not reliable we just try and if we don't get - # anything we will come back around - try: - proc_name, pid, exception = self._exception_queue.get_nowait() - except pyq.Empty: - pass + # create the multiprocessing context to use for spawning + # processes here + self._mp_ctx = mp.get_context(method=self._proc_start_method) - else: - logger.error( - "Exception occured in process {}; pid {}.".format(proc_name, pid) + # if the number of walkers was given for this init() call use + # that, otherwise we use the default that was specified when + # the object was created + if num_workers is not None and self.num_workers is None: + self._num_workers = num_workers + elif num_workers is not None and self.num_workers is not None: + logger.warning( + "Both num_workers and self.num_workers are set, using default from self" + ) + + # the number of workers must be given here or set as an object attribute + if self.num_workers is None: + raise ValueError( + "The number of workers must be given, received {}".format( + self.num_workers ) + ) - # we can handle Task and Worker exceptions differently - if type(exception) == TaskException: - logger.critical( - "Exception encountered in a task which is unrecoverable." - "You will need to reconfigure your components in a stable manner." - ) + # update the worker segment times + self._worker_segment_times = {i: [] for i in range(self.num_workers)} - self.force_shutdown() + def cleanup(self, **kwargs): + # ALERT: is this all we need to do? I have a hunch there is + # more caveats, but these context objects are not really + # documented - logger.critical("Shutdown complete.") - raise exception + # make sure the context for this work mapper is destroyed + del self._mp_ctx - elif type(exception) == WorkerException: - # we make just an error message to say that errors - # in the worker may be due to the network or - # something and could recover - logger.error( - "Exception encountered in the work mapper worker process." - "Recovery possible, see further messages." - ) + @property + def num_workers(self): + """The number of worker processes.""" + return self._num_workers - # However, the current implementation doesn't - # support retries or whatever so we issue a - # critical log informing that it has been elevated - # to critical and will force shutdown - logger.critical( - "Worker error mode resiliency not supported at this time." - "Performing force shutdown and simulation ending." - ) + @property + def worker_segment_times(self): + """The run timings for each segment for each walker. - self.force_shutdown() + Returns + ------- + worker_seg_times : dict of int : list of float + Dictionary mapping worker indices to a list of times in + seconds for each segment run. - logger.critical("Shutdown complete.") - raise exception + """ + return self._worker_segment_times - else: - logger.critical("Unknown exception encountered.") + def _make_task(self, *args, **kwargs): + """Generate a task from 'segment_func' attribute. - self.force_shutdown() + Similar to partial evaluation (or currying). - logger.critical("Shutdown complete.") + Args will be eventually used as the arguments to the call of + 'segment_func' by the worker processes when they receive the + task from the queue. - raise exception + Returns + ------- + task : Task object - # attempt to get something off of the results queue - try: - result = self._result_queue.get_nowait() - except pyq.Empty: - pass + """ + return Task(self._func, *args, **kwargs) - # if we get something handle it - else: - logger.info("Retrieved result: {}".format(result)) - results.append(result) - # reduce the counter so we know when we are done - n_results_left -= 1 +# ---------------------------------- +# everything below this logically belongs in worker.py and should be imported from there - # sort the results according to their task_idx - results.sort() - # save the task run times, so they can be accessed if desired, - # after clearing the task times from the last mapping +class WorkerException(WrapperException): + pass - # DEBUG: removing this because it should be set on init() - # self._worker_segment_times = {i : [] for i in range(self.num_workers)} - for task_idx, worker_idx, task_time, result in results: - self._worker_segment_times[worker_idx].append(task_time) +class WorkerKilledError(ChildProcessError): + pass - # then just return the values of the function - return [result for task_idx, worker_idx, task_time, result in results] +# TODO: move this class to the wepy.work_mapper.worker class where it +# belongs. It shouldn't be in this namespace, but we will leave it +# here. Furthermore I would like to rename it since we now have +# different worker mapper implementations with different concurrency +# models +class WorkerMapper(ABCWorkerMapper): + """Work mapper implementation using multiple worker processes and task + queue. -# same for the worker in terms of refactoring -class Worker(mp.Process): - """Worker process. + Uses the python multiprocessing module to spawn multiple worker + processes which watch a task queue of walker segments. + """ + + def __init__( + self, + num_workers: int | None = None, + worker_type: type[Worker] | None = None, + worker_attributes: dict | None = None, + segment_func: SegmentFunc | None = None, + **kwargs, + ): + """Constructor for WorkerMapper. + + Parameters + ---------- + num_workers : int + The number of worker processes to spawn. + + worker_type : callable, optional + Callable that generates an object implementing the Worker + interface, typically a type from a Worker class. + + worker_attributes : dictionary + A dictionary of values that are passed to the worker + constructor as key-word arguments. + + segment_func : callable, optional + Set a default segment_func. Typically set at runtime. + + """ + + super().__init__( + num_workers=num_workers, + segment_func=segment_func, + **kwargs, + ) + + # since the workers will be their own process classes we + # handle this data - This is a subclass of process with an overriden `__init__` - constructor that will automatically generate the Process. + # attributes that will be passed to the worker constructors + if worker_attributes is not None: + self._worker_attributes = worker_attributes + else: + self._worker_attributes = {} - When this class is constructed a new process will be formed. + # choose the type of the worker + if worker_type is None: + self._worker_type = Worker + warn("worker_type not given using the default base class", stacklevel=1) + logger.warning("worker_type not given using the default base class") + else: + self._worker_type = worker_type - """ + @property + def worker_type(self): + """The callable that generates a worker object. - NAME_TEMPLATE = "Worker-{}" - """A string formatting template to identify worker processes in - logs. The field will be filled with the worker index.""" + Typically this is just the type from the class definition of + the Worker where the constructor is called. - def __init__( + """ + return self._worker_type + + def init( self, - worker_idx, - task_queue, - result_queue, - exception_queue, - interrupt_connection, - mapper_attributes=None, - log_level="INFO", + segment_func: SegmentFunc | None = None, + num_workers: int | None = None, **kwargs, ): - """Constructor for the Worker class. + """Runtime initialization and setting of function to map over walkers. Parameters ---------- - worker_idx : int - The index of the worker. Should be unique. - - task_queue : multiprocessing.JoinableQueue - The shared task queue the worker will watch for new tasks to complete. - - result_queue : multiprocessing.Queue - The shared queue that completed task results will be placed on. - - interrupt_connection : multiprocessing.Connection - One end of a pipe to listen for messages specific to this worker. - - mapper_attributes : None or dict - A dictionary of the attributes of the mapper for reference in workers. + num_workers : int + The number of worker processes to spawn - kwargs : - The worker specific attributes + segment_func : callable implementing the Runner.run_segment interface """ - # call the Process constructor - mp.Process.__init__(self, name=self.NAME_TEMPLATE.format(worker_idx)) - - self._exception_queue = exception_queue - self._exception = None - self._traceback = None - - # the queue that will trigger a shutdown in the event of failure - self._irq_channel = interrupt_connection + super().init( + num_workers=num_workers, + segment_func=segment_func, + **kwargs, + ) - # also register the SIGTERM signal handler for graceful - # shutdown with reporting to mapper - signal.signal(signal.SIGTERM, self._sigterm_shutdown) + if self.num_workers is None: + raise ValueError("num_workers is not set") - self._worker_idx = worker_idx + manager = self._mp_ctx.Manager() - self._mapper_attributes = mapper_attributes + # Establish communication queues - # set all the kwargs into an attributes dictionary - self._attributes = kwargs + # A queue for errors + self._exception_queue = manager.Queue() - # the queues for work to be done and work done - self._task_queue = task_queue - self._result_queue = result_queue + # queue for the tasks we know the batch size so we don't need + # a JoinableQueue + self._task_queue = manager.Queue() - logger.debug("{} process created".format(self.name)) + # results queue + self._result_queue = manager.Queue() - @property - def worker_idx(self): - """Dictionary of attributes of the worker.""" - return self._worker_idx + # use pipes for communication channels between this parent + # process and the children for sending specific interrupts + # such as the signal to kill them. Note that the clean way to + # end the process is to send poison pills on the task queue, + # this is for other stuff. IRQ is a common abbreviation for + # interrupts + self._irq_parent_conns = [] - @property - def attributes(self): - """Dictionary of attributes of the worker.""" - return self._attributes + # Start workers, giving them all the queues + self._workers = [] + for i in range(self.num_workers): + # make a pipe to communicate with this worker for the int + parent_conn, child_conn = self._mp_ctx.Pipe() + self._irq_parent_conns.append(parent_conn) - @property - def mapper_attributes(self): - """Dictionary of attributes of the worker.""" - return self._mapper_attributes + # create the worker giving it all of the communication + # channels + worker = self.worker_type( + i, + self._task_queue, + self._result_queue, + self._exception_queue, + child_conn, + mapper_attributes=self._attributes, + **self._worker_attributes, + ) + self._workers.append(worker) - def run(self): - logger.debug("{}: starting to run".format(self.name)) + # start the worker processes + for worker in self._workers: + worker.start() - # try to run the worker and it's task, except either class of - # error that can come from it either from the worker - # (WorkerException) or the task (TaskException) and communicate it - # back to the main process + logger.info( + "Worker process started as name: {}; PID: {}".format( + worker.name, worker.pid + ) + ) - # if we get an exception there is some cleanup logic - run_exception = None + # now that we have started the processes register the handler + # for SIGTERM signals that will clean up our children cleanly + signal.signal(signal.SIGTERM, self._sigterm_shutdown) - try: - # run the worker, which will retrieve its task from the - # queue attempt to run the task, and if it succeeds will - # put the results on the result queue, if the task fails - # it will catch it and wrap it as a task exception - self._run_worker() + def _sigterm_shutdown(self, signum, frame): + logger.critical("Received external SIGTERM, forcing shutdown.") - except TaskException as task_exception: - logger.error("{}: TaskException caught".format(self.name)) + self.force_shutdown() - run_exception = task_exception + def force_shutdown(self, **kwargs): + logger.critical("Forcing shutdown") - # anything else is considered a WorkerException so take the - # original exception and generate a worker exception from that - except Exception as exception: - logger.debug("{}: WorkerError caught".format(self.name)) + # our primary job is to shut down all of the running processes + # without just shutting down the queues and breaking the pipes - # get the traceback - tb = sys.exc_info()[2] + # to do this we send the kill signals to them on the kill + # channel. - msg = "Exception '{}({})' caught in a worker.".format( - type(exception).__name__, exception - ) - traceback_log_msg = """Traceback: --------------------------------------------------------------------------------- -{} --------------------------------------------------------------------------------- - """.format( - "".join(traceback.format_exception(type(exception), exception, tb)), + for worker_idx, worker in enumerate(self._workers): + logger.critical( + "Sending SIGTERM message on {} to worker {}".format( + self._irq_parent_conns[worker_idx].fileno(), worker_idx + ) ) - logger.error("{}:".format(self.name) + msg + "\n" + traceback_log_msg) + # send a kill message to the worker + self._irq_parent_conns[worker_idx].send(signal.SIGTERM) - # raise a TaskError to distinguish it from the worker - # errors with the metadata about the original exception + logger.critical("All kill messages sent to workers") - worker_exception = WorkerException( - "Error occured during worker execution.", - wrapped_exception=exception, - tb=tb, - ) + # check that all have exited + alive_workers = [worker.is_alive() for worker in self._workers] + worker_acks = {} + worker_exitcodes = {} + premature_exit = False + while any(alive_workers) and not premature_exit: + for worker_idx, worker in enumerate(self._workers): + # ignore already known dead workers + if not alive_workers[worker_idx]: + continue - run_exception = worker_exception + if worker.is_alive(): + # if it is still alive and we have an ack from it + # just terminate. There is a bug in the code and + # is out of our control + if worker_idx in worker_acks: + logger.debug( + "Ack received from {} but has not shut down".format( + worker.name + ) + ) + premature_exit = True - # raise worker_exception - if run_exception is not None: - logger.debug("{}: Putting exception on exception queue".format(self.name)) + # otherwise we need to try and receive the ack + elif self._irq_parent_conns[worker_idx].poll(1): + # receive the acknowledgement + ack = self._irq_parent_conns[worker_idx].recv() - # then put the exception and the traceback onto the queue - # so we can communicate back to the parent process - try: - self._exception_queue.put((self.name, self.pid, run_exception)) - except BrokenPipeError as exc: - logger.error( - "Pipe is broken indicating the root process has already exited:\n{}".format( - exc - ) - ) + logger.debug( + "Received {} acknowledgement from {}".format( + ack, worker.name + ) + ) - # TODO: not sure if this is good or not - # then reraise the exception so it can be caught - # raise run_exception + # make sure the ack is affirmative + if ack is True: + worker_acks[worker_idx] = ack - def _sigterm_shutdown(self, signum, frame): - logger.debug("Received external SIGTERM kill command.") + # if it is an exeption wrap it as a worker + # error and use the os to kill the process + elif issubclass(type(ack), Exception): + # wrap it as a worker exception + exception = WorkerException(wrapped_exception=ack) + worker_acks[worker_idx] = exception - logger.debug("Alerting mapper that this will be honored.") + logger.critical( + "{} not responding, terminating with SIGTERM".format( + worker.name + ) + ) - # send an error to the mapper that the worker has been killed - self._irq_channel.send( - WorkerKilledError( - "{} (pid: {}) killed by external SIGTERM signal".format( - self.name, self.pid + worker.terminate() + + else: + alive_workers[worker_idx] = False + worker_exitcodes[worker_idx] = worker.exitcode + + if any(alive_workers): + logger.critical( + "Terminating main process with running workers {}".format( + ",".join([ + str(worker_idx) + for worker_idx in range(len(self._workers)) + if alive_workers[worker_idx] + ]) ) ) - ) - logger.debug("Acknowledgment sent") + def cleanup(self, **kwargs): + """Runtime post-simulation tasks. - logger.debug("Shutting down process") + This is run either at the end of a successful simulation or + upon an error in the main process of the simulation manager + call to `run_cycle`. - def _shutdown(self): - logger.debug("Received SIGTERM kill command from mapper") + The Mapper class performs no actions here and all arguments + are ignored. - logger.debug("Acknowledging kill request will be honored") + """ - # report back that we are shutting down with a True - self._irq_channel.send(True) + super().cleanup(**kwargs) - logger.debug("Acknowledgment sent") + if self.num_workers is None: + raise ValueError("num_workers is not set") - logger.debug("Shutting down process") + # send poison pills (Stop signals) to the queues to stop them in a nice way + # and let them finish up + for _ in range(self.num_workers): + self._task_queue.put((None, None)) - def _run_worker(self): - # run the logic associated with communication and liveness of - # the worker process itself, this is not necessarily a fatal - # (critical) error and restarting a worker might resolve the - # problem. This calls the _run_task method though which is - # always critical since the logic in the code cannot be - # disputed + # delete the queues and workers + self._task_queue = None + self._result_queue = None + self._workers = None - # TODO remove when confirmed that this works - # worker_process = mp.current_process() + def map(self, *args, **kwargs): + # docstring in superclass + + map_process = self._mp_ctx.current_process() logger.info( - "{}: Worker process started as name: {}; PID: {}".format( - self.name, self.name, self.pid - ) + "Mapping from process {}; PID {}".format(map_process.name, map_process.pid) ) - while True: - # check to see if there is any signals in the interrupt channel - if self._irq_channel.poll(): - # get the message - message = self._irq_channel.recv() - - logger.debug( - "{}: Received message from mapper on filehandle {}: {}".format( - self.name, self._irq_channel.fileno(), message - ) - ) - - # handle the message + # make tuples for the arguments to each function call + task_args = zip(*args, strict=False) + kwargs = {key: list(kwarg) for key, kwarg in kwargs.items()} - # a SIGTERM is a signal to kill the process - # unconditionally - if message is signal.SIGTERM: - self._shutdown() + num_tasks = len(args[0]) + # Enqueue the jobs + for task_idx, task_arg in enumerate(task_args): + task_kwargs = {key: value[task_idx] for key, value in kwargs.items()} - # break from the event (while) loop and shut down - break + # a task will be the actual task and its task idx so we can + # sort them later + self._task_queue.put((task_idx, self._make_task(*task_arg, **task_kwargs))) - # anything is not recognized and we will continue and - # report back that we don't recognize the message with - # a ValueError object - else: - logger.error( - "{}: Message not recognized, continuing operations and" - " sending error to mapper".format(self.name) - ) - self._irq_channel.send( - ValueError( - "Message: {} not recognized continuing operations".format( - message - ) - ) - ) + logger.info("Waiting for tasks to be run") - # get the next task + # poll the exception and result queues for results + n_results_left = num_tasks + results = [] + while n_results_left > 0: + # first check if any errors came back but don't wait, + # since the methods for querying whether it is empty or + # not are not reliable we just try and if we don't get + # anything we will come back around try: - task_idx, next_task = self._task_queue.get(block=False, timeout=None) - - logger.debug("{}: Got task {}".format(self.name, task_idx)) - + proc_name, pid, exception = self._exception_queue.get_nowait() except pyq.Empty: - task_idx = None - next_task = Ellipsis + pass - # # check for the poison pill which is the signal to stop - if next_task is None: - logger.info( - "{}: received {} {}: FINISHED".format( - self.name, task_idx, next_task - ) + else: + logger.error( + "Exception occured in process {}; pid {}.".format(proc_name, pid) ) - # TODO remove since we aren't using joinble queue anymore - # mark the poison pill task as done - # self.task_queue.task_done() - - # and exit the loop - break - - # only execute this if a task was actually receieved from - # the queue; an Ellipsis indicates continue the loop - elif next_task is not Ellipsis: - logger.info( - "{}; task_idx : {}; args : {} ".format( - self.name, task_idx, next_task.args + # we can handle Task and Worker exceptions differently + if type(exception) == TaskException: + logger.critical( + "Exception encountered in a task which is unrecoverable." + "You will need to reconfigure your components in a stable manner." ) - ) - - # run the task - start = time.time() - answer = self._run_task(next_task) + self.force_shutdown() - end = time.time() - task_time = end - start + logger.critical("Shutdown complete.") + raise exception - logger.info( - "{}: task_idx : {}; COMPLETED in {} s".format( - self.name, task_idx, task_time + elif type(exception) == WorkerException: + # we make just an error message to say that errors + # in the worker may be due to the network or + # something and could recover + logger.error( + "Exception encountered in the work mapper worker process." + "Recovery possible, see further messages." ) - ) - - # put the results into the results queue with it's task - # index so we can sort them later - self._result_queue.put((task_idx, self.worker_idx, task_time, answer)) - - def run_task(self, task): - """Actually executes the task. - - This default runner simply executes the task thunk. - - This can be customized by subclasses in order to allow for - injection of worker specific data. - Parameters - ---------- - task : Task object - The partially evaluated task; function plus arguments + # However, the current implementation doesn't + # support retries or whatever so we issue a + # critical log informing that it has been elevated + # to critical and will force shutdown + logger.critical( + "Worker error mode resiliency not supported at this time." + "Performing force shutdown and simulation ending." + ) - Returns - ------- - task_result - Results of running the task. + self.force_shutdown() - """ + logger.critical("Shutdown complete.") + raise exception - return task() + else: + logger.critical("Unknown exception encountered.") - def _run_task(self, task): - """Runs the given task and returns the results. + self.force_shutdown() - This manages handling exceptions and tracebacks from the - actual `run_task` function which is intended to be specialized - by different workers to inject worker specific arguments to - tasks. Such as node and device identification. + logger.critical("Shutdown complete.") - Parameters - ---------- - task : Task object - The partially evaluated task; function plus arguments + raise exception - Returns - ------- - task_result - Results of running the task. + # attempt to get something off of the results queue + try: + result = self._result_queue.get_nowait() + except pyq.Empty: + pass - """ + # if we get something handle it + else: + logger.info("Retrieved result: {}".format(result)) + results.append(result) - logger.info("Running task") - try: - return self.run_task(task) + # reduce the counter so we know when we are done + n_results_left -= 1 - except Exception as task_exception: - # get the traceback for the exception - tb = sys.exc_info()[2] + # sort the results according to their task_idx + results.sort() - msg = "Exception '{}({})' caught in a task.".format( - type(task_exception).__name__, task_exception - ) - traceback_log_msg = """Traceback: --------------------------------------------------------------------------------- -{} --------------------------------------------------------------------------------- - """.format( - "".join( - traceback.format_exception(type(task_exception), task_exception, tb) - ), - ) + # save the task run times, so they can be accessed if desired, + # after clearing the task times from the last mapping - logger.critical(msg + "\n" + traceback_log_msg) + # DEBUG: removing this because it should be set on init() + # self._worker_segment_times = {i : [] for i in range(self.num_workers)} - # raise a TaskException to distinguish it from the worker - # errors with the metadata about the original exception + for task_idx, worker_idx, task_time, result in results: + self._worker_segment_times[worker_idx].append(task_time) - raise TaskException( - "Error occured during task execution, recovery not possible.", - wrapped_exception=task_exception, - tb=tb, - ) + # then just return the values of the function + return [result for task_idx, worker_idx, task_time, result in results] diff --git a/src/wepy/work_mapper/task_mapper.py b/src/wepy/work_mapper/task_mapper.py index faeb7c5c..31395e92 100644 --- a/src/wepy/work_mapper/task_mapper.py +++ b/src/wepy/work_mapper/task_mapper.py @@ -1,7 +1,6 @@ # Standard Library import logging -logger = logging.getLogger(__name__) # Standard Library import multiprocessing as mp import pickle @@ -15,11 +14,14 @@ # First Party Library from wepy.work_mapper.mapper import ( ABCWorkerMapper, + SegmentFunc, Task, TaskException, WrapperException, ) +logger = logging.getLogger(__name__) + class TaskProcessException(WrapperException): pass @@ -29,455 +31,170 @@ class TaskProcessKilledError(ChildProcessError): pass -class TaskMapper(ABCWorkerMapper): - """Process-per-task mapper. - - This method of work mapper starts new processes for each runner - segment task that needs to be run. This allows cheap copying of - shared state using the operating system primitives. On linux this - would be either 'fork' (default) or 'spawn'. Fork is cheap but - doesn't initialize certain process namespace things, whereas spawn - is much more expensive but properly cleans things up. Fork should - be sufficient in most cases, however spawn may be needed when you - have some special contexts in the parent process. This is the case - with starting CUDA contexts in the main parent process and then - forking new processes from it. We suggest using fork and avoiding - making these kinds of contexts in the main process. - - This method avoids using shared memory or sending objects through - interprocess communication (that has a serialization and - deserialization cost associated with them) by using OS copying - mechanism. However, a new process will be created each cycle for - each walker in the simulation. So if you want a large number of - walkers you may experience a large overhead. If your walker states - are very small or a very fast serializer is available you may also - not benefit from full process address space copies. Instead the - WorkerMapper may be better suited. - - """ +class WalkerTaskProcess(mp.Process): + NAME_TEMPLATE = "Walker-{}" def __init__( - self, walker_task_type=None, num_workers=None, segment_func=None, **kwargs + self, + walker_idx: int, + mapper_attributes: dict, + func, + task_args, + task_kwargs, + worker_queue, + results_list, + worker_segment_times, + interrupt_connection, + **kwargs, ): - super().__init__(num_workers=num_workers, segment_func=segment_func, **kwargs) + # initialize the process customizing the name + super().__init__( + name=self.NAME_TEMPLATE.format(walker_idx), + **kwargs, + ) - # choose the type of the worker - if walker_task_type is None: - self._walker_task_type = WalkerTaskProcess - warn("walker_task_type not given using the default base class") - logger.warning("walker_task_type not given using the default base class") - else: - self._walker_task_type = walker_task_type + # the idea with this TaskProcess thing is that we pass in all + # the data to the constructor to create a "thunk" (a closure + # that is ready to be run without arguments) and then when run + # is called there will be no arguments to be passed. This + # simplifies the flow of data and underscores that the task is + # the process. - # initialize a list to put results in - self.results = None + # task arguments + self._func = func + self._task_args = task_args + self._task_kwargs = task_kwargs - # this is meant to be a transient variable, will be initialized and deinitialized - self._walker_processes = None + self.walker_idx = walker_idx + self._worker_idx = None + self.mapper_attributes = mapper_attributes - def init(self, **kwargs): - super().init(**kwargs) + # set the managed datastructure proxies as an attribute so we + self._worker_queue = worker_queue + self._results_list = results_list + self._worker_segment_times = worker_segment_times + self._irq_channel = interrupt_connection - # now that we have started the processes register the handler - # for SIGTERM signals that will clean up our children cleanly - signal.signal(signal.SIGTERM, self._sigterm_shutdown) + # also register the SIGTERM signal handler for graceful + # shutdown with reporting to mapper + signal.signal(signal.SIGTERM, self._external_sigterm_shutdown) - def _sigterm_shutdown(self, signum, frame): - logger.critical("Received external SIGTERM, forcing shutdown.") + def _external_sigterm_shutdown(self, signum, frame): + logger.debug("Received external SIGTERM kill command.") - self.force_shutdown() + logger.debug("Alerting mapper that this will be honored.") - logger.critical("Shutdown complete.") + # send an error to the mapper that the worker has been killed + self._irq_channel.send( + TaskProcessKilledError( + "{} (pid: {}) killed by external SIGTERM signal".format( + self.name, self.pid + ) + ) + ) - @property - def walker_task_type(self): - """The callable that generates a worker object. + logger.debug("Acknowledgment sent") - Typically this is just the type from the class definition of - the Worker where the constructor is called. + logger.debug("Shutting down process") - """ - return self._walker_task_type + def _shutdown(self): + """The normal shutdown which can be ordered by the work mapper.""" - def force_shutdown(self): - # send sigterm signals to processes to kill them - for walker_idx, walker_process in enumerate(self._walker_processes): - logger.critical( - "Sending SIGTERM message on {} to worker {}".format( - self._irq_parent_conns[walker_idx].fileno(), walker_idx - ) - ) + logger.debug("Received SIGTERM kill command from mapper") - # send a kill message to the worker - self._irq_parent_conns[walker_idx].send(signal.SIGTERM) + logger.debug("Acknowledging kill request will be honored") - logger.critical("All kill messages sent to workers") + # report back that we are shutting down with a True + self._irq_channel.send(True) - # wait for the walkers to finish and handle errors in them - # appropriately - alive_walkers = [walker.is_alive() for walker in self._walker_processes] - walker_exitcodes = {} - premature_exit = False - while any(alive_walkers): - for walker_idx, walker in enumerate(self._walker_processes): - if not alive_walkers[walker_idx]: - continue + logger.debug("Acknowledgment sent") - if walker.is_alive(): - pass + logger.debug("Shutting down process") - # otherwise the walker is done - else: - alive_walkers[walker_idx] = False - walker_exitcodes[walker_idx] = walker.exitcode + @property + def attributes(self, key): + return self._attributes - def map(self, *args, **kwargs): - # run computations in a Manager context - with self._mp_ctx.Manager() as manager: - num_walkers = len(args[0]) + @attributes.getter + def attributes(self, key): + return self._attributes[key] - # to manage access to worker resources we use a queue with - # the index of the worker - worker_queue = manager.Queue() + def _run_task(self, task): + # run the task thunk + logger.info("{}: Running task".format(self.name)) + try: + result = self.run_task(task) + except Exception as task_exception: + # get the traceback for the exception + tb = sys.exc_info()[2] - # put the workers onto the queue - for worker_idx in range(self.num_workers): - worker_queue.put(worker_idx) + msg = "Exception '{}({})' caught in a task.".format( + type(task_exception).__name__, task_exception + ) + traceback_log_msg = """Traceback: +-------------------------------------------------------------------------------- +{} +-------------------------------------------------------------------------------- + """.format( + "".join( + traceback.format_exception(type(task_exception), task_exception, tb) + ), + ) - # initialize segment times for workers to - # fill in - worker_segment_times = manager.dict() + logger.critical("{}: ".format(self.name) + msg + "\n" + traceback_log_msg) - # initialize for the number of workers, since these will be - # the slots to put timing results in - for i in range(self.num_workers): - worker_segment_times[i] = [] + # raise a TaskException to distinguish it from the worker + # errors with the metadata about the original exception - # make a shared list for the walker results - results = manager.list() + raise TaskException( + "Error occured during task execution, recovery not possible.", + wrapped_exception=task_exception, + tb=tb, + ) - # since this will be indexed by walker index initialize the - # length of the array - for walker in range(num_walkers): - results.append(None) + return result - # use pipes for communication channels between this parent - # process and the children for sending specific interrupts - # such as the signal to kill them. Note that the clean way to - # end the process is to send poison pills on the task queue, - # this is for other stuff. IRQ is a common abbreviation for - # interrupts - self._irq_parent_conns = [] + def run_task(self, task): + logger.info("Running an unspecialized task") - # unpack the generator for the kwargs - kwargs = {key: list(kwarg) for key, kwarg in kwargs.items()} + task_kwargs = dict() - # create the task based processes - self._walker_processes = [] - for walker_idx, task_args in enumerate(zip(*args)): - task_kwargs = {key: value[walker_idx] for key, value in kwargs.items()} + if self.mapper_attributes.get("platform") is not None: + task_kwargs["platform"] = self.mapper_attributes.get("platform") - # make the interrupt pipe - parent_conn, child_conn = self._mp_ctx.Pipe() - self._irq_parent_conns.append(parent_conn) + if task_kwargs["platform"] == "CUDA": + device_ids = self.mapper_attributes.get("device_ids") + if device_ids is not None and isinstance(device_ids, list): + if self._worker_idx >= len(device_ids): + raise ValueError( + "Worker index {} is greater than the number of device ids {} given to mapper. Set the 'device_ids' attribute in TaskMapper appropriately.".format( + self._worker_idx, len(device_ids) + ) + ) - # start a process for this walker - walker_process = self.walker_task_type( - walker_idx, - self._attributes, - self._func, - task_args, - task_kwargs, - worker_queue, - results, - worker_segment_times, - child_conn, - ) + task_kwargs["platform_kwargs"] = { + "DeviceIndex": str(device_ids[self._worker_idx]) + } - walker_process.start() + else: + raise ValueError( + "CUDA platform selected but no device ids given to mapper. Set the 'device_ids' attribute in TaskMapper." + ) - self._walker_processes.append(walker_process) + return task( + **task_kwargs, + ) - new_walkers = [None for _ in range(num_walkers)] - results_found = [False for _ in range(num_walkers)] - while not all(results_found): - # go through the results list and handle the values that may be there - for walker_idx, result in enumerate(results): - if results_found[walker_idx]: - continue + def run(self): + logger.debug("{}: starting to run".format(self.name)) - # logger.info("Checking for walker {}".format(walker_idx)) + # try to run the worker and it's task, except either class of + # error that can come from it either from the worker + # (WorkerException) or the task (TaskException) and communicate it + # back to the main process - # first check to see if any of the task processes were - # terminated from the system - if self._irq_parent_conns[walker_idx].poll(): - irq = self._irq_parent_conns[walker_idx].recv() - - if issubclass(type(irq), TaskProcessKilledError): - # just terminate if a worker goes down. We - # could handle this better but it is not implemented now - logger.critical( - "Process {} was killed by sigterm, shutting down.".format( - walker_process[walker_idx].name - ) - ) - - logger.info( - "Recovery is possible here, but is not implemented " - "so we opt to fail fast and let you know a problem exists." - "Please use checkpointing to avoid lost data." - ) - - self.force_shutdown() - logger.critical("Shutdown complete.") - - logger.debug( - "Received {} acknowledgement from {}".format( - ack, worker.name - ) - ) - - # if no interrupts were handled we continue - - # if it is None no response has been made at all - # yet, this is the initialized value - if result is None: - pass - - # walker results are returned serialized as - # pickles, they are packed into a tuple so that we - # can associate them with an explicit marker, if - # we have a tuple then we can handle that - # appropriately - elif type(result) == tuple: - logger.debug("Received a results tuple") - - assert ( - len(result) == 2 - ), "Result tuples should be only be (ID, pickle)" - - result_id, payload = result - - # there was a walker successfully returned - if result_id == "Walker": - logger.debug("Received a serialized results walker") - - # deserialize - logger.debug("deserializing") - new_walker = pickle.loads(payload) - - logger.info("Got result for walker {}".format(walker_idx)) - - new_walkers[walker_idx] = new_walker - results_found[walker_idx] = True - - else: - raise ValueError("Unkown result ID: {}".format(result_id)) - - elif issubclass(type(result), TaskException): - logger.critical( - "Exception encountered in a task which is unrecoverable." - "You will need to reconfigure your components in a stable manner." - ) - - self.force_shutdown() - - logger.critical("Shutdown complete.") - raise result - - elif issubclass(type(result), TaskProcessException): - # we make just an error message to say that errors - # in the worker may be due to the network or - # something and could recover - logger.error( - "Exception encountered in the work mapper task process." - "Recovery possible, see further messages." - ) - - # However, the current implementation doesn't - # support retries or whatever so we issue a - # critical log informing that it has been elevated - # to critical and will force shutdown - logger.critical( - "Task process error mode resiliency not supported at this time." - "Performing force shutdown and simulation ending." - ) - - self.force_shutdown() - - logger.critical("Shutdown complete.") - raise result - - elif issubclass(type(result), Exception): - logger.critical( - "Unknown exception {} encountered.".format(result) - ) - - self.force_shutdown() - - logger.critical("Shutdown complete.") - - raise result - - else: - logger.critical( - "Unknown result value {} encountered.".format(result) - ) - - self.force_shutdown() - - logger.critical("Shutdown complete.") - - # save the managed list of the recorded worker times locally - for key, val in worker_segment_times.items(): - self._worker_segment_times[key] = val - - # wait for the processes to end - # for walker in self._walker_processes: - # walker.join() - # logger.info("Joined {}".format(walker.name)) - - # deinitialize the current walker processes - self._walker_processes = None - - return new_walkers - - -class WalkerTaskProcess(mp.Process): - NAME_TEMPLATE = "Walker-{}" - - def __init__( - self, - walker_idx, - mapper_attributes, - func, - task_args, - task_kwargs, - worker_queue, - results_list, - worker_segment_times, - interrupt_connection, - **kwargs, - ): - # initialize the process customizing the name - mp.Process.__init__(self, name=self.NAME_TEMPLATE.format(walker_idx), **kwargs) - - # the idea with this TaskProcess thing is that we pass in all - # the data to the constructor to create a "thunk" (a closure - # that is ready to be run without arguments) and then when run - # is called there will be no arguments to be passed. This - # simplifies the flow of data and underscores that the task is - # the process. - - # task arguments - self._func = func - self._task_args = task_args - self._task_kwargs = task_kwargs - - self.walker_idx = walker_idx - self._worker_idx = None - self.mapper_attributes = mapper_attributes - - # set the managed datastructure proxies as an attribute so we - self._worker_queue = worker_queue - self._results_list = results_list - self._worker_segment_times = worker_segment_times - self._irq_channel = interrupt_connection - - # also register the SIGTERM signal handler for graceful - # shutdown with reporting to mapper - signal.signal(signal.SIGTERM, self._external_sigterm_shutdown) - - def _external_sigterm_shutdown(self, signum, frame): - logger.debug("Received external SIGTERM kill command.") - - logger.debug("Alerting mapper that this will be honored.") - - # send an error to the mapper that the worker has been killed - self._irq_channel.send( - TaskProcessKilledError( - "{} (pid: {}) killed by external SIGTERM signal".format( - self.name, self.pid - ) - ) - ) - - logger.debug("Acknowledgment sent") - - logger.debug("Shutting down process") - - def _shutdown(self): - """The normal shutdown which can be ordered by the work mapper.""" - - logger.debug("Received SIGTERM kill command from mapper") - - logger.debug("Acknowledging kill request will be honored") - - # report back that we are shutting down with a True - self._irq_channel.send(True) - - logger.debug("Acknowledgment sent") - - logger.debug("Shutting down process") - - @property - def attributes(self, key): - return self._attributes - - @attributes.getter - def attributes(self, key): - return self._attributes[key] - - def _run_task(self, task): - # run the task thunk - logger.info("{}: Running task".format(self.name)) - try: - result = self.run_task(task) - except Exception as task_exception: - # get the traceback for the exception - tb = sys.exc_info()[2] - - msg = "Exception '{}({})' caught in a task.".format( - type(task_exception).__name__, task_exception - ) - traceback_log_msg = """Traceback: --------------------------------------------------------------------------------- -{} --------------------------------------------------------------------------------- - """.format( - "".join( - traceback.format_exception(type(task_exception), task_exception, tb) - ), - ) - - logger.critical("{}: ".format(self.name) + msg + "\n" + traceback_log_msg) - - # raise a TaskException to distinguish it from the worker - # errors with the metadata about the original exception - - raise TaskException( - "Error occured during task execution, recovery not possible.", - wrapped_exception=task_exception, - tb=tb, - ) - - return result - - def run_task(self, task): - logger.info("Running an unspecialized task") - - return task() - - def run(self): - logger.debug("{}: starting to run".format(self.name)) - - # try to run the worker and it's task, except either class of - # error that can come from it either from the worker - # (WorkerException) or the task (TaskException) and communicate it - # back to the main process - - # if we get an exception there is some cleanup logic - run_exception = None + # if we get an exception there is some cleanup logic + run_exception = None try: # run the worker, which will retrieve its task from the @@ -664,3 +381,333 @@ def _run_walker(self): self._worker_segment_times[worker_idx] = seg_times logger.info("{}: Exiting normally having completed the task".format(self.name)) + + +class TaskMapper(ABCWorkerMapper): + """Process-per-task mapper. + + This method of work mapper starts new processes for each runner + segment task that needs to be run. This allows cheap copying of + shared state using the operating system primitives. On linux this + would be either 'fork' (default) or 'spawn'. Fork is cheap but + doesn't initialize certain process namespace things, whereas spawn + is much more expensive but properly cleans things up. Fork should + be sufficient in most cases, however spawn may be needed when you + have some special contexts in the parent process. This is the case + with starting CUDA contexts in the main parent process and then + forking new processes from it. We suggest using fork and avoiding + making these kinds of contexts in the main process. + + This method avoids using shared memory or sending objects through + interprocess communication (that has a serialization and + deserialization cost associated with them) by using OS copying + mechanism. However, a new process will be created each cycle for + each walker in the simulation. So if you want a large number of + walkers you may experience a large overhead. If your walker states + are very small or a very fast serializer is available you may also + not benefit from full process address space copies. Instead the + WorkerMapper may be better suited. + + """ + + def __init__( + self, + num_workers: int | None = None, + walker_task_type: type[WalkerTaskProcess] | None = None, + segment_func: SegmentFunc | None = None, + **kwargs, + ): + super().__init__(num_workers=num_workers, segment_func=segment_func, **kwargs) + + # choose the type of the worker + if walker_task_type is None: + self._walker_task_type = WalkerTaskProcess + warn("walker_task_type not given using the default base class") + logger.warning("walker_task_type not given using the default base class") + else: + self._walker_task_type = walker_task_type + + # initialize a list to put results in + self.results = None + + # this is meant to be a transient variable, will be initialized and deinitialized + self._walker_processes = None + + def init( + self, + segment_func: SegmentFunc | None = None, + num_workers: int | None = None, + **kwargs, + ): + super().init( + segment_func=segment_func, + num_workers=num_workers, + **kwargs, + ) + + # now that we have started the processes register the handler + # for SIGTERM signals that will clean up our children cleanly + signal.signal(signal.SIGTERM, self._sigterm_shutdown) + + def _sigterm_shutdown(self, signum, frame): + logger.critical("Received external SIGTERM, forcing shutdown.") + + self.force_shutdown() + + logger.critical("Shutdown complete.") + + @property + def walker_task_type(self): + """The callable that generates a worker object. + + Typically this is just the type from the class definition of + the Worker where the constructor is called. + + """ + return self._walker_task_type + + def force_shutdown(self): + # send sigterm signals to processes to kill them + for walker_idx, walker_process in enumerate(self._walker_processes): + logger.critical( + "Sending SIGTERM message on {} to worker {}".format( + self._irq_parent_conns[walker_idx].fileno(), walker_idx + ) + ) + + # send a kill message to the worker + self._irq_parent_conns[walker_idx].send(signal.SIGTERM) + + logger.critical("All kill messages sent to workers") + + # wait for the walkers to finish and handle errors in them + # appropriately + alive_walkers = [walker.is_alive() for walker in self._walker_processes] + walker_exitcodes = {} + premature_exit = False + while any(alive_walkers): + for walker_idx, walker in enumerate(self._walker_processes): + if not alive_walkers[walker_idx]: + continue + + if walker.is_alive(): + pass + + # otherwise the walker is done + else: + alive_walkers[walker_idx] = False + walker_exitcodes[walker_idx] = walker.exitcode + + def map(self, *args, **kwargs): + if self.num_workers is None: + raise ValueError("num_workers is not set") + + # run computations in a Manager context + with self._mp_ctx.Manager() as manager: + num_walkers = len(args[0]) + + # to manage access to worker resources we use a queue with + # the index of the worker + worker_queue = manager.Queue() + + # put the workers onto the queue + for worker_idx in range(self.num_workers): + worker_queue.put(worker_idx) + + # initialize segment times for workers to + # fill in + worker_segment_times = manager.dict() + + # initialize for the number of workers, since these will be + # the slots to put timing results in + for i in range(self.num_workers): + worker_segment_times[i] = [] + + # make a shared list for the walker results + results = manager.list() + + # since this will be indexed by walker index initialize the + # length of the array + for _ in range(num_walkers): + results.append(None) + + # use pipes for communication channels between this parent + # process and the children for sending specific interrupts + # such as the signal to kill them. Note that the clean way to + # end the process is to send poison pills on the task queue, + # this is for other stuff. IRQ is a common abbreviation for + # interrupts + self._irq_parent_conns = [] + + # unpack the generator for the kwargs + kwargs = {key: list(kwarg) for key, kwarg in kwargs.items()} + + # create the task based processes + self._walker_processes = [] + for walker_idx, task_args in enumerate(zip(*args)): + task_kwargs = {key: value[walker_idx] for key, value in kwargs.items()} + + # make the interrupt pipe + parent_conn, child_conn = self._mp_ctx.Pipe() + self._irq_parent_conns.append(parent_conn) + + # start a process for this walker + walker_process = self.walker_task_type( + walker_idx, + self._attributes, + self._func, + task_args, + task_kwargs, + worker_queue, + results, + worker_segment_times, + child_conn, + ) + + walker_process.start() + + self._walker_processes.append(walker_process) + + new_walkers = [None for _ in range(num_walkers)] + results_found = [False for _ in range(num_walkers)] + while not all(results_found): + # go through the results list and handle the values that may be there + for walker_idx, result in enumerate(results): + if results_found[walker_idx]: + continue + + # logger.info("Checking for walker {}".format(walker_idx)) + + # first check to see if any of the task processes were + # terminated from the system + if self._irq_parent_conns[walker_idx].poll(): + irq = self._irq_parent_conns[walker_idx].recv() + + if issubclass(type(irq), TaskProcessKilledError): + # just terminate if a worker goes down. We + # could handle this better but it is not implemented now + logger.critical( + "Process {} was killed by sigterm, shutting down.".format( + walker_process[walker_idx].name + ) + ) + + logger.info( + "Recovery is possible here, but is not implemented " + "so we opt to fail fast and let you know a problem exists." + "Please use checkpointing to avoid lost data." + ) + + self.force_shutdown() + logger.critical("Shutdown complete.") + + logger.debug( + "Received {} acknowledgement from {}".format( + ack, worker.name + ) + ) + + # if no interrupts were handled we continue + + # if it is None no response has been made at all + # yet, this is the initialized value + if result is None: + pass + + # walker results are returned serialized as + # pickles, they are packed into a tuple so that we + # can associate them with an explicit marker, if + # we have a tuple then we can handle that + # appropriately + elif type(result) == tuple: + logger.debug("Received a results tuple") + + assert len(result) == 2, ( + "Result tuples should be only be (ID, pickle)" + ) + + result_id, payload = result + + # there was a walker successfully returned + if result_id == "Walker": + logger.debug("Received a serialized results walker") + + # deserialize + logger.debug("deserializing") + new_walker = pickle.loads(payload) + + logger.info("Got result for walker {}".format(walker_idx)) + + new_walkers[walker_idx] = new_walker + results_found[walker_idx] = True + + else: + raise ValueError("Unkown result ID: {}".format(result_id)) + + elif issubclass(type(result), TaskException): + logger.critical( + "Exception encountered in a task which is unrecoverable." + "You will need to reconfigure your components in a stable manner." + ) + + self.force_shutdown() + + logger.critical("Shutdown complete.") + raise result + + elif issubclass(type(result), TaskProcessException): + # we make just an error message to say that errors + # in the worker may be due to the network or + # something and could recover + logger.error( + "Exception encountered in the work mapper task process." + "Recovery possible, see further messages." + ) + + # However, the current implementation doesn't + # support retries or whatever so we issue a + # critical log informing that it has been elevated + # to critical and will force shutdown + logger.critical( + "Task process error mode resiliency not supported at this time." + "Performing force shutdown and simulation ending." + ) + + self.force_shutdown() + + logger.critical("Shutdown complete.") + raise result + + elif issubclass(type(result), Exception): + logger.critical( + "Unknown exception {} encountered.".format(result) + ) + + self.force_shutdown() + + logger.critical("Shutdown complete.") + + raise result + + else: + logger.critical( + "Unknown result value {} encountered.".format(result) + ) + + self.force_shutdown() + + logger.critical("Shutdown complete.") + + # save the managed list of the recorded worker times locally + for key, val in worker_segment_times.items(): + self._worker_segment_times[key] = val + + # wait for the processes to end + # for walker in self._walker_processes: + # walker.join() + # logger.info("Joined {}".format(walker.name)) + + # deinitialize the current walker processes + self._walker_processes = None + + return new_walkers diff --git a/src/wepy_test_drive.py b/src/wepy_test_drive.py index 50578dcd..67d37ec1 100644 --- a/src/wepy_test_drive.py +++ b/src/wepy_test_drive.py @@ -5,7 +5,7 @@ # Third Party Library import click -import simtk.unit as unit +import openmm.unit as unit from multiprocessing_logging import install_mp_handler # First Party Library