import argparse
from schrodinger.models import parameters
from schrodinger.tasks.hosts import Host
from schrodinger.ui.qt.appframework2 import application
[docs]def start_task(task):
"""
Wrapper around task.start that runs preprocessing, handles any pre-
processing errors or warnings, starts the task, and prints a message if the
task fails during start.
"""
def callback(result):
if result.message:
if result.passed:
print(f'Warning: {result.message}')
else:
print(f'Error: {result.message}')
return result.passed
results = task.runPreprocessing(callback, calling_context=task.CMDLINE)
if all(results):
task.start(skip_preprocessing=True)
[docs]def run_task_from_cmdline(TaskClass, args=None):
"""
Given a TaskClass, create a task and pass it arguments from the command-
line. Arguments take the form of `--foo.bar.x VALUE', where `foo` is a
subparam on the input of the task, `bar` is a subparam on the `foo`,
and `x` is a subparam on `bar`. Any omitted arguments will take the
default value as specified by their corresponding params.
:param args: A list of strings describing the arguments to the task.
If `None`, the args will be parsed from the command-line.
:type args: list[str]
:return: The finished task.
:rtype: TaskClass
"""
application.get_application(create=True, use_qtcore_app=True)
task = build_task_from_args(TaskClass, args)
start_task(task)
task.wait() # OK: commandline
return task
[docs]def build_task_from_args(TaskClass, args):
parser = _build_parser_from_task(TaskClass)
parsed_args = parser.parse_args(args)
if parsed_args.task_json:
task = TaskClass.fromJsonFilename(parsed_args.task_json)
else:
arg_dict = vars(parsed_args)
arg_dict.pop('task_json')
arg_dict = _flat_dict_to_nested_dict(arg_dict)
task = TaskClass()
if arg_dict:
task.input.setValue(**arg_dict)
return task
[docs]def run_jobtask_from_cmdline(JobTaskClass, args=None):
application.get_application(create=True, use_qtcore_app=True)
task = build_jobtask_from_args(JobTaskClass, args)
start_task(task)
return task
[docs]def build_jobtask_from_args(JobTaskClass, args=None):
parser = _build_parser_from_jobtask(JobTaskClass)
parsed_args, unparsed_args = parser.parse_known_args(args)
task = build_task_from_args(JobTaskClass, unparsed_args)
if parsed_args.jobname:
task.name = parsed_args.jobname
if parsed_args.run_as_backend is not None:
task.run_as_backend = parsed_args.run_as_backend
if parsed_args.host:
task.job_config.host_settings.host = Host(parsed_args.host)
else:
task.job_config.host_settings.host = None
if task.job_config.host_settings.num_subjobs is not None and parsed_args.nsubjobs:
task.job_config.host_settings.num_subjobs = parsed_args.nsubjobs
return task
def _build_parser_from_jobtask(JobTaskClass):
task = JobTaskClass()
parser = argparse.ArgumentParser()
parser.add_argument('--host', type=str)
parser.add_argument('--jobname', type=str, default='')
parser.add_argument('--run_as_backend', type=bool, default=None)
if task.job_config.host_settings.num_subjobs is not None:
parser.add_argument('--nsubjobs', type=int, default=0)
return parser
def _build_parser_from_task(TaskClass):
"""
Given a TaskClass, return a parser that takes arguments for all
input parameters. For example, if we have a task::
class MyTask(tasks.AbstractTask):
class Input(parameters.CompoundParam):
coord = Coord()
label = parameters.StringParam()
input = Input()
The returned parser will take the arguments '--coord.x', '--coord.y', and
'--z'.
"""
task = TaskClass()
parser = argparse.ArgumentParser()
# TODO: what to do about non-compound param inputs?
param_names_to_datatypes = _get_param_names_to_datatypes(task.input)
for name, data_type in param_names_to_datatypes.items():
parser.add_argument('--' + name, type=data_type)
parser.add_argument('--task_json', type=str, default='')
return parser
def _get_param_names_to_datatypes(param):
"""
Given a CompoundParam, return a flat dictionary mapping every subparam
to its data type. See `nested_dict_to_flat_dict` for additional
documentation on the structure of the return value.
"""
return _nested_dict_to_flat_dict(param.toDict())
def _flat_dict_to_nested_dict(flat_dict):
"""
Convert a dictionary mapping qualified keys into a nested dictionary.
For example::
{'foo.bar.kelp':10,
'foo.bar.foop':'a',
'foo.car':True}
would be converted into ::
{'foo':
{'bar':
{'kelp':10,
'foop':'a'},
'car':True
}
}
Any value mapped to `None` will be ignored.
"""
nested_dict = {}
for flat_key, v in flat_dict.items():
if v is not None:
d = nested_dict
all_keys = flat_key.split('.')
for key in all_keys[:-1]:
d = d.setdefault(key, dict())
d[all_keys[-1]] = v
return nested_dict
def _nested_dict_to_flat_dict(param_dict):
"""
Convert a dict representation of a param into a single-level dictionary
mapping qualified names to the param value type. For example, a param
dictionary ::
{'atom':
{'coord':
{'x':1,
'y':2
}
}
}
will be converted to ::
{'atom.coord.x':int,
'atom.coord.y':int}
"""
flat_dict = {}
for prefix_key, v in param_dict.items():
if isinstance(v, parameters.ParamDict):
converted_v = _nested_dict_to_flat_dict(v)
for suffix_key in converted_v:
flat_dict[prefix_key + '.' +
suffix_key] = converted_v[suffix_key]
else:
flat_dict[prefix_key] = type(v)
return flat_dict