diff --git a/olive/cli/base.py b/olive/cli/base.py index a64a9bf473..7ca05fe36a 100644 --- a/olive/cli/base.py +++ b/olive/cli/base.py @@ -617,10 +617,11 @@ def update_input_model_options(args, config): config["input_model"] = get_input_model_config(args) -def add_logging_options(sub_parser: ArgumentParser, default: int = 3): +def add_logging_options(sub_parser: ArgumentParser, default: int | None = 3): """Add logging options to the sub_parser.""" sub_parser.add_argument( "--log_level", + "--log-level", type=int, default=default, help=f"Logging level. Default is {default}. level 0: DEBUG, 1: INFO, 2: WARNING, 3: ERROR, 4: CRITICAL", @@ -632,12 +633,14 @@ def add_save_config_file_options(sub_parser: ArgumentParser): """Add save config file options to the sub_parser.""" sub_parser.add_argument( "--save_config_file", + "--save-config-file", action="store_true", help="Generate and save the config file for the command.", ) sub_parser.add_argument( "--dry_run", + "--dry-run", action="store_true", help="Enable dry run mode. This will not perform any actual optimization but will validate the configuration.", ) @@ -672,6 +675,7 @@ def add_input_model_options( model_group.add_argument( "-m", "--model_name_or_path", + "--model-name-or-path", type=str, # only pytorch model doesn't require model_name_or_path required=required and not enable_pt, @@ -689,7 +693,10 @@ def add_input_model_options( help=f"Task for which the huggingface model is used. Default task is {DEFAULT_HF_TASK}.", ) model_group.add_argument( - "--trust_remote_code", action="store_true", help="Trust remote code when loading a huggingface model." + "--trust_remote_code", + "--trust-remote-code", + action="store_true", + help="Trust remote code when loading a huggingface model.", ) model_group.add_argument( "--test", @@ -701,6 +708,7 @@ def add_input_model_options( ) model_group.add_argument( "--test_metrics", + "--test-metrics", type=_parse_test_metrics, nargs="+", help=( @@ -715,6 +723,7 @@ def add_input_model_options( ) model_group.add_argument( "--test_llama_path", + "--test-llama-path", type=str, default=None, help=( @@ -729,12 +738,14 @@ def add_input_model_options( model_group.add_argument( "-a", "--adapter_path", + "--adapter-path", type=str, help="Path to the adapters weights saved after peft fine-tuning. Local folder or huggingface id.", ) if enable_diffusers: model_group.add_argument( "--model_variant", + "--model-variant", type=DiffusersModelVariant, choices=[ DiffusersModelVariant.AUTO, @@ -752,17 +763,20 @@ def add_input_model_options( model_group.add_argument( "-a", "--adapter_path", + "--adapter-path", type=str, help="Path to the LoRA adapter weights. Local folder or huggingface id.", ) if enable_pt: model_group.add_argument( "--model_script", + "--model-script", type=str, help="The script file containing the model definition. Required for the local PyTorch model.", ) model_group.add_argument( "--script_dir", + "--script-dir", type=str, help=( "The directory containing the local PyTorch model script file." @@ -773,6 +787,7 @@ def add_input_model_options( model_group.add_argument( "-o", "--output_path", + "--output-path", type=output_path_type if directory_output else str, required=required and default_output_path is None, default=default_output_path, @@ -797,18 +812,19 @@ def add_dataset_options(sub_parser, required=True, include_train=True, include_e dataset_group.add_argument( "-d", "--data_name", + "--data-name", type=str, required=required, help="The dataset name.", ) if include_train: - dataset_group.add_argument("--train_subset", type=str, help="The subset to use for training.") - dataset_group.add_argument("--train_split", type=str, default="train", help="The split to use for training.") + dataset_group.add_argument("--train_subset", "--train-subset", type=str, help="The subset to use for training.") + dataset_group.add_argument("--train_split", "--train-split", type=str, default="train", help="The split to use for training.") if include_eval: - dataset_group.add_argument("--eval_subset", type=str, help="The subset to use for evaluation.") - dataset_group.add_argument("--eval_split", default="", help="The dataset split to evaluate on.") + dataset_group.add_argument("--eval_subset", "--eval-subset", type=str, help="The subset to use for evaluation.") + dataset_group.add_argument("--eval_split", "--eval-split", type=str, default="", help="The dataset split to evaluate on.") if not (include_train and include_eval): dataset_group.add_argument("--subset", type=str, help="The subset of the dataset to use.") @@ -816,48 +832,55 @@ def add_dataset_options(sub_parser, required=True, include_train=True, include_e # TODO(jambayk): currently only supports single file or list of files, support mapping dataset_group.add_argument( - "--data_files", type=str, help="The dataset files. If multiple files, separate by comma." + "--data_files", "--data-files", type=str, help="The dataset files. If multiple files, separate by comma." ) text_group = dataset_group.add_mutually_exclusive_group(required=False) text_group.add_argument( "--text_field", + "--text-field", type=str, help="The text field to use for fine-tuning.", ) text_group.add_argument( "--text_template", + "--text-template", # using special string type to allow for escaped characters like \n type=unescaped_str, help=r"Template to generate text field from. E.g. '### Question: {prompt} \n### Answer: {response}'", ) - text_group.add_argument("--use_chat_template", action="store_true", help="Use chat template for text field.") + text_group.add_argument("--use_chat_template", "--use-chat-template", action="store_true", help="Use chat template for text field.") dataset_group.add_argument( "--max_seq_len", + "--max-seq-len", type=int, default=1024, help="Maximum sequence length for the data.", ) dataset_group.add_argument( "--add_special_tokens", + "--add-special-tokens", type=bool, default=False, help="Whether to add special tokens during preprocessing.", ) dataset_group.add_argument( "--max_samples", + "--max-samples", type=int, default=256, help="Maximum samples to select from the dataset.", ) dataset_group.add_argument( "--batch_size", + "--batch-size", type=int, default=1, help="Batch size.", ) dataset_group.add_argument( "--input_cols", + "--input-cols", type=str, nargs="+", help=( @@ -902,11 +925,13 @@ def add_shared_cache_options(sub_parser: ArgumentParser): shared_cache_group = sub_parser shared_cache_group.add_argument( "--account_name", + "--account-name", type=str, help="Azure storage account name for shared cache.", ) shared_cache_group.add_argument( "--container_name", + "--container-name", type=str, help="Azure storage container name for shared cache.", ) @@ -945,6 +970,7 @@ def add_accelerator_options(sub_parser, single_provider: bool = True): else: accelerator_group.add_argument( "--providers_list", + "--providers-list", type=str, nargs="*", choices=execution_providers, @@ -975,46 +1001,12 @@ def update_accelerator_options(args, config, single_provider: bool = True): set_nested_dict_value(config, k, v) -def add_search_options(sub_parser: ArgumentParser): - search_strategy_group = sub_parser - search_strategy_group.add_argument( - "--enable_search", - type=str, - default=None, - const="sequential", - nargs="?", - choices=["random", "sequential", "tpe"], - help=( - "Enable search to produce optimal model for the given criteria. " - "Optionally provide sampler from available choices. " - "By default, uses sequential sampler." - ), - ) - search_strategy_group.add_argument("--seed", type=int, default=0, help="Random seed for search sampler") - - def add_telemetry_options(sub_parser: ArgumentParser): """Add telemetry options to the sub_parser.""" - sub_parser.add_argument("--disable_telemetry", action="store_true", help="Disable telemetry for this command.") - return sub_parser - - -def update_search_options(args, config): - to_replace = [] - to_replace.extend( - [ - ( - "search_strategy", - { - "execution_order": "joint", - "sampler": args.enable_search, - "seed": args.seed, - }, - ), - ] + sub_parser.add_argument( + "--disable_telemetry", + "--disable-telemetry", + action="store_true", + help="Disable telemetry for this command.", ) - - for keys, value in to_replace: - if value is None: - continue - set_nested_dict_value(config, keys, value) + return sub_parser diff --git a/olive/cli/launcher.py b/olive/cli/launcher.py index 55e6ffdeb4..da670d0db5 100644 --- a/olive/cli/launcher.py +++ b/olive/cli/launcher.py @@ -22,6 +22,7 @@ from olive.cli.quantize import QuantizeCommand from olive.cli.run import WorkflowRunCommand from olive.cli.run_pass import RunPassCommand +from olive.cli.search import WorkflowSearchCommand from olive.cli.session_params_tuning import SessionParamsTuningCommand from olive.cli.shared_cache import SharedCacheCommand from olive.telemetry import Telemetry @@ -41,6 +42,7 @@ def get_cli_parser(called_as_console_script: bool = True) -> ArgumentParser: # NOTE: The order of the commands is to organize the documentation better. InitCommand.register_subcommand(commands_parser) WorkflowRunCommand.register_subcommand(commands_parser) + WorkflowSearchCommand.register_subcommand(commands_parser) RunPassCommand.register_subcommand(commands_parser) AutoOptCommand.register_subcommand(commands_parser) OptimizeCommand.register_subcommand(commands_parser) diff --git a/olive/cli/search.py b/olive/cli/search.py new file mode 100644 index 0000000000..d5d2e174f0 --- /dev/null +++ b/olive/cli/search.py @@ -0,0 +1,545 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import json +import re +from argparse import ArgumentParser, ArgumentTypeError, Namespace +from collections import OrderedDict +from pathlib import Path +from typing import Any + +from olive.cli.base import ( + BaseOliveCLICommand, + add_dataset_options, + add_logging_options, + add_save_config_file_options, + add_telemetry_options, + update_dataset_options, +) +from olive.telemetry import action + +# Endpoint used to submit remote Olive-as-a-Service jobs. +OLIVE_AAS_RUN_URL = "https://oliveaas.azure-api.net/v1/run" + + +def _hex_string(value: str) -> str: + """Validate that the value is a 32-character hexadecimal string.""" + if not value or not re.fullmatch(r"[0-9a-fA-F]{32}", value): + raise ArgumentTypeError("must be a 32-character hexadecimal string") + return value + + +class WorkflowSearchCommand(BaseOliveCLICommand): + @staticmethod + def register_subcommand(parser: ArgumentParser): + sub_parser = parser.add_parser( + "search", + help="Run an olive workflow with search strategy configuration", + ) + + # ------------------------------------------------------------------ # + # Mode selection: mutually exclusive — config file OR build from args # + # ------------------------------------------------------------------ # + mode_group = sub_parser.add_mutually_exclusive_group(required=True) + mode_group.add_argument( + "--run-config", + "--config", + type=str, + dest="run_config", + default=None, + help="Path to a JSON workflow config file. Mutually exclusive with --model_name_or_path.", + ) + mode_group.add_argument( + "-m", + "--model_name_or_path", + type=str, + dest="model_name_or_path", + default=None, + help=( + "Model to search over (HuggingFace id, local path, or AzureML URI). " + "Mutually exclusive with --run-config." + ), + ) + + # ------------------------------------------------------------------ # + # Common options (both modes) # + # ------------------------------------------------------------------ # + 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." + ), + ) + sub_parser.add_argument( + "--az_apim_subscription_key", + type=_hex_string, + default=None, + help=( + "Azure APIM subscription key (hexadecimal string). When provided, each generated config " + "is submitted as a remote Olive-as-a-Service job instead of running locally, and the " + "returned operation id (op_xxxxxxxxx) is printed." + ), + ) + + # ------------------------------------------------------------------ # + # Search strategy options (both modes) # + # ------------------------------------------------------------------ # + search_group = sub_parser.add_argument_group("Search strategy options") + search_group.add_argument( + "--execution-order", + type=str, + choices=["joint", "pass-by-pass"], + default=None, + help=( + "Order in which passes are searched. " + "'joint' searches all passes together; 'pass-by-pass' searches each pass sequentially." + ), + ) + search_group.add_argument( + "--sampler", + type=str, + choices=["sequential", "random", "tpe"], + default=None, + help="Sampler algorithm to use for exploring the search space.", + ) + search_group.add_argument( + "--max-iter", + type=int, + default=None, + help="Maximum number of search iterations. Only applies to joint execution order.", + ) + search_group.add_argument( + "--max-time", + type=int, + default=None, + help="Maximum search time in seconds. Only applies to joint execution order.", + ) + search_group.add_argument( + "--output-model-num", + type=int, + default=None, + help="Number of output models to produce from the search.", + ) + search_group.add_argument( + "--stop-when-goals-met", + action="store_true", + default=False, + help="Stop searching once all metric goals are met. Only applies to joint execution order.", + ) + search_group.add_argument( + "--include-pass-params", + action="store_true", + default=None, + help="Include individual pass parameters in the search space.", + ) + + # ------------------------------------------------------------------ # + # Evaluation options (lm-eval based, used to score searched models) # + # ------------------------------------------------------------------ # + eval_group = sub_parser.add_argument_group( + "Evaluation options", + description=( + "lm-eval based evaluation used to score models produced during the search. " + "Provide --eval-tasks to enable evaluation when building a config from arguments." + ), + ) + eval_group.add_argument( + "--eval-tasks", + type=str, + nargs="*", + default=None, + help="List of lm-eval tasks to evaluate searched models on (e.g. hellaswag winogrande).", + ) + eval_group.add_argument( + "--eval-batch-size", + type=int, + default=1, + help="Batch size for evaluation. Default is 1.", + ) + eval_group.add_argument( + "--eval-max-length", + type=int, + default=1024, + help="Maximum length of input + output for evaluation. Default is 1024.", + ) + eval_group.add_argument( + "--eval-limit", + type=float, + default=1, + help="Number (or percentage of dataset) of samples to use for evaluation. Default is 1.", + ) + eval_group.add_argument( + "--eval-backend", + type=str, + default="auto", + choices=["auto", "ort", "ortgenai"], + help=( + "Backend for lm-eval model evaluation. 'ort' and 'ortgenai' require ONNX input; " + "'ortgenai' additionally requires GenAI-packaged model assets (e.g., genai_config.json). " + "'auto' infers backend from model type." + ), + ) + + # ------------------------------------------------------------------ # + # Build-mode options (only used when --model_name_or_path is given) # + # ------------------------------------------------------------------ # + build_group = sub_parser.add_argument_group( + "Build-mode options", + description=( + "These options are used to construct a run config on-the-fly " + "when --model_name_or_path is provided instead of --run-config." + ), + ) + build_group.add_argument( + "--passes", + type=str, + nargs="+", + metavar="PASSES", + default=None, + help=( + "Ordered pass groups to search over. Each group generates a separate config file named " + "'config_%%02d.json'. Wrap a multi-pass group in square brackets, e.g. " + "--passes [OnnxConversion, OrtTransformersOptimization] [OnnxConversion, OnnxQuantization]. " + "Unbracketed entries are each treated as their own single-pass group, e.g. " + "--passes OnnxConversion OnnxQuantization generates two configs. " + "Each value must be a valid Olive pass type name and passes use their default configuration." + ), + ) + build_group.add_argument( + "--device", + type=str, + default="cpu", + choices=["cpu", "gpu", "npu"], + help="Target device for the search. Default is cpu.", + ) + build_group.add_argument( + "--provider", + type=str, + default="CPUExecutionProvider", + choices=[ + "CPUExecutionProvider", + "CUDAExecutionProvider", + "QNNExecutionProvider", + "VitisAIExecutionProvider", + "OpenVINOExecutionProvider", + "WebGpuExecutionProvider", + "NvTensorRTRTXExecutionProvider", + ], + help="Execution provider to target during search. Default is CPUExecutionProvider.", + ) + build_group.add_argument( + "-t", + "--task", + type=str, + default=None, + help="HuggingFace task for the model (e.g. text-generation).", + ) + build_group.add_argument( + "--trust_remote_code", + action="store_true", + help="Trust remote code when loading a HuggingFace model.", + ) + build_group.add_argument( + "-o", + "--output_path", + type=str, + default="search-output", + help="Directory to save the search output. Default is 'search-output'.", + ) + + # ------------------------------------------------------------------ # + # Dataset options (optional, used by evaluation/data-dependent passes) # + # ------------------------------------------------------------------ # + add_dataset_options(sub_parser, required=False, include_train=True, include_eval=True) + + add_logging_options(sub_parser, default=None) + add_save_config_file_options(sub_parser) + add_telemetry_options(sub_parser) + sub_parser.set_defaults(func=WorkflowSearchCommand) + + # ---------------------------------------------------------------------- # + # Helpers # + # ---------------------------------------------------------------------- # + + @staticmethod + def _build_search_strategy_overrides(args: Namespace) -> dict: + overrides = {} + if args.execution_order is not None: + overrides["execution_order"] = args.execution_order + if args.sampler is not None: + overrides["sampler"] = args.sampler + if args.max_iter is not None: + overrides["max_iter"] = args.max_iter + if args.max_time is not None: + overrides["max_time"] = args.max_time + if args.output_model_num is not None: + overrides["output_model_num"] = args.output_model_num + if args.stop_when_goals_met: + overrides["stop_when_goals_met"] = True + if args.include_pass_params is not None: + overrides["include_pass_params"] = args.include_pass_params + return overrides + + @staticmethod + def _apply_search_strategy(run_config: dict, overrides: dict) -> None: + if not overrides: + return + engine_config = run_config.setdefault("engine", {}) + existing = engine_config.get("search_strategy") or {} + if not isinstance(existing, dict): + existing = {} + existing.update(overrides) + engine_config["search_strategy"] = existing + print(f"Applying search strategy overrides: {overrides}") + + def _build_run_config(self, passes: list[str], output_dir: str) -> dict[str, Any]: + """Construct a run config dict for a single ordered group of passes.""" + args = self.args + + # Build passes section using default configuration for each pass + passes_config: dict[str, Any] = OrderedDict() + for pass_type in passes: + passes_config[pass_type.lower()] = {"type": pass_type} + + # Input model + input_model: dict[str, Any] = {"type": "HfModel", "model_path": args.model_name_or_path} + if getattr(args, "task", None): + input_model["task"] = args.task + if getattr(args, "trust_remote_code", False): + input_model["load_kwargs"] = {"trust_remote_code": True} + + # Accelerator / system + run_config: dict[str, Any] = { + "input_model": input_model, + "passes": passes_config, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": args.device, + "execution_providers": [args.provider], + } + ], + } + }, + "target": "local_system", + "output_dir": output_dir, + "no_artifacts": True, + } + + # Add lm-eval based evaluator when evaluation tasks are provided. + self._add_evaluator(run_config) + + # Add data config when a dataset is provided. + self._add_data_config(run_config) + + if args.log_level is not None: + run_config["log_severity_level"] = args.log_level + + return run_config + + def _add_data_config(self, run_config: dict[str, Any]) -> None: + """Add a data config to the run config when --data_name is provided.""" + if not getattr(self.args, "data_name", None): + return + + run_config["data_configs"] = [ + { + "name": "default_data_config", + "type": "HuggingfaceContainer", + "load_dataset_config": {}, + "pre_process_data_config": {}, + "dataloader_config": {}, + "post_process_data_config": {}, + } + ] + update_dataset_options(self.args, run_config) + + def _add_evaluator(self, run_config: dict[str, Any]) -> None: + """Add an lm-eval based evaluator to the run config when --eval-tasks is provided.""" + args = self.args + if not getattr(args, "eval_tasks", None): + return + + run_config["evaluators"] = { + "evaluator": { + "type": "LMEvaluator", + "tasks": args.eval_tasks, + "batch_size": args.eval_batch_size, + "max_length": args.eval_max_length, + "device": args.device, + "limit": args.eval_limit, + } + } + if args.eval_backend != "auto": + run_config["evaluators"]["evaluator"]["model_class"] = args.eval_backend + run_config["evaluator"] = "evaluator" + run_config.setdefault("host", "local_system") + + @staticmethod + def _parse_pass_groups(raw: list[str]) -> list[list[str]]: + """Parse the --passes tokens into a list of ordered pass groups. + + Bracketed and unbracketed entries can be freely mixed, and order is preserved: + - A bracketed segment is a single multi-pass group: + ``[A, B, C]`` -> ``["A", "B", "C"]`` + - Each unbracketed entry is its own single-pass group: + ``A B`` -> ``["A"], ["B"]`` + So ``[A, B] C [D, E]`` -> ``[["A", "B"], ["C"], ["D", "E"]]``. + Commas and spaces are both accepted as separators within a bracketed group. + """ + joined = " ".join(raw) + + groups: list[list[str]] = [] + pos = 0 + for match in re.finditer(r"\[([^\]]*)\]", joined): + # Unbracketed tokens appearing before this bracketed segment each form their own group. + groups.extend([token] for token in joined[pos : match.start()].replace(",", " ").split()) + # The bracketed segment forms a single multi-pass group. + passes = [p.strip() for p in match.group(1).replace(",", " ").split() if p.strip()] + if passes: + groups.append(passes) + pos = match.end() + + # Any trailing unbracketed tokens after the last bracketed segment. + groups.extend([token] for token in joined[pos:].replace(",", " ").split()) + + if not groups: + raise ValueError("No valid pass groups parsed from --passes.") + return groups + + @staticmethod + def _submit_remote_job(run_config: dict[str, Any], subscription_key: str) -> str: + """Submit a run config as a remote Olive-as-a-Service job and return the operation id.""" + import requests + + response = requests.post( + OLIVE_AAS_RUN_URL, + headers={ + "Content-Type": "application/json", + "Ocp-Apim-Subscription-Key": subscription_key, + }, + json={"run_config": run_config}, + timeout=60, + ) + response.raise_for_status() + + # The service returns the operation id, either as a bare string or wrapped in JSON. + try: + payload = response.json() + except ValueError: + payload = response.text + op_id = payload if isinstance(payload, str) else (payload.get("op_id") or payload.get("id")) + if not op_id: + raise RuntimeError(f"Unexpected response from remote job submission: {response.text!r}") + return op_id.strip().strip('"') + + # ---------------------------------------------------------------------- # + # Entry point # + # ---------------------------------------------------------------------- # + + @action + def run(self): + from olive.workflows import run as olive_run + + if self.args.run_config is not None: + # ---- Config-file mode ---------------------------------------- # + from olive.common.config_utils import load_config_file + + run_config = self.args.run_config + if not isinstance(run_config, dict): + run_config = load_config_file(run_config) + + for arg_key, rc_key in [("log_level", "log_severity_level")]: + if (arg_value := getattr(self.args, arg_key, None)) is not None: + print(f"Replacing {rc_key} in run config with {arg_value}") + run_config.get("engine", {}).pop(rc_key, None) + run_config[rc_key] = arg_value + + # Apply search strategy overrides + overrides = self._build_search_strategy_overrides(self.args) + self._apply_search_strategy(run_config, overrides) + + if self.args.save_config_file or self.args.dry_run: + self._save_config_file(run_config) + + if self.args.dry_run: + print("Dry run mode enabled. Configuration file is generated but no workflow is executed.") + return + + if self.args.az_apim_subscription_key: + op_id = self._submit_remote_job(run_config, self.args.az_apim_subscription_key) + print(f"Remote job submitted. Operation id: {op_id}") + return + + olive_run( + run_config, + list_required_packages=self.args.list_required_packages, + tempdir=self.args.tempdir, + package_config=self.args.package_config, + ) + + if self.args.list_required_packages is True: + print("Required packages listed!") + return + + # ---- Build mode -------------------------------------------------- # + if not self.args.passes: + raise ValueError( + "At least one pass group must be specified with --passes when using build mode (--model_name_or_path)." + ) + + groups = self._parse_pass_groups(self.args.passes) + overrides = self._build_search_strategy_overrides(self.args) + + output_root = Path(self.args.output_path) + output_root.mkdir(parents=True, exist_ok=True) + + # Generate one config per ordered pass group, saved as config_%02d.json. + run_configs: list[tuple[str, dict[str, Any]]] = [] + for idx, group in enumerate(groups): + config_name = f"config_{idx:02d}" + run_config = self._build_run_config(group, (output_root / config_name).as_posix()) + self._apply_search_strategy(run_config, overrides) + + config_path = output_root / f"{config_name}.json" + with open(config_path, "w") as f: + json.dump(run_config, f, indent=4) + print(f"Config file saved at {config_path} for pass group {group}") + run_configs.append((config_name, run_config)) + + if self.args.dry_run: + print( + f"Dry run mode enabled. {len(run_configs)} configuration file(s) generated but no workflow is executed." + ) + return + + if self.args.az_apim_subscription_key: + for config_name, run_config in run_configs: + print(f"Submitting remote job for {config_name} ...") + op_id = self._submit_remote_job(run_config, self.args.az_apim_subscription_key) + print(f"Remote job submitted for {config_name}. Operation id: {op_id}") + return + + for config_name, run_config in run_configs: + print(f"Running search workflow for {config_name} ...") + olive_run( + run_config, + list_required_packages=self.args.list_required_packages, + tempdir=self.args.tempdir, + package_config=self.args.package_config, + ) + + if self.args.list_required_packages is True: + print("Required packages listed!") diff --git a/test/cli/test_search.py b/test/cli/test_search.py new file mode 100644 index 0000000000..fa79ac059f --- /dev/null +++ b/test/cli/test_search.py @@ -0,0 +1,600 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import json +import subprocess +import sys +from argparse import ArgumentParser, ArgumentTypeError +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from olive.cli.launcher import main as cli_main +from olive.cli.search import OLIVE_AAS_RUN_URL, WorkflowSearchCommand, _hex_string + +# pylint: disable=W0212 + + +def _build_parser(): + parser = ArgumentParser() + sub_parsers = parser.add_subparsers() + WorkflowSearchCommand.register_subcommand(sub_parsers) + return parser + + +# --------------------------------------------------------------------------- # +# Help / registration # +# --------------------------------------------------------------------------- # + + +def test_search_command_help(): + # setup + command_args = [sys.executable, "-m", "olive", "search", "--help"] + + # execute + out = subprocess.run(command_args, check=True, capture_output=True) + + # assert + help_text = out.stdout.decode("utf-8") + assert "usage:" in help_text + assert "search" in help_text + assert "--run-config" in help_text + assert "--passes" in help_text + + +# --------------------------------------------------------------------------- # +# _hex_string validation # +# --------------------------------------------------------------------------- # + + +def test_hex_string_accepts_valid_32_char_hex(): + value = "0123456789abcdef0123456789ABCDEF" + assert _hex_string(value) == value + + +@pytest.mark.parametrize( + "value", + [ + "", + "0123456789abcdef", # too short + "0123456789abcdef0123456789abcdef0", # too long + "0123456789abcdef0123456789abcdeg", # invalid char 'g' + "0123456789abcdef 123456789abcdef", # space + ], +) +def test_hex_string_rejects_invalid_values(value): + with pytest.raises(ArgumentTypeError): + _hex_string(value) + + +# --------------------------------------------------------------------------- # +# Argument parsing # +# --------------------------------------------------------------------------- # + + +def test_search_mutually_exclusive_mode_required(): + parser = _build_parser() + with pytest.raises(SystemExit): + parser.parse_args(["search"]) + + +def test_search_run_config_and_model_are_mutually_exclusive(): + parser = _build_parser() + with pytest.raises(SystemExit): + parser.parse_args(["search", "--run-config", "config.json", "-m", "some-model"]) + + +def test_search_parses_build_mode_args(): + parser = _build_parser() + args = parser.parse_args( + [ + "search", + "-m", + "microsoft/phi-2", + "--passes", + "OnnxConversion", + "OnnxQuantization", + "--device", + "gpu", + "--provider", + "CUDAExecutionProvider", + "--output_path", + "out", + ] + ) + assert args.model_name_or_path == "microsoft/phi-2" + assert args.run_config is None + assert args.passes == ["OnnxConversion", "OnnxQuantization"] + assert args.device == "gpu" + assert args.provider == "CUDAExecutionProvider" + assert args.output_path == "out" + + +def test_search_parses_config_mode_args(): + parser = _build_parser() + args = parser.parse_args(["search", "--config", "config.json"]) + assert args.run_config == "config.json" + assert args.model_name_or_path is None + + +def test_search_rejects_invalid_subscription_key(): + parser = _build_parser() + with pytest.raises(SystemExit): + parser.parse_args(["search", "-m", "model", "--az_apim_subscription_key", "not-hex"]) + + +def test_search_parses_search_strategy_args(): + parser = _build_parser() + args = parser.parse_args( + [ + "search", + "-m", + "model", + "--execution-order", + "pass-by-pass", + "--sampler", + "tpe", + "--max-iter", + "10", + "--max-time", + "60", + "--output-model-num", + "3", + "--stop-when-goals-met", + "--include-pass-params", + ] + ) + assert args.execution_order == "pass-by-pass" + assert args.sampler == "tpe" + assert args.max_iter == 10 + assert args.max_time == 60 + assert args.output_model_num == 3 + assert args.stop_when_goals_met is True + assert args.include_pass_params is True + + +def test_search_rejects_invalid_sampler(): + parser = _build_parser() + with pytest.raises(SystemExit): + parser.parse_args(["search", "-m", "model", "--sampler", "invalid"]) + + +# --------------------------------------------------------------------------- # +# _parse_pass_groups # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + # single unbracketed pass -> one single-pass group + (["OnnxConversion"], [["OnnxConversion"]]), + # multiple unbracketed passes -> each its own group + (["OnnxConversion", "OnnxQuantization"], [["OnnxConversion"], ["OnnxQuantization"]]), + # single bracketed multi-pass group + (["[OnnxConversion, OnnxQuantization]"], [["OnnxConversion", "OnnxQuantization"]]), + # two bracketed groups + ( + ["[OnnxConversion, OrtTransformersOptimization]", "[OnnxConversion, OnnxQuantization]"], + [["OnnxConversion", "OrtTransformersOptimization"], ["OnnxConversion", "OnnxQuantization"]], + ), + # mixed bracketed and unbracketed, order preserved + ( + [ + "[OnnxConversion, OrtTransformersOptimization]", + "OnnxQuantization", + "[OnnxConversion, OnnxFloatToFloat16]", + ], + [ + ["OnnxConversion", "OrtTransformersOptimization"], + ["OnnxQuantization"], + ["OnnxConversion", "OnnxFloatToFloat16"], + ], + ), + # space-separated passes inside brackets + (["[OnnxConversion OnnxQuantization]"], [["OnnxConversion", "OnnxQuantization"]]), + ], +) +def test_parse_pass_groups(raw, expected): + assert WorkflowSearchCommand._parse_pass_groups(raw) == expected + + +def test_parse_pass_groups_empty_raises(): + with pytest.raises(ValueError, match="No valid pass groups"): + WorkflowSearchCommand._parse_pass_groups([""]) + + +# --------------------------------------------------------------------------- # +# _build_search_strategy_overrides # +# --------------------------------------------------------------------------- # + + +def test_build_search_strategy_overrides_all_set(): + args = SimpleNamespace( + execution_order="joint", + sampler="tpe", + max_iter=10, + max_time=60, + output_model_num=2, + stop_when_goals_met=True, + include_pass_params=True, + ) + overrides = WorkflowSearchCommand._build_search_strategy_overrides(args) + assert overrides == { + "execution_order": "joint", + "sampler": "tpe", + "max_iter": 10, + "max_time": 60, + "output_model_num": 2, + "stop_when_goals_met": True, + "include_pass_params": True, + } + + +def test_build_search_strategy_overrides_none_set(): + args = SimpleNamespace( + execution_order=None, + sampler=None, + max_iter=None, + max_time=None, + output_model_num=None, + stop_when_goals_met=False, + include_pass_params=None, + ) + assert WorkflowSearchCommand._build_search_strategy_overrides(args) == {} + + +def test_build_search_strategy_overrides_include_pass_params_false(): + args = SimpleNamespace( + execution_order=None, + sampler=None, + max_iter=None, + max_time=None, + output_model_num=None, + stop_when_goals_met=False, + include_pass_params=False, + ) + assert WorkflowSearchCommand._build_search_strategy_overrides(args) == {"include_pass_params": False} + + +# --------------------------------------------------------------------------- # +# _apply_search_strategy # +# --------------------------------------------------------------------------- # + + +def test_apply_search_strategy_no_overrides_leaves_config_untouched(): + run_config = {"input_model": {}} + WorkflowSearchCommand._apply_search_strategy(run_config, {}) + assert run_config == {"input_model": {}} + + +def test_apply_search_strategy_creates_engine_section(): + run_config = {} + WorkflowSearchCommand._apply_search_strategy(run_config, {"sampler": "tpe"}) + assert run_config["engine"]["search_strategy"] == {"sampler": "tpe"} + + +def test_apply_search_strategy_merges_existing(): + run_config = {"engine": {"search_strategy": {"max_iter": 5, "sampler": "random"}}} + WorkflowSearchCommand._apply_search_strategy(run_config, {"sampler": "tpe"}) + assert run_config["engine"]["search_strategy"] == {"max_iter": 5, "sampler": "tpe"} + + +def test_apply_search_strategy_replaces_non_dict_existing(): + run_config = {"engine": {"search_strategy": "not-a-dict"}} + WorkflowSearchCommand._apply_search_strategy(run_config, {"sampler": "tpe"}) + assert run_config["engine"]["search_strategy"] == {"sampler": "tpe"} + + +# --------------------------------------------------------------------------- # +# _build_run_config / _add_evaluator / _add_data_config # +# --------------------------------------------------------------------------- # + + +def _make_command(**overrides): + """Create a WorkflowSearchCommand instance with a fake args namespace.""" + defaults = { + "model_name_or_path": "microsoft/phi-2", + "task": None, + "trust_remote_code": False, + "device": "cpu", + "provider": "CPUExecutionProvider", + "eval_tasks": None, + "eval_batch_size": 1, + "eval_max_length": 1024, + "eval_limit": 1, + "eval_backend": "auto", + "data_name": None, + "log_level": None, + } + defaults.update(overrides) + command = WorkflowSearchCommand.__new__(WorkflowSearchCommand) + command.args = SimpleNamespace(**defaults) + return command + + +def test_build_run_config_basic(): + command = _make_command() + config = command._build_run_config(["OnnxConversion"], "out/config_00") + assert config["input_model"] == {"type": "HfModel", "model_path": "microsoft/phi-2"} + assert config["passes"] == {"onnxconversion": {"type": "OnnxConversion"}} + assert config["target"] == "local_system" + assert config["output_dir"] == "out/config_00" + assert config["no_artifacts"] is True + assert config["systems"]["local_system"]["accelerators"][0] == { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"], + } + # no evaluator / data config / log level by default + assert "evaluator" not in config + assert "data_configs" not in config + assert "log_severity_level" not in config + + +def test_build_run_config_with_task_and_trust_remote_code(): + command = _make_command(task="text-generation", trust_remote_code=True) + config = command._build_run_config(["OnnxConversion"], "out") + assert config["input_model"]["task"] == "text-generation" + assert config["input_model"]["load_kwargs"] == {"trust_remote_code": True} + + +def test_build_run_config_preserves_pass_order(): + command = _make_command() + config = command._build_run_config(["OnnxConversion", "OrtTransformersOptimization"], "out") + assert list(config["passes"].keys()) == ["onnxconversion", "orttransformersoptimization"] + + +def test_build_run_config_sets_log_level(): + command = _make_command(log_level=2) + config = command._build_run_config(["OnnxConversion"], "out") + assert config["log_severity_level"] == 2 + + +def test_add_evaluator_when_eval_tasks_provided(): + command = _make_command(eval_tasks=["hellaswag", "winogrande"], eval_limit=0.1, device="gpu", eval_backend="ort") + config = command._build_run_config(["OnnxConversion"], "out") + evaluator = config["evaluators"]["evaluator"] + assert evaluator["type"] == "LMEvaluator" + assert evaluator["tasks"] == ["hellaswag", "winogrande"] + assert evaluator["device"] == "gpu" + assert evaluator["limit"] == 0.1 + assert evaluator["model_class"] == "ort" + assert config["evaluator"] == "evaluator" + assert config["host"] == "local_system" + + +def test_add_evaluator_auto_backend_has_no_model_class(): + command = _make_command(eval_tasks=["hellaswag"], eval_backend="auto") + config = command._build_run_config(["OnnxConversion"], "out") + assert "model_class" not in config["evaluators"]["evaluator"] + + +def test_add_data_config_when_data_name_provided(): + command = _make_command(data_name="wikitext") + with patch("olive.cli.search.update_dataset_options") as mock_update: + config = command._build_run_config(["OnnxStaticQuantization"], "out") + assert config["data_configs"][0]["name"] == "default_data_config" + assert config["data_configs"][0]["type"] == "HuggingfaceContainer" + mock_update.assert_called_once() + + +def test_add_data_config_absent_without_data_name(): + command = _make_command(data_name=None) + config = command._build_run_config(["OnnxConversion"], "out") + assert "data_configs" not in config + + +# --------------------------------------------------------------------------- # +# _submit_remote_job # +# --------------------------------------------------------------------------- # + + +def test_submit_remote_job_returns_op_id_from_string(): + mock_response = MagicMock() + mock_response.json.return_value = "op_12345" + with patch("requests.post", return_value=mock_response) as mock_post: + op_id = WorkflowSearchCommand._submit_remote_job({"key": "value"}, "a" * 32) + assert op_id == "op_12345" + mock_post.assert_called_once() + _, kwargs = mock_post.call_args + assert kwargs["headers"]["Ocp-Apim-Subscription-Key"] == "a" * 32 + assert kwargs["json"] == {"run_config": {"key": "value"}} + + +def test_submit_remote_job_returns_op_id_from_json(): + mock_response = MagicMock() + mock_response.json.return_value = {"op_id": "op_67890"} + with patch("requests.post", return_value=mock_response): + op_id = WorkflowSearchCommand._submit_remote_job({"key": "value"}, "b" * 32) + assert op_id == "op_67890" + + +def test_submit_remote_job_strips_quotes(): + mock_response = MagicMock() + mock_response.json.return_value = '"op_quoted"' + with patch("requests.post", return_value=mock_response): + op_id = WorkflowSearchCommand._submit_remote_job({}, "c" * 32) + assert op_id == "op_quoted" + + +def test_submit_remote_job_raises_when_no_op_id(): + mock_response = MagicMock() + mock_response.json.return_value = {"unexpected": "field"} + mock_response.text = '{"unexpected": "field"}' + with patch("requests.post", return_value=mock_response), pytest.raises(RuntimeError, match="Unexpected response"): + WorkflowSearchCommand._submit_remote_job({}, "d" * 32) + + +def test_submit_remote_job_uses_correct_url(): + mock_response = MagicMock() + mock_response.json.return_value = "op_1" + with patch("requests.post", return_value=mock_response) as mock_post: + WorkflowSearchCommand._submit_remote_job({}, "e" * 32) + args, _ = mock_post.call_args + assert args[0] == OLIVE_AAS_RUN_URL + + +# --------------------------------------------------------------------------- # +# Build-mode end-to-end (dry run / file generation) # +# --------------------------------------------------------------------------- # + + +def test_build_mode_dry_run_generates_config_files(tmp_path): + output_dir = tmp_path / "search-out" + command_args = [ + "search", + "-m", + "microsoft/phi-2", + "--passes", + "OnnxConversion", + "OnnxQuantization", + "--dry_run", + "-o", + str(output_dir), + ] + + cli_main(command_args) + + config_0 = output_dir / "config_00.json" + config_1 = output_dir / "config_01.json" + assert config_0.exists() + assert config_1.exists() + + data_0 = json.loads(config_0.read_text()) + data_1 = json.loads(config_1.read_text()) + assert data_0["passes"] == {"onnxconversion": {"type": "OnnxConversion"}} + assert data_1["passes"] == {"onnxquantization": {"type": "OnnxQuantization"}} + # output_dir must use posix separators regardless of OS + assert "\\" not in data_0["output_dir"] + assert data_0["output_dir"].endswith("config_00") + + +def test_build_mode_dry_run_bracketed_group_single_config(tmp_path): + output_dir = tmp_path / "search-out" + command_args = [ + "search", + "-m", + "microsoft/phi-2", + "--passes", + "[OnnxConversion, OnnxQuantization]", + "--dry_run", + "-o", + str(output_dir), + ] + + cli_main(command_args) + + config_0 = output_dir / "config_00.json" + assert config_0.exists() + assert not (output_dir / "config_01.json").exists() + data = json.loads(config_0.read_text()) + assert list(data["passes"].keys()) == ["onnxconversion", "onnxquantization"] + + +@patch("olive.workflows.run") +def test_build_mode_runs_olive_for_each_group(mock_run, tmp_path): + output_dir = tmp_path / "search-out" + command_args = [ + "search", + "-m", + "microsoft/phi-2", + "--passes", + "OnnxConversion", + "OnnxQuantization", + "-o", + str(output_dir), + ] + + cli_main(command_args) + + assert mock_run.call_count == 2 + + +@patch("olive.cli.search.WorkflowSearchCommand._submit_remote_job", return_value="op_123") +def test_build_mode_submits_remote_job_when_key_provided(mock_submit, tmp_path): + output_dir = tmp_path / "search-out" + command_args = [ + "search", + "-m", + "microsoft/phi-2", + "--passes", + "OnnxConversion", + "--az_apim_subscription_key", + "f" * 32, + "-o", + str(output_dir), + ] + + cli_main(command_args) + + mock_submit.assert_called_once() + + +# --------------------------------------------------------------------------- # +# Config-file mode # +# --------------------------------------------------------------------------- # + + +@patch("olive.workflows.run") +def test_config_mode_runs_olive(mock_run, tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"input_model": {"type": "HfModel"}, "output_dir": "out"})) + command_args = ["search", "--run-config", str(config_path)] + + cli_main(command_args) + + mock_run.assert_called_once() + + +@patch("olive.workflows.run") +def test_config_mode_applies_search_strategy(mock_run, tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"input_model": {"type": "HfModel"}, "output_dir": "out"})) + command_args = [ + "search", + "--run-config", + str(config_path), + "--sampler", + "tpe", + "--max-iter", + "10", + ] + + cli_main(command_args) + + passed_config = mock_run.call_args[0][0] + assert passed_config["engine"]["search_strategy"] == {"sampler": "tpe", "max_iter": 10} + + +def test_config_mode_dry_run_does_not_run(tmp_path): + config_path = tmp_path / "config.json" + output_dir = tmp_path / "out" + output_dir.mkdir() + config_path.write_text(json.dumps({"input_model": {"type": "HfModel"}, "output_dir": str(output_dir)})) + command_args = ["search", "--run-config", str(config_path), "--dry_run"] + + with patch("olive.workflows.run") as mock_run: + cli_main(command_args) + + mock_run.assert_not_called() + + +@patch("olive.cli.search.WorkflowSearchCommand._submit_remote_job", return_value="op_abc") +def test_config_mode_submits_remote_job_when_key_provided(mock_submit, tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"input_model": {"type": "HfModel"}, "output_dir": "out"})) + command_args = [ + "search", + "--run-config", + str(config_path), + "--az_apim_subscription_key", + "1" * 32, + ] + + with patch("olive.workflows.run") as mock_run: + cli_main(command_args) + + mock_submit.assert_called_once() + mock_run.assert_not_called()