"""
Common command line arguments
Lightweight framework for defining new command line arguments.
Copyright Schrodinger, LLC. All rights reserved.
"""
import argparse
import os
import sys
import time
import warnings
from typing import Dict
from typing import List
from typing import NamedTuple
from typing import Optional
from typing import Set
from typing import Tuple
from typing import Union
from schrodinger.application.desmond import util
from schrodinger.infra import mm
from schrodinger.utils.cmdline import DEBUG
from schrodinger.utils.cmdline import HOST
from schrodinger.utils.cmdline import LOCAL
from schrodinger.utils.cmdline import NICE
from schrodinger.utils.cmdline import OPLSDIR
from schrodinger.utils.cmdline import RESTART
from schrodinger.utils.cmdline import RETRIES
from schrodinger.utils.cmdline import SAVE
from schrodinger.utils.cmdline import SUBHOST
from schrodinger.utils.cmdline import TMPDIR
from schrodinger.utils.cmdline import WAIT
from schrodinger.utils.cmdline import add_jobcontrol_options
from schrodinger.utils.cmdline import add_standard_options
Destination = Union[None, str, Dict[str, str]]
ERROR_MISSING_INPUT = 'ERROR: Must provide <pv-file> for running a new job.'
[docs]class Option(NamedTuple):
name: Union[str, List]
default: object
help: str
dest: Dict = None
[docs]def define_options(parser: argparse.ArgumentParser, options: List[Option]):
"""
Define the options on a specified parser.
:param parser: Add options to this parser.
:param options: List of options in the format (name, default, help, dest).
"""
for opt in options:
default = opt.default
kwarg = {
"default": default,
"help": opt.help,
}
if isinstance(default, str):
kwarg["metavar"] = "<string>"
elif isinstance(default, bool):
kwarg["action"] = "store_%s" % ("false" if default else "true")
elif isinstance(default, int):
kwarg["metavar"], kwarg["type"] = "<integer>", int
elif isinstance(default, float):
kwarg["metavar"], kwarg["type"] = "<real>", float
elif isinstance(default, list):
if len(default) >= 1:
if (isinstance(default[0], str)):
kwarg["metavar"] = "<string>"
elif (isinstance(default[0], int)):
kwarg["metavar"], kwarg["type"] = "<integer>", int
elif (isinstance(default[0], float)):
kwarg["metavar"], kwarg["type"] = "<real>", float
kwarg["nargs"] = '*'
kwarg["action"] = "append"
dest = opt.dest
if isinstance(dest, str):
kwarg["dest"] = dest
elif isinstance(dest, dict):
kwarg.update(dest)
opt_names = isinstance(opt.name, list) and opt.name or [opt.name]
parser.add_argument(*opt_names, **kwarg)
[docs]def auto_int(string):
if string == "auto":
return string
else:
try:
return int(string)
except TypeError:
msg = f'{string} is not an integer or the word \"auto\"'
raise argparse.ArgumentTypeError(msg)
[docs]def get_common_options() -> List[Option]:
"""
Return list of options common to all fep scripts.
"""
ffld_names = mm.opls_names()
return [
Option(
"-ff",
ffld_names[-1],
f"Specify the forcefield to use. Default: {ffld_names[-1]}.",
{
"dest": "forcefield",
"metavar": "{%s}" % "|".join(ffld_names)
},
),
Option(
"-seed",
2014,
"Specify seed of pseudorandom number generator for initial atom"
" velocities. Default: 2014",
),
Option(
"-ppj",
0,
"Specify number of processors per job. Default: 4.",
{"metavar": "PPJ"},
),
Option("-mps", 0, argparse.SUPPRESS, {"type": auto_int}),
Option(
"-checkpoint",
"",
"Specify the multisim checkpoint file.",
{"metavar": "<multisim-checkpiont-file>"},
),
Option(
"-prepare",
False,
"Do not run job. Only prepare multisim input files.",
),
Option(
"-skip_traj",
False,
argparse.SUPPRESS,
),
# "Do not copy trajectories to host machine on restart (saves bandwidth,
# "but may affect the cleanup stage)."
Option(
"-JOBNAME",
"",
"Specify the job name.",
),
Option(
"-buffer",
5.0,
"Specify a larger buffer size (in Angstroms). Defaults: 5 in"
" complex leg; 5 in solvent leg of protein-residue-mutation FEP;"
" 10 in solubility FEP; 10 in solvent leg of other types of FEP."
" The custom value will be used only if it's greater than the"
" corresponding default values.",
),
Option(
"-maxjob",
0,
"Maximum number of simultaneous subjobs. Default: 0 (unlimited)",
),
Option(
["-lambda-windows", "-lambda_windows"],
12,
"Number of lambda windows for the default protocol. Default: 12",
),
Option(
"-no_concat",
False,
argparse.SUPPRESS,
),
Option(
"-max-walltime",
0,
argparse.SUPPRESS,
# "Specify the maximum number of seconds to run the subjobs "
# "before automatically checkpointing and requeuing them. "
# "The default of 0 means to run all subjobs to completion. ""
),
Option("-ffbuilder", False, "Run the ffbuilder workflow."),
Option(
"-ff-host", "",
"Host for the ffbuilder jobs specified as HOST:MAX_FF_BUILDER_JOBs. "
"This must be set if using -ffbuilder. ")
]
[docs]def suppress_options(options: List[Option], excluded: Set[str]):
"""
Modify the options as specified by `excluded` by replacing the help text
with `argparse.SUPPRESS`, which will effectively hide the specified options
in the command line.
No effects if either `options` or `excluded` is an empty container.
"""
for index, (name, default, _, dest) in enumerate(options):
name = isinstance(name, list) and name[0] or name
if name in excluded:
suppressed_opt = Option(name, default, argparse.SUPPRESS, dest)
options[index] = suppressed_opt
[docs]def get_parser(usage: str, options: List[Option],
add_help: bool = True, add_subhost: bool = True) \
-> argparse.ArgumentParser:
"""
Return a command-line parser with the given options.
:param usage: Usage to display if no arguments given.
:param options: List of options
:param add_help: Whether to add help flag to the parser
:return: Configured command-line parser
"""
parser = argparse.ArgumentParser(usage=usage,
add_help=add_help,
allow_abbrev=False)
define_options(parser, options)
add_jobcontrol_options(parser,
options=(
HOST,
WAIT,
LOCAL,
DEBUG,
TMPDIR,
NICE,
SAVE,
OPLSDIR,
))
standard_options = (SUBHOST, RETRIES, RESTART) if add_subhost else \
(RETRIES, RESTART)
add_standard_options(parser, options=standard_options)
return parser
[docs]def parse_known_options(
usage: str,
options: List[Option],
argv: List[str],
add_subhost: bool = True
) -> Tuple[argparse.Namespace, List[str], argparse.ArgumentParser]:
"""
Parse and return the parsed options.
:param usage: Usage to display if no arguments given.
:param options: List of options in the format (name, default, help, destination).
:param argv: List of input arguments.
:return: (Known parsed options, unknown options)
:raise SystemExit: If no arguments given, show the usage and exit.
"""
parser = get_parser(usage, options, add_subhost=add_subhost)
if len(argv) < 1:
parser.print_help()
sys.exit(0)
opts, other_args = parser.parse_known_args(argv)
return opts, other_args, parser
[docs]def parse_options(usage: str, options: List[Option],
argv: List[str], add_subhost: bool =True) \
-> argparse.Namespace:
"""
Parse and return the parsed options.
:param usage: Usage to display if no arguments given.
:param options: List of options in the format (name, default, help, dest).
:param argv: List of input arguments
:param add_subhost: Add the -SUBHOST option?
:return: Parsed options
:raise SystemExit: If no arguments given or if there is unknown arguments,
show the usage and exit.
"""
check_discontinued_args(argv)
opts, other_args, parser = parse_known_options(usage,
options,
argv,
add_subhost=add_subhost)
for other_arg in other_args:
if other_arg.startswith('-'):
parser.print_help()
print("ERROR: Unrecognized option {}".format(other_arg))
sys.exit(1)
return opts
[docs]def check_discontinued_args(args: List[str]):
"""
Check for the presence of arguments that have been removed and exit if any
are used.
"""
for arg in args:
if arg == "-m":
sys.exit(
"ERROR: The -m option has been removed in favor of using the "
"'-prepare' option and then directly editing the msj before "
"running multisim.")
[docs]def check_jobname(jobname: str) -> Optional[str]:
"""
Check whether the given job name contains problematic characters. It
cannot start with a "-" or contain a "=".
"""
if jobname.startswith("-"):
return f'ERROR: The job name {jobname} cannot start with "-"'
if "=" in jobname:
return f'ERROR: The job name {jobname} cannot contain "="'
return None
[docs]class BaseArgs:
"""
Base class for arguments.
A subclass of this class contains all the command-line arguments. The
jobcontrol-related arguments, which are consumed by the toplevel script,
are recovered.
"""
[docs] def __init__(self, opt: argparse.Namespace):
"""
:param opt: Command line options with corresponding values.
"""
self.msj = None
self.forcefield = opt.forcefield
self.seed = opt.seed
# FIXME: Remove `self.ppr`?
# Was processors per replica, deprecated option
self.ppr = 1
self.ppj_set_by_user = bool(opt.ppj)
self.ppj = opt.ppj or 4
self.mps_factor = 0
self.set_mps_factor(opt.mps)
self.checkpoint = opt.checkpoint
self.prepare = opt.prepare
self.skip_traj = opt.skip_traj
self.buffer = opt.buffer
self.maxjob = opt.maxjob
self.lambda_windows = opt.lambda_windows
self.concat = not opt.no_concat
self.max_walltime = opt.max_walltime
self.ffbuilder = opt.ffbuilder
self.ff_host = opt.ff_host
self.inp_file = opt.inp_file
# Jobcontrol options
self.JOBNAME = opt.JOBNAME
self.RESTART = opt.restart
self.RETRIES = "1" if opt.retries is None else str(opt.retries)
# NOTE: JOBHOST and SCHRODINGER_NODELIST may be ""
# in which case localhost should be used.
self.HOST = os.getenv("JOBHOST") or "localhost"
self.SUBHOST = os.getenv("SCHRODINGER_NODELIST") or "localhost"
self.WAIT = opt.wait
self.LOCAL = opt.local
self.DEBUG = opt.debug
self.TMPDIR = os.getenv("SCHRODINGER_TMPDIR")
# NOTE: -OPLSDIR option will be stripped from the command line by
# toplevel script and it will set OPLS_DIR environment.
self.OPLSDIR = os.getenv("OPLS_DIR")
self.NICE = os.getenv("SCHRODINGER_NICE")
self.SAVE = opt.save
self.set_restart()
self.generate_jobname()
self.validate()
def __str__(self):
s = [f"{k} = {v}" for k, v in vars(self).items()]
return "\n".join(s)
[docs] def validate(self):
"""
Validate the parameters.
:raise SystemExit: For invalid parameters.
"""
try:
ff_int = mm.opls_name_to_version(self.forcefield)
if mm.mmffld_license_ok(ff_int) != mm.TRUE:
sys.exit(
f"ERROR: The {self.forcefield} forcefield requires a license.\n"
f"The {mm.OPLS_NAME_F14} forcefield does not require a license.\n"
f"You can specify this with \"-ff {mm.OPLS_NAME_F14}\" in the command."
)
except IndexError:
pass # CHARMM, AMBER, etc. don't require license check
if self.ffbuilder and not self.ff_host:
sys.exit(
"ERROR: The -ffbuilder option requires the -ff-host HOSTNAME:MAX_FFBUILDER_SUBJOBS argument."
)
err = check_jobname(self.JOBNAME)
if err:
sys.exit(err)
[docs] def copy_parser_attributes(self, opt: argparse.ArgumentParser):
"""
Copy parser options (e.g: time, buffer, ...) from `opt` to `self`.
Subclass needs to call this method in __init__
"""
# FIXME: It would be better if this method is called in base class?
# How?
for attr in dir(opt):
if not attr.startswith('_'):
setattr(self, attr, getattr(opt, attr))
[docs] def check_ppj(self):
"""
Raise a warning if restarting and trying to set ppj.
:raise UserWarning: If ppj set for a restarted job.
"""
if self.ppj_set_by_user and (self.RESTART or self.checkpoint):
warnings.warn(
'WARNING: When restarting, the "-ppj" option is not supported'
' and will be ignored.')
[docs] def set_restart(self):
"""
Set the RESTART flag if only the checkpoint is given.
"""
# FIXME: `self.RESTART` seems determined by the state of `inp_file` and
# `checkpoint` parameters. Is the `-RESTART` option redundant
# then?
if self.inp_file is None and self.checkpoint:
self.RESTART = True
[docs] def generate_jobname(self):
"""
If the JOBNAME was not set and not RESTART'ing,
automatically generate a job name.
"""
if not self.JOBNAME:
self.JOBNAME = os.environ.get("SCHRODINGER_JOBNAME")
if not self.JOBNAME and not self.RESTART:
self.JOBNAME = util.getlogin() + time.strftime("%Y%m%dT%H%M%S")
print(
f"Using an automatically-generated job name: {self.JOBNAME}"
)
[docs] def set_mps_factor(self, val):
"""
Set the mps oversubcription factor. If val is `auto`, the mps factor
will be determined automatically. Otherwise it is set directly, and
should have an `int` value. `0` is treated equivalently to the value
`1`.
"""
if val != "auto":
if not (0 <= val < 9):
raise ValueError("explicit MPS value must be an integer "
"between 0 and 8")
# interpret 0 as mps_factor == 1
if not val:
val = 1
if val == "auto" or val > 1:
warnings.warn(
'WARNING: MPS is enabled - this may cause undefined behavior '
'if the queue being used does not support MPS')
self.mps_factor = val