Using SearchParameters in OptunaRCG

model_analyzer/config/generate/optuna_plus_concurrency_sweep_run_config_generator.py

@@ -16,7 +16,7 @@
 
 import logging
 from copy import deepcopy
-from typing import Generator, List, Optional
+from typing import Dict, Generator, List, Optional
 
 from model_analyzer.config.generate.model_profile_spec import ModelProfileSpec
 from model_analyzer.config.generate.model_variant_name_manager import (
@@ -52,7 +52,7 @@ def __init__(
         models: List[ModelProfileSpec],
         result_manager: ResultManager,
         model_variant_name_manager: ModelVariantNameManager,
-        search_parameters: SearchParameters,
+        search_parameters: Dict[str, SearchParameters],
     ):
         """
         Parameters

model_analyzer/config/generate/optuna_run_config_generator.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 import logging
-from typing import Any, Dict, Generator, List, Optional, Union
+from typing import Any, Dict, Generator, List, Optional, TypeAlias, Union
 
 import optuna
 
@@ -29,9 +29,16 @@
 from model_analyzer.config.generate.model_variant_name_manager import (
     ModelVariantNameManager,
 )
+from model_analyzer.config.generate.search_parameter import (
+    ParameterCategory,
+    SearchParameter,
+)
 from model_analyzer.config.generate.search_parameters import SearchParameters
 from model_analyzer.config.input.config_command_profile import ConfigCommandProfile
-from model_analyzer.config.input.config_defaults import DEFAULT_BATCH_SIZES
+from model_analyzer.config.input.config_defaults import (
+    DEFAULT_BATCH_SIZES,
+    DEFAULT_RUN_CONFIG_MAX_CONCURRENCY,
+)
 from model_analyzer.config.run.model_run_config import ModelRunConfig
 from model_analyzer.config.run.run_config import RunConfig
 from model_analyzer.constants import LOGGER_NAME
@@ -44,20 +51,31 @@
 
 logger = logging.getLogger(LOGGER_NAME)
 
+TrialObjectives: TypeAlias = Dict[str, Any]
+ParameterCombo: TypeAlias = Dict[str, Any]
+
 
 class OptunaRunConfigGenerator(ConfigGeneratorInterface):
     """
     Use Optuna algorithm to create RunConfigs
     """
 
+    # This list represents all possible parameters Optuna can currently search for
+    optuna_parameter_list = [
+        "batch_sizes",
+        "instance_group",
+        "concurrency",
+        "max_queue_delay_microseconds",
+    ]
+
     def __init__(
         self,
         config: ConfigCommandProfile,
         gpu_count: int,
         models: List[ModelProfileSpec],
         model_variant_name_manager: ModelVariantNameManager,
-        search_parameters: SearchParameters,
-        seed: Optional[int] = 0,
+        search_parameters: Dict[str, SearchParameters],
+        seed: Optional[int] = None,
     ):
         """
         Parameters
@@ -74,7 +92,8 @@ def __init__(
         self._config = config
         self._gpu_count = gpu_count
         self._models = models
-        self._search_parameters = search_parameters
+        # TODO: TMA-1927: Add support for multi-model
+        self._search_parameters = search_parameters[models[0].model_name()]
 
         self._model_variant_name_manager = model_variant_name_manager
 
@@ -126,24 +145,70 @@ def get_configs(self) -> Generator[RunConfig, None, None]:
         # TODO: TMA-1885: Need an early exit strategy
         for _ in range(n_trials):
             trial = self._study.ask()
-            self._create_trial_objectives(trial)
-            run_config = self._create_objective_based_run_config()
+            trial_objectives = self._create_trial_objectives(trial)
+            run_config = self._create_objective_based_run_config(trial_objectives)
             yield run_config
             score = self._calculate_score()
             self._study.tell(trial, score)
 
-    def _create_trial_objectives(self, trial) -> None:
-        # TODO: TMA-1925: Use SearchParameters here
-        self._instance_count = trial.suggest_int("instance_count", 1, 8)
-        self._batch_size = int(2 ** trial.suggest_int("batch_size", 1, 10))
+    def _create_trial_objectives(self, trial: optuna.Trial) -> TrialObjectives:
+        trial_objectives: TrialObjectives = {}
+        for parameter_name in OptunaRunConfigGenerator.optuna_parameter_list:
+            parameter = self._search_parameters.get_parameter(parameter_name)
+
+            if parameter:
+                trial_objectives = self._add_trial_objective(
+                    trial, parameter_name, parameter, trial_objectives
+                )
 
         # TODO: TMA-1884: Need an option to choose btw. concurrency formula and optuna searching
-        self._concurrency = 2 * self._instance_count * self._batch_size
-        if self._concurrency > 1024:
-            self._concurrency = 1024
+        trial_objectives = self._set_objective_concurrency(trial_objectives)
+
+        return trial_objectives
+
+    def _add_trial_objective(
+        self,
+        trial: optuna.Trial,
+        name: str,
+        parameter: SearchParameter,
+        trial_objectives: TrialObjectives,
+    ) -> TrialObjectives:
+        if parameter.category is ParameterCategory.INTEGER:
+            objective = trial.suggest_int(
+                name, parameter.min_range, parameter.max_range
+            )
+        elif parameter.category is ParameterCategory.EXPONENTIAL:
+            objective = int(
+                2 ** trial.suggest_int(name, parameter.min_range, parameter.max_range)
+            )
+        elif parameter.category is ParameterCategory.INT_LIST:
+            objective = int(trial.suggest_categorical(name, parameter.enumerated_list))
+        elif parameter.category is ParameterCategory.STR_LIST:
+            objective = trial.suggest_categorical(name, parameter.enumerated_list)
+
+        trial_objectives[name] = objective
 
-    def _create_objective_based_run_config(self) -> RunConfig:
-        param_combo = self._create_parameter_combo()
+        return trial_objectives
+
+    def _set_objective_concurrency(
+        self, trial_objectives: TrialObjectives
+    ) -> TrialObjectives:
+        concurrency = (
+            2 * trial_objectives["instance_group"] * trial_objectives["batch_sizes"]
+        )
+        concurrency = (
+            DEFAULT_RUN_CONFIG_MAX_CONCURRENCY
+            if concurrency > DEFAULT_RUN_CONFIG_MAX_CONCURRENCY
+            else concurrency
+        )
+        trial_objectives["concurrency"] = concurrency
+
+        return trial_objectives
+
+    def _create_objective_based_run_config(
+        self, trial_objectives: TrialObjectives
+    ) -> RunConfig:
+        param_combo = self._create_parameter_combo(trial_objectives)
 
         # TODO: TMA-1927: Add support for multi-model
         run_config = RunConfig(self._triton_env)
@@ -159,27 +224,38 @@ def _create_objective_based_run_config(self) -> RunConfig:
         model_run_config = self._create_model_run_config(
             model=self._models[0],
             model_config_variant=model_config_variant,
+            trial_objectives=trial_objectives,
         )
 
         run_config.add_model_run_config(model_run_config=model_run_config)
 
         return run_config
 
-    def _create_parameter_combo(self) -> Dict[str, Any]:
-        # TODO: TMA-1925: Use SearchParameters here
-        param_combo: Dict["str", Any] = {}
-        param_combo["dynamic_batching"] = {}
+    def _create_parameter_combo(
+        self, trial_objectives: TrialObjectives
+    ) -> ParameterCombo:
+        param_combo: ParameterCombo = {}
+        param_combo["dynamic_batching"] = []
 
         # TODO: TMA-1927: Add support for multi-model
-        kind = "KIND_CPU" if self._models[0].cpu_only() else "KIND_GPU"
-        param_combo["instance_group"] = [
-            {
-                "count": self._instance_count,
-                "kind": kind,
+        if trial_objectives["instance_group"]:
+            kind = "KIND_CPU" if self._models[0].cpu_only() else "KIND_GPU"
+            param_combo["instance_group"] = [
+                {
+                    "count": trial_objectives["instance_group"],
+                    "kind": kind,
+                }
+            ]
+
+        if trial_objectives["batch_sizes"]:
+            param_combo["max_batch_size"] = trial_objectives["batch_sizes"]
+
+        if trial_objectives["max_queue_delay_microseconds"]:
+            param_combo["dynamic_batching"] = {
+                "max_queue_delay_microseconds": trial_objectives[
+                    "max_queue_delay_microseconds"
+                ]
             }
-        ]
-
-        param_combo["max_batch_size"] = self._batch_size
 
         return param_combo
 
@@ -261,9 +337,10 @@ def _create_model_run_config(
         self,
         model: ModelProfileSpec,
         model_config_variant: ModelConfigVariant,
+        trial_objectives: TrialObjectives,
     ) -> ModelRunConfig:
         perf_analyzer_config = self._create_perf_analyzer_config(
-            model.model_name(), model, self._concurrency
+            model.model_name(), model, trial_objectives["concurrency"]
         )
         model_run_config = ModelRunConfig(
             model.model_name(), model_config_variant, perf_analyzer_config

model_analyzer/config/generate/run_config_generator_factory.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List
+from typing import Dict, List
 
 from model_analyzer.config.generate.model_profile_spec import ModelProfileSpec
 from model_analyzer.config.generate.model_variant_name_manager import (
@@ -60,7 +60,7 @@ def create_run_config_generator(
         client: TritonClient,
         result_manager: ResultManager,
         model_variant_name_manager: ModelVariantNameManager,
-        search_parameters: SearchParameters,
+        search_parameters: Dict[str, SearchParameters],
     ) -> ConfigGeneratorInterface:
         """
         Parameters
@@ -149,7 +149,7 @@ def _create_optuna_plus_concurrency_sweep_run_config_generator(
         models: List[ModelProfileSpec],
         result_manager: ResultManager,
         model_variant_name_manager: ModelVariantNameManager,
-        search_parameters: SearchParameters,
+        search_parameters: Dict[str, SearchParameters],
     ) -> ConfigGeneratorInterface:
         return OptunaPlusConcurrencySweepRunConfigGenerator(
             config=command_config,

model_analyzer/config/generate/search_parameter.py

@@ -28,7 +28,8 @@ class ParameterUsage(Enum):
 class ParameterCategory(Enum):
     INTEGER = auto()
     EXPONENTIAL = auto()
-    LIST = auto()
+    STR_LIST = auto()
+    INT_LIST = auto()
 
 
 @dataclass
@@ -40,7 +41,7 @@ class SearchParameter:
     usage: ParameterUsage
     category: ParameterCategory
 
-    # This is only applicable to LIST category
+    # This is only applicable to the LIST categories
     enumerated_list: Optional[List[Any]] = None
 
     # These are only applicable to INTEGER and EXPONENTIAL categories

model_analyzer/config/generate/search_parameters.py

@@ -49,11 +49,13 @@ def __init__(
 
         self._populate_search_parameters()
 
-    def get_parameters(self) -> List[SearchParameter]:
-        return [v for v in self._search_parameters.values()]
+    def get_parameter(self, name: str) -> Optional[SearchParameter]:
+        try:
+            parameter = self._search_parameters[name]
+        except KeyError:
+            return None
 
-    def get_parameter(self, name: str) -> SearchParameter:
-        return self._search_parameters[name]
+        return parameter
 
     def get_type(self, name: str) -> ParameterUsage:
         return self._search_parameters[name].usage
@@ -73,6 +75,9 @@ def get_list(self, name: str) -> Optional[List[Any]]:
     def _populate_search_parameters(self) -> None:
         if self._parameters:
             self._populate_parameters()
+        else:
+            # Always populate batch size if nothing is specified
+            self._populate_batch_sizes()
 
         self._populate_model_config_parameters()
 
@@ -86,10 +91,11 @@ def _populate_model_config_parameters(self) -> None:
         self._populate_max_queue_delay_microseconds()
 
     def _populate_batch_sizes(self) -> None:
-        if self._parameters["batch_sizes"]:
+        if self._parameters and self._parameters["batch_sizes"]:
             self._populate_list_parameter(
                 parameter_name="batch_sizes",
                 parameter_list=self._parameters["batch_sizes"],
+                parameter_category=ParameterCategory.INT_LIST,
             )
         else:
             self._populate_rcs_parameter(
@@ -103,6 +109,7 @@ def _populate_concurrency(self) -> None:
             self._populate_list_parameter(
                 parameter_name="concurrency",
                 parameter_list=self._parameters["concurrency"],
+                parameter_category=ParameterCategory.INT_LIST,
             )
         else:
             self._populate_rcs_parameter(
@@ -135,6 +142,7 @@ def _populate_instance_group(self) -> None:
             self._populate_list_parameter(
                 parameter_name="instance_group",
                 parameter_list=parameter_list,
+                parameter_category=ParameterCategory.INT_LIST,
             )
         else:
             self._populate_rcs_parameter(
@@ -157,6 +165,7 @@ def _populate_max_queue_delay_microseconds(self) -> None:
                 parameter_list=self._model_config_parameters["dynamic_batching"][0][
                     "max_queue_delay_microseconds"
                 ],
+                parameter_category=ParameterCategory.INT_LIST,
             )
 
     def _is_max_queue_delay_in_model_config_parameters(self) -> bool:
@@ -176,14 +185,15 @@ def _is_max_queue_delay_in_model_config_parameters(self) -> bool:
     def _populate_list_parameter(
         self,
         parameter_name: str,
-        parameter_list: List[int],
+        parameter_list: List[int | str],
+        parameter_category: ParameterCategory,
     ) -> None:
         usage = self._determine_parameter_usage(parameter_name)
 
         self._add_search_parameter(
             name=parameter_name,
             usage=usage,
-            category=ParameterCategory.LIST,
+            category=parameter_category,
             enumerated_list=parameter_list,
         )
 
@@ -257,7 +267,10 @@ def _check_for_illegal_input(
         max_range: Optional[int],
         enumerated_list: List[Any],
     ) -> None:
-        if category is ParameterCategory.LIST:
+        if (
+            category is ParameterCategory.INT_LIST
+            or category is ParameterCategory.STR_LIST
+        ):
             self._check_for_illegal_list_input(min_range, max_range, enumerated_list)
         else:
             if min_range is None or max_range is None:

model_analyzer/model_manager.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 import logging
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from model_analyzer.config.generate.model_variant_name_manager import (
     ModelVariantNameManager,
@@ -60,7 +60,7 @@ def __init__(
         result_manager: ResultManager,
         state_manager: AnalyzerStateManager,
         constraint_manager: ConstraintManager,
-        search_parameters: SearchParameters,
+        search_parameters: Dict[str, SearchParameters],
     ):
         """
         Parameters