Source code for schrodinger.application.desmond.stage.run_method_parallel

import pickle
import sys
from io import BytesIO

from schrodinger.application.desmond import cmj
from schrodinger.utils import sea

# Certain keys are expected not to match, so ignore them
IGNORE_KEYS = [
    "_PREV_STAGE", "_NEXT_STAGE", "_ID", "_is_shown", "_is_packed", "_pre_job",
    "_pre_jobre", "_cap_job", "_rls_job", "_pack_fname", 'NAME',
    '_final_status', '_is_called_prestage', '_start_time', '_param',
    '_job_manager'
]


class InconsistentAttributeError(Exception):
    pass


def check_stage_attributes(title, ref_attr_dict, stage_attr_dict):
    err = ''
    if stage_attr_dict != ref_attr_dict:
        for k in ref_attr_dict.keys():
            if k in IGNORE_KEYS:
                continue
            if k not in stage_attr_dict:
                err += f'missing key: {k}\n'
            else:
                v0 = ref_attr_dict[k]
                v1 = stage_attr_dict[k]
                if k == 'param':
                    check_stage_attributes(f'{title}: param', v0, v1)
                    continue
                if v0 != v1:
                    err += f'different values for key "{k}": {v0} != {v1}\n'
    if err:
        raise InconsistentAttributeError(
            f'Inconsistent stage attributes: {title} {err}')


def main(pkl_fname=None):
    # Read in the input
    with open(pkl_fname, 'rb') as f:
        input_args = pickle.load(f)

    # Set up the multisim environment
    sea.set_macro_dict(input_args['macro_dict'])
    cmj.ENGINE = cmj.Engine.deserialize(fh=BytesIO(input_args['engine']))
    cmj.ENGINE.loglevel = "silent"
    cmj.ENGINE.boot_setup(base_dir=cmj.ENGINE.base_dir)
    # Only use a single process for any code run here
    cmj.ENGINE.maxjob = 1
    cmj.ENGINE.run_set_family(input_args['stage_idx'])
    cmj.GENERAL_LOGLEVEL = input_args['loglevel']

    # Execute the method
    stage = cmj.ENGINE.stage[input_args['stage_idx']]

    # Set up the stage attributes and
    # copy the tags for GCMC.
    for k in input_args['param']:
        if k in input_args['setbyuser_keys']:
            input_args['param'][k].add_tag('setbyuser')
    input_args['param'].add_tag('setbyuser', propagate=False)
    stage.param = input_args['param']
    stage.check_param()

    check_stage_attributes('initial', stage.__dict__,
                           input_args['stage_attributes'])
    method = getattr(stage, input_args['method'])
    results = method(input_args['batch'], *input_args['args'],
                     **input_args['kwargs'])
    check_stage_attributes('final', stage.__dict__,
                           input_args['stage_attributes'])

    # Write out the results pkl
    with open(input_args['out_pkl'], 'wb') as f:
        pickle.dump({'results': results}, f)


if __name__ == '__main__':
    main(pkl_fname=sys.argv[-1])