Save output model to output_dir

examples/whisper/test_transcription.py

@@ -79,7 +79,7 @@ def main(raw_args=None):
     ep = config["systems"]["local_system"]["accelerators"][0]["execution_providers"][0]
 
     # load output model json
-    output_model_json_path = Path(config["output_dir"]) / "output_model" / "model_config.json"
+    output_model_json_path = Path(config["output_dir"]) / "model_config.json"
     with output_model_json_path.open() as f:
         output_model_json = json.load(f)
 

olive/cli/base.py

@@ -416,7 +416,7 @@ def save_output_model(config: Dict, output_model_dir: Union[str, Path]):
 
     This assumes a single accelerator workflow.
     """
-    run_output_path = Path(config["output_dir"]) / "output_model"
+    run_output_path = Path(config["output_dir"])
     if not any(run_output_path.rglob("model_config.json")):
         # there must be an run_output_path with at least one model_config.json
         print("Command failed. Please set the log_level to 1 for more detailed logs.")

olive/engine/engine.py

@@ -15,7 +15,7 @@
 from olive.common.config_utils import validate_config
 from olive.common.constants import DEFAULT_WORKFLOW_ID, LOCAL_INPUT_MODEL_ID
 from olive.engine.config import FAILED_CONFIG, INVALID_CONFIG, PRUNED_CONFIGS
-from olive.engine.footprint import Footprint, FootprintNodeMetric
+from olive.engine.footprint import Footprint, FootprintNode, FootprintNodeMetric, get_best_candidate_node
 from olive.engine.packaging.packaging_generator import generate_output_artifacts
 from olive.evaluator.metric import Metric
 from olive.evaluator.metric_result import MetricResult, joint_metric_key
@@ -220,12 +220,12 @@ def run(
                     output_dir/output_footprints.json: footprint of the output models
 
                     A. One pass flow:
-                        output_dir/output_model/metrics.json: evaluation results of the output model
-                        output_dir/output_model/model_config.json: output model configuration
-                        output_dir/output_model/...: output model files
+                        output_dir/metrics.json: evaluation results of the output model
+                        output_dir/model_config.json: output model configuration
+                        output_dir/...: output model files
 
                     B. Multiple pass flows:
-                        output_dir/output_model/{pass_flow}/...: Same as A but for each pass flow
+                        output_dir/{pass_flow}/...: Same as A but for each pass flow
 
                 2. Multiple accelerator specs
                     output_dir/{acclerator_spec}/...: Same as 1 but for each accelerator spec
@@ -279,6 +279,14 @@ def run(
         else:
             logger.debug("No packaging config provided, skip packaging artifacts")
 
+        # TODO(team): refactor output structure
+        # Do not change condition order. For no search, values of outputs are MetricResult
+        # Consolidate the output structure for search and no search mode
+        if outputs and self.passes and not next(iter(outputs.values())).check_empty_nodes():
+            best_node: FootprintNode = get_best_candidate_node(outputs, self.footprints)
+            self.cache.save_model(model_id=best_node.model_id, output_dir=output_dir, overwrite=True)
+            logger.info("Saved output model to %s", outputs)
+
         return outputs
 
     def run_accelerator(
@@ -387,8 +395,7 @@ def run_no_search(
                 pass_name = pass_item["name"]
                 raise ValueError(f"Pass {pass_name} has search space but search strategy is None")
 
-        # output models will be saved in output_dir/output_model
-        output_model_dir = Path(output_dir) / "output_model"
+        output_model_dir = Path(output_dir)
 
         output_model_ids = []
         for pass_flow in self.pass_flows:
@@ -422,9 +429,6 @@ def run_no_search(
                     json.dump(signal.to_json(), f, indent=4)
                 logger.info("Saved evaluation results of output model to %s", results_path)
 
-            self.cache.save_model(model_id=model_ids[-1], output_dir=flow_output_dir, overwrite=True)
-            logger.info("Saved output model to %s", flow_output_dir)
-
             output_model_ids.append(model_ids[-1])
 
         output_footprints = self.footprints[accelerator_spec].create_footprints_by_model_ids(output_model_ids)

olive/engine/footprint.py

@@ -6,11 +6,15 @@
 import logging
 from collections import OrderedDict, defaultdict
 from copy import deepcopy
-from typing import DefaultDict, Dict, List, NamedTuple, Optional
+from typing import TYPE_CHECKING, DefaultDict, Dict, List, NamedTuple, Optional
 
 from olive.common.config_utils import ConfigBase, config_json_dumps, config_json_loads
 from olive.evaluator.metric_result import MetricResult
 
+if TYPE_CHECKING:
+    from olive.hardware import AcceleratorSpec
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -159,6 +163,9 @@ def trace_back_run_history(self, model_id) -> Dict[str, Dict]:
             model_id = self.nodes[model_id].parent_model_id
         return rls
 
+    def check_empty_nodes(self):
+        return self.nodes is None or len(self.nodes) == 0
+
     def to_df(self):
         # to pandas.DataFrame
         raise NotImplementedError
@@ -422,3 +429,46 @@ def _plot_pareto_frontier(self, ranks=None, save_path=None, is_show=True, save_f
 
         if is_show:
             fig.show()
+
+
+def get_best_candidate_node(
+    pf_footprints: Dict["AcceleratorSpec", Footprint], footprints: Dict["AcceleratorSpec", Footprint]
+):
+    """Select the best candidate node from the pareto frontier footprints.
+
+    This function evaluates nodes from the given pareto frontier footprints and selects the top-ranked node
+    based on specified objective metrics. It compares nodes from two dictionaries of footprints and
+    ranks them according to their metrics.
+
+    Args:
+        pf_footprints (Dict["AcceleratorSpec", Footprint]): A dictionary mapping accelerator specifications
+            to their corresponding pareto frontier footprints, which contain nodes and their metrics.
+        footprints (Dict["AcceleratorSpec", Footprint"]): A dictionary mapping accelerator specifications
+            to their corresponding footprints, which contain nodes and their metrics.
+
+    Returns:
+        Node: The top-ranked node based on the specified objective metrics.
+
+    """
+    objective_dict = next(iter(pf_footprints.values())).objective_dict
+    top_nodes = []
+    for accelerator_spec, pf_footprint in pf_footprints.items():
+        footprint = footprints[accelerator_spec]
+        if pf_footprint.nodes and footprint.nodes:
+            top_nodes.append(next(iter(pf_footprint.get_top_ranked_nodes(1))))
+    return next(
+        iter(
+            sorted(
+                top_nodes,
+                key=lambda x: tuple(
+                    (
+                        x.metrics.value[metric].value
+                        if x.metrics.cmp_direction[metric] == 1
+                        else -x.metrics.value[metric].value
+                    )
+                    for metric in objective_dict
+                ),
+                reverse=True,
+            )
+        )
+    )

olive/engine/packaging/packaging_generator.py

@@ -18,6 +18,7 @@
 
 from olive.common.constants import OS
 from olive.common.utils import retry_func, run_subprocess
+from olive.engine.footprint import get_best_candidate_node
 from olive.engine.packaging.packaging_config import (
     AzureMLDeploymentPackagingConfig,
     DockerfilePackagingConfig,
@@ -68,7 +69,7 @@ def _package_dockerfile(
     config: DockerfilePackagingConfig = packaging_config.config
     logger.info("Packaging output models to Dockerfile")
     base_image = config.base_image
-    best_node = _get_best_candidate_node(pf_footprints, footprints)
+    best_node = get_best_candidate_node(pf_footprints, footprints)
 
     docker_context_path = "docker_content"
     content_path = output_dir / docker_context_path
@@ -133,7 +134,7 @@ def _package_azureml_deployment(
 
     try:
         # Get best model from footprint
-        best_node = _get_best_candidate_node(pf_footprints, footprints)
+        best_node = get_best_candidate_node(pf_footprints, footprints)
 
         with tempfile.TemporaryDirectory() as temp_dir:
             tempdir = Path(temp_dir)
@@ -303,33 +304,6 @@ def _package_azureml_deployment(
         raise
 
 
-def _get_best_candidate_node(
-    pf_footprints: Dict["AcceleratorSpec", "Footprint"], footprints: Dict["AcceleratorSpec", "Footprint"]
-):
-    objective_dict = next(iter(pf_footprints.values())).objective_dict
-    top_nodes = []
-    for accelerator_spec, pf_footprint in pf_footprints.items():
-        footprint = footprints[accelerator_spec]
-        if pf_footprint.nodes and footprint.nodes:
-            top_nodes.append(next(iter(pf_footprint.get_top_ranked_nodes(1))))
-    return next(
-        iter(
-            sorted(
-                top_nodes,
-                key=lambda x: tuple(
-                    (
-                        x.metrics.value[metric].value
-                        if x.metrics.cmp_direction[metric] == 1
-                        else -x.metrics.value[metric].value
-                    )
-                    for metric in objective_dict
-                ),
-                reverse=True,
-            )
-        )
-    )
-
-
 def _is_generative_model(config: Dict[str, Any]) -> bool:
     model_attributes = config.get("model_attributes") or {}
     return model_attributes.get("generative", False)
@@ -353,7 +327,7 @@ def _package_candidate_models(
         tempdir = Path(temp_dir)
 
         if packaging_type == PackagingType.Zipfile:
-            best_node: FootprintNode = _get_best_candidate_node(pf_footprints, footprints)
+            best_node: FootprintNode = get_best_candidate_node(pf_footprints, footprints)
             is_generative = _is_generative_model(best_node.model_config["config"])
 
             if packaging_config.include_runtime_packages:

test/unit_test/cli/test_cli.py

@@ -126,8 +126,7 @@ def test_finetune_command(_, mock_tempdir, mock_run, tmp_path):
 
     # setup
     mock_tempdir.return_value = tmpdir.resolve()
-    workflow_output_dir = tmpdir / "output_model"
-    workflow_output_dir.mkdir(parents=True)
+    workflow_output_dir = tmpdir
     dummy_output = workflow_output_dir / "model_config.json"
     with open(dummy_output, "w") as f:
         json.dump({"dummy": "output"}, f)
@@ -196,8 +195,7 @@ def test_capture_onnx_command(_, mock_tempdir, mock_run, use_model_builder, tmp_
 
     # setup
     mock_tempdir.return_value = tmpdir.resolve()
-    workflow_output_dir = tmpdir / "output_model"
-    workflow_output_dir.mkdir(parents=True)
+    workflow_output_dir = tmpdir
     dummy_output = workflow_output_dir / "model_config.json"
     with open(dummy_output, "w") as f:
         json.dump({"config": {"inference_settings": {"dummy-key": "dummy-value"}}}, f)
@@ -284,9 +282,7 @@ def test_quantize_command(mock_repo_exists, mock_tempdir, mock_run, algorithm_na
     mock_tempdir.return_value = tmpdir.resolve()
     mock_run.return_value = {}
 
-    workflow_output_dir = tmpdir / "output_model" / algorithm_name
-    workflow_output_dir.mkdir(parents=True)
-    model_config_path = workflow_output_dir / "model_config.json"
+    model_config_path = tmpdir / "model_config.json"
     with model_config_path.open("w") as f:
         f.write("{}")
 
@@ -309,7 +305,7 @@ def test_quantize_command(mock_repo_exists, mock_tempdir, mock_run, algorithm_na
 
     config = mock_run.call_args[0][0]
     assert config["input_model"]["model_path"] == "dummy_model"
-    assert {el.name for el in output_dir.iterdir()} == {algorithm_name}
+    assert {el.name for el in output_dir.iterdir()} == {"model_config.json"}
 
 
 # TODO(anyone): Add tests for ManageAMLComputeCommand

test/unit_test/engine/test_engine.py

@@ -285,7 +285,7 @@ def test_run_no_search(self, mock_local_system_init, tmp_path):
         # output model to output_dir
         output_dir = tmp_path / "output_dir"
         expected_metrics = MetricResult.parse_obj(metric_result_dict)
-        expected_saved_model_config = get_onnx_model_config(model_path=output_dir / "output_model" / "model.onnx")
+        expected_saved_model_config = get_onnx_model_config(model_path=output_dir / "model.onnx")
 
         # execute
         footprint = engine.run(
@@ -299,17 +299,56 @@ def test_run_no_search(self, mock_local_system_init, tmp_path):
         assert output_node.model_config == onnx_model_config
         assert expected_metrics == output_node.metrics.value
 
-        output_model_dir = output_dir / "output_model"
-        model_json_path = output_model_dir / "model_config.json"
+        model_json_path = output_dir / "model_config.json"
         assert model_json_path.is_file()
         with model_json_path.open() as f:
             assert json.load(f) == expected_saved_model_config.to_json()
 
-        result_json_path = output_model_dir / "metrics.json"
+        result_json_path = output_dir / "metrics.json"
         assert result_json_path.is_file()
         with result_json_path.open() as f:
             assert json.load(f) == expected_metrics.__root__
 
+    @pytest.mark.parametrize(
+        "search_strategy",
+        [
+            {
+                "execution_order": "joint",
+                "search_algorithm": "random",
+            },
+            None,
+        ],
+    )
+    def test_run_output_model(self, search_strategy, tmp_path):
+        # setup
+        model_config = get_pytorch_model_config()
+        metric = get_accuracy_metric(AccuracySubType.ACCURACY_SCORE)
+        evaluator_config = OliveEvaluatorConfig(metrics=[metric])
+        options = {
+            "cache_config": {
+                "cache_dir": tmp_path,
+                "clean_cache": True,
+                "clean_evaluation_cache": True,
+            },
+            "search_strategy": search_strategy,
+            "evaluator": evaluator_config,
+        }
+        engine = Engine(**options)
+        _, p_config = get_onnxconversion_pass(ignore_pass_config=False, target_opset=13)
+        engine.register(OnnxConversion, config=p_config)
+        # output model to output_dir
+        output_dir = tmp_path / "output_dir"
+
+        # execute
+        engine.run(
+            model_config,
+            [DEFAULT_CPU_ACCELERATOR],
+            output_dir=output_dir,
+        )
+
+        # assert
+        assert Path(output_dir / "model.onnx").is_file()
+
     def test_pass_exception(self, caplog, tmpdir):
         # Need explicitly set the propagate to allow the message to be logged into caplog
         # setup