Creating optuna RCG factory

model_analyzer/analyzer.py

@@ -204,6 +204,7 @@ def _create_model_manager(self, client, gpus):
             metrics_manager=self._metrics_manager,
             state_manager=self._state_manager,
             constraint_manager=self._constraint_manager,
+            search_parameters=self._search_parameters,
         )
 
     def _get_server_only_metrics(self, client, gpus):

model_analyzer/config/generate/optuna_plus_concurrency_sweep_run_config_generator.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+
+# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from copy import deepcopy
+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 (
+    ModelVariantNameManager,
+)
+from model_analyzer.config.generate.optuna_run_config_generator import (
+    OptunaRunConfigGenerator,
+)
+from model_analyzer.config.generate.search_parameters import SearchParameters
+from model_analyzer.config.input.config_command_profile import ConfigCommandProfile
+from model_analyzer.config.run.run_config import RunConfig
+from model_analyzer.constants import LOGGER_NAME
+from model_analyzer.device.gpu_device import GPUDevice
+from model_analyzer.record.metrics_manager import MetricsManager
+from model_analyzer.result.parameter_search import ParameterSearch
+from model_analyzer.result.result_manager import ResultManager
+from model_analyzer.result.run_config_measurement import RunConfigMeasurement
+from model_analyzer.triton.client.client import TritonClient
+
+from .config_generator_interface import ConfigGeneratorInterface
+
+logger = logging.getLogger(LOGGER_NAME)
+
+
+class OptunaPlusConcurrencySweepRunConfigGenerator(ConfigGeneratorInterface):
+    """
+    First run OptunaConfigGenerator for an Optuna search, then use
+    ParameterSearch for a concurrency sweep + binary search of the default
+    and Top N results
+    """
+
+    def __init__(
+        self,
+        config: ConfigCommandProfile,
+        gpus: List[GPUDevice],
+        models: List[ModelProfileSpec],
+        client: TritonClient,
+        result_manager: ResultManager,
+        metrics_manager: MetricsManager,
+        model_variant_name_manager: ModelVariantNameManager,
+        search_parameters: SearchParameters,
+    ):
+        """
+        Parameters
+        ----------
+        config: ConfigCommandProfile
+            Profile configuration information
+        gpus: List of GPUDevices
+        models: List of ModelProfileSpec
+            List of models to profile
+        client: TritonClient
+        result_manager: ResultManager
+            The object that handles storing and sorting the results from the perf analyzer
+        model_variant_name_manager: ModelVariantNameManager
+            Maps model variants to config names
+        search_parameters: SearchParameters
+            The object that handles the users configuration search parameters
+        """
+        self._config = config
+        self._gpus = gpus
+        self._models = models
+        self._client = client
+        self._result_manager = result_manager
+        self._metrics_manager = metrics_manager
+        self._model_variant_name_manager = model_variant_name_manager
+        self._search_parameters = search_parameters
+
+    def set_last_results(
+        self, measurements: List[Optional[RunConfigMeasurement]]
+    ) -> None:
+        self._last_measurement = measurements[-1]
+        self._rcg.set_last_results(measurements)
+
+    def get_configs(self) -> Generator[RunConfig, None, None]:
+        """
+        Returns
+        -------
+        RunConfig
+            The next RunConfig generated by this class
+        """
+
+        logger.info("")
+        logger.info("Starting Optuna mode search to find optimal configs")
+        logger.info("")
+        yield from self._execute_optuna_search()
+        logger.info("")
+        logger.info(
+            "Done with Optuna mode search. Gathering concurrency sweep measurements for reports"
+        )
+        logger.info("")
+        yield from self._sweep_concurrency_over_top_results()
+        logger.info("")
+        logger.info("Done gathering concurrency sweep measurements for reports")
+        logger.info("")
+
+    def _execute_optuna_search(self) -> Generator[RunConfig, None, None]:
+        self._rcg: ConfigGeneratorInterface = self._create_optuna_run_config_generator()
+
+        yield from self._rcg.get_configs()
+
+    def _create_optuna_run_config_generator(self) -> OptunaRunConfigGenerator:
+        return OptunaRunConfigGenerator(
+            config=self._config,
+            gpus=self._gpus,
+            models=self._models,
+            client=self._client,
+            model_variant_name_manager=self._model_variant_name_manager,
+            search_parameters=self._search_parameters,
+            metrics_manager=self._metrics_manager,
+        )
+
+    def _sweep_concurrency_over_top_results(self) -> Generator[RunConfig, None, None]:
+        for model_name in self._result_manager.get_model_names():
+            top_results = self._result_manager.top_n_results(
+                model_name=model_name,
+                n=self._config.num_configs_per_model,
+                include_default=True,
+            )
+
+            for result in top_results:
+                run_config = deepcopy(result.run_config())
+                parameter_search = ParameterSearch(self._config)
+                for concurrency in parameter_search.search_parameters():
+                    run_config = self._set_concurrency(run_config, concurrency)
+                    yield run_config
+                    parameter_search.add_run_config_measurement(self._last_measurement)
+
+    def _set_concurrency(self, run_config: RunConfig, concurrency: int) -> RunConfig:
+        for model_run_config in run_config.model_run_configs():
+            perf_config = model_run_config.perf_config()
+            perf_config.update_config({"concurrency-range": concurrency})
+
+        return run_config

model_analyzer/config/generate/quick_plus_concurrency_sweep_run_config_generator.py

@@ -26,6 +26,7 @@
     QuickRunConfigGenerator,
 )
 from model_analyzer.config.generate.search_config import SearchConfig
+from model_analyzer.config.generate.search_parameters import SearchParameters
 from model_analyzer.config.input.config_command_profile import ConfigCommandProfile
 from model_analyzer.config.run.run_config import RunConfig
 from model_analyzer.constants import LOGGER_NAME

model_analyzer/config/generate/run_config_generator_factory.py

@@ -20,13 +20,15 @@
 from model_analyzer.config.generate.model_variant_name_manager import (
     ModelVariantNameManager,
 )
+from model_analyzer.config.generate.search_parameters import SearchParameters
 from model_analyzer.config.input.config_command_profile import ConfigCommandProfile
 from model_analyzer.config.input.objects.config_model_profile_spec import (
     ConfigModelProfileSpec,
 )
 from model_analyzer.constants import MIN_INITIALIZED, RADIUS
 from model_analyzer.device.gpu_device import GPUDevice
 from model_analyzer.model_analyzer_exceptions import TritonModelAnalyzerException
+from model_analyzer.record.metrics_manager import MetricsManager
 from model_analyzer.result.result_manager import ResultManager
 from model_analyzer.triton.client.client import TritonClient
 from model_analyzer.triton.model.model_config import ModelConfig
@@ -35,6 +37,9 @@
     BrutePlusBinaryParameterSearchRunConfigGenerator,
 )
 from .config_generator_interface import ConfigGeneratorInterface
+from .optuna_plus_concurrency_sweep_run_config_generator import (
+    OptunaPlusConcurrencySweepRunConfigGenerator,
+)
 from .quick_plus_concurrency_sweep_run_config_generator import (
     QuickPlusConcurrencySweepRunConfigGenerator,
 )
@@ -55,7 +60,9 @@ def create_run_config_generator(
         models: List[ConfigModelProfileSpec],
         client: TritonClient,
         result_manager: ResultManager,
+        metrics_manager: MetricsManager,
         model_variant_name_manager: ModelVariantNameManager,
+        search_parameters: SearchParameters,
     ) -> ConfigGeneratorInterface:
         """
         Parameters
@@ -71,6 +78,8 @@ def create_run_config_generator(
             The object that handles storing and sorting the results from the perf analyzer
         model_variant_name_manager: ModelVariantNameManager
             Maps model variants to config names
+        search_parameters: SearchParameters
+            The object that handles the users configuration search parameters
 
         Returns
         -------
@@ -85,7 +94,18 @@ def create_run_config_generator(
             new_models, command_config, client, gpus
         )
 
-        if command_config.run_config_search_mode == "quick" or composing_models:
+        if command_config.run_config_search_mode == "optuna":
+            return RunConfigGeneratorFactory._create_optuna_plus_concurrency_sweep_run_config_generator(
+                command_config=command_config,
+                gpus=gpus,
+                models=new_models,
+                client=client,
+                result_manager=result_manager,
+                metrics_manager=metrics_manager,
+                search_parameters=search_parameters,
+                model_variant_name_manager=model_variant_name_manager,
+            )
+        elif command_config.run_config_search_mode == "quick" or composing_models:
             return RunConfigGeneratorFactory._create_quick_plus_concurrency_sweep_run_config_generator(
                 command_config=command_config,
                 gpus=gpus,
@@ -127,6 +147,28 @@ def _create_brute_plus_binary_parameter_search_run_config_generator(
             model_variant_name_manager=model_variant_name_manager,
         )
 
+    @staticmethod
+    def _create_optuna_plus_concurrency_sweep_run_config_generator(
+        command_config: ConfigCommandProfile,
+        gpus: List[GPUDevice],
+        models: List[ModelProfileSpec],
+        client: TritonClient,
+        result_manager: ResultManager,
+        metrics_manager: MetricsManager,
+        model_variant_name_manager: ModelVariantNameManager,
+        search_parameters: SearchParameters,
+    ) -> ConfigGeneratorInterface:
+        return OptunaPlusConcurrencySweepRunConfigGenerator(
+            config=command_config,
+            gpus=gpus,
+            models=models,
+            client=client,
+            result_manager=result_manager,
+            metrics_manager=metrics_manager,
+            model_variant_name_manager=model_variant_name_manager,
+            search_parameters=search_parameters,
+        )
+
     @staticmethod
     def _create_quick_plus_concurrency_sweep_run_config_generator(
         command_config: ConfigCommandProfile,

model_analyzer/config/generate/search_parameters.py

@@ -49,6 +49,9 @@ 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) -> SearchParameter:
         return self._search_parameters[name]
 

model_analyzer/model_manager.py

@@ -23,6 +23,7 @@
 from model_analyzer.config.generate.run_config_generator_factory import (
     RunConfigGeneratorFactory,
 )
+from model_analyzer.config.generate.search_parameters import SearchParameters
 from model_analyzer.config.input.config_command_profile import ConfigCommandProfile
 from model_analyzer.config.input.objects.config_model_profile_spec import (
     ConfigModelProfileSpec,
@@ -59,6 +60,7 @@ def __init__(
         result_manager: ResultManager,
         state_manager: AnalyzerStateManager,
         constraint_manager: ConstraintManager,
+        search_parameters: SearchParameters,
     ):
         """
         Parameters
@@ -79,6 +81,8 @@ def __init__(
         constraint_manager: ConstraintManager
             The object that handles processing and applying
             constraints on a given measurements
+        search_parameters: SearchParameters
+            The object that handles the users configuration search parameters
         """
 
         self._config = config
@@ -89,6 +93,7 @@ def __init__(
         self._result_manager = result_manager
         self._state_manager = state_manager
         self._constraint_manager = constraint_manager
+        self._search_parameters = search_parameters
 
         if state_manager.starting_fresh_run():
             self._init_state()
@@ -131,6 +136,8 @@ def run_models(self, models: List[ConfigModelProfileSpec]) -> None:
             models=models,
             client=self._client,
             result_manager=self._result_manager,
+            metrics_manager=self._metrics_manager,
+            search_parameters=self._search_parameters,
             model_variant_name_manager=self._model_variant_name_manager,
         )
 

tests/test_model_manager.py

@@ -1160,6 +1160,7 @@ def test_ensemble_illegal_checks(self, *args):
             MagicMock(),
             state_manager,
             MagicMock(),
+            MagicMock(),
         )
 
         # Multiple model check
@@ -1210,6 +1211,7 @@ def test_ensemble_makes_quick_default(self, *args):
             MagicMock(),
             state_manager,
             MagicMock(),
+            MagicMock(),
         )
 
         models = [
@@ -1253,6 +1255,7 @@ def test_cpu_only_composing_models_error(self, *args):
             MagicMock(),
             state_manager,
             MagicMock(),
+            MagicMock(),
         )
 
         # RunConfigSearch check
@@ -1292,6 +1295,7 @@ def _test_model_manager(self, yaml_content, expected_ranges, args=None):
             MagicMock(),
             state_manager,
             MagicMock(),
+            MagicMock(),
         )
 
         model_manager.run_models([config.profile_models[0]])