add --yes to avoid prompt.

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

Author: luarss

tools/AutoTuner/src/autotuner/distributed.py

@@ -306,6 +306,9 @@ def parse_arguments():
         default=None,
         help="Time limit (in hours) for the whole Autotuning process.",
     )
+    parser.add_argument(
+        "-y", "--yes", action="store_true", help="Skip confirmation prompt."
+    )
     tune_parser.add_argument(
         "--resume",
         action="store_true",
@@ -514,9 +517,11 @@ def parse_arguments():
         args.timeout = args.cpu_budget / os.cpu_count()
         template = calculate_expected_numbers(args.timeout, args.samples)
         print(template)
-        ans = input("Do you want to continue? [Y/n]")
-        if ans.lower() != "y":
-            sys.exit(0)
+        if not args.yes:
+            ans = input("Are you sure you want to proceed? (y/n): ")
+            if ans.lower() != "y":
+                print("Exiting AutoTuner.")
+                sys.exit(0)
 
     return args
 

tools/AutoTuner/test/smoke_test_sweep.py

@@ -38,6 +38,7 @@ def setUp(self):
             f" --experiment {self.experiment}"
             f" --config {self.config}"
             f" --cpu_budget {self.cpu_budget}"
+            f" --yes"
             f" sweep --resources_per_trial {self.resources_per_trial}"
         )
 

tools/AutoTuner/test/smoke_test_timeout.py

@@ -26,6 +26,7 @@ def setUp(self):
             f" --platform {self.platform}"
             f" --experiment {self.experiment}"
             f" --config {self.config}"
+            f" --yes"
             f" {flag}"
             f" tune --samples 1"
             for flag in timeout_flags

tools/AutoTuner/test/smoke_test_tune.py

@@ -21,7 +21,8 @@ def setUp(self):
             f" --platform {self.platform}"
             f" --experiment {self.experiment}"
             f" --config {self.config}"
-            f" tune --samples 5"
+            f" --yes"
+            f" tune --samples 1"
         )
 
     def test_tune(self):