From 01668b2e5172ee2e75cd92d58e72f984b0afcd8b Mon Sep 17 00:00:00 2001
From: Sukrit Kalra <sukritkalra@gmail.com>
Date: Wed, 27 Dec 2023 22:17:50 -0800
Subject: [PATCH] Enable definition of multiple workload profiles in main.py

---
 main.py | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/main.py b/main.py
index 35acc87d..3a6176bc 100644
--- a/main.py
+++ b/main.py
@@ -77,8 +77,30 @@
 flags.DEFINE_string("log_level", "debug", "Level of logging.")
 flags.DEFINE_string(
     "workload_profile_path",
-    "./profiles/workload/pylot-workload-profile.json",
-    "Path of the description of the Workload to schedule.",
+    None,
+    "Path of the description of the Workload to schedule. The path can be a file, or a "
+    "directory, and the behavior is defined by the workload loader used. Use "
+    "`workload_profile_paths` if specific files are to be loaded.",
+)
+flags.DEFINE_list(
+    "workload_profile_paths",
+    [],
+    "A list of paths of the description of the Workload to schedule.",
+)
+flags.register_multi_flags_validator(
+    ["workload_profile_path", "workload_profile_paths"],
+    lambda flags: not (
+        len(flags["workload_profile_paths"]) > 0
+        and flags["workload_profile_path"] is not None
+    ),
+    message="Only one of workload_profile_path and workload_profile_paths must be set.",
+)
+flags.register_multi_flags_validator(
+    ["workload_profile_path", "workload_profile_paths"],
+    lambda flags: len(flags["workload_profile_paths"]) > 0
+    or flags["workload_profile_path"] is not None,
+    message="At least one of workload_profile_path and workload_profile_paths must be "
+    "set.",
 )
 flags.DEFINE_string(
     "worker_profile_path",
@@ -137,7 +159,7 @@
     "benchmark_num_cpus", 10, "Number of CPUs available for benchmarking."
 )
 
-# Task related flags.
+# AlibabaLoader related flags.
 flags.DEFINE_integer(
     "alibaba_loader_task_cpu_divisor",
     25,
@@ -161,6 +183,8 @@
     False,
     "If true, we use heterogeneous resource types with %difference in runtime.",
 )
+
+# Task related flags.
 flags.DEFINE_integer(
     "max_timestamp",
     None,