diff --git a/.lintrunner.toml b/.lintrunner.toml index 3c5b386afc..34e39e68d0 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -62,7 +62,7 @@ init_command = [ is_formatter = true [[linter]] -code = 'BLACK-ISORT' +code = 'RUFF-FORMAT' include_patterns = [ '**/*.py' ] @@ -74,7 +74,7 @@ command = [ '-m', 'lintrunner_adapters', 'run', - 'black_isort_linter', + 'ruff_format_linter', '--', '@{{PATHSFILE}}' ] diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5b53b967b4..7e0c4ab2b6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -7,29 +7,11 @@ repos: args: [--markdown-linebreak-ext=md] - id: check-yaml - id: requirements-txt-fixer - - - repo: https://github.com/psf/black - rev: 24.1.0 - hooks: - - id: black - name: Format code - exclude: "hadamard_utils.py" - - repo: https://github.com/pycqa/isort - rev: 5.11.5 - hooks: - - id: isort - name: Format imports - repo: https://github.com/MarcoGorelli/absolufy-imports rev: v0.3.1 hooks: - id: absolufy-imports exclude: examples/ - - repo: https://github.com/astral-sh/ruff-pre-commit - # Ruff version. - rev: v0.6.0 - hooks: - - id: ruff - args: [ --fix ] - repo: local hooks: - id: format-json diff --git a/docs/source/exts/auto_config_doc/__init__.py b/docs/source/exts/auto_config_doc/__init__.py index eec7368a80..45948ada9e 100644 --- a/docs/source/exts/auto_config_doc/__init__.py +++ b/docs/source/exts/auto_config_doc/__init__.py @@ -73,9 +73,9 @@ def run(self): (class_name,) = self.arguments package_config = OlivePackageConfig.load_default_config() auto_config_class = import_class(class_name, package_config) - assert issubclass(auto_config_class, AutoConfigClass) or issubclass( - auto_config_class, Pass - ), f"{class_name} is not a subclass of AutoConfigClass or Pass" + assert issubclass(auto_config_class, AutoConfigClass) or issubclass(auto_config_class, Pass), ( + f"{class_name} is not a subclass of AutoConfigClass or Pass" + ) node = nodes.section() node.document = self.state.document diff --git a/docs/source/exts/gallery_directive.py b/docs/source/exts/gallery_directive.py index cd197a7483..e935b9c9f3 100644 --- a/docs/source/exts/gallery_directive.py +++ b/docs/source/exts/gallery_directive.py @@ -107,7 +107,7 @@ def run(self) -> Optional[List[nodes.Node]]: # Parse the template with Sphinx Design to create an output container # Prep the options for the template grid - class_ = "gallery-directive" + f' {self.options.get("class-container", "")}' + class_ = "gallery-directive" + f" {self.options.get('class-container', '')}" options = {"gutter": 2, "class-container": class_} options_str = "\n".join(f":{k}: {v}" for k, v in options.items()) diff --git a/examples/directml/llm/chat_app/interface/hddr_llm_onnx_dml_interface.py b/examples/directml/llm/chat_app/interface/hddr_llm_onnx_dml_interface.py index 6a6956c4e4..9a895f080a 100644 --- a/examples/directml/llm/chat_app/interface/hddr_llm_onnx_dml_interface.py +++ b/examples/directml/llm/chat_app/interface/hddr_llm_onnx_dml_interface.py @@ -288,10 +288,13 @@ def predict( sentence = sentence[: sentence.index(ai_token)].strip() break sentence = sentence.strip() - a, b = [[y[0], convert_to_markdown(y[1])] for y in history] + [[text, convert_to_markdown(sentence)]], [ - *history, - [text, sentence], - ] + a, b = ( + [[y[0], convert_to_markdown(y[1])] for y in history] + [[text, convert_to_markdown(sentence)]], + [ + *history, + [text, sentence], + ], + ) yield a, b, "Generating..." if shared_state.interrupted: diff --git a/examples/inception/prepare_config.py b/examples/inception/prepare_config.py index 03b39de67e..1141ec8235 100644 --- a/examples/inception/prepare_config.py +++ b/examples/inception/prepare_config.py @@ -10,7 +10,6 @@ def resolve_windows_config(): - with Path("inception_config.json").open() as f: snpe_windows_config = json.load(f) diff --git a/examples/llama2/tensor_parallel_generate.py b/examples/llama2/tensor_parallel_generate.py index 1e70298062..600f056776 100644 --- a/examples/llama2/tensor_parallel_generate.py +++ b/examples/llama2/tensor_parallel_generate.py @@ -42,8 +42,8 @@ torch.cuda.empty_cache() model_id = "meta-llama/Llama-2-7b-hf" -model_path = "models/tensor_parallel/tensor_parallel-conversion-transformers_optimization_fp16/gpu-cuda_model/model_{:02d}".format( # noqa: E501 - rank +model_path = ( + f"models/tensor_parallel/tensor_parallel-conversion-transformers_optimization_fp16/gpu-cuda_model/model_{rank:02d}" ) prompt = "What is an apple?" # prompt = "Is it normal to have a dark ring around the iris of my eye?" diff --git a/examples/open_llama/user_script.py b/examples/open_llama/user_script.py index d235cf97cb..91e3e74916 100644 --- a/examples/open_llama/user_script.py +++ b/examples/open_llama/user_script.py @@ -54,10 +54,13 @@ def __iter__(self): inp = trainenc["input_ids"][i:j].unsqueeze(0) mask = torch.ones(inp.shape) if self.sess is None: - yield { - "input_ids": inp.detach().cpu().numpy().astype("int64"), - "attention_mask": mask.detach().cpu().numpy().astype("int64"), - }, 0 + yield ( + { + "input_ids": inp.detach().cpu().numpy().astype("int64"), + "attention_mask": mask.detach().cpu().numpy().astype("int64"), + }, + 0, + ) else: outputs = self.sess.run( None, diff --git a/examples/phi2/generate.py b/examples/phi2/generate.py index 115cfdd387..31de092c7b 100644 --- a/examples/phi2/generate.py +++ b/examples/phi2/generate.py @@ -2,7 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -# copied from https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/phi2/inference_example.py # noqa: E501 +# copied from https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/phi2/inference_example.py import numpy as np import onnxruntime as ort import torch @@ -40,7 +40,6 @@ def __init__(self, decoder_path): self.tokenizer = None def get_initial_inputs_and_outputs(self, encodings_dict): - input_ids = torch.tensor(encodings_dict["input_ids"], device=self.device, dtype=torch.int32) attention_mask = torch.tensor(encodings_dict["attention_mask"], device=self.device, dtype=torch.int32) step = torch.tensor([0], device=self.device, dtype=torch.int64) @@ -278,7 +277,7 @@ def genai_run(prompts, model_path, max_length=200): app_started_timestamp = time.time() model = og.Model(model_path) model_loaded_timestamp = time.time() - print("Model loaded in {:.2f} seconds".format(model_loaded_timestamp - app_started_timestamp)) + print(f"Model loaded in {model_loaded_timestamp - app_started_timestamp:.2f} seconds") tokenizer = og.Tokenizer(model) input_tokens = tokenizer.encode_batch(prompts) @@ -299,7 +298,7 @@ def genai_run(prompts, model_path, max_length=200): output_token_count = 0 for i, prompt in enumerate(prompts): - print(f"Prompt #{i+1:02d}: {prompt}") + print(f"Prompt #{i + 1:02d}: {prompt}") print(tokenizer.decode(generator.get_sequence(i))) output_token_count = sum(len(generator.get_sequence(i)) for i in range(len(prompts))) diff --git a/examples/phi3/phi3.py b/examples/phi3/phi3.py index cca41fa834..52e66fd0ea 100644 --- a/examples/phi3/phi3.py +++ b/examples/phi3/phi3.py @@ -199,7 +199,6 @@ def use_passes(template_json, *passes): def generate_config(args): - json_file_template = "phi3_template.json" with open(json_file_template) as f: template_json = json.load(f) @@ -261,14 +260,13 @@ def generate_config(args): def genai_run(prompt, model_path, max_length): - print("\nModel inference starts...") print("Loading model...") app_started_timestamp = time.time() model = og.Model(model_path) model_loaded_timestamp = time.time() - print("Model loaded in {:.2f} seconds".format(model_loaded_timestamp - app_started_timestamp)) + print(f"Model loaded in {model_loaded_timestamp - app_started_timestamp:.2f} seconds") print("Creating tokenizer...") tokenizer = og.Tokenizer(model) @@ -320,7 +318,7 @@ def genai_run(prompt, model_path, max_length): "\n\n" f"Prompt tokens: {len(input_tokens)}, New tokens: {len(new_tokens)}," f" Time to first: {(first_token_timestamp - started_timestamp):.2f}s," - f" New tokens per second: {len(new_tokens)/run_time:.2f} tps" + f" New tokens per second: {len(new_tokens) / run_time:.2f} tps" ) diff --git a/examples/red_pajama/user_script.py b/examples/red_pajama/user_script.py index ca3e33062d..1cdde67fd4 100644 --- a/examples/red_pajama/user_script.py +++ b/examples/red_pajama/user_script.py @@ -7,6 +7,6 @@ MIN_TRANSFORMERS_VERSION = "4.30.2" # check transformers version -assert ( - transformers.__version__ >= MIN_TRANSFORMERS_VERSION -), f"Please upgrade transformers to version {MIN_TRANSFORMERS_VERSION} or higher." +assert transformers.__version__ >= MIN_TRANSFORMERS_VERSION, ( + f"Please upgrade transformers to version {MIN_TRANSFORMERS_VERSION} or higher." +) diff --git a/examples/resnet/prepare_model_data.py b/examples/resnet/prepare_model_data.py index 63c18faaeb..e2cd428a86 100644 --- a/examples/resnet/prepare_model_data.py +++ b/examples/resnet/prepare_model_data.py @@ -105,11 +105,7 @@ def prepare_model(num_epochs=1, models_dir="models", data_dir="data"): loss.backward() optimizer.step() if (i + 1) % 100 == 0: - print( - "Epoch [{}/{}], Step [{}/{}] Loss: {:.4f}".format( - epoch + 1, num_epochs, i + 1, total_step, loss.item() - ) - ) + print(f"Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{total_step}] Loss: {loss.item():.4f}") # Decay learning rate if (epoch + 1) % 20 == 0: curr_lr /= 3 @@ -129,7 +125,7 @@ def prepare_model(num_epochs=1, models_dir="models", data_dir="data"): total += labels.size(0) correct += (predicted == labels).sum().item() - print("Accuracy of the model on the test images: {} %".format(100 * correct / total)) + print(f"Accuracy of the model on the test images: {100 * correct / total} %") # Save the model model.to("cpu") diff --git a/examples/utils/generator.py b/examples/utils/generator.py index 1b56425b6b..7ec8fbc0c2 100644 --- a/examples/utils/generator.py +++ b/examples/utils/generator.py @@ -428,9 +428,9 @@ def get_initial_inputs( else [apply_template(template, p) for p in prompt] ) else: - assert isinstance(prompt, str) or ( - isinstance(prompt, list) and all(isinstance(p, str) for p in prompt) - ), "tuple prompts require a template" + assert isinstance(prompt, str) or (isinstance(prompt, list) and all(isinstance(p, str) for p in prompt)), ( + "tuple prompts require a template" + ) encodings_dict = self.tokenizer(prompt, return_tensors="np", padding=True) input_ids = encodings_dict["input_ids"].astype(self.input_info["input_ids"]["dtype"]) diff --git a/examples/utils/kv_cache_utils.py b/examples/utils/kv_cache_utils.py index cfcacf10e8..ec16f71e3d 100644 --- a/examples/utils/kv_cache_utils.py +++ b/examples/utils/kv_cache_utils.py @@ -160,18 +160,18 @@ def update(self, present_kvs: List["NDArray"]): assert present_len > 0, "present_kvs cannot be empty" if self.seen_len == 0: - assert ( - present_len <= self.max_cache_len - ), "present_kvs is longer than max_cache_len during prompt processing" + assert present_len <= self.max_cache_len, ( + "present_kvs is longer than max_cache_len during prompt processing" + ) # prompt processing for k, v in zip(self.past_names, present_kvs): self.cache[k][:, :, :present_len] = v self.seen_len = present_len return - assert ( - present_len == self.max_cache_len + 1 - ), "present_kvs must be one step longer than max_cache_len in token generation" + assert present_len == self.max_cache_len + 1, ( + "present_kvs must be one step longer than max_cache_len in token generation" + ) for k, v in zip(self.past_names, present_kvs): self.cache[k][:, :, self.seen_len] = v[:, :, -1] self.seen_len += 1 @@ -414,9 +414,9 @@ def update(self, present_kvs: List[OrtValue]): # prompt processing present_len = present_kvs[0].shape()[2] assert present_len > 0, "present_kvs cannot be empty" - assert ( - present_len <= self.max_cache_len - ), "present_kvs is longer than max_cache_len during prompt processing" + assert present_len <= self.max_cache_len, ( + "present_kvs is longer than max_cache_len during prompt processing" + ) for k, v in zip(self.past_names, present_kvs): if self.backend == "torch": import torch @@ -431,9 +431,9 @@ def update(self, present_kvs: List[OrtValue]): # token generation for past_k, present_k, v in zip(self.past_names, self.present_names, present_kvs): - assert ( - self.output_cache[present_k].data_ptr() == v.data_ptr() - ), "out cache ortvalue should be same as present ortvalue" + assert self.output_cache[present_k].data_ptr() == v.data_ptr(), ( + "out cache ortvalue should be same as present ortvalue" + ) if self.backend == "torch": self.cache[past_k][:, :, self.seen_len] = self.output_cache[present_k][:, :, -1] else: diff --git a/examples/vgg/prepare_config.py b/examples/vgg/prepare_config.py index 78aa223a6d..4861462bfa 100644 --- a/examples/vgg/prepare_config.py +++ b/examples/vgg/prepare_config.py @@ -10,7 +10,6 @@ def resolve_windows_config(): - with Path("vgg_config.json").open() as f: snpe_windows_config = json.load(f) diff --git a/examples/vit/val_tiny_imagenet/val_tiny_imagenet.py b/examples/vit/val_tiny_imagenet/val_tiny_imagenet.py index 866b221281..aac7546694 100644 --- a/examples/vit/val_tiny_imagenet/val_tiny_imagenet.py +++ b/examples/vit/val_tiny_imagenet/val_tiny_imagenet.py @@ -99,7 +99,7 @@ def evaluate_onnx_model(session, dataloader): ground_truth = val_idx_to_name[label[0]] pred_label = vit_id2label["id2label"][str(top1_pred)] - print(f"Image {i+1}: {img_name[0]}") + print(f"Image {i + 1}: {img_name[0]}") print(f" Ground Truth: {ground_truth}") print(f" Top-1 Prediction: {pred_label}") print(f" Top-5 Predictions: {[vit_id2label['id2label'][str(pred)] for pred in top5_preds]}\n") diff --git a/olive/auto_optimizer/regulate_mixins.py b/olive/auto_optimizer/regulate_mixins.py index 5801ea2ccf..d5238ba796 100644 --- a/olive/auto_optimizer/regulate_mixins.py +++ b/olive/auto_optimizer/regulate_mixins.py @@ -69,9 +69,9 @@ def _regulate_precision(self, pass_config, pass_flows): is_cuda_ep = self.accelerator_spec.execution_provider != "TensorrtExecutionProvider" is_trt_ep = self.accelerator_spec.execution_provider == "TensorrtExecutionProvider" - assert ( - not is_cuda_ep or not is_trt_ep - ), "can not support CUDA/DmlExecutionProvider and TensorrtExecutionProvider at the same time" + assert not is_cuda_ep or not is_trt_ep, ( + "can not support CUDA/DmlExecutionProvider and TensorrtExecutionProvider at the same time" + ) customized_fp16 = self._allow_precision("fp16") cuda_fp16 = customized_fp16 and is_cuda_ep diff --git a/olive/cli/auto_opt.py b/olive/cli/auto_opt.py index 78693db561..cbec4019d6 100644 --- a/olive/cli/auto_opt.py +++ b/olive/cli/auto_opt.py @@ -27,7 +27,6 @@ class AutoOptCommand(BaseOliveCLICommand): - @staticmethod def register_subcommand(parser: ArgumentParser): sub_parser = parser.add_parser( diff --git a/olive/cli/generate_adapter.py b/olive/cli/generate_adapter.py index f28555c533..c0a0fec730 100644 --- a/olive/cli/generate_adapter.py +++ b/olive/cli/generate_adapter.py @@ -49,9 +49,9 @@ def run(self): def _get_run_config(self, tempdir: str) -> Dict: input_model_config = get_input_model_config(self.args) - assert ( - input_model_config["type"].lower() == "onnxmodel" - ), "Only ONNX models are supported in generate-adapter command." + assert input_model_config["type"].lower() == "onnxmodel", ( + "Only ONNX models are supported in generate-adapter command." + ) to_replace = [ ("input_model", input_model_config), diff --git a/olive/cli/quantize.py b/olive/cli/quantize.py index c38e638160..50e6b349c6 100644 --- a/olive/cli/quantize.py +++ b/olive/cli/quantize.py @@ -31,7 +31,6 @@ class QuantizeCommand(BaseOliveCLICommand): - @staticmethod def register_subcommand(parser: ArgumentParser): sub_parser = parser.add_parser( @@ -130,8 +129,7 @@ def _get_pass_list(self, precision, algo, impl, is_hf_model): if not pass_list: raise ValueError( - f"Quantiation for precision {precision}, algorithm {algo} " - f"and implementation {impl} is not supported" + f"Quantiation for precision {precision}, algorithm {algo} and implementation {impl} is not supported" ) logger.info("pass list: %s", pass_list) return pass_list diff --git a/olive/cli/session_params_tuning.py b/olive/cli/session_params_tuning.py index 147a642927..907ecd5560 100644 --- a/olive/cli/session_params_tuning.py +++ b/olive/cli/session_params_tuning.py @@ -26,7 +26,6 @@ class SessionParamsTuningCommand(BaseOliveCLICommand): - @staticmethod def register_subcommand(parser: ArgumentParser): sub_parser = parser.add_parser( diff --git a/olive/data/component/dataset.py b/olive/data/component/dataset.py index df003c369a..06ee68dcd0 100644 --- a/olive/data/component/dataset.py +++ b/olive/data/component/dataset.py @@ -157,9 +157,9 @@ def __init__( self.annotations = None if annotations_file is not None: self.annotations = np.load(self.data_dir / annotations_file) - assert len(self.annotations) == len( - self.input_files - ), "Number of annotations should be equal to number of input files." + assert len(self.annotations) == len(self.input_files), ( + "Number of annotations should be equal to number of input files." + ) def __len__(self): return len(self.input_files) diff --git a/olive/data/component/load_dataset.py b/olive/data/component/load_dataset.py index cda95e739a..0c4b21663f 100644 --- a/olive/data/component/load_dataset.py +++ b/olive/data/component/load_dataset.py @@ -26,7 +26,7 @@ def huggingface_dataset( split: Optional[str] = "validation", data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None, col_filters: Optional[Mapping[str, Union[str, int, float]]] = None, - **kwargs + **kwargs, ): """Create a dataset from huggingface datasets.""" from datasets.utils.logging import disable_progress_bar, set_verbosity_error diff --git a/olive/data/component/pre_process_data.py b/olive/data/component/pre_process_data.py index c0f407567b..b77f233963 100644 --- a/olive/data/component/pre_process_data.py +++ b/olive/data/component/pre_process_data.py @@ -189,7 +189,7 @@ def audio_classification_pre_process( max_samples: Optional[int] = None, trust_remote_code: Optional[bool] = None, feature_extractor_args: Optional[Dict[str, Any]] = None, - **kwargs + **kwargs, ): """Pre-process data for audio classification task. diff --git a/olive/data/component/text_generation.py b/olive/data/component/text_generation.py index 1e99dfc9d6..d5c9fce31b 100644 --- a/olive/data/component/text_generation.py +++ b/olive/data/component/text_generation.py @@ -229,7 +229,7 @@ def text_gen_pre_process(dataset, tokenizer, all_kwargs): examples_to_get = min(args.processing_batch_size, total_examples - example_idx) # batch tokenize batched_input_ids = tokenizer( - text_list[example_idx : example_idx + examples_to_get], # noqa: E203, RUF100 + text_list[example_idx : example_idx + examples_to_get], add_special_tokens=False, truncation=False, )["input_ids"] @@ -307,7 +307,9 @@ def text_gen_pre_process(dataset, tokenizer, all_kwargs): examples_to_get = min(args.max_samples - num_samples, total_examples - example_idx) # batch tokenize tokenized_texts = batch_tokenize_text( - text_list[example_idx : example_idx + examples_to_get], tokenizer, args # noqa: E203, RUF100 + text_list[example_idx : example_idx + examples_to_get], + tokenizer, + args, ) for native_input_ids, native_attention_mask in tokenized_texts: append_text_gen_input_ids( diff --git a/olive/evaluator/olive_evaluator.py b/olive/evaluator/olive_evaluator.py index 28df68ceeb..786a0c7dbc 100644 --- a/olive/evaluator/olive_evaluator.py +++ b/olive/evaluator/olive_evaluator.py @@ -330,7 +330,6 @@ def evaluate( class OnnxEvaluatorMixin: - @staticmethod def get_inference_settings(metric: Metric, model: ONNXModelHandler) -> Dict[str, Any]: # user.config.inference_settings > model.inference_settings > default inference_settings @@ -351,7 +350,6 @@ def get_inference_settings(metric: Metric, model: ONNXModelHandler) -> Dict[str, @Registry.register(str(Framework.ONNX)) @Registry.register("OnnxEvaluator") class OnnxEvaluator(_OliveEvaluator, OnnxEvaluatorMixin): - @staticmethod def get_session_wrapper( model: ONNXModelHandler, @@ -825,7 +823,6 @@ def _evaluate_raw_latency( @Registry.register(str(Framework.SNPE)) @Registry.register("SNPEEvaluator") class SNPEEvaluator(_OliveEvaluator): - def _inference( self, model: "SNPEModelHandler", @@ -915,7 +912,6 @@ def _prepare_dataloader( @Registry.register(str(Framework.OPENVINO)) @Registry.register("OpenVINOEvaluator") class OpenVINOEvaluator(_OliveEvaluator): - def _inference( self, model: "OpenVINOModelHandler", @@ -979,7 +975,6 @@ def _evaluate_raw_latency( @Registry.register(str(Framework.QNN)) @Registry.register("QNNEvaluator") class QNNEvaluator(_OliveEvaluator): - def _inference( self, model: "QNNModelHandler", diff --git a/olive/model/handler/composite.py b/olive/model/handler/composite.py index 64c7ff1cae..f3000e90b7 100644 --- a/olive/model/handler/composite.py +++ b/olive/model/handler/composite.py @@ -47,9 +47,9 @@ def __init__( self._model_components = [ validate_config(m, ModelConfig).create_model() if isinstance(m, dict) else m for m in model_components ] - assert all( - isinstance(m, OliveModelHandler) for m in self._model_components - ), "All components must be OliveModelHandler or dict" + assert all(isinstance(m, OliveModelHandler) for m in self._model_components), ( + "All components must be OliveModelHandler or dict" + ) assert len(self._model_components) == len(model_component_names), "Number of components and names must match" self.model_component_names = model_component_names diff --git a/olive/model/handler/mixin/resource.py b/olive/model/handler/mixin/resource.py index 6b8b56dd74..1e81dd2507 100644 --- a/olive/model/handler/mixin/resource.py +++ b/olive/model/handler/mixin/resource.py @@ -35,9 +35,9 @@ def set_resource(self, resource_name: str, resource_path: Union[Path, str, Resou if resource_path is not None: resolved_resource_path = create_resource_path(resource_path) - assert ( - resolved_resource_path.is_local_resource_or_string_name() - ), f"{resource_name} must be local path or string name." + assert resolved_resource_path.is_local_resource_or_string_name(), ( + f"{resource_name} must be local path or string name." + ) resource_path = resolved_resource_path.get_path() self.resource_paths[resource_name] = resource_path @@ -65,9 +65,9 @@ def _add_resources(self, resources: Dict[str, OLIVE_RESOURCE_ANNOTATIONS]): for resource_name, resource_path in resources.items(): if resource_path is not None: resolved_resource_path = create_resource_path(resource_path) - assert ( - resolved_resource_path.is_local_resource_or_string_name() - ), f"{resource_name} must be local path or string name." + assert resolved_resource_path.is_local_resource_or_string_name(), ( + f"{resource_name} must be local path or string name." + ) self.resource_paths[resource_name] = resolved_resource_path.get_path() else: self.resource_paths[resource_name] = None diff --git a/olive/model/handler/pytorch.py b/olive/model/handler/pytorch.py index cb4b53c4b0..1288e01abf 100644 --- a/olive/model/handler/pytorch.py +++ b/olive/model/handler/pytorch.py @@ -22,9 +22,7 @@ logger = logging.getLogger(__name__) -class PyTorchModelHandlerBase( - OliveModelHandler, DummyInputsMixin, PytorchKvCacheMixin -): # pylint: disable=too-many-ancestors +class PyTorchModelHandlerBase(OliveModelHandler, DummyInputsMixin, PytorchKvCacheMixin): # pylint: disable=too-many-ancestors """Base class for PyTorch model handler.""" def prepare_session( diff --git a/olive/passes/olive_pass.py b/olive/passes/olive_pass.py index de395fd99b..526e1c264f 100644 --- a/olive/passes/olive_pass.py +++ b/olive/passes/olive_pass.py @@ -134,9 +134,9 @@ def generate_config( point = point or {} config_class, fixed_values, search_params = cls.get_config_params(accelerator_spec, config, disable_search) - assert ( - set(point.keys()).intersection(set(search_params.keys())) == point.keys() - ), "Search point is not in the search space." + assert set(point.keys()).intersection(set(search_params.keys())) == point.keys(), ( + "Search point is not in the search space." + ) return config_class.parse_obj({**fixed_values, **search_params, **point}) @classmethod @@ -188,9 +188,9 @@ def default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassConf for param, param_config in config.items(): if param.endswith("data_config"): param_type = param_config.type_ - assert param_type == DataConfig or DataConfig in get_args( - param_type - ), f"{param} ending with data_config must be of type DataConfig." + assert param_type == DataConfig or DataConfig in get_args(param_type), ( + f"{param} ending with data_config must be of type DataConfig." + ) return config @classmethod diff --git a/olive/passes/onnx/append_pre_post_processing_ops.py b/olive/passes/onnx/append_pre_post_processing_ops.py index 4e03b65865..1ef97decd8 100644 --- a/olive/passes/onnx/append_pre_post_processing_ops.py +++ b/olive/passes/onnx/append_pre_post_processing_ops.py @@ -92,9 +92,9 @@ def _run_for_config( if tool_command == "whisper": from onnxruntime_extensions import __version__ as ortext_version - assert version.parse(ortext_version) >= version.parse( - "0.9.0" - ), "Whisper pre-post processing requires onnxruntime_extensions>=0.9.0" + assert version.parse(ortext_version) >= version.parse("0.9.0"), ( + "Whisper pre-post processing requires onnxruntime_extensions>=0.9.0" + ) from olive.passes.utils.whisper_prepost import add_pre_post_processing_to_model diff --git a/olive/passes/onnx/bnb_quantization.py b/olive/passes/onnx/bnb_quantization.py index f28c09feb2..cc90b2dfda 100644 --- a/olive/passes/onnx/bnb_quantization.py +++ b/olive/passes/onnx/bnb_quantization.py @@ -52,9 +52,9 @@ def _run_for_config( ) -> ONNXModelHandler: from onnxruntime import __version__ as OrtVersion - assert version.parse(OrtVersion) >= version.parse( - "1.16.2" - ), "MatMulBnb4Quantizer is only supported in onnxruntime >= 1.16.2" + assert version.parse(OrtVersion) >= version.parse("1.16.2"), ( + "MatMulBnb4Quantizer is only supported in onnxruntime >= 1.16.2" + ) from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer diff --git a/olive/passes/onnx/compose.py b/olive/passes/onnx/compose.py index c520211371..126d0069f2 100644 --- a/olive/passes/onnx/compose.py +++ b/olive/passes/onnx/compose.py @@ -52,9 +52,9 @@ def _run_for_config( output_model_path: str, ) -> Union[ONNXModelHandler, CompositeModelHandler]: assert isinstance(model, CompositeModelHandler), "ComposeOnnxModels pass only supports CompositeModelHandler" - assert all( - isinstance(m, ONNXModelHandler) for m in model.model_components - ), "All components must be ONNXModelHandler" + assert all(isinstance(m, ONNXModelHandler) for m in model.model_components), ( + "All components must be ONNXModelHandler" + ) if pipeline := (model.model_attributes or {}).get("llm_pipeline"): output_model_path = Path(output_model_path).with_suffix("") @@ -117,13 +117,13 @@ def _get_composed_model( dag_inputs = set(dag.get_input_names()) dag_outputs = set(dag.get_output_names()) # avoid circular connection, model_2 output cannot be model_1 input - assert dag_outputs.isdisjoint( - seen_inputs - ), f"Output names {dag_outputs.intersection(seen_inputs)} are already used as input names." + assert dag_outputs.isdisjoint(seen_inputs), ( + f"Output names {dag_outputs.intersection(seen_inputs)} are already used as input names." + ) # avoid reused output name - assert dag_outputs.isdisjoint( - seen_outputs - ), f"Output names {dag_outputs.intersection(seen_outputs)} are already used as output names." + assert dag_outputs.isdisjoint(seen_outputs), ( + f"Output names {dag_outputs.intersection(seen_outputs)} are already used as output names." + ) # update seen inputs and outputs seen_inputs.update(dag_inputs) @@ -143,12 +143,12 @@ def _get_composed_model( cd_initializer_names = set(composed_dag.get_initializer_names()) for input_name in dag.get_input_names(): if input_name in cd_input_names | cd_output_names: - assert dag.get_io_shape(input_name) == composed_dag.get_io_shape( - input_name - ), f"Input shape mismatch: {input_name}" - assert dag.get_io_dtype(input_name) == composed_dag.get_io_dtype( - input_name - ), f"Input dtype mismatch: {input_name}" + assert dag.get_io_shape(input_name) == composed_dag.get_io_shape(input_name), ( + f"Input shape mismatch: {input_name}" + ) + assert dag.get_io_dtype(input_name) == composed_dag.get_io_dtype(input_name), ( + f"Input dtype mismatch: {input_name}" + ) continue # will add to graph 0 for now @@ -160,9 +160,12 @@ def _get_composed_model( for init_name in dag.get_initializer_names(): if init_name in cd_initializer_names: - np.testing.assert_array_equal( - dag.get_initializer_np_array(init_name), composed_dag.get_initializer_np_array(init_name) - ), f"Initializer mismatch: {init_name}" + ( + np.testing.assert_array_equal( + dag.get_initializer_np_array(init_name), composed_dag.get_initializer_np_array(init_name) + ), + f"Initializer mismatch: {init_name}", + ) continue composed_dag.add_initializer(dag.get_initializer_proto(init_name), 0) diff --git a/olive/passes/onnx/context_binary.py b/olive/passes/onnx/context_binary.py index 384c37157f..d20f84076c 100644 --- a/olive/passes/onnx/context_binary.py +++ b/olive/passes/onnx/context_binary.py @@ -73,9 +73,9 @@ def _run_for_config( from onnxruntime import get_available_providers # TODO(jambayk): validate and support other NPU EPs - assert ( - self.accelerator_spec.execution_provider == "QNNExecutionProvider" - ), "Only QNNExecutionProvider is supported for now." + assert self.accelerator_spec.execution_provider == "QNNExecutionProvider", ( + "Only QNNExecutionProvider is supported for now." + ) assert self.accelerator_spec.execution_provider in get_available_providers(), ( f"Execution provider {self.accelerator_spec.execution_provider} is not available. Available providers:" f" {get_available_providers()}" diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py index cd95f9eea9..625aebe766 100644 --- a/olive/passes/onnx/dynamic_to_fixed_shape.py +++ b/olive/passes/onnx/dynamic_to_fixed_shape.py @@ -71,7 +71,6 @@ def _run_for_config( config: Type[BasePassConfig], output_model_path: str, ) -> ONNXModelHandler: - onnx_model = model.load_model() output_model_path = resolve_onnx_path(output_model_path) diff --git a/olive/passes/onnx/graph_surgeries.py b/olive/passes/onnx/graph_surgeries.py index 8dbc788c02..2a82ba629a 100644 --- a/olive/passes/onnx/graph_surgeries.py +++ b/olive/passes/onnx/graph_surgeries.py @@ -173,7 +173,6 @@ def __call__(self, model: ModelProto): class ReplaceErfWithTanh(Surgeon): - DTYPE_MAP = { TensorProto.FLOAT: np.float32, TensorProto.FLOAT16: np.float16, diff --git a/olive/passes/onnx/mixed_precision_overrides.py b/olive/passes/onnx/mixed_precision_overrides.py index 3976cc26e5..16d6161364 100644 --- a/olive/passes/onnx/mixed_precision_overrides.py +++ b/olive/passes/onnx/mixed_precision_overrides.py @@ -139,7 +139,6 @@ def handle_conflict(tensor_name, node) -> bool: element_wise_binary_ops = config.element_wise_binary_ops or ["Add", "Sub", "Mul", "Div"] for node in onnx_model.graph.node: - if node.op_type in element_wise_binary_ops: # For ElementWiseBinaryOps inputs and outputs should be of same type if node.output[0] in activations_16bit: diff --git a/olive/passes/onnx/nvmo_quantization.py b/olive/passes/onnx/nvmo_quantization.py index a98f0655bd..4a3a4a3ce0 100644 --- a/olive/passes/onnx/nvmo_quantization.py +++ b/olive/passes/onnx/nvmo_quantization.py @@ -382,7 +382,6 @@ def convert_opset_to_21_proto(self, model_proto: ModelProto) -> ModelProto: def _run_for_config( self, model: OliveModelHandler, config: Type[BasePassConfig], output_model_path: str ) -> OliveModelHandler: - try: logger.debug("Loading the original ONNX model from %s.", model.model_path) quant_config = self.initialize_quant_config(config) diff --git a/olive/passes/onnx/onnx_dag.py b/olive/passes/onnx/onnx_dag.py index f2ec6670ad..276f755310 100644 --- a/olive/passes/onnx/onnx_dag.py +++ b/olive/passes/onnx/onnx_dag.py @@ -278,9 +278,9 @@ def add_value_info(self, value_info: ValueInfoProto, graph_idx: int, overwrite: self.ios[name] = OnnxIO(proto=[value_info], graph_idx=graph_idx) return - assert ( - overwrite or not self.ios[name].proto - ), f"Value info for {name} already exists in the graph but overwrite is False." + assert overwrite or not self.ios[name].proto, ( + f"Value info for {name} already exists in the graph but overwrite is False." + ) self.ios[name].proto = [value_info] def is_io(self, io_name: str) -> bool: diff --git a/olive/passes/onnx/optimum_conversion.py b/olive/passes/onnx/optimum_conversion.py index b5c420f8f0..4ddf0f4993 100644 --- a/olive/passes/onnx/optimum_conversion.py +++ b/olive/passes/onnx/optimum_conversion.py @@ -101,9 +101,9 @@ def _run_for_config( # check the exported components exported_models = [name.stem for name in Path(output_model_path).iterdir() if name.suffix == ".onnx"] if config.components: - assert all( - component in exported_models for component in config.components - ), f"Components {config['components']} are not exported. Only {exported_models} are exported." + assert all(component in exported_models for component in config.components), ( + f"Components {config['components']} are not exported. Only {exported_models} are exported." + ) components = config.components or exported_models logger.debug("Exported models are: %s. Returning components: %s.", exported_models, components) diff --git a/olive/passes/onnx/pipeline/step_utils.py b/olive/passes/onnx/pipeline/step_utils.py index b19f010d86..ccec0fe968 100644 --- a/olive/passes/onnx/pipeline/step_utils.py +++ b/olive/passes/onnx/pipeline/step_utils.py @@ -8,13 +8,11 @@ import onnx # pylint: disable=wildcard-import -from onnxruntime_extensions.tools.pre_post_processing import * # noqa: F401, F403, RUF100 +from onnxruntime_extensions.tools.pre_post_processing import * # noqa: F403 from onnxruntime_extensions.tools.pre_post_processing.utils import create_named_value from olive.passes.onnx.pipeline import resolve_placeholder -# ruff: noqa: RUF100, PLW2901 - def parse_steps(model: onnx.ModelProto, config: List[Dict]): """Parse the config and return a dictionary of step name and its parameters. @@ -181,7 +179,7 @@ def parse_step_params(model: onnx.ModelProto, step_config: Dict): param_cls = get_customized_class(param_type) params[param_name] = param_cls(**param_args) elif param_type in ("tuple", "list"): - param_value = param_value.get("value") + param_value = param_value.get("value") # noqa: PLW2901 # explicitly list or tuple type is specified assert isinstance(param_value, list) diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py index 5a76ae30e4..51a1775eb3 100644 --- a/olive/passes/onnx/static_llm.py +++ b/olive/passes/onnx/static_llm.py @@ -57,9 +57,9 @@ def _run_for_config( assert isinstance(model, CompositeModelHandler), "StaticLLM pass only supports CompositeModelHandler" model_components = list(model.model_components) assert all(isinstance(m, ONNXModelHandler) for m in model_components), "All components must be ONNXModelHandler" - assert ( - len(model_components) >= 3 - ), "There should be at least 3 components in the model: embedding, transformer, and lm_head." + assert len(model_components) >= 3, ( + "There should be at least 3 components in the model: embedding, transformer, and lm_head." + ) # only gqa models are supported for now assert ( @@ -175,7 +175,7 @@ def fix_shape(model_proto: onnx.ModelProto, param_mapping: Dict[str, int]): for old_dim, new_dim in zip(original_shape, new_shape): if isinstance(old_dim, str) and isinstance(new_dim, int): if old_dim in param_mapping: - assert ( - param_mapping[old_dim] == new_dim - ), f"Param {old_dim} already exists with different value. Something is wrong." + assert param_mapping[old_dim] == new_dim, ( + f"Param {old_dim} already exists with different value. Something is wrong." + ) param_mapping[old_dim] = new_dim diff --git a/olive/passes/onnx/vitis_ai/quantizer.py b/olive/passes/onnx/vitis_ai/quantizer.py index 6853c0d0e1..4ce219e464 100644 --- a/olive/passes/onnx/vitis_ai/quantizer.py +++ b/olive/passes/onnx/vitis_ai/quantizer.py @@ -1265,9 +1265,9 @@ def quantize_initializer_impl(self, weight, qType, reduce_range=False, keep_floa scale = np.array(quant_overrides["scale"]) q_weight_data = quantize_nparray(qType, weight_data.flatten(), scale, zero_point) assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}" - assert ( - zero_point.dtype != np.float32 and zero_point.dtype != np.float16 - ), f"Unexpected dtype {zero_point.dtype}" + assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, ( + f"Unexpected dtype {zero_point.dtype}" + ) assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}" else: @@ -1280,9 +1280,9 @@ def quantize_initializer_impl(self, weight, qType, reduce_range=False, keep_floa ) assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}" - assert ( - zero_point.dtype != np.float32 and zero_point.dtype != np.float16 - ), f"Unexpected dtype {zero_point.dtype}" + assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, ( + f"Unexpected dtype {zero_point.dtype}" + ) assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}" scale_dtype = weight.data_type @@ -1383,13 +1383,13 @@ def quantize_weight_per_channel_impl( weight_qType, per_channel_data.flatten(), scale, zero_point ) assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}" - assert ( - zero_point.dtype != np.float32 and zero_point.dtype != np.float16 - ), f"Unexpected dtype {zero_point.dtype}" + assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, ( + f"Unexpected dtype {zero_point.dtype}" + ) assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}" - assert isinstance( - quantized_per_channel_data, np.ndarray - ), f"Unexpected type {type(quantized_per_channel_data)}" + assert isinstance(quantized_per_channel_data, np.ndarray), ( + f"Unexpected type {type(quantized_per_channel_data)}" + ) else: _, _, zero_point, scale, quantized_per_channel_data = quantize_data_pof2s( @@ -1401,13 +1401,13 @@ def quantize_weight_per_channel_impl( ) assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}" - assert ( - zero_point.dtype != np.float32 and zero_point.dtype != np.float16 - ), f"Unexpected dtype {zero_point.dtype}" + assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, ( + f"Unexpected dtype {zero_point.dtype}" + ) assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}" - assert isinstance( - quantized_per_channel_data, np.ndarray - ), f"Unexpected type {type(quantized_per_channel_data)}" + assert isinstance(quantized_per_channel_data, np.ndarray), ( + f"Unexpected type {type(quantized_per_channel_data)}" + ) zero_point_list.append(zero_point) scale_list.append(scale) diff --git a/olive/passes/onnx/vitis_ai/refine.py b/olive/passes/onnx/vitis_ai/refine.py index d1f0df5bbb..3fb48eb9b4 100644 --- a/olive/passes/onnx/vitis_ai/refine.py +++ b/olive/passes/onnx/vitis_ai/refine.py @@ -174,8 +174,7 @@ def adjust_shift_cut(self): new_wpos = new_sc + opos - ipos self.set_pos(wpos_node, new_wpos) logger.info( - "Shift cut of layer {} is {}. It exceeds range [{}, {}]. " - "Modify wpos from {} to {}.".format( + "Shift cut of layer {} is {}. It exceeds range [{}, {}]. Modify wpos from {} to {}.".format( node.input[1], int(sc), int(min_sc), int(max_sc), int(wpos), int(new_wpos) ) ) @@ -222,8 +221,7 @@ def adjust_shift_bias(self): new_bpos = wpos + ipos - new_sb self.set_pos(self.get_node_by_name(node.input[2].strip(postfix)), new_bpos) logger.info( - "Shift bias of layer {} is {}. It exceeds range [{}, {}]. " - "Modify bpos from {} to {}.".format( + "Shift bias of layer {} is {}. It exceeds range [{}, {}]. Modify bpos from {} to {}.".format( node.input[2], int(shift_bias), int(min_sb), int(max_sb), int(bpos), int(new_bpos) ) ) @@ -263,8 +261,9 @@ def adjust_vitis_sigmoid(self): if new_ipos != ipos: self.set_pos(self.get_node_by_name(node.input[0].strip(postfix)), new_ipos) logger.info( - "Input quantize pos of VitisSimoid layer {} is {}, modify it to {} " - "to meet the DPU constraints.".format(node.input[0], int(ipos), int(new_ipos)) + "Input quantize pos of VitisSimoid layer {} is {}, modify it to {} to meet the DPU constraints.".format( + node.input[0], int(ipos), int(new_ipos) + ) ) if new_opos != opos: @@ -296,8 +295,9 @@ def adjust_shift_read(self): ipos, _ = self.get_pos_by_name(i) if ipos is None: logger.info( - "Fail to get quantize position for layer {}, " - "skip adjust_shift_read for it.".format(ipos_layers[i]) + "Fail to get quantize position for layer {}, skip adjust_shift_read for it.".format( + ipos_layers[i] + ) ) skip = True iposes.append(ipos) @@ -318,8 +318,7 @@ def adjust_shift_read(self): new_ipos_max = iposes[id_min] + new_sr self.set_pos(self.get_node_by_name(ipos_layers[id_max].strip(postfix)), new_ipos_max) logger.info( - "Shift read of layer {} is {}({}-{}). It exceeds range [{}, {}]. " - "Modify ipos from {} to {}.".format( + "Shift read of layer {} is {}({}-{}). It exceeds range [{}, {}]. Modify ipos from {} to {}.".format( node.name, int(sr), int(iposes[id_max]), @@ -362,8 +361,9 @@ def adjust_shift_write(self): opos, _ = self.get_pos_by_name(opos_name) if opos is None: logger.info( - "Fail to get quantize position for layer {}(output:0), " - "skip adjust_shift_write for it.".format(node.name) + "Fail to get quantize position for layer {}(output:0), skip adjust_shift_write for it.".format( + node.name + ) ) if skip: continue @@ -417,8 +417,9 @@ def align_concat(self): self.set_pos(self.get_node_by_name(self.find_o_name(node.output[0])), min_pos) logger.info( ( - "Output pos of concat node {} is {}, min_pos is {}. " - "Modify opos from {} to {}.".format(node.name, int(opos), int(min_pos), int(opos), int(min_pos)) + "Output pos of concat node {} is {}, min_pos is {}. Modify opos from {} to {}.".format( + node.name, int(opos), int(min_pos), int(opos), int(min_pos) + ) ) ) for name in ipos_layers: @@ -426,8 +427,9 @@ def align_concat(self): if ipos is not None and ipos != min_pos: self.set_pos(ipos_node, min_pos) logger.info( - "Input pos of concat node {} is {}, min_pos is {}. " - "Modify ipos from {} to {}.".format(node.name, int(ipos), int(min_pos), int(ipos), int(min_pos)) + "Input pos of concat node {} is {}, min_pos is {}. Modify ipos from {} to {}.".format( + node.name, int(ipos), int(min_pos), int(ipos), int(min_pos) + ) ) def align_pool(self): diff --git a/olive/passes/openvino/quantization.py b/olive/passes/openvino/quantization.py index f489051a7c..0a459f3485 100644 --- a/olive/passes/openvino/quantization.py +++ b/olive/passes/openvino/quantization.py @@ -186,7 +186,6 @@ def _get_extra_params(config): class OpenVINOQuantization(OpenVINOQuantizationBase): - def _run_for_config( self, model: OpenVINOModelHandler, config: Type[BasePassConfig], output_model_path: str ) -> OpenVINOModelHandler: @@ -275,7 +274,7 @@ def _run_for_config( validation_fn=validate_func, max_drop=config.max_drop, drop_type=drop_type, - **extra_params + **extra_params, ) model_name = "ov_model" diff --git a/olive/passes/pass_config.py b/olive/passes/pass_config.py index 757bf1354c..70f4fcfb38 100644 --- a/olive/passes/pass_config.py +++ b/olive/passes/pass_config.py @@ -94,7 +94,6 @@ def get_user_script_data_config( class BasePassConfig(ConfigBase): - @validator("*", pre=True) def _validate_default_str(cls, v, field): if not isinstance(v, (str, PassParamDefault)) or v not in DEFAULT_SET: @@ -167,7 +166,6 @@ def create_config_class( class PassModuleConfig(ConfigBase): - ACCELERATORS: ClassVar[Set[str]] = {v.value for v in Device} PRECISIONS: ClassVar[Set[str]] = {v.value for v in Precision} QUANT_ALGORITHMS: ClassVar[Set[str]] = {v.value for v in QuantAlgorithm} diff --git a/olive/passes/pytorch/qat_utils.py b/olive/passes/pytorch/qat_utils.py index bcd53d82b1..c89943cbdb 100644 --- a/olive/passes/pytorch/qat_utils.py +++ b/olive/passes/pytorch/qat_utils.py @@ -92,7 +92,7 @@ def execute_local(self) -> PyTorchModelHandler: max_steps=self.config.num_steps, logger=self.config.logger, default_root_dir=self.config.checkpoint_path, - **kwargs + **kwargs, ) trainer.fit(ptl_module, datamodule=ptl_data_module) diff --git a/olive/passes/pytorch/sgdg.py b/olive/passes/pytorch/sgdg.py index 38a8fe3bdf..c74ecd97a3 100644 --- a/olive/passes/pytorch/sgdg.py +++ b/olive/passes/pytorch/sgdg.py @@ -134,7 +134,6 @@ def step(self, closure=None): unity, _ = unit(p.data.view(p.size()[0], -1)) if stiefel and unity.size()[0] <= unity.size()[1]: - weight_decay = group["weight_decay"] dampening = group["dampening"] nesterov = group["nesterov"] diff --git a/olive/passes/pytorch/sparsegpt_utils.py b/olive/passes/pytorch/sparsegpt_utils.py index 4ff7d0d665..03b701160f 100644 --- a/olive/passes/pytorch/sparsegpt_utils.py +++ b/olive/passes/pytorch/sparsegpt_utils.py @@ -14,7 +14,7 @@ logger = logging.getLogger(__name__) -# ruff: noqa: N802, N806, RUF100 +# ruff: noqa: N802, N806 # model types supported by SparseGPT supported_models = ("bloom", "gpt2", "gpt_neox", "llama", "opt") @@ -236,10 +236,7 @@ def prune(self, mode, sparsity=None, n=None, m=None, blocksize=128, percdamp=0.0 if mode == "structured" and col % m == 0: # every mth column, set bottom n weights to True (prune) - magnitude = ( - W1[:, col : (col + m)] ** 2 # noqa: E203, RUF100 - / (torch.diag(Hinv1)[col : (col + m)].reshape((1, -1))) ** 2 # noqa: E203, RUF100 - ) + magnitude = W1[:, col : (col + m)] ** 2 / (torch.diag(Hinv1)[col : (col + m)].reshape((1, -1))) ** 2 mask1.scatter_(1, col + torch.topk(magnitude, n, dim=1, largest=False)[1], True) # freeze weights in current column diff --git a/olive/passes/pytorch/tensor_parallel_layers.py b/olive/passes/pytorch/tensor_parallel_layers.py index b078995d48..42805ede4b 100644 --- a/olive/passes/pytorch/tensor_parallel_layers.py +++ b/olive/passes/pytorch/tensor_parallel_layers.py @@ -81,9 +81,7 @@ def reset_parameters(self) -> None: def extra_repr(self) -> str: # From `torch.nn.Linear` - return "in_features={}, out_features={}, bias={}".format( - self.in_features, self.out_features, self.bias is not None - ) + return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}" def forward(self, ip: torch.Tensor) -> torch.Tensor: return F.linear(ip, weight=self.weight, bias=self.bias) # pylint: disable=not-callable @@ -136,6 +134,4 @@ def forward(self, ip: torch.Tensor) -> torch.Tensor: def extra_repr(self) -> str: # From `torch.nn.Linear` - return "in_features={}, out_features={}, bias={}".format( - self.in_features, self.out_features, self.bias is not None - ) + return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}" diff --git a/olive/passes/pytorch/trt_utils.py b/olive/passes/pytorch/trt_utils.py index e7bb1c06be..bde1291f81 100644 --- a/olive/passes/pytorch/trt_utils.py +++ b/olive/passes/pytorch/trt_utils.py @@ -13,8 +13,12 @@ from typing import TYPE_CHECKING import tensorrt as trt -from torch_tensorrt.fx import compile # noqa: A004 # pylint: disable=redefined-builtin -from torch_tensorrt.fx import InputTensorSpec, TRTInterpreter, TRTModule +from torch_tensorrt.fx import ( # pylint: disable=redefined-builtin + InputTensorSpec, + TRTInterpreter, + TRTModule, + compile, # noqa: A004 +) from torch_tensorrt.fx.tracer.acc_tracer import acc_tracer if TYPE_CHECKING: diff --git a/olive/platform_sdk/qualcomm/utils/data_loader.py b/olive/platform_sdk/qualcomm/utils/data_loader.py index 12a71ce480..67301ac1ba 100644 --- a/olive/platform_sdk/qualcomm/utils/data_loader.py +++ b/olive/platform_sdk/qualcomm/utils/data_loader.py @@ -75,21 +75,19 @@ def prepare_batches(self): self.batches = [] for i in range(0, len(self.batch_input_list_metadata), self.batch_size): - self.batches.append(self.batch_input_list_metadata[i : i + self.batch_size]) # noqa: E203, RUF100 + self.batches.append(self.batch_input_list_metadata[i : i + self.batch_size]) self.num_batches = len(self.batches) def get_batch(self, batch_id): if batch_id >= self.num_batches: - raise ValueError("batch_id should be less than {}".format(self.num_batches)) + raise ValueError(f"batch_id should be less than {self.num_batches}") if self.batch_size is None: return self.data_dir, self.input_list, self.annotation else: annotation = None if self.annotation is not None: - annotation = self.annotation[ - self.batch_size * batch_id : self.batch_size * (batch_id + 1) # noqa: E203, RUF100 - ] + annotation = self.annotation[self.batch_size * batch_id : self.batch_size * (batch_id + 1)] batch = self.batches[batch_id] # empty batch directory and copy current batch data @@ -279,9 +277,9 @@ def __init__(self, dataloader: Any, io_config: dict, batch_size: int = None): # get permutation from source shape to target shape target_shape = input_spec["target_shape"] - assert len(source_shape) == len( - target_shape - ), f"Source shape {source_shape} and target shape {target_shape} must have the same length" + assert len(source_shape) == len(target_shape), ( + f"Source shape {source_shape} and target shape {target_shape} must have the same length" + ) # find the permutation of the source shape that matches the target shape # e.g. source_shape = [1, 3, 224, 224], target_shape = [1, 224, 224, 3] @@ -321,9 +319,9 @@ def __init__(self, dataloader: Any, io_config: dict, batch_size: int = None): input_data = dict(zip(input_specs.keys(), [input_data_i])) else: input_data = input_data_i - assert isinstance( - input_data, dict - ), f"Input data must be a tuple, torch.Tensor, np.ndarray, or dict. Got {type(input_data)}" + assert isinstance(input_data, dict), ( + f"Input data must be a tuple, torch.Tensor, np.ndarray, or dict. Got {type(input_data)}" + ) input_file_name = f"{i}.bin".zfill(sample_digits + 4) input_order.append(input_file_name) diff --git a/olive/search/search_parameter.py b/olive/search/search_parameter.py index f3aa2ef486..3ca9845583 100644 --- a/olive/search/search_parameter.py +++ b/olive/search/search_parameter.py @@ -159,9 +159,9 @@ def condition(self, parent_values: Dict[str, Any]) -> SearchParameter: parent_value = parent_values[parent] parent_idx = i break - new_parents = self.parents[:parent_idx] + self.parents[parent_idx + 1 :] # noqa: E203, RUF100 + new_parents = self.parents[:parent_idx] + self.parents[parent_idx + 1 :] new_support = { - key[:parent_idx] + key[parent_idx + 1 :]: value # noqa: E203, RUF100 + key[:parent_idx] + key[parent_idx + 1 :]: value for key, value in self.support.items() if key[parent_idx] == parent_value } diff --git a/olive/search/search_space.py b/olive/search/search_space.py index caa7fae173..b886fe06a7 100644 --- a/olive/search/search_space.py +++ b/olive/search/search_space.py @@ -90,9 +90,9 @@ class SearchSpace: """ def __init__(self, parameters: List[Tuple[str, Union[SearchParameter, "SearchSpace"]]]): - assert len(parameters) == len( - {name for name, _ in parameters} - ), "Parameter name in search space should be unique." + assert len(parameters) == len({name for name, _ in parameters}), ( + "Parameter name in search space should be unique." + ) self._parameters = self._order_search_space(parameters) diff --git a/pyproject.toml b/pyproject.toml index 90f04478f2..01ec259455 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,3 @@ -[tool.black] -line-length = 120 - -[tool.isort] -profile = "black" -line_length = 120 - [tool.pylint.BASIC] good-names = [ "a", @@ -77,7 +70,8 @@ disable = [ "missing-docstring", "fixme", "unspecified-encoding", - "unused-argument" + "unused-argument", + "wrong-import-order" # import order is handled by the formatter (ruff) ] [tool.ruff] @@ -95,6 +89,7 @@ select = [ "F", # Pyflakes "FLY", # flake8-flynt "G", # flake8-logging-format + "I", # isort "ICN", # flake8-import-conventions "INP", # flake8-no-pep420 "INT", # flake8-gettext @@ -143,6 +138,7 @@ ignore = [ "D107", # Ignore missing docstring in __init__ "D406", # Ignore new line after section name "D407", # Ignore dashed-underline-after-section + "E501", # Ignore line too long. The formatter will handle it. "N803", # Argument casing "N812", # Allow import torch.nn.functional as F "N999", # Module names @@ -172,14 +168,7 @@ ignore = [ "TRY002", # Ignore create custom exception "TRY003", # Ignore check message not defined in the exception class "TRY004", # Ignore prefer TypeError over ValueError - "TRY300", # Ignore check if return in try block - "UP032", # Ignore string format calls - "UP038" # Ignore old isinstance hints -] -ignore-init-module-imports = true -unfixable = [ - "F401", # Unused imports - "SIM112" # Use upper case for env vars + "TRY300" # Ignore check if return in try block ] [tool.ruff.lint.flake8-tidy-imports] @@ -198,4 +187,3 @@ classmethod-decorators = ["classmethod", "olive.common.pydantic_v1.validator", " "scripts/**" = ["INP001"] "examples/directml/llm/chat_app/**" = ["TID252", "UP006", "T201"] "olive/cli/**" = ["T201"] -"olive/**/hadamard_utils.py" = ["E501"] diff --git a/requirements-dev.txt b/requirements-dev.txt index 0de2b3a271..2bd5e1988a 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,9 +1,6 @@ -r requirements.txt -black editorconfig-checker -flake8 -isort lintrunner lintrunner-adapters pylint<3.2.4 # https://github.com/pylint-dev/pylint/issues/9751 -ruff>=0.2.0 +ruff==0.11.4 diff --git a/test/integ_test/aml_model_test/test_aml_model.py b/test/integ_test/aml_model_test/test_aml_model.py index 04462f7daf..f5d02917c6 100644 --- a/test/integ_test/aml_model_test/test_aml_model.py +++ b/test/integ_test/aml_model_test/test_aml_model.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.integ_test.utils import get_olive_workspace_config from olive.azureml.azureml_client import AzureMLClientConfig from olive.model import ModelConfig @@ -11,6 +10,7 @@ from olive.passes.onnx.conversion import OnnxConversion from olive.resource_path import ResourcePath from olive.systems.azureml import AzureMLDockerConfig, AzureMLSystem +from test.integ_test.utils import get_olive_workspace_config def test_aml_model_pass_run(tmp_path): diff --git a/test/integ_test/aml_resource_path/test_aml_resource_path.py b/test/integ_test/aml_resource_path/test_aml_resource_path.py index 45cb3c434a..8b44ccaee1 100644 --- a/test/integ_test/aml_resource_path/test_aml_resource_path.py +++ b/test/integ_test/aml_resource_path/test_aml_resource_path.py @@ -3,11 +3,11 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.integ_test.utils import get_olive_workspace_config import pytest from olive.resource_path import ResourceType, create_resource_path +from test.integ_test.utils import get_olive_workspace_config # pylint: disable=attribute-defined-outside-init diff --git a/test/integ_test/evaluator/azureml_eval/test_aml_evaluation.py b/test/integ_test/evaluator/azureml_eval/test_aml_evaluation.py index b599ef2b59..d8afe4ecbc 100644 --- a/test/integ_test/evaluator/azureml_eval/test_aml_evaluation.py +++ b/test/integ_test/evaluator/azureml_eval/test_aml_evaluation.py @@ -2,6 +2,14 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +from typing import ClassVar, List + +import pytest + +from olive.evaluator.metric_result import joint_metric_key +from olive.evaluator.olive_evaluator import OliveEvaluatorConfig +from olive.hardware import DEFAULT_CPU_ACCELERATOR +from olive.model import ModelConfig from test.integ_test.evaluator.azureml_eval.utils import ( delete_directories, download_data, @@ -13,14 +21,6 @@ get_onnx_model, get_pytorch_model, ) -from typing import ClassVar, List - -import pytest - -from olive.evaluator.metric_result import joint_metric_key -from olive.evaluator.olive_evaluator import OliveEvaluatorConfig -from olive.hardware import DEFAULT_CPU_ACCELERATOR -from olive.model import ModelConfig class TestAMLEvaluation: diff --git a/test/integ_test/evaluator/azureml_eval/utils.py b/test/integ_test/evaluator/azureml_eval/utils.py index ca2d9bc54a..aee30214cc 100644 --- a/test/integ_test/evaluator/azureml_eval/utils.py +++ b/test/integ_test/evaluator/azureml_eval/utils.py @@ -4,7 +4,6 @@ # -------------------------------------------------------------------------- import shutil from pathlib import Path -from test.integ_test.utils import download_azure_blob, get_olive_workspace_config from torchvision import datasets from torchvision.transforms import ToTensor @@ -14,6 +13,7 @@ from olive.data.config import DataComponentConfig, DataConfig from olive.evaluator.metric import AccuracySubType, LatencySubType, Metric, MetricType from olive.systems.azureml import AzureMLDockerConfig, AzureMLSystem +from test.integ_test.utils import download_azure_blob, get_olive_workspace_config # pylint: disable=redefined-outer-name diff --git a/test/integ_test/evaluator/docker_eval/test_docker_evaluation.py b/test/integ_test/evaluator/docker_eval/test_docker_evaluation.py index 89ca0e646a..21ff52ca3a 100644 --- a/test/integ_test/evaluator/docker_eval/test_docker_evaluation.py +++ b/test/integ_test/evaluator/docker_eval/test_docker_evaluation.py @@ -4,6 +4,15 @@ # -------------------------------------------------------------------------- import platform from functools import partial +from typing import ClassVar, List + +import pytest + +from olive.common.constants import OS +from olive.evaluator.metric_result import joint_metric_key +from olive.evaluator.olive_evaluator import OliveEvaluatorConfig +from olive.hardware import DEFAULT_CPU_ACCELERATOR +from olive.model import ModelConfig from test.integ_test.evaluator.docker_eval.utils import ( delete_directories, download_data, @@ -17,15 +26,6 @@ get_openvino_model, get_pytorch_model, ) -from typing import ClassVar, List - -import pytest - -from olive.common.constants import OS -from olive.evaluator.metric_result import joint_metric_key -from olive.evaluator.olive_evaluator import OliveEvaluatorConfig -from olive.hardware import DEFAULT_CPU_ACCELERATOR -from olive.model import ModelConfig class TestDockerEvaluation: diff --git a/test/integ_test/evaluator/docker_eval/utils.py b/test/integ_test/evaluator/docker_eval/utils.py index 36517abf71..9807a7ffd9 100644 --- a/test/integ_test/evaluator/docker_eval/utils.py +++ b/test/integ_test/evaluator/docker_eval/utils.py @@ -4,7 +4,6 @@ # -------------------------------------------------------------------------- import shutil from pathlib import Path -from test.integ_test.utils import download_azure_blob from zipfile import ZipFile from torchvision import datasets @@ -14,6 +13,7 @@ from olive.data.config import DataComponentConfig, DataConfig from olive.evaluator.metric import AccuracySubType, LatencySubType, Metric, MetricType from olive.systems.docker import DockerSystem, LocalDockerConfig +from test.integ_test.utils import download_azure_blob # pylint: disable=redefined-outer-name diff --git a/test/integ_test/evaluator/local_eval/test_local_evaluation.py b/test/integ_test/evaluator/local_eval/test_local_evaluation.py index 76432ee6fb..5f7bfef8a6 100644 --- a/test/integ_test/evaluator/local_eval/test_local_evaluation.py +++ b/test/integ_test/evaluator/local_eval/test_local_evaluation.py @@ -3,6 +3,15 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from functools import partial +from typing import ClassVar, List + +import pytest + +from olive.evaluator.metric_result import joint_metric_key +from olive.evaluator.olive_evaluator import OliveEvaluatorConfig +from olive.hardware import DEFAULT_CPU_ACCELERATOR +from olive.model import ModelConfig +from olive.systems.local import LocalSystem from test.integ_test.evaluator.local_eval.utils import ( delete_directories, get_accuracy_metric, @@ -15,15 +24,6 @@ get_openvino_model, get_pytorch_model, ) -from typing import ClassVar, List - -import pytest - -from olive.evaluator.metric_result import joint_metric_key -from olive.evaluator.olive_evaluator import OliveEvaluatorConfig -from olive.hardware import DEFAULT_CPU_ACCELERATOR -from olive.model import ModelConfig -from olive.systems.local import LocalSystem # pylint: disable=redefined-builtin diff --git a/test/integ_test/evaluator/local_eval/utils.py b/test/integ_test/evaluator/local_eval/utils.py index df5a0a1bdd..576dbda1b8 100644 --- a/test/integ_test/evaluator/local_eval/utils.py +++ b/test/integ_test/evaluator/local_eval/utils.py @@ -4,12 +4,12 @@ # -------------------------------------------------------------------------- import shutil from pathlib import Path -from test.integ_test.utils import download_azure_blob from zipfile import ZipFile from olive.common.config_utils import validate_config from olive.data.config import DataComponentConfig, DataConfig from olive.evaluator.metric import AccuracySubType, LatencySubType, Metric, MetricType +from test.integ_test.utils import download_azure_blob # pylint: disable=redefined-outer-name diff --git a/test/integ_test/pass_runner/test_docker_system.py b/test/integ_test/pass_runner/test_docker_system.py index 7e5454cdd2..be24c2a564 100644 --- a/test/integ_test/pass_runner/test_docker_system.py +++ b/test/integ_test/pass_runner/test_docker_system.py @@ -1,11 +1,4 @@ import platform -from test.integ_test.evaluator.docker_eval.utils import ( - delete_directories, - download_models, - get_directories, - get_docker_target, - get_onnx_model, -) import pytest @@ -15,6 +8,13 @@ from olive.model.config.model_config import ModelConfig from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.session_params_tuning import OrtSessionParamsTuning +from test.integ_test.evaluator.docker_eval.utils import ( + delete_directories, + download_models, + get_directories, + get_docker_target, + get_onnx_model, +) @pytest.mark.skipif(platform.system() == OS.WINDOWS, reason="Docker target does not support windows") diff --git a/test/multiple_ep/test_aml_system.py b/test/multiple_ep/test_aml_system.py index 60a755fac8..3f0fcb1101 100644 --- a/test/multiple_ep/test_aml_system.py +++ b/test/multiple_ep/test_aml_system.py @@ -19,11 +19,10 @@ class TestOliveAzureMLSystem: @pytest.fixture(autouse=True) def setup(self): # use the olive managed AzureML system as the test environment - from test.integ_test.utils import get_olive_workspace_config - from test.multiple_ep.utils import download_data, download_models, get_onnx_model - from olive.azureml.azureml_client import AzureMLClientConfig from olive.systems.system_config import AzureMLTargetUserConfig, SystemConfig + from test.integ_test.utils import get_olive_workspace_config + from test.multiple_ep.utils import download_data, download_models, get_onnx_model aml_compute = "cpu-cluster" azureml_client_config = AzureMLClientConfig(**get_olive_workspace_config()) diff --git a/test/multiple_ep/test_docker_system.py b/test/multiple_ep/test_docker_system.py index 27a9a0d5f5..4f8ee004e7 100644 --- a/test/multiple_ep/test_docker_system.py +++ b/test/multiple_ep/test_docker_system.py @@ -5,13 +5,13 @@ import logging import platform from pathlib import Path -from test.multiple_ep.utils import get_directories import pytest from olive.common.constants import OS from olive.logging import set_default_logger_severity from olive.model import ModelConfig +from test.multiple_ep.utils import get_directories # pylint: disable=attribute-defined-outside-init @@ -20,9 +20,8 @@ class TestOliveManagedDockerSystem: @pytest.fixture(autouse=True) def setup(self): - from test.multiple_ep.utils import download_data, download_models, get_onnx_model - from olive.systems.system_config import DockerTargetUserConfig, SystemConfig + from test.multiple_ep.utils import download_data, download_models, get_onnx_model # use the olive managed Docker system as the test environment self.system_config = SystemConfig( diff --git a/test/multiple_ep/test_python_env_system.py b/test/multiple_ep/test_python_env_system.py index ab6df03431..ca1bd49f41 100644 --- a/test/multiple_ep/test_python_env_system.py +++ b/test/multiple_ep/test_python_env_system.py @@ -3,12 +3,12 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import platform -from test.unit_test.utils import create_onnx_model_file, get_custom_metric, get_onnx_model_config import pytest from olive.common.constants import OS from olive.systems.system_config import PythonEnvironmentTargetUserConfig, SystemConfig +from test.unit_test.utils import create_onnx_model_file, get_custom_metric, get_onnx_model_config # pylint: disable=attribute-defined-outside-init diff --git a/test/multiple_ep/utils.py b/test/multiple_ep/utils.py index 51243a416b..212086100b 100644 --- a/test/multiple_ep/utils.py +++ b/test/multiple_ep/utils.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.integ_test.utils import download_azure_blob from torchvision import datasets from torchvision.transforms import ToTensor @@ -14,6 +13,7 @@ from olive.evaluator.olive_evaluator import OliveEvaluatorConfig from olive.passes.onnx.session_params_tuning import OrtSessionParamsTuning from olive.systems.accelerator_creator import create_accelerators +from test.integ_test.utils import download_azure_blob # pylint: disable=redefined-outer-name diff --git a/test/unit_test/auto_optimizer/test_auto_optimizer.py b/test/unit_test/auto_optimizer/test_auto_optimizer.py index c7972363f5..f39f74827b 100644 --- a/test/unit_test/auto_optimizer/test_auto_optimizer.py +++ b/test/unit_test/auto_optimizer/test_auto_optimizer.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import get_accuracy_metric, get_glue_huggingface_data_config import pytest import yaml @@ -14,6 +13,7 @@ from olive.evaluator.olive_evaluator import OliveEvaluatorConfig from olive.hardware import DEFAULT_CPU_ACCELERATOR, DEFAULT_GPU_CUDA_ACCELERATOR, DEFAULT_GPU_TRT_ACCELERATOR from olive.model import ModelConfig +from test.unit_test.utils import get_accuracy_metric, get_glue_huggingface_data_config # pylint: disable=attribute-defined-outside-init diff --git a/test/unit_test/conftest.py b/test/unit_test/conftest.py index 672741aed6..95fc71cfe7 100644 --- a/test/unit_test/conftest.py +++ b/test/unit_test/conftest.py @@ -3,10 +3,11 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import shutil -from test.unit_test.utils import create_onnx_model_file, delete_onnx_model_files import pytest +from test.unit_test.utils import create_onnx_model_file, delete_onnx_model_files + @pytest.fixture(scope="session", autouse=True) def setup_onnx_model(request, tmp_path_factory): diff --git a/test/unit_test/data_container/test_data_config.py b/test/unit_test/data_container/test_data_config.py index dd4e864e80..3627e739de 100644 --- a/test/unit_test/data_container/test_data_config.py +++ b/test/unit_test/data_container/test_data_config.py @@ -3,13 +3,12 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import get_data_config - import pytest from olive.common.config_utils import validate_config from olive.data.config import DataConfig from olive.data.registry import Registry +from test.unit_test.utils import get_data_config # pylint: disable=attribute-defined-outside-init diff --git a/test/unit_test/data_container/test_data_container.py b/test/unit_test/data_container/test_data_container.py index 18e0f9a8a3..8a81b9ce63 100644 --- a/test/unit_test/data_container/test_data_container.py +++ b/test/unit_test/data_container/test_data_container.py @@ -3,6 +3,11 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +import numpy as np +import pytest + +from olive.data.config import DataConfig +from olive.data.container.data_container import DataContainer from test.unit_test.utils import ( create_raw_data, get_data_config, @@ -10,12 +15,6 @@ get_transformer_dummy_input_data_config, ) -import numpy as np -import pytest - -from olive.data.config import DataConfig -from olive.data.container.data_container import DataContainer - # pylint: disable=attribute-defined-outside-init diff --git a/test/unit_test/data_container/test_dataloader.py b/test/unit_test/data_container/test_dataloader.py index b49cf667b6..2336dd2c41 100644 --- a/test/unit_test/data_container/test_dataloader.py +++ b/test/unit_test/data_container/test_dataloader.py @@ -3,13 +3,12 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import make_local_tiny_llama - import pytest from olive.data.component.dataloader import LLMAugmentedDataLoader from olive.data.template import huggingface_data_config_template from olive.passes.olive_pass import create_pass_from_dict +from test.unit_test.utils import make_local_tiny_llama @pytest.mark.parametrize("use_gqa", [True, False]) diff --git a/test/unit_test/data_container/test_template.py b/test/unit_test/data_container/test_template.py index 2f957ea3a9..340f49bfe9 100644 --- a/test/unit_test/data_container/test_template.py +++ b/test/unit_test/data_container/test_template.py @@ -3,13 +3,13 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import create_raw_data from unittest.mock import patch import pytest import olive.data.template as data_config_template from olive.data.config import DataComponentConfig +from test.unit_test.utils import create_raw_data class TestDataConfigTemplate: diff --git a/test/unit_test/engine/packaging/test_packaging_generator.py b/test/unit_test/engine/packaging/test_packaging_generator.py index 9e108fea62..42388161a7 100644 --- a/test/unit_test/engine/packaging/test_packaging_generator.py +++ b/test/unit_test/engine/packaging/test_packaging_generator.py @@ -6,7 +6,6 @@ import shutil import zipfile from pathlib import Path -from test.unit_test.utils import get_accuracy_metric, get_pytorch_model_config from unittest.mock import Mock, patch import mlflow @@ -32,6 +31,7 @@ from olive.hardware import DEFAULT_CPU_ACCELERATOR from olive.hardware.accelerator import AcceleratorSpec from olive.passes.onnx.conversion import OnnxConversion +from test.unit_test.utils import get_accuracy_metric, get_pytorch_model_config # TODO(team): no engine API envolved, use generate_output_artifacts API directly @@ -536,12 +536,12 @@ def verify_models_rank_json_file(output_dir, file_path, save_as_external_data=Fa model_path = output_dir / Path(model_data["model_config"]["config"]["model_path"]) assert model_path.exists(), "Model path in model rank file does not exist." if export_in_mlflow_format: - assert mlflow.onnx.load_model( - str(model_path) - ), "Model path in model rank file is not a valid MLflow model path." + assert mlflow.onnx.load_model(str(model_path)), ( + "Model path in model rank file is not a valid MLflow model path." + ) elif save_as_external_data: - assert onnx.load( - str(model_path / "model.onnx") - ), "With external data, model path in model rank file is not a valid ONNX model path." + assert onnx.load(str(model_path / "model.onnx")), ( + "With external data, model path in model rank file is not a valid ONNX model path." + ) else: assert onnx.load(str(model_path)), "Model path in model rank file is not a valid ONNX model path." diff --git a/test/unit_test/engine/test_engine.py b/test/unit_test/engine/test_engine.py index 85ad2ea61f..e8be0766e8 100644 --- a/test/unit_test/engine/test_engine.py +++ b/test/unit_test/engine/test_engine.py @@ -5,14 +5,6 @@ import json import logging from pathlib import Path -from test.unit_test.utils import ( - get_accuracy_metric, - get_composite_onnx_model_config, - get_onnx_model_config, - get_onnxconversion_pass, - get_pytorch_model_config, - get_pytorch_model_io_config, -) from unittest.mock import MagicMock, patch import pytest @@ -30,6 +22,14 @@ from olive.systems.accelerator_creator import create_accelerators from olive.systems.common import SystemType from olive.systems.system_config import LocalTargetUserConfig, SystemConfig +from test.unit_test.utils import ( + get_accuracy_metric, + get_composite_onnx_model_config, + get_onnx_model_config, + get_onnxconversion_pass, + get_pytorch_model_config, + get_pytorch_model_io_config, +) # pylint: disable=protected-access diff --git a/test/unit_test/evaluator/test_metric_backend.py b/test/unit_test/evaluator/test_metric_backend.py index 2a376d4dac..d93b59aae2 100644 --- a/test/unit_test/evaluator/test_metric_backend.py +++ b/test/unit_test/evaluator/test_metric_backend.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from functools import partial -from test.unit_test.utils import get_accuracy_metric, get_onnx_model_config, get_pytorch_model_config from typing import ClassVar, List from unittest.mock import patch @@ -14,6 +13,7 @@ from olive.evaluator.olive_evaluator import OliveEvaluatorConfig from olive.hardware import DEFAULT_CPU_ACCELERATOR from olive.systems.local import LocalSystem +from test.unit_test.utils import get_accuracy_metric, get_onnx_model_config, get_pytorch_model_config # pylint: disable=attribute-defined-outside-init, redefined-outer-name diff --git a/test/unit_test/evaluator/test_olive_evaluator.py b/test/unit_test/evaluator/test_olive_evaluator.py index 21ed6f1616..0600e6af3a 100644 --- a/test/unit_test/evaluator/test_olive_evaluator.py +++ b/test/unit_test/evaluator/test_olive_evaluator.py @@ -3,17 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from functools import partial -from test.unit_test.utils import ( - get_accuracy_metric, - get_custom_metric, - get_custom_metric_no_eval, - get_latency_metric, - get_mock_openvino_model, - get_mock_snpe_model, - get_onnx_model, - get_pytorch_model, - get_throughput_metric, -) from types import FunctionType from typing import ClassVar, List from unittest.mock import MagicMock, patch @@ -32,6 +21,17 @@ ) from olive.exception import OliveEvaluationError from olive.hardware.accelerator import Device +from test.unit_test.utils import ( + get_accuracy_metric, + get_custom_metric, + get_custom_metric_no_eval, + get_latency_metric, + get_mock_openvino_model, + get_mock_snpe_model, + get_onnx_model, + get_pytorch_model, + get_throughput_metric, +) class TestOliveEvaluator: @@ -308,7 +308,7 @@ def test_evaluate_latency_with_tunable_op(self, inference_session_mock): { "ep": "ROCMExecutionProvider", "results": { - "onnxruntime::rocm::tunable::blas::internal::GemmTunableOp<__half, ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::RowMajor>": { # noqa: E501 + "onnxruntime::rocm::tunable::blas::internal::GemmTunableOp<__half, ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::RowMajor>": { "NN_992_4096_4096": 300, "NN_992_4096_11008": 664, "NN_984_4096_4096": 1295, diff --git a/test/unit_test/model/test_composite_model.py b/test/unit_test/model/test_composite_model.py index 8acf353c7b..66f3dad384 100644 --- a/test/unit_test/model/test_composite_model.py +++ b/test/unit_test/model/test_composite_model.py @@ -2,13 +2,12 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import get_onnx_model - import pytest from olive.model.config.model_config import ModelConfig from olive.model.handler.composite import CompositeModelHandler from olive.model.handler.onnx import ONNXModelHandler +from test.unit_test.utils import get_onnx_model @pytest.mark.parametrize("as_handler", [True, False]) diff --git a/test/unit_test/model/test_onnx_model.py b/test/unit_test/model/test_onnx_model.py index d5f4af4c37..ef5dad1407 100644 --- a/test/unit_test/model/test_onnx_model.py +++ b/test/unit_test/model/test_onnx_model.py @@ -1,4 +1,3 @@ -from test.unit_test.utils import get_onnx_model from unittest.mock import ANY, MagicMock, patch import numpy as np @@ -8,6 +7,7 @@ from olive.exception import OliveEvaluationError from olive.hardware.accelerator import Device from olive.model import ONNXModelHandler +from test.unit_test.utils import get_onnx_model @patch("onnxruntime.InferenceSession") diff --git a/test/unit_test/passes/onnx/pipeline/test_step_utils.py b/test/unit_test/passes/onnx/pipeline/test_step_utils.py index 603bd11f81..34aa2662e8 100644 --- a/test/unit_test/passes/onnx/pipeline/test_step_utils.py +++ b/test/unit_test/passes/onnx/pipeline/test_step_utils.py @@ -4,13 +4,13 @@ # -------------------------------------------------------------------------- import json from pathlib import Path + +from olive.passes.onnx.pipeline.step_utils import parse_steps from test.unit_test.passes.onnx.test_pre_post_processing_op import ( convert_superresolution_model, get_superresolution_model, ) -from olive.passes.onnx.pipeline.step_utils import parse_steps - class CustomizedParam: def __init__(self, params: dict): diff --git a/test/unit_test/passes/onnx/test_bnb_quantization.py b/test/unit_test/passes/onnx/test_bnb_quantization.py index 979e75f2d6..549b1ba121 100644 --- a/test/unit_test/passes/onnx/test_bnb_quantization.py +++ b/test/unit_test/passes/onnx/test_bnb_quantization.py @@ -2,8 +2,6 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import get_onnx_model, pytorch_model_loader - import onnx import onnxruntime import pytest @@ -13,6 +11,7 @@ from olive.model import ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.bnb_quantization import OnnxBnb4Quantization +from test.unit_test.utils import get_onnx_model, pytorch_model_loader # pylint: disable=protected-access diff --git a/test/unit_test/passes/onnx/test_common.py b/test/unit_test/passes/onnx/test_common.py index f7934eab92..ce4bdf72c4 100644 --- a/test/unit_test/passes/onnx/test_common.py +++ b/test/unit_test/passes/onnx/test_common.py @@ -3,14 +3,13 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import ONNX_MODEL_PATH, get_hf_model - import onnx import pytest from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.common import model_proto_to_olive_model, resave_model from olive.passes.onnx.conversion import OnnxConversion +from test.unit_test.utils import ONNX_MODEL_PATH, get_hf_model @pytest.mark.parametrize( diff --git a/test/unit_test/passes/onnx/test_compose.py b/test/unit_test/passes/onnx/test_compose.py index 26c8de6409..9bd1c8cc2c 100644 --- a/test/unit_test/passes/onnx/test_compose.py +++ b/test/unit_test/passes/onnx/test_compose.py @@ -4,7 +4,6 @@ # -------------------------------------------------------------------------- import json from pathlib import Path -from test.unit_test.utils import make_local_tiny_llama import pytest @@ -15,6 +14,7 @@ from olive.passes.onnx.model_builder import ModelBuilder from olive.passes.onnx.split import SplitModel from olive.passes.onnx.static_llm import StaticLLM +from test.unit_test.utils import make_local_tiny_llama @pytest.mark.parametrize("use_mb", [True, False]) diff --git a/test/unit_test/passes/onnx/test_context_binary.py b/test/unit_test/passes/onnx/test_context_binary.py index 68f4bba906..9a6325e81a 100644 --- a/test/unit_test/passes/onnx/test_context_binary.py +++ b/test/unit_test/passes/onnx/test_context_binary.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import json -from test.unit_test.utils import get_onnx_model import onnxruntime import pytest @@ -13,6 +12,7 @@ from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.common import resave_model from olive.passes.onnx.context_binary import EPContextBinaryGenerator +from test.unit_test.utils import get_onnx_model @pytest.mark.skipif( diff --git a/test/unit_test/passes/onnx/test_conversion.py b/test/unit_test/passes/onnx/test_conversion.py index 8a266e9cd1..6171656a4a 100644 --- a/test/unit_test/passes/onnx/test_conversion.py +++ b/test/unit_test/passes/onnx/test_conversion.py @@ -6,7 +6,6 @@ import shutil from itertools import chain from pathlib import Path -from test.unit_test.utils import ONNX_MODEL_PATH, get_hf_model, get_onnx_model, get_pytorch_model, pytorch_model_loader from typing import Dict, Tuple from unittest.mock import patch @@ -20,6 +19,7 @@ from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.conversion import OnnxConversion, OnnxOpVersionConversion from olive.passes.pytorch.gptq import GptqQuantizer +from test.unit_test.utils import ONNX_MODEL_PATH, get_hf_model, get_onnx_model, get_pytorch_model, pytorch_model_loader @pytest.mark.parametrize( diff --git a/test/unit_test/passes/onnx/test_dynamic_to_fixed_shape.py b/test/unit_test/passes/onnx/test_dynamic_to_fixed_shape.py index d8b4c0dbb8..22cadf5602 100644 --- a/test/unit_test/passes/onnx/test_dynamic_to_fixed_shape.py +++ b/test/unit_test/passes/onnx/test_dynamic_to_fixed_shape.py @@ -1,10 +1,9 @@ -from test.unit_test.utils import create_onnx_model_with_dynamic_axis - import pytest from olive.model import ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.dynamic_to_fixed_shape import DynamicToFixedShape +from test.unit_test.utils import create_onnx_model_with_dynamic_axis @pytest.mark.parametrize( diff --git a/test/unit_test/passes/onnx/test_extract_adapters.py b/test/unit_test/passes/onnx/test_extract_adapters.py index 0194807901..e18834c272 100644 --- a/test/unit_test/passes/onnx/test_extract_adapters.py +++ b/test/unit_test/passes/onnx/test_extract_adapters.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import get_onnx_model import numpy as np import onnx @@ -22,6 +21,7 @@ from olive.passes.onnx.conversion import OnnxConversion from olive.passes.onnx.extract_adapters import ExtractAdapters from olive.passes.onnx.quantization import OnnxMatMul4Quantizer +from test.unit_test.utils import get_onnx_model class LlamaCalibrationDataLoader(CalibrationDataReader): diff --git a/test/unit_test/passes/onnx/test_float16_conversion.py b/test/unit_test/passes/onnx/test_float16_conversion.py index 97083e8ad9..c64038fb30 100644 --- a/test/unit_test/passes/onnx/test_float16_conversion.py +++ b/test/unit_test/passes/onnx/test_float16_conversion.py @@ -2,13 +2,12 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import get_onnx_model - import onnx import pytest from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.float16_conversion import OnnxFloatToFloat16 +from test.unit_test.utils import get_onnx_model @pytest.mark.parametrize("keep_io_types", [True, False]) diff --git a/test/unit_test/passes/onnx/test_graph_surgeries.py b/test/unit_test/passes/onnx/test_graph_surgeries.py index 188e11cbe8..9e78de68fb 100644 --- a/test/unit_test/passes/onnx/test_graph_surgeries.py +++ b/test/unit_test/passes/onnx/test_graph_surgeries.py @@ -4,7 +4,6 @@ # -------------------------------------------------------------------------- import json from pathlib import Path -from test.unit_test.utils import make_local_tiny_llama import numpy as np import onnx @@ -18,6 +17,7 @@ from olive.passes.onnx.graph_surgeries import GraphSurgeries from olive.passes.onnx.model_builder import ModelBuilder from olive.passes.onnx.onnx_dag import OnnxDAG +from test.unit_test.utils import make_local_tiny_llama def get_onnx_model(model_path): @@ -346,9 +346,9 @@ def test_expose_quantized_output(tmp_path): # Validate that the scale node and its initializer exist in the modified model assert any(node.name == scale_node_name for node in output_model.graph.node), "Scale node not added." scale_initializer = next(init for init in output_model.graph.initializer if init.name == scale_initializer_name) - assert np.allclose( - numpy_helper.to_array(scale_initializer), np.array([original_scale_value], dtype=np.float32) - ), "Scale value mismatch." + assert np.allclose(numpy_helper.to_array(scale_initializer), np.array([original_scale_value], dtype=np.float32)), ( + "Scale value mismatch." + ) # Validate that the zero_point node and its initializer exist in the modified model assert any(node.name == zero_point_node_name for node in output_model.graph.node), "Zero point node not added." diff --git a/test/unit_test/passes/onnx/test_insert_beam_search.py b/test/unit_test/passes/onnx/test_insert_beam_search.py index efaf0db974..32c6d45a85 100644 --- a/test/unit_test/passes/onnx/test_insert_beam_search.py +++ b/test/unit_test/passes/onnx/test_insert_beam_search.py @@ -3,13 +3,12 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import get_onnx_model - from transformers import AutoConfig from olive.model import CompositeModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.insert_beam_search import InsertBeamSearch +from test.unit_test.utils import get_onnx_model def test_insert_beam_search_pass(tmp_path): diff --git a/test/unit_test/passes/onnx/test_mixed_precision.py b/test/unit_test/passes/onnx/test_mixed_precision.py index 8e7dd24535..4aa44ce287 100644 --- a/test/unit_test/passes/onnx/test_mixed_precision.py +++ b/test/unit_test/passes/onnx/test_mixed_precision.py @@ -3,10 +3,9 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import get_onnx_model - from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.mixed_precision import OrtMixedPrecision +from test.unit_test.utils import get_onnx_model def test_ort_mixed_precision_pass(tmp_path): diff --git a/test/unit_test/passes/onnx/test_model_builder.py b/test/unit_test/passes/onnx/test_model_builder.py index 1d5c3f2b53..cdd0413357 100644 --- a/test/unit_test/passes/onnx/test_model_builder.py +++ b/test/unit_test/passes/onnx/test_model_builder.py @@ -3,13 +3,13 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import make_local_tiny_llama import pytest from olive.model import ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.model_builder import ModelBuilder +from test.unit_test.utils import make_local_tiny_llama @pytest.mark.parametrize("metadata_only", [True, False]) diff --git a/test/unit_test/passes/onnx/test_nvmo_quantization.py b/test/unit_test/passes/onnx/test_nvmo_quantization.py index 8ea5cd0c91..723a736630 100644 --- a/test/unit_test/passes/onnx/test_nvmo_quantization.py +++ b/test/unit_test/passes/onnx/test_nvmo_quantization.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import get_onnx_model import onnxruntime import pytest @@ -11,6 +10,7 @@ from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.nvmo_quantization import NVModelOptQuantization +from test.unit_test.utils import get_onnx_model @pytest.mark.skipif( diff --git a/test/unit_test/passes/onnx/test_optimum_conversion.py b/test/unit_test/passes/onnx/test_optimum_conversion.py index 0a0b53e54a..fbdec27dc4 100644 --- a/test/unit_test/passes/onnx/test_optimum_conversion.py +++ b/test/unit_test/passes/onnx/test_optimum_conversion.py @@ -4,13 +4,13 @@ # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import get_hf_model import pytest from olive.model import CompositeModelHandler, HfModelHandler, ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.optimum_conversion import OptimumConversion +from test.unit_test.utils import get_hf_model @pytest.mark.parametrize("extra_args", [{"atol": 0.1}, {"atol": None}]) diff --git a/test/unit_test/passes/onnx/test_peephole_optimizer.py b/test/unit_test/passes/onnx/test_peephole_optimizer.py index e27232e509..b4a323e428 100644 --- a/test/unit_test/passes/onnx/test_peephole_optimizer.py +++ b/test/unit_test/passes/onnx/test_peephole_optimizer.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import get_onnx_model from typing import TYPE_CHECKING, Any, Dict from unittest.mock import patch @@ -14,6 +13,7 @@ from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.common import model_proto_to_olive_model from olive.passes.onnx.peephole_optimizer import OnnxPeepholeOptimizer +from test.unit_test.utils import get_onnx_model if TYPE_CHECKING: from olive.model import ONNXModelHandler diff --git a/test/unit_test/passes/onnx/test_qnn_mixed_precision_overrides.py b/test/unit_test/passes/onnx/test_qnn_mixed_precision_overrides.py index 411340a76e..224de39264 100644 --- a/test/unit_test/passes/onnx/test_qnn_mixed_precision_overrides.py +++ b/test/unit_test/passes/onnx/test_qnn_mixed_precision_overrides.py @@ -1,7 +1,6 @@ -from test.unit_test.utils import get_onnx_model - from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.mixed_precision_overrides import MixedPrecisionOverrides +from test.unit_test.utils import get_onnx_model def test_qnn_mixed_precision_overrides(tmp_path): diff --git a/test/unit_test/passes/onnx/test_qnn_preprocess.py b/test/unit_test/passes/onnx/test_qnn_preprocess.py index 0cdbbd4440..3551d6cd2b 100644 --- a/test/unit_test/passes/onnx/test_qnn_preprocess.py +++ b/test/unit_test/passes/onnx/test_qnn_preprocess.py @@ -1,5 +1,4 @@ import shutil -from test.unit_test.utils import get_onnx_model from unittest.mock import patch import pytest @@ -8,6 +7,7 @@ from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.qnn.qnn_preprocess import QNNPreprocess +from test.unit_test.utils import get_onnx_model @pytest.mark.skipif( diff --git a/test/unit_test/passes/onnx/test_quantization.py b/test/unit_test/passes/onnx/test_quantization.py index 96b7ce524d..6d50d53646 100644 --- a/test/unit_test/passes/onnx/test_quantization.py +++ b/test/unit_test/passes/onnx/test_quantization.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import logging -from test.unit_test.utils import get_onnx_model, get_pytorch_model_dummy_input from unittest.mock import patch import onnx @@ -23,6 +22,7 @@ OnnxQuantizationPreprocess, OnnxStaticQuantization, ) +from test.unit_test.utils import get_onnx_model, get_pytorch_model_dummy_input class DummyCalibrationDataReader(CalibrationDataReader): diff --git a/test/unit_test/passes/onnx/test_session_params_tuning.py b/test/unit_test/passes/onnx/test_session_params_tuning.py index fb3b58f023..18da9ea72d 100644 --- a/test/unit_test/passes/onnx/test_session_params_tuning.py +++ b/test/unit_test/passes/onnx/test_session_params_tuning.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import logging -from test.unit_test.utils import get_onnx_model from unittest.mock import PropertyMock, patch import pytest @@ -15,6 +14,7 @@ from olive.hardware.accelerator import DEFAULT_GPU_CUDA_ACCELERATOR from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.session_params_tuning import OrtSessionParamsTuning +from test.unit_test.utils import get_onnx_model def _get_tuning_data_config(input_shapes, input_names=None): @@ -68,13 +68,13 @@ def test_ort_session_params_tuning_with_customized_configs(mock_run, config): # assert if "providers_list" not in config: - assert ( - mock_run.call_args.args[1].providers_list == "CPUExecutionProvider" - ), "providers_list is not set correctly as ['CPUExecutionProvider'] by default when user does not specify it" + assert mock_run.call_args.args[1].providers_list == "CPUExecutionProvider", ( + "providers_list is not set correctly as ['CPUExecutionProvider'] by default when user does not specify it" + ) if "device" not in config: - assert ( - mock_run.call_args.args[1].device == "cpu" - ), "device is not set correctly as cpu by default when user does not specify it" + assert mock_run.call_args.args[1].device == "cpu", ( + "device is not set correctly as cpu by default when user does not specify it" + ) for k, v in config.items(): assert getattr(mock_run.call_args.args[1], k) == v, f"{k} is not set correctly as {v}" @@ -145,9 +145,9 @@ def mock_evaluate_method(model, metrics, device, execution_providers): assert "io_bind" in result.inference_settings assert acutal_eps == [execution_provider] if execution_provider == "CUDAExecutionProvider": - assert result.inference_settings["provider_options"][0][ - "enable_cuda_graph" - ], "enable_cuda_graph should be overridden to True" + assert result.inference_settings["provider_options"][0]["enable_cuda_graph"], ( + "enable_cuda_graph should be overridden to True" + ) assert result.inference_settings["provider_options"][0]["arena_extend_strategy"] == "kNextPowerOfTwo" else: assert "enable_cuda_graph" not in result.inference_settings["provider_options"][0] diff --git a/test/unit_test/passes/onnx/test_static_llm.py b/test/unit_test/passes/onnx/test_static_llm.py index 6d7b6199a4..16e4492244 100644 --- a/test/unit_test/passes/onnx/test_static_llm.py +++ b/test/unit_test/passes/onnx/test_static_llm.py @@ -3,11 +3,11 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import json -from test.unit_test.utils import make_local_tiny_llama from olive.model import CompositeModelHandler, ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.static_llm import StaticLLM +from test.unit_test.utils import make_local_tiny_llama def test_static_llm(tmp_path): diff --git a/test/unit_test/passes/onnx/test_transformer_optimization.py b/test/unit_test/passes/onnx/test_transformer_optimization.py index 423a4b5653..284965bf56 100644 --- a/test/unit_test/passes/onnx/test_transformer_optimization.py +++ b/test/unit_test/passes/onnx/test_transformer_optimization.py @@ -4,7 +4,6 @@ # -------------------------------------------------------------------------- import logging import shutil -from test.unit_test.utils import ONNX_MODEL_PATH, get_onnx_model from unittest.mock import MagicMock, patch import pytest @@ -14,6 +13,7 @@ from olive.hardware.accelerator import AcceleratorSpec, Device from olive.passes.onnx.common import get_external_data_config from olive.passes.onnx.transformer_optimization import OrtTransformersOptimization +from test.unit_test.utils import ONNX_MODEL_PATH, get_onnx_model # pylint: disable=redefined-outer-name, abstract-method, protected-access diff --git a/test/unit_test/passes/openvino/test_openvino_conversion.py b/test/unit_test/passes/openvino/test_openvino_conversion.py index 5a145558c6..b0375c9601 100644 --- a/test/unit_test/passes/openvino/test_openvino_conversion.py +++ b/test/unit_test/passes/openvino/test_openvino_conversion.py @@ -4,10 +4,10 @@ # -------------------------------------------------------------------------- import shutil from pathlib import Path -from test.unit_test.utils import get_pytorch_model, get_pytorch_model_dummy_input from olive.passes.olive_pass import create_pass_from_dict from olive.passes.openvino.conversion import OpenVINOConversion +from test.unit_test.utils import get_pytorch_model, get_pytorch_model_dummy_input def test_openvino_conversion_pass(tmp_path): diff --git a/test/unit_test/passes/pytorch/test_gptq.py b/test/unit_test/passes/pytorch/test_gptq.py index 72b26028ff..79de2e1e83 100644 --- a/test/unit_test/passes/pytorch/test_gptq.py +++ b/test/unit_test/passes/pytorch/test_gptq.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import make_local_tiny_llama import pytest import torch @@ -13,6 +12,7 @@ from olive.model import HfModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.pytorch.gptq import GptqQuantizer +from test.unit_test.utils import make_local_tiny_llama test_gptq_dc_config = DataConfig( name="test_gptq_dc_config", diff --git a/test/unit_test/passes/pytorch/test_quantization_aware_training.py b/test/unit_test/passes/pytorch/test_quantization_aware_training.py index abe21cfeca..d50d99fbdb 100644 --- a/test/unit_test/passes/pytorch/test_quantization_aware_training.py +++ b/test/unit_test/passes/pytorch/test_quantization_aware_training.py @@ -2,8 +2,6 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import get_pytorch_model - from torch.utils.data import DataLoader from olive.data.component.dataset import DummyDataset @@ -11,6 +9,7 @@ from olive.hardware.accelerator import AcceleratorSpec from olive.passes.olive_pass import FullPassConfig, create_pass_from_dict from olive.passes.pytorch.quantization_aware_training import QuantizationAwareTraining +from test.unit_test.utils import get_pytorch_model @Registry.register_dataloader() diff --git a/test/unit_test/passes/pytorch/test_rotate.py b/test/unit_test/passes/pytorch/test_rotate.py index 0f6b7be29f..99f9178ed0 100644 --- a/test/unit_test/passes/pytorch/test_rotate.py +++ b/test/unit_test/passes/pytorch/test_rotate.py @@ -2,7 +2,6 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -from test.unit_test.utils import make_local_tiny_llama from unittest.mock import patch import pytest @@ -12,6 +11,7 @@ from olive.model import HfModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.pytorch.rotate import QuaRot, SpinQuant +from test.unit_test.utils import make_local_tiny_llama def common_test_rotate(rotate_pass, tmp_path, model_path, rotate_mode, atol, **config_kwargs): diff --git a/test/unit_test/passes/qnn/test_qnn_conversion.py b/test/unit_test/passes/qnn/test_qnn_conversion.py index 6703810c84..11c45253a8 100644 --- a/test/unit_test/passes/qnn/test_qnn_conversion.py +++ b/test/unit_test/passes/qnn/test_qnn_conversion.py @@ -1,5 +1,4 @@ import platform -from test.unit_test.utils import get_onnx_model from unittest.mock import patch import pytest @@ -7,6 +6,7 @@ from olive.common.constants import OS from olive.passes.olive_pass import create_pass_from_dict from olive.passes.qnn.conversion import QNNConversion +from test.unit_test.utils import get_onnx_model @patch("olive.passes.qnn.conversion.QNNSDKRunner") diff --git a/test/unit_test/passes/vitis_ai/test_vitis_ai_quantization.py b/test/unit_test/passes/vitis_ai/test_vitis_ai_quantization.py index 623c3b0344..371fe030bc 100644 --- a/test/unit_test/passes/vitis_ai/test_vitis_ai_quantization.py +++ b/test/unit_test/passes/vitis_ai/test_vitis_ai_quantization.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from pathlib import Path -from test.unit_test.utils import get_onnx_model import onnxruntime import pytest @@ -15,6 +14,7 @@ from olive.data.registry import Registry from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.vitis_ai_quantization import VitisAIQuantization +from test.unit_test.utils import get_onnx_model class RandomDataReader(CalibrationDataReader): diff --git a/test/unit_test/systems/azureml/test_aml_system.py b/test/unit_test/systems/azureml/test_aml_system.py index e947ae5489..1ded524515 100644 --- a/test/unit_test/systems/azureml/test_aml_system.py +++ b/test/unit_test/systems/azureml/test_aml_system.py @@ -7,16 +7,6 @@ import tempfile from functools import partial from pathlib import Path -from test.unit_test.utils import ( - ONNX_MODEL_PATH, - get_accuracy_metric, - get_custom_metric, - get_glue_latency_metric, - get_hf_model_config, - get_latency_metric, - get_onnx_model_config, - get_onnxconversion_pass, -) from typing import ClassVar, List from unittest.mock import ANY, MagicMock, Mock, patch @@ -41,6 +31,16 @@ from olive.systems.azureml.aml_pass_runner import main as aml_pass_runner_main from olive.systems.azureml.aml_system import AzureMLSystem from olive.systems.common import AzureMLDockerConfig, AzureMLEnvironmentConfig +from test.unit_test.utils import ( + ONNX_MODEL_PATH, + get_accuracy_metric, + get_custom_metric, + get_glue_latency_metric, + get_hf_model_config, + get_latency_metric, + get_onnx_model_config, + get_onnxconversion_pass, +) # pylint: disable=attribute-defined-outside-init, protected-access diff --git a/test/unit_test/systems/docker/test_docker_system.py b/test/unit_test/systems/docker/test_docker_system.py index b1eaca8dd9..25a95f67e7 100644 --- a/test/unit_test/systems/docker/test_docker_system.py +++ b/test/unit_test/systems/docker/test_docker_system.py @@ -6,7 +6,6 @@ import json import shutil from pathlib import Path -from test.unit_test.utils import ONNX_MODEL_PATH, get_accuracy_metric, get_onnx_model_config, get_pytorch_model_config from unittest.mock import ANY, MagicMock, patch import pytest @@ -21,6 +20,7 @@ from olive.systems.docker.docker_system import DockerSystem from olive.systems.system_config import DockerTargetUserConfig, SystemConfig from olive.systems.utils import create_managed_system +from test.unit_test.utils import ONNX_MODEL_PATH, get_accuracy_metric, get_onnx_model_config, get_pytorch_model_config # pylint: disable=attribute-defined-outside-init,protected-access diff --git a/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py b/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py index c3e86cd9cc..7d9e7815a3 100644 --- a/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py +++ b/test/unit_test/systems/isolated_ort/test_isolated_ort_system.py @@ -9,7 +9,6 @@ import sys import venv from pathlib import Path -from test.unit_test.utils import get_accuracy_metric, get_latency_metric, get_onnx_model_config from unittest import mock from unittest.mock import MagicMock, patch @@ -27,6 +26,7 @@ from olive.systems.isolated_ort.inference_runner import main as inference_runner_main from olive.systems.isolated_ort.isolated_ort_system import IsolatedORTEvaluator from olive.systems.system_config import IsolatedORTTargetUserConfig, SystemConfig +from test.unit_test.utils import get_accuracy_metric, get_latency_metric, get_onnx_model_config # pylint: disable=attribute-defined-outside-init, protected-access diff --git a/test/unit_test/systems/python_environment/test_python_environment_system.py b/test/unit_test/systems/python_environment/test_python_environment_system.py index 4087c62476..ffed9786b3 100644 --- a/test/unit_test/systems/python_environment/test_python_environment_system.py +++ b/test/unit_test/systems/python_environment/test_python_environment_system.py @@ -8,14 +8,6 @@ import tempfile import venv from pathlib import Path -from test.unit_test.utils import ( - get_glue_accuracy_metric, - get_glue_latency_metric, - get_hf_model_config, - get_onnx_model, - get_onnx_model_config, - get_onnxconversion_pass, -) from unittest.mock import MagicMock, patch import pytest @@ -30,6 +22,14 @@ from olive.systems.python_environment.pass_runner import main as pass_runner_main from olive.systems.system_config import PythonEnvironmentTargetUserConfig, SystemConfig from olive.systems.utils import create_managed_system, create_managed_system_with_cache +from test.unit_test.utils import ( + get_glue_accuracy_metric, + get_glue_latency_metric, + get_hf_model_config, + get_onnx_model, + get_onnx_model_config, + get_onnxconversion_pass, +) # pylint: disable=no-value-for-parameter, attribute-defined-outside-init, protected-access diff --git a/test/unit_test/systems/test_local.py b/test/unit_test/systems/test_local.py index cde0f2918b..d9c0760ee2 100644 --- a/test/unit_test/systems/test_local.py +++ b/test/unit_test/systems/test_local.py @@ -3,7 +3,6 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- from functools import partial -from test.unit_test.utils import get_accuracy_metric, get_custom_metric, get_latency_metric from typing import ClassVar, List from unittest.mock import MagicMock, patch @@ -16,6 +15,7 @@ from olive.hardware import DEFAULT_CPU_ACCELERATOR from olive.model import PyTorchModelHandler from olive.systems.local import LocalSystem +from test.unit_test.utils import get_accuracy_metric, get_custom_metric, get_latency_metric # pylint: disable=attribute-defined-outside-init diff --git a/test/unit_test/workflows/test_workflow_run.py b/test/unit_test/workflows/test_workflow_run.py index f4a4ce6e83..cc8de8d23e 100644 --- a/test/unit_test/workflows/test_workflow_run.py +++ b/test/unit_test/workflows/test_workflow_run.py @@ -1,11 +1,5 @@ from copy import deepcopy from pathlib import Path -from test.unit_test.utils import ( - get_pytorch_model, - get_pytorch_model_config, - get_pytorch_model_io_config, - pytorch_model_loader, -) from unittest.mock import patch import pytest @@ -13,6 +7,12 @@ from olive.data.registry import Registry from olive.hardware.accelerator import AcceleratorSpec from olive.workflows import run as olive_run +from test.unit_test.utils import ( + get_pytorch_model, + get_pytorch_model_config, + get_pytorch_model_io_config, + pytorch_model_loader, +) @Registry.register_dataloader()