Add custom onnx model name support for output dir#2235
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR adds support for saving ONNX models with custom output file paths, including the ability to specify a custom .onnx filename rather than defaulting to model.onnx. The changes enable better handling of models with external data files and provide more flexibility in output file naming.
- Extended
save_to_dirmethod to support aflattenparameter for copying folder contents directly - Modified
save_modeland_save_modelmethods in cache to handle file paths with.onnxsuffix as output file paths - Updated engine to create parent directory when output_dir has a suffix
- Added comprehensive test coverage for the new custom filename functionality
Reviewed Changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| olive/resource_path.py | Added flatten parameter to save_to_dir method to control whether folder contents are copied directly or nested |
| olive/cache.py | Refactored save_model and _save_model to detect and handle custom ONNX filenames when output path has .onnx suffix |
| olive/engine/engine.py | Modified directory creation logic to handle output paths with suffixes by creating parent directory instead |
| olive/passes/onnx/common.py | Updated file deletion logic to handle directories, removed extraneous blank lines |
| test/test_cache.py | Added test for custom ONNX filenames and updated existing tests to verify external data file handling |
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
olive/resource_path.py
Outdated
| @@ -237,14 +239,13 @@ def save_to_dir(self, dir_path: Union[Path, str], name: str = None, overwrite: b | |||
| new_path_name = self.config.path.name | |||
| new_path = dir_path / new_path_name | |||
| _overwrite_helper(new_path, overwrite) | |||
There was a problem hiding this comment.
i think the is_file check and flatten conditions needs to be handled before the _overwrite_helper call so that we clean up the destination directory?
but this is interesting. should we be cleaning the dir_path if we are flattening or we just call copy_dir with dirs_exist_ok=True instead
There was a problem hiding this comment.
i don't think we should clean up the dir_path if flatten. The dir_path may have engine artifacts.
Describe your changes
Add custom onnx model name support for output dir
Checklist before requesting a review
lintrunner -a(Optional) Issue link