source/sphenixprod/argparsing.py

0001 import argparse
0002
0003 def parse_and_set_loglevel(parser) -> argparse.Namespace:
0004     args = parser.parse_args()
0005     if args.verbose == 1:
0006         args.loglevel = 'INFO'
0007     if args.debug or args.verbose == 2:
0008         args.loglevel = 'DEBUG'
0009     if args.chatty or args.verbose == 3:
0010         args.loglevel = 'CHATTY'
0011
0012     return args
0013
0014 def _base_arguments(parser):
0015     """Add common arguments to the parser."""
0016     # General arguments
0017     parser.add_argument('--dryrun', '--no-submit', '-n',
0018                         help="Job will not be submitted, DBs not updated. Just print things", dest="dryrun", action="store_true")
0019     parser.add_argument('--test-mode', dest="test_mode", default=False,
0020                         help="Sets testing mode, which will mangle DST names and directory paths.", action="store_true")
0021     parser.add_argument('--profile', help="Enable profiling", action="store_true")
0022
0023     # sPHENIX files have specific names and locations. Override for testing or special purposes.
0024     parser.add_argument('--mangle-dirpath', dest='mangle_dirpath',
0025                         help="Inserts string after sphnxpro/ (or tmp/) in the directory structure", default=None,
0026                         type=int)
0027
0028     vgroup = parser.add_argument_group('Logging level')
0029     exclusive_vgroup = vgroup.add_mutually_exclusive_group()
0030     exclusive_vgroup.add_argument('-v', '--verbose', help="Prints more information per repetition", action='count', default=0)
0031     exclusive_vgroup.add_argument('-d', '--debug', help="Prints even more information", action="store_true")
0032     exclusive_vgroup.add_argument('-c', '--chatty', help="Prints the most information", action="store_true")
0033     exclusive_vgroup.add_argument('--loglevel', dest='loglevel', default='INFO',
0034                                   help="Specific logging level (CHATTY, DEBUG, INFO, WARN, ERROR, CRITICAL)")
0035
0036     # Input-specific
0037     rgroup = parser.add_argument_group('Run selection')
0038     exclusive_rgroup = rgroup.add_mutually_exclusive_group()
0039     exclusive_rgroup.add_argument('--runs', nargs='*',
0040                                   help="One argument for a specific run.  Two arguments an inclusive range.  Three or more, a list",
0041                                   default=None)
0042     exclusive_rgroup.add_argument('--runlist',
0043                                   help="Flat text file containing list of runs to process, separated by whitespace / newlines.",
0044                                   default=None)
0045     parser.add_argument('--physics-mode', '--experiment-mode', dest="physicsmode",
0046                         help="Specifies the experiment mode (cosmics, commissioning, physics) for direct lookup of input files.",
0047                         default=None)
0048
0049     parser.add_argument('--submitdir', dest='submitdir', default='./tosubmit', help="Directory for condor submission files")
0050     parser.add_argument('--sublogdir', dest='sublogdir', default=None,
0051                         help="Directory for submission script logging (defaults under /tmp)")
0052
0053     return parser
0054
0055 # ============================================================================================
0056 def submission_args():
0057     """Handle command line tedium for submitting jobs."""
0058     parser = argparse.ArgumentParser(prog='create_submission.py',
0059                                      description='"Production script to submit jobs to the batch system for sPHENIX."')
0060     parser = _base_arguments(parser)
0061
0062     # General arguments
0063     parser.add_argument("--force", "-f", dest="force", help="Override existing output in file and prod db.",
0064                         action="store_true")
0065     parser.add_argument("--force-delete", "-fd", dest="force_delete", help="Set --force and delete existing files that are reproduced.",
0066                         action="store_true")
0067
0068     parser.add_argument('--print-query', dest='printquery', help="Print the query after parameter substitution and exit",
0069                         action="store_true")
0070     parser.add_argument('--andgo', dest='andgo', help="Submit condor jobs at the end", action="store_true")
0071
0072     # Job description arguments
0073     parser.add_argument('--config', dest='config', required=True,
0074                         help="Name of the YAML file containing production rules.",
0075                         default="DST_STREAMING_run3auau_new_2024p012.yaml")
0076     parser.add_argument('--rulename', dest='rulename', required=True, help="Name of submission rule",
0077                         default="DST_EVENT")
0078     parser.add_argument('-N', '--nevents', default=0, dest='nevents', help='Number of events to process.  0=all.',
0079                         type=int)
0080
0081     # Copy additional file to the job work directory
0082     parser.add_argument('--append-to-rsync', dest='append2rsync', default=None,
0083                         help="Appends the argument to the list of rsync files to copy to the worker node")
0084
0085     # Input file selection arguments
0086     parser.add_argument('--onlyseg0', help='Combine only segment 0 files.', action=argparse.BooleanOptionalAction)
0087     parser.add_argument('--choose20', help='Randomly choose 20%% of available files for combining only (no effect downstream)',
0088                         action="store_true")
0089     parser.add_argument('--cut-segment', dest='cut_segment', default=1, type=int,
0090                         help="Only submit jobs for segments where segment %% cut_segment == 0.")
0091
0092     # parser.add_argument("--dbinput", default=True, action="store_true",
0093     #                     help="Passes input filelist through the production status db rather than the argument list of the production script.")
0094     # parser.add_argument("--no-dbinput", dest="dbinput", action="store_false", help="Unsets dbinput flag.")
0095
0096     # Queue-related constraints
0097     parser.add_argument('--mem', help="Override memory allocated for a job", default=None)
0098     parser.add_argument('--priority', help="Override condor priority for this job (more is higher)", default=None)
0099     parser.add_argument('--maxjobs', dest="maxjobs", help="Maximum number of jobs to create in one pass", default=0, type=int)
0100     parser.add_argument('--maxqueued', dest="maxqueued", help="Maximum number of jobs to have waiting for submission", default=70000, type=int)
0101     parser.add_argument('--chunk-size', dest="chunk_size", help="Process runs in chunks of this size (0 = process all at once)", default=0, type=int)
0102     parser.add_argument('--docstring', default=None, help="Appends a documentation string to the log entry")
0103
0104     return parse_and_set_loglevel(parser)
0105
0106 # ============================================================================================
0107 def monitor_args():
0108     """Handle command line tedium for monitoring jobs."""
0109     parser = argparse.ArgumentParser(prog='altmonitor.py',
0110                                      description='"Production script to monitor jobs in the batch system for sPHENIX."')
0111     parser = _base_arguments(parser)
0112
0113     # Job description arguments - here, they are optional
0114     parser.add_argument('--config', dest='config', required=False, help="Name of the YAML file containing production rules.")
0115     parser.add_argument('--rulename', dest='rulename', required=False, help="Name of submission rule")
0116
0117     # Can be used to query/manipulate the queue directly
0118     parser.add_argument('--base_batchname', default=None, help="Select a specific condor batch by name.")
0119
0120     parser.add_argument('-p', '--plot', dest='plot', default=True, action='store_true',
0121                         help='Create plots for held jobs')
0122     parser.add_argument('--memory', dest='memory', default=None, type=int,
0123                         help='Exact memory (MB) to request for resubmitted held jobs (default: None)')
0124     parser.add_argument('--scale-memory', dest='memory_scale_factor', default=1.5, type=float,
0125                         help='Factor to scale memory for resubmitted jobs (default: 1.5)')
0126     parser.add_argument('--max-memory', dest='max_memory', default=13000, type=int,
0127                         help='Maximum memory (MB) to request for resubmitted held jobs (default: 12000)')
0128     parser.add_argument('-r', '--resubmit', dest='resubmit', default=False, action='store_true',
0129                         help='Held jobs are killed and resubmitted with adjusted memory requests')
0130     parser.add_argument('-k', '--kill', dest='kill', default=False, action='store_true',
0131                         help='Held jobs above max-mem are serialized, then killed and not resubmitted')
0132
0133     return parse_and_set_loglevel(parser)
0134
0135 # ============================================================================================
0136