reintroduce args.jobs, and recalculate timeouts accordingly

Signed-off-by: Jack Luar <[email protected]>

Author: luarss

tools/AutoTuner/src/autotuner/distributed.py

@@ -446,6 +446,13 @@ def parse_arguments():
         default=int(np.floor(cpu_count() / 2)),
         help="CPU Hours",
     )
+    parser.add_argument(
+        "--jobs",
+        type=int,
+        metavar="<int>",
+        default=int(np.floor(cpu_count() / 2)),
+        help="Max number of concurrent jobs.",
+    )
     parser.add_argument(
         "--openroad_threads",
         type=int,
@@ -510,14 +517,17 @@ def parse_arguments():
     if args.timeout is not None:
         args.timeout = round(args.timeout * 3600)
 
-    # Calculate jobs
-    jobs = int(args.cpu_budget / args.resources_per_trial)
-    args.jobs = jobs
-
     # Calculate timeout based on cpu_budget
     if args.cpu_budget is not None:
         args.timeout = round(args.cpu_budget / os.cpu_count() * 3600)
-        template = calculate_expected_numbers(args.timeout, args.samples)
+        args.timeout_per_trial = round(
+            args.cpu_budget / (args.jobs * args.resources_per_trial) * 3600
+        )
+        overall_timeout = min(args.timeout, args.timeout_per_trial)
+        if args.mode == "tune":
+            template = calculate_expected_numbers(overall_timeout, args.samples)
+        else:
+            template = calculate_expected_numbers(overall_timeout, 1)
         print(template)
         if not args.yes:
             ans = input("Are you sure you want to proceed? (y/n): ")

tools/AutoTuner/test/smoke_test_sweep.py

@@ -29,7 +29,8 @@ def setUp(self):
         # limit jobs because ray.get() does not terminate if jobs > number of samples
         core = os.cpu_count()
         self.resources_per_trial = 1
-        self.cpu_budget = 4 if core >= 4 else core
+        self.jobs = 4
+        self.cpu_budget = 1  # test should terminate in 1/(4*1) = 0.25 hours
         self.experiment = f"smoke-test-sweep-{self.platform}"
         self.command = (
             "python3 -m autotuner.distributed"
@@ -38,6 +39,7 @@ def setUp(self):
             f" --experiment {self.experiment}"
             f" --config {self.config}"
             f" --cpu_budget {self.cpu_budget}"
+            f" --jobs {self.jobs}"
             f" --yes"
             f" sweep --resources_per_trial {self.resources_per_trial}"
         )