Olive/olive/cli/run.py at 35d355255c7224f8a5edbd569210228367da2216 · microsoft/Olive · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from argparse import ArgumentParser

from olive.cli.base import (
    BaseOliveCLICommand,
    add_discrepancy_check_pass,
    add_hf_test_model_config,
    add_input_model_options,
    add_logging_options,
    add_telemetry_options,
    get_input_model_config,
    mark_test_output_path,
    validate_test_output_path,
)
from olive.telemetry import action


class WorkflowRunCommand(BaseOliveCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        sub_parser = parser.add_parser("run", help="Run an olive workflow")
        sub_parser.add_argument("--run-config", "--config", type=str, help="Path to json config file", required=True)
        sub_parser.add_argument(
            "--list_required_packages", help="List packages required to run the workflow", action="store_true"
        )
        sub_parser.add_argument(
            "--tempdir", type=str, help="Root directory for tempfile directories and files", required=False
        )
        sub_parser.add_argument(
            "--package-config",
            type=str,
            required=False,
            help=(
                "For advanced users. Path to optional package (json) config file with location "
                "of individual pass module implementation and corresponding dependencies. "
                "Configuration might also include user owned/proprietary/private pass implementations."
            ),
        )
        add_logging_options(sub_parser, default=None)
        add_input_model_options(
            sub_parser.add_argument_group("Model options (not required)"),
            enable_hf=True,
            enable_hf_adapter=True,
            enable_pt=True,
            enable_onnx=True,
            required=False,
        )
        add_telemetry_options(sub_parser)
        sub_parser.set_defaults(func=WorkflowRunCommand)

    @action
    def run(self):
        from copy import deepcopy
        from pathlib import Path

        from olive.common.config_utils import load_config_file
        from olive.workflows import run as olive_run

        # allow the run_config to be a dict already (for api use)
        run_config_input = self.args.run_config
        run_config = (
            deepcopy(run_config_input) if isinstance(run_config_input, dict) else load_config_file(run_config_input)
        )
        config_overrides = {}
        if input_model_config := get_input_model_config(self.args, required=False):
            print("Replacing input model config in run config")
            run_config["input_model"] = input_model_config
            config_overrides["input_model"] = input_model_config
        elif self.args.test not in (None, False):
            input_model = run_config.get("input_model")
            if not isinstance(input_model, dict) or input_model.get("type", "").lower() != "hfmodel":
                raise ValueError("--test for olive run requires a Hugging Face input_model in the run config.")
            output_path = (
                self.args.output_path or run_config.get("output_dir") or run_config.get("engine", {}).get("output_dir")
            )
            run_config["input_model"] = add_hf_test_model_config(input_model, self.args.test, output_path)

        for arg_key, rc_key in [("output_path", "output_dir"), ("log_level", "log_severity_level")]:
            if (arg_value := getattr(self.args, arg_key)) is not None:
                print(f"Replacing {rc_key} in run config with {arg_value}")
                # remove value from engine config if it exists
                run_config.get("engine", {}).pop(rc_key, None)
                # add value to run config directly
                run_config[rc_key] = arg_value
                config_overrides[rc_key] = arg_value

        recipe_telemetry_metadata = {
            "recipe_command": "WorkflowRun",
            "recipe_source": "config_dict" if isinstance(run_config_input, dict) else "config_file",
            "recipe_format": "dict"
            if isinstance(run_config_input, dict)
            else Path(run_config_input).suffix.lstrip(".").lower() or "unknown",
            "execution_mode": "list_required_packages" if self.args.list_required_packages else "run",
            "package_config_provided": bool(self.args.package_config),
        }
        if config_overrides:
            recipe_telemetry_metadata["config_overrides"] = config_overrides

        output_path = run_config.get("output_dir") or run_config.get("engine", {}).get("output_dir")
        validate_test_output_path(output_path, self.args.test)
        if self.args.test not in (None, False):
            run_config = add_discrepancy_check_pass(run_config)
        workflow_output = olive_run(
            run_config,
            list_required_packages=self.args.list_required_packages,
            tempdir=self.args.tempdir,
            package_config=self.args.package_config,
            recipe_telemetry_metadata=recipe_telemetry_metadata,
        )
        if self.args.test not in (None, False):
            mark_test_output_path(output_path)

        if self.args.list_required_packages is True:
            print("Required packages listed!")

        return workflow_output