casm-calc submit maximum flag

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."
+        ),
+    )