only use one command for criteo run

experiments/runner.cli.py

@@ -139,17 +139,25 @@ def microbenchmark_varying_epoch_granularity(ray_session_dir):
 ## ----------------- CRITEO ----------------- ##
 
 
-def _criteo_run(ray_session_dir, impressions_path, conversions_path, logs_dir):
+def criteo_run(ray_session_dir):
     dataset = "criteo"
+    impressions_path = f"{dataset}/{dataset}_query_pool_impressions.csv"
+    conversions_path = f"{dataset}/{dataset}_query_pool_conversions.csv"
+    logs_dir = f"{dataset}/bias_varying_epoch_size"
 
     workload_generation = OmegaConf.load("data/criteo/config.json")
 
+    augment_rate = workload_generation.get("augment_rate")
+
+    epoch_first_batch = [1, 7, 14, 21]
+    epoch_second_batch = [30, 60, 90]
+
     config = {
         "baseline": ["ipa", "cookiemonster_base", "cookiemonster"],
         "dataset_name": f"{dataset}",
         "impressions_path": impressions_path,
         "conversions_path": conversions_path,
-        "num_days_per_epoch": [1, 7, 14, 21],
+        "num_days_per_epoch": epoch_first_batch,
         "num_days_attribution_window": [30],
         "workload_size": [1_000],  # force a high number so that we run on all queries
         "max_scheduling_batch_size_per_query": workload_generation.max_batch_size,
@@ -162,25 +170,17 @@ def _criteo_run(ray_session_dir, impressions_path, conversions_path, logs_dir):
     }
 
     grid_run(**config)
-    config["num_days_per_epoch"] = [30, 60, 90]
+    config["num_days_per_epoch"] = epoch_second_batch
     config["ray_init"] = False
     grid_run(**config)
 
-
-def criteo_varying_epoch_size(ray_session_dir):
-    dataset = "criteo"
-    impressions_path = f"{dataset}/{dataset}_query_pool_impressions.csv"
-    conversions_path = f"{dataset}/{dataset}_query_pool_conversions.csv"
-    logs_dir = f"{dataset}/bias_varying_epoch_size"
-    _criteo_run(ray_session_dir, impressions_path, conversions_path, logs_dir)
-
-
-def criteo_augmented_varying_epoch_size(ray_session_dir):
-    dataset = "criteo"
-    impressions_path = f"{dataset}/{dataset}_query_pool_augmented_impressions.csv"
-    conversions_path = f"{dataset}/{dataset}_query_pool_conversions.csv"
-    logs_dir = f"{dataset}/augmented_bias_varying_epoch_size"
-    _criteo_run(ray_session_dir, impressions_path, conversions_path, logs_dir)
+    if augment_rate:
+        config["impressions_path"] = f"{dataset}/{dataset}_query_pool_augmented_impressions.csv"
+        config["num_days_per_epoch"] = epoch_first_batch
+        config["logs_dir"] = f"{dataset}/augmented_bias_varying_epoch_size"
+        grid_run(**config)
+        config["num_days_per_epoch"] = epoch_second_batch
+        grid_run(**config)
 
 
 ## ----------------- PATCG ----------------- ##