Add Percentage Search Space to Optuna (#882)

* Added method for calculating total possible configurations

* Added min/max percentage of search space to CLI

* Connected up in optuna RCG

model_analyzer/config/generate/optuna_run_config_generator.py

@@ -141,10 +141,17 @@ def get_configs(self) -> Generator[RunConfig, None, None]:
         yield default_run_config
         self._default_measurement = self._last_measurement
 
-        # TODO: TMA-1884: Need a default + config option for trial number
-        n_trials = 20
+        total_num_of_possible_configs = (
+            self._search_parameters.number_of_total_possible_configurations()
+        )
+        max_configs_to_search = int(
+            total_num_of_possible_configs
+            * self._config.max_percentage_of_search_space
+            / 100
+        )
+
         # TODO: TMA-1885: Need an early exit strategy
-        for _ in range(n_trials):
+        for _ in range(max_configs_to_search):
             trial = self._study.ask()
             trial_objectives = self._create_trial_objectives(trial)
             run_config = self._create_objective_based_run_config(trial_objectives)

model_analyzer/config/generate/search_parameters.py

@@ -67,6 +67,28 @@ def get_range(self, name: str) -> Tuple[Optional[int], Optional[int]]:
     def get_list(self, name: str) -> Optional[List[Any]]:
         return self._search_parameters[name].enumerated_list
 
+    def number_of_total_possible_configurations(self) -> int:
+        total_number_of_configs = 1
+        for parameter in self._search_parameters.values():
+            total_number_of_configs *= self._number_of_configurations_for_parameter(
+                parameter
+            )
+
+        return total_number_of_configs
+
+    def _number_of_configurations_for_parameter(
+        self, parameter: SearchParameter
+    ) -> int:
+        if (
+            parameter.category is ParameterCategory.INTEGER
+            or parameter.category is ParameterCategory.EXPONENTIAL
+        ):
+            number_of_parameter_configs = parameter.max_range - parameter.min_range + 1  # type: ignore
+        else:
+            number_of_parameter_configs = len(parameter.enumerated_list)  # type: ignore
+
+        return number_of_parameter_configs
+
     def _populate_search_parameters(self) -> None:
         if self._parameters:
             self._populate_parameters()

model_analyzer/config/input/config_command_profile.py

@@ -62,6 +62,8 @@
     DEFAULT_OFFLINE_PLOTS,
     DEFAULT_ONLINE_OBJECTIVES,
     DEFAULT_ONLINE_PLOTS,
+    DEFAULT_OPTUNA_MAX_PERCENTAGE_OF_SEARCH_SPACE,
+    DEFAULT_OPTUNA_MIN_PERCENTAGE_OF_SEARCH_SPACE,
     DEFAULT_OUTPUT_MODEL_REPOSITORY,
     DEFAULT_OVERRIDE_OUTPUT_REPOSITORY_FLAG,
     DEFAULT_PERF_ANALYZER_CPU_UTIL,
@@ -916,6 +918,24 @@ def _add_run_search_configs(self):
                 description="Maximum number of steps take during the binary concurrency search.",
             )
         )
+        self._add_config(
+            ConfigField(
+                "min_percentage_of_search_space",
+                flags=["--min_percentage_of_search_space"],
+                field_type=ConfigPrimitive(int),
+                default_value=DEFAULT_OPTUNA_MIN_PERCENTAGE_OF_SEARCH_SPACE,
+                description="Minimum percentage of the search space to profile when using Optuna",
+            )
+        )
+        self._add_config(
+            ConfigField(
+                "max_percentage_of_search_space",
+                flags=["--max_percentage_of_search_space"],
+                field_type=ConfigPrimitive(int),
+                default_value=DEFAULT_OPTUNA_MAX_PERCENTAGE_OF_SEARCH_SPACE,
+                description="Maximum percentage of the search space to profile when using Optuna",
+            )
+        )
         self._add_config(
             ConfigField(
                 "run_config_search_mode",

model_analyzer/config/input/config_defaults.py

@@ -54,6 +54,8 @@
 DEFAULT_RUN_CONFIG_SEARCH_DISABLE = False
 DEFAULT_RUN_CONFIG_SEARCH_MODE = "brute"
 DEFAULT_RUN_CONFIG_PROFILE_MODELS_CONCURRENTLY_ENABLE = False
+DEFAULT_OPTUNA_MIN_PERCENTAGE_OF_SEARCH_SPACE = 5
+DEFAULT_OPTUNA_MAX_PERCENTAGE_OF_SEARCH_SPACE = 10
 DEFAULT_REQUEST_RATE_SEARCH_ENABLE = False
 DEFAULT_TRITON_LAUNCH_MODE = "local"
 DEFAULT_TRITON_DOCKER_IMAGE = "nvcr.io/nvidia/tritonserver:24.05-py3"

tests/test_cli.py

@@ -85,6 +85,8 @@ def get_test_options():
         OptionStruct("int", "profile", "--run-config-search-min-instance-count", None, "2", "1"),
         OptionStruct("int", "profile", "--run-config-search-max-instance-count", None, "10", "5"),
         OptionStruct("int", "profile", "--run-config-search-max-binary-search-steps", None, "10", "5"),
+        OptionStruct("int", "profile", "--min_percentage_of_search_space", None, "10", "5"),
+        OptionStruct("int", "profile", "--max_percentage_of_search_space", None, "5", "10"),
         OptionStruct("float", "profile", "--monitoring-interval", "-i", "10.0", "1.0"),
         OptionStruct("float", "profile", "--perf-analyzer-cpu-util", None, "10.0", str(psutil.cpu_count() * 80.0)),
         OptionStruct("int", "profile", "--num-configs-per-model", None, "10", "3"),

tests/test_search_parameters.py

@@ -68,7 +68,7 @@ def setUp(self):
         )
 
         self.search_parameters._add_search_parameter(
-            name="instance_count",
+            name="instance_group",
             usage=ParameterUsage.MODEL,
             category=ParameterCategory.INTEGER,
             min_range=1,
@@ -105,13 +105,13 @@ def test_integer_parameter(self):
 
         self.assertEqual(
             ParameterUsage.MODEL,
-            self.search_parameters.get_type("instance_count"),
+            self.search_parameters.get_type("instance_group"),
         )
         self.assertEqual(
             ParameterCategory.INTEGER,
-            self.search_parameters.get_category("instance_count"),
+            self.search_parameters.get_category("instance_group"),
         )
-        self.assertEqual((1, 8), self.search_parameters.get_range("instance_count"))
+        self.assertEqual((1, 8), self.search_parameters.get_range("instance_group"))
 
     def test_list_parameter(self):
         """
@@ -371,6 +371,46 @@ def test_search_parameter_creation_multi_model_non_default(self):
             default.DEFAULT_RUN_CONFIG_MAX_INSTANCE_COUNT, instance_group.max_range
         )
 
+    def test_number_of_configs_range(self):
+        """
+        Test number of configs for a range (INTEGER/EXPONENTIAL)
+        """
+
+        # INTEGER
+        # =====================================================================
+        num_of_configs = self.search_parameters._number_of_configurations_for_parameter(
+            self.search_parameters.get_parameter("instance_group")
+        )
+        self.assertEqual(8, num_of_configs)
+
+        # EXPONENTIAL
+        # =====================================================================
+        num_of_configs = self.search_parameters._number_of_configurations_for_parameter(
+            self.search_parameters.get_parameter("concurrency")
+        )
+        self.assertEqual(11, num_of_configs)
+
+    def test_number_of_configs_list(self):
+        """
+        Test number of configs for a list
+        """
+
+        num_of_configs = self.search_parameters._number_of_configurations_for_parameter(
+            self.search_parameters.get_parameter("size")
+        )
+        self.assertEqual(3, num_of_configs)
+
+    def test_total_possible_configurations(self):
+        """
+        Test number of total possible configurations
+        """
+        total_num_of_possible_configurations = (
+            self.search_parameters.number_of_total_possible_configurations()
+        )
+
+        # batch_sizes (8) * instance group (8) * concurrency (11) * size (3)
+        self.assertEqual(8 * 8 * 11 * 3, total_num_of_possible_configurations)
+
 
 if __name__ == "__main__":
     unittest.main()