Adding cli option for optuna search

triton-inference-server · Apr 26, 2024 · c3d3f3c · c3d3f3c
1 parent eefea55
commit c3d3f3c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/model_analyzer/config/input/config_command_profile.py b/model_analyzer/config/input/config_command_profile.py
@@ -920,14 +920,15 @@ def _add_run_search_configs(self):
             ConfigField(
                 "run_config_search_mode",
                 flags=["--run-config-search-mode"],
-                choices=["brute", "quick"],
+                choices=["brute", "quick", "optuna"],
                 field_type=ConfigPrimitive(str),
                 default_value=DEFAULT_RUN_CONFIG_SEARCH_MODE,
                 description="The search mode for Model Analyzer to find and evaluate"
                 " model configurations. 'brute' will brute force all combinations of"
                 " configuration options.  'quick' will attempt to find a near-optimal"
                 " configuration as fast as possible, but isn't guaranteed to find the"
-                " best.",
+                " best. 'optuna' is a more generalized version of quick search allowing "
+                " the user to quick search over any parameter.",
             )
         )
         self._add_config(

diff --git a/tests/test_config.py b/tests/test_config.py
@@ -1979,7 +1979,7 @@ def test_copy(self):
 
     def test_multi_model_search_mode(self):
         """
-        Test that multi-model is only run in quick
+        Test that multi-model is only run in quick/optuna
         """
         args = [
             "model-analyzer",
@@ -2011,6 +2011,13 @@ def test_multi_model_search_mode(self):
 
         self._evaluate_config(new_args, yaml_content, subcommand="profile")
 
+        # Optuna should pass
+        new_args = list(args)
+        new_args.append("--run-config-search-mode")
+        new_args.append("optuna")
+
+        self._evaluate_config(new_args, yaml_content, subcommand="profile")
+
     def test_quick_search_mode(self):
         """
         Test that only legal options are specified in quick search