diff --git a/casm/tools/calc/commands/submit.py b/casm/tools/calc/commands/submit.py index fc66e0b..50a0a9a 100644 --- a/casm/tools/calc/commands/submit.py +++ b/casm/tools/calc/commands/submit.py @@ -22,6 +22,8 @@ def run_submit(args): if not specified. - `args.dry_run`: If set, prints configurations that would be submitted and their current calc status, but does not submit. + - `args.max`: int, optional, Maximum number of jobs to submit (or display in + dry-run mode). No limit if not specified. Returns ------- @@ -79,6 +81,7 @@ def run_submit(args): print("-" * 78) n_jobs_not_ready = 0 + jobs_submitted = 0 for record in config_selection: @@ -119,10 +122,15 @@ def run_submit(args): n_jobs_not_ready += 1 msg = "skipping" else: + # --- Check Limit --- + if args.max is not None and jobs_submitted >= args.max: + print(f"Reached dry-run display limit of {args.max}.") + break + jobs_submitted += 1 msg = dry_run_msg jobid = record.calc_jobid - # --- Printstatus --- + # --- Print status --- print(f"{record.name:36}{record.calc_status:12}{jobid:12}{msg:18}") elif record.calc_status != "setup" or record.calc_jobid != "none": n_jobs_not_ready += 1 @@ -134,6 +142,11 @@ def run_submit(args): else: + # --- Check Limit --- + if args.max is not None and jobs_submitted >= args.max: + print(f"Reached submission limit of {args.max}.") + break + # --- Submit job w/ hold --- completed_processes = record.run_subprocess( args=["sbatch", "--hold", "submit.sh"], @@ -170,6 +183,9 @@ def run_submit(args): capture_output=True, text=True, ) + + # --- Update Counter --- + jobs_submitted += 1 else: print("~" * 40) print(f"Error submitting job for {record.name}: ") @@ -279,3 +295,12 @@ def make_submit_subparser(c): action="store_true", help=("If given, list enumeration, calctype, clex, and selection options."), ) + submit.add_argument( + "--max", + type=int, + help=( + "Maximum number of jobs to submit in this execution. " + "In dry-run mode, limits the number of configurations displayed. " + "Jobs are counted only on successful submission." + ), + )