[Core] Tying wandb CLI Args in Capacity Search (#6)

* wandb CLI args for capacity search

* benchmark config bug fix

* benchmark config bug fix

* make format

* Adding "no-"" to bool CLI args

metron/capacity_search/capacity_search.py

@@ -241,9 +241,10 @@ def is_under_sla(self, qps: float) -> Tuple[bool, float, float, float, float, st
             ),
             qps=qps,
             tbt_deadline=self.args.tbt_slo,
-            wandb_project="Metron",
-            wandb_group="gatech-sysml",
-            wandb_run_name=f"qps_{qps}_model_{self.job_config.model_config.name}_ttftslack_{self.args.ttft_slack_slo}_tbt_{self.args.tbt_slo}_tpot_{self.args.tpot_slo}_ttft_{self.args.ttft_slo}_trace_{self.job_config.request_generator_config.trace_file_name}",
+            wandb_project=self.args.wandb_project,
+            wandb_group=self.args.wandb_group,
+            wandb_run_name=f"qps_{qps}_model_{self.job_config.model_config.name}_engine_{self.job_config.server_config.openai_server_engine}",
+            should_write_metrics=self.args.should_write_metrics_to_wandb,
         )
 
         run_dir = benchmark_config.get_run_dir()
@@ -347,7 +348,7 @@ def search(self):
             f"Best Run ID: {best_run_id}",
         )
 
-        if self.args.wandb_project is not None:
+        if self.args.wandb_project is not None and self.args.enable_wandb_sweep:
             best_run = wandb.Api().run(f"{self.args.wandb_project}/{best_run_id}")
             best_run.tags.append("BEST_CONFIG")
             best_run.update()

metron/capacity_search/config/config.py

@@ -354,6 +354,7 @@ class BenchmarkConfig:
     wandb_group: Optional[str] = None
     wandb_project: Optional[str] = None
     wandb_run_name: Optional[str] = None
+    should_write_metrics: Optional[bool] = True
 
     def to_config_dict(self):
         return {
@@ -366,14 +367,21 @@ def to_config_dict(self):
             "wandb-group": self.wandb_group,
             "wandb-project": self.wandb_project,
             "wandb-run-name": self.wandb_run_name,
+            "should-write-metrics": self.should_write_metrics,
         }
 
     def to_args(self):
         config_dict = self.to_config_dict()
         args = []
         for key, value in config_dict.items():
             if value is not None:
-                args.append(f"--{key} {value}")
+                if isinstance(value, bool):
+                    if value:
+                        args.append(f"--{key}")
+                    else:
+                        args.append(f"--no-{key}")
+                else:
+                    args.append(f"--{key} {value}")
         return " ".join(args)
 
     def get_run_id(self):

metron/capacity_search/main.py

@@ -49,12 +49,25 @@ def get_args():
         "--debug", action="store_true", help="Print debug logs and commands"
     )
     parser.add_argument("--wandb-project", type=str, default=None)
+    parser.add_argument("--wandb-group", type=str, default=None)
+    parser.add_argument(
+        "--should-write-metrics-to-wandb",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
     parser.add_argument("--wandb-sweep-name", type=str, default=None)
     parser.add_argument("--wandb-sweep-id", type=str, default=None)
+    parser.add_argument(
+        "--enable-wandb-sweep",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
 
     args = parser.parse_args()
 
-    if args.wandb_project:
+    if args.wandb_project and args.enable_wandb_sweep:
         assert (
             args.wandb_sweep_name or args.wandb_sweep_id
         ), "wandb-sweep-name/id is required with wandb-project"
@@ -80,7 +93,7 @@ def get_args():
     # store the config and args
     json.dump(config, open(f"{args.output_dir}/config.json", "w"))
 
-    if args.wandb_project and not args.wandb_sweep_id:
+    if args.wandb_project and args.enable_wandb_sweep and not args.wandb_sweep_id:
         config["name"] = args.wandb_sweep_name
         config["method"] = "custom"
 

metron/run_benchmark.py

@@ -530,6 +530,7 @@ def parse_args():
         "--should-use-given-dir",  # added to prevent the creation of a new directories for the capacity search
         type=bool,
         default=True,
+        action=argparse.BooleanOptionalAction,
         help=(
             "Whether to add directly use --output-dir directory or create new directories for the results. (default: %(default)s)"
         ),