Add request rate parameter for optuna search (#915)

* Add reeuest rate parameter for optuna search

* Fix pre-commit error

* Remove unused variable

* Fix comments

* Modify TestOptunaRunConfigGenerator to support search based on request-rate

* Add search by request-rate to OptunRunConfigGenerator

* Changes

* Fix codeql errors

* Modification to disable concurrency sweep when request rate is specified

* Modify test to check if concurrency-range is in perf_config

* Fix codeql errors

* Fix codeql errors

* Fix comments

* Fix error

Author: lkomali

model_analyzer/config/generate/optuna_run_config_generator.py

@@ -37,7 +37,10 @@
 )
 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_MIN_REQUEST_RATE,
+)
 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
@@ -76,6 +79,7 @@ class OptunaRunConfigGenerator(ConfigGeneratorInterface):
         "instance_group",
         "concurrency",
         "max_queue_delay_microseconds",
+        "request_rate",
     ]
 
     # TODO: TMA-1927: Figure out the correct value for this
@@ -380,7 +384,6 @@ def _create_trial_objectives(self, trial: optuna.Trial) -> AllTrialObjectives:
                 parameter = self._search_parameters[model_name].get_parameter(
                     parameter_name
                 )
-
                 if parameter:
                     objective_name = self._create_trial_objective_name(
                         model_name=model_name, parameter_name=parameter_name
@@ -668,14 +671,21 @@ def _create_default_perf_analyzer_config(
             model_config.get_field("name"), self._config
         )
 
-        default_concurrency = self._calculate_default_concurrency(model_config)
-
-        perf_config_params = {
-            "batch-size": DEFAULT_BATCH_SIZES,
-            "concurrency-range": default_concurrency,
-        }
+        if self._search_parameters[model_config.get_field("name")].get_parameter(
+            "request_rate"
+        ):
+            perf_config_params = {
+                "batch-size": DEFAULT_BATCH_SIZES,
+                "request-rate-range": DEFAULT_RUN_CONFIG_MIN_REQUEST_RATE,
+            }
+            self._config.concurrency_sweep_disable = True
+        else:
+            default_concurrency = self._calculate_default_concurrency(model_config)
+            perf_config_params = {
+                "batch-size": DEFAULT_BATCH_SIZES,
+                "concurrency-range": default_concurrency,
+            }
         default_perf_analyzer_config.update_config(perf_config_params)
-
         default_perf_analyzer_config.update_config(model.perf_analyzer_flags())
 
         return default_perf_analyzer_config
@@ -712,16 +722,10 @@ def _create_model_run_config(
         composing_model_config_variants: List[ModelConfigVariant],
         trial_objectives: ModelTrialObjectives,
     ) -> ModelRunConfig:
-        trial_batch_sizes = (
-            int(trial_objectives["batch_sizes"])
-            if "batch_sizes" in trial_objectives
-            else DEFAULT_BATCH_SIZES
-        )
         perf_analyzer_config = self._create_perf_analyzer_config(
             model_name=model.model_name(),
             model=model,
-            concurrency=int(trial_objectives["concurrency"]),
-            batch_sizes=trial_batch_sizes,
+            trial_objectives=trial_objectives,
         )
         model_run_config = ModelRunConfig(
             model.model_name(), model_config_variant, perf_analyzer_config
@@ -738,17 +742,30 @@ def _create_perf_analyzer_config(
         self,
         model_name: str,
         model: ModelProfileSpec,
-        concurrency: int,
-        batch_sizes: int,
+        trial_objectives: ModelTrialObjectives,
     ) -> PerfAnalyzerConfig:
         perf_analyzer_config = PerfAnalyzerConfig()
 
         perf_analyzer_config.update_config_from_profile_config(model_name, self._config)
 
-        perf_config_params = {
-            "batch-size": batch_sizes,
-            "concurrency-range": concurrency,
-        }
+        batch_sizes = (
+            int(trial_objectives["batch_sizes"])
+            if "batch_sizes" in trial_objectives
+            else DEFAULT_BATCH_SIZES
+        )
+
+        perf_config_params = {"batch-size": batch_sizes}
+
+        if "concurrency" in trial_objectives:
+            perf_config_params["concurrency-range"] = int(
+                trial_objectives["concurrency"]
+            )
+        elif "request_rate" in trial_objectives:
+            perf_config_params["request-rate-range"] = int(
+                trial_objectives["request_rate"]
+            )
+            self._config.concurrency_sweep_disable = True
+
         perf_analyzer_config.update_config(perf_config_params)
 
         perf_analyzer_config.update_config(model.perf_analyzer_flags())

model_analyzer/config/generate/search_parameters.py

@@ -31,15 +31,20 @@ class SearchParameters:
 
     # These map to the run-config-search fields
     # See github.com/triton-inference-server/model_analyzer/blob/main/docs/config.md
-    exponential_rcs_parameters = ["max_batch_size", "batch_sizes", "concurrency"]
+    exponential_rcs_parameters = [
+        "max_batch_size",
+        "batch_sizes",
+        "concurrency",
+        "request_rate",
+    ]
     linear_rcs_parameters = ["instance_group"]
 
     model_parameters = [
         "max_batch_size",
         "instance_group",
         "max_queue_delay_microseconds",
     ]
-    runtime_parameters = ["batch_sizes", "concurrency"]
+    runtime_parameters = ["batch_sizes", "concurrency", "request_rate"]
 
     def __init__(
         self,
@@ -129,8 +134,10 @@ def _populate_parameters(self) -> None:
         self._populate_batch_sizes()
 
         if not self._is_composing_model:
-            self._populate_concurrency()
-            # TODO: Populate request rate - TMA-1903
+            if self._config.is_request_rate_specified(self._parameters):
+                self._populate_request_rate()
+            else:
+                self._populate_concurrency()
 
     def _populate_model_config_parameters(self) -> None:
         self._populate_max_batch_size()
@@ -161,6 +168,20 @@ def _populate_concurrency(self) -> None:
                 rcs_parameter_max_value=self._config.run_config_search_max_concurrency,
             )
 
+    def _populate_request_rate(self) -> None:
+        if self._parameters and self._parameters["request_rate"]:
+            self._populate_list_parameter(
+                parameter_name="request_rate",
+                parameter_list=self._parameters["request_rate"],
+                parameter_category=ParameterCategory.INT_LIST,
+            )
+        else:
+            self._populate_rcs_parameter(
+                parameter_name="request_rate",
+                rcs_parameter_min_value=self._config.run_config_search_min_request_rate,
+                rcs_parameter_max_value=self._config.run_config_search_max_request_rate,
+            )
+
     def _populate_max_batch_size(self) -> None:
         # Example config format:
         # model_config_parameters:

tests/test_optuna_run_config_generator.py

@@ -25,7 +25,10 @@
     OptunaRunConfigGenerator,
 )
 from model_analyzer.config.generate.search_parameters import SearchParameters
-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_MIN_REQUEST_RATE,
+)
 from model_analyzer.config.input.objects.config_model_profile_spec import (
     ConfigModelProfileSpec,
 )
@@ -174,9 +177,9 @@ def test_min_number_of_configs_to_search_both(self):
         # Since both are specified we will use the larger of the two (trials=6)
         self.assertEqual(min_configs_to_search, 6)
 
-    def test_create_default_run_config(self):
+    def test_create_default_run_config_with_concurrency(self):
         """
-        Test that a default run config is properly created
+        Test that a default run config with concurrency is properly created
         """
         default_run_config = self._rcg._create_default_run_config()
 
@@ -190,9 +193,50 @@ def test_create_default_run_config(self):
             perf_config["concurrency-range"], 2 * self._default_max_batch_size
         )
 
-    def test_create_objective_based_run_config(self):
+    def test_create_default_run_config_with_request_rate(self):
+        """
+        Test that a default run config with request rate is properly created
+        """
+        config = self._create_config(["--request-rate-search-enable"])
+        mock_model_config = MockModelConfig()
+        mock_model_config.start()
+        model = ModelProfileSpec(
+            config.profile_models[0], config, MagicMock(), MagicMock()
+        )
+        mock_model_config.stop()
+        search_parameters = SearchParameters(
+            model=model,
+            config=config,
+        )
+
+        rcg = OptunaRunConfigGenerator(
+            config=config,
+            state_manager=MagicMock(),
+            gpu_count=1,
+            models=self._mock_models,
+            composing_models=[],
+            model_variant_name_manager=ModelVariantNameManager(),
+            search_parameters={"add_sub": search_parameters},
+            composing_search_parameters={},
+            user_seed=100,
+        )
+
+        default_run_config = rcg._create_default_run_config()
+        self.assertEqual(len(default_run_config.model_run_configs()), 1)
+
+        model_config = default_run_config.model_run_configs()[0].model_config()
+        perf_config = default_run_config.model_run_configs()[0].perf_config()
+
+        self.assertEqual(model_config.to_dict()["name"], self._test_config_dict["name"])
+        self.assertEqual(perf_config["batch-size"], DEFAULT_BATCH_SIZES)
+        self.assertEqual(
+            perf_config["request-rate-range"], DEFAULT_RUN_CONFIG_MIN_REQUEST_RATE
+        )
+        self.assertEqual(perf_config["concurrency-range"], None)
+
+    def test_create_objective_based_run_config_with_concurrency(self):
         """
-        Test that an objective based run config is properly created
+        Test that an objective based run config with concurrency is properly created
         """
         trial = self._rcg._study.ask()
         trial_objectives = self._rcg._create_trial_objectives(trial)
@@ -215,6 +259,47 @@ def test_create_objective_based_run_config(self):
         self.assertEqual(perf_config["batch-size"], DEFAULT_BATCH_SIZES)
         self.assertEqual(perf_config["concurrency-range"], 64)
 
+    def test_create_objective_based_run_config_with_request_rate(self):
+        """
+        Test that an objective based run config with request rate is properly created
+        """
+        config = self._create_config(["--request-rate", "1024,2048"])
+        mock_model_config = MockModelConfig()
+        mock_model_config.start()
+        model = ModelProfileSpec(
+            config.profile_models[0], config, MagicMock(), MagicMock()
+        )
+        mock_model_config.stop()
+        search_parameters = SearchParameters(
+            model=model,
+            config=config,
+        )
+
+        rcg = OptunaRunConfigGenerator(
+            config=config,
+            state_manager=MagicMock(),
+            gpu_count=1,
+            models=self._mock_models,
+            composing_models=[],
+            model_variant_name_manager=ModelVariantNameManager(),
+            search_parameters={"add_sub": search_parameters},
+            composing_search_parameters={},
+            user_seed=100,
+        )
+
+        trial = rcg._study.ask()
+        trial_objectives = rcg._create_trial_objectives(trial)
+        run_config = rcg._create_objective_based_run_config(trial_objectives, None)
+
+        model_config = run_config.model_run_configs()[0].model_config()
+        perf_config = run_config.model_run_configs()[0].perf_config()
+
+        # These values are the result of using a fixed seed of 100
+        self.assertEqual(model_config.to_dict()["name"], self._test_config_dict["name"])
+        self.assertEqual(perf_config["batch-size"], DEFAULT_BATCH_SIZES)
+        self.assertEqual(perf_config["request-rate-range"], 2048)
+        self.assertEqual(perf_config["concurrency-range"], None)
+
     def test_create_run_config_with_concurrency_formula(self):
         """
         Tests that the concurrency formula option is used correctly
@@ -284,12 +369,10 @@ def test_create_run_bls_config(self):
             config=config,
         )
         add_search_parameters = SearchParameters(
-            model=add_model,
-            config=config,
+            model=add_model, config=config, is_composing_model=True
         )
         sub_search_parameters = SearchParameters(
-            model=sub_model,
-            config=config,
+            model=sub_model, config=config, is_composing_model=True
         )
         rcg = OptunaRunConfigGenerator(
             config=config,
@@ -331,16 +414,16 @@ def test_create_run_bls_config(self):
         self.assertEqual(add_model_config.to_dict()["instanceGroup"][0]["count"], 3)
         self.assertEqual(
             add_model_config.to_dict()["dynamicBatching"]["maxQueueDelayMicroseconds"],
-            "300",
+            "100",
         )
 
         # SUB (composing model)
         # =====================================================================
         self.assertEqual(sub_model_config.to_dict()["name"], "sub")
-        self.assertEqual(sub_model_config.to_dict()["instanceGroup"][0]["count"], 5)
+        self.assertEqual(sub_model_config.to_dict()["instanceGroup"][0]["count"], 4)
         self.assertEqual(
             sub_model_config.to_dict()["dynamicBatching"]["maxQueueDelayMicroseconds"],
-            "500",
+            "400",
         )
 
     def test_create_run_multi_model_config(self):

tests/test_search_parameters.py

@@ -98,6 +98,7 @@ def test_exponential_parameter(self):
         Test exponential parameter, accessing dataclass directly
         """
 
+        # concurrency
         parameter = self.search_parameters.get_parameter("concurrency")
 
         self.assertEqual(ParameterUsage.RUNTIME, parameter.usage)
@@ -322,10 +323,10 @@ def test_search_parameter_creation_multi_model_non_default(self):
             mult_div:
                 parameters:
                     concurrency: [1, 8, 64, 256]
+
         """
 
         config = TestConfig()._evaluate_config(args, yaml_content)
-
         analyzer = Analyzer(config, MagicMock(), MagicMock(), MagicMock())
         mock_model_config = MockModelConfig()
         mock_model_config.start()
@@ -418,6 +419,48 @@ def test_search_parameter_creation_multi_model_non_default(self):
             default.DEFAULT_RUN_CONFIG_MAX_INSTANCE_COUNT, instance_group.max_range
         )
 
+    def test_search_parameter_request_rate(self):
+        """
+        Test that request rate is correctly set in
+        a non-default optuna case
+        """
+
+        args = [
+            "model-analyzer",
+            "profile",
+            "--model-repository",
+            "cli-repository",
+            "-f",
+            "path-to-config-file",
+            "--run-config-search-mode",
+            "optuna",
+        ]
+
+        yaml_content = """
+        run_config_search_mode: optuna
+        profile_models:
+            mult_div:
+                parameters:
+                    request_rate: [1, 8, 64, 256]
+
+        """
+        config = TestConfig()._evaluate_config(args, yaml_content)
+        analyzer = Analyzer(config, MagicMock(), MagicMock(), MagicMock())
+        mock_model_config = MockModelConfig()
+        mock_model_config.start()
+        analyzer._populate_search_parameters(MagicMock(), MagicMock())
+        mock_model_config.stop()
+
+        # request_rate
+        # ===================================================================
+
+        request_rate = analyzer._search_parameters["mult_div"].get_parameter(
+            "request_rate"
+        )
+        self.assertEqual(ParameterUsage.RUNTIME, request_rate.usage)
+        self.assertEqual(ParameterCategory.INT_LIST, request_rate.category)
+        self.assertEqual([1, 8, 64, 256], request_rate.enumerated_list)
+
     def test_number_of_configs_range(self):
         """
         Test number of configs for a range (INTEGER/EXPONENTIAL)