Merge branch 'main' of https://github.com/microsoft/olive into rel-0.5.1

README.md

@@ -66,8 +66,8 @@ pip install olive-ai[directml]
 ```
 
 ### Optional Dependencies
-Olive has optional dependencies that can be installed to enable additional features. Please refer to [extra dependencies](./olive/extra_dependencies.json) for
-the list of extras and their dependencies.
+Olive has optional dependencies that can be installed to enable additional features. Please refer to
+[Olive package config](./olive/olive_config.json) for the list of extras and their dependencies.
 
 ## Pipeline Status
 

docs/architecture.md

@@ -35,7 +35,7 @@ Passes are the building blocks of an Olive workflow. Olive uses multiple Passes
 The base class for Pass:
 ```python
 class Pass(ABC):
-    def __init__(self, accelerator_spec: AcceleratorSpec, config: Union[Dict[str, Any], BaseModel], disable_search: Optional[bool] = False):
+    def __init__(self, accelerator_spec: AcceleratorSpec, config: Union[Dict[str, Any], BaseModel], disable_search: Optional[bool] = False, host_device: Optional[str] = None):
         ...
 
     @classmethod

docs/source/api/systems.rst

@@ -82,7 +82,7 @@ PythonEnvironmentSystem
 
 IsolatedORTSystem
 ^^^^^^^^^^^^^^^^^^^^^^^^
-.. autoclass:: olive.systems.ort_evironment.IsolatedORTSystem
+.. autoclass:: olive.systems.isolated_ort.IsolatedORTSystem
 
 .. _olive_system_alias:
 

docs/source/exts/auto_config_doc/__init__.py

@@ -15,16 +15,19 @@
 
 from olive.common.auto_config import AutoConfigClass
 from olive.hardware import DEFAULT_CPU_ACCELERATOR
+from olive.package_config import OlivePackageConfig
 from olive.passes import Pass
 
 # pylint: skip-file
 
 
-def import_class(class_name: str):
-    module_name = ".".join(class_name.split(".")[:-1])
-    class_name = class_name.split(".")[-1]
-    module = import_module(module_name)
-    return getattr(module, class_name)
+def import_class(class_name: str, package_config: OlivePackageConfig):
+    module_path, module_name = class_name.rsplit(".", 1)
+    if module_name in package_config.passes:
+        return package_config.import_pass_module(module_name)
+
+    module = import_module(module_path)
+    return getattr(module, module_name)
 
 
 class AutoConfigDirective(Directive):
@@ -68,7 +71,8 @@ def make_doc(self, auto_config_class: Union[AutoConfigClass, Pass]):
 
     def run(self):
         (class_name,) = self.arguments
-        auto_config_class = import_class(class_name)
+        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"

docs/source/getstarted/installation.md

@@ -62,5 +62,6 @@ pip install -e .
 ```
 
 ## Optional Dependencies
-Olive has optional dependencies that can be installed to enable additional features. Please refer to [extra dependencies](https://github.com/microsoft/Olive/blob/main/olive/extra_dependencies.json)
-for the list of extras and their dependencies.
+Olive has optional dependencies that can be installed to enable additional features. Please refer to
+[Olive package config](https://github.com/microsoft/Olive/blob/main/olive/olive_config.json) for the list of extras
+and their dependencies.

docs/source/overview/quicktour.md

@@ -31,6 +31,9 @@ olive_run("my_model_acceleration_description.json")
 
 You can use setup mode `python -m olive.workflows.run --config my_model_acceleration_description.json --setup` to identify list of additional packages you may need to install for your workflow.
 
+To include user implemented (or proprietary, or private) passes as part of workflow, clone olive_config.json and update it.
+Provide the path to the cloned _olive_config.json_ file at launch using the '--package-config' command line option.
+
 You can also change the default directory for temporary files and directories using `--tempdir` option.
 Set this to a local directory if you want to avoid using the default tempdir for reasons such as disk space and permissions.
 

examples/phi2/README.md

@@ -76,6 +76,8 @@ For better generation experience, here is the way to run inference with the opti
 python phi2.py --model_type cpu_fp32 --inference --optimum_optimization --prompt "Write a extremely long story starting with once upon a time"
 ```
 
+## Export output models in MLFlow format
+If you want to output the optimized models to a zip file in MLFlow format, add `--export_mlflow_format` argument. The MLFlow model will be packaged in a zip file named `mlflow_model` in the output folder.
 
 ## Limitations
 1. The latest ONNXRuntime implements specific fusion patterns for better performance but only works for ONNX model from TorchDynamo-based ONNX Exporter. And the TorchDynamo-based ONNX Exporter is only available on Linux.

examples/phi2/phi2.py

@@ -78,7 +78,12 @@ def get_args(raw_args):
     parser.add_argument(
         "--optimum_optimization",
         action="store_true",
-        help="Run inference with optimized model",
+        help="Use optimum optimization",
+    )
+    parser.add_argument(
+        "--export_mlflow_format",
+        action="store_true",
+        help="Export the model in mlflow format.",
     )
     parser.add_argument(
         "--prompt",
@@ -173,6 +178,15 @@ def main(raw_args=None):
             del template_json["passes"][pass_name]
             continue
 
+    if args.export_mlflow_format:
+        template_json["engine"]["packaging_config"] = [
+            {
+                "type": "Zipfile",
+                "name": "mlflow_model",
+                "config": {"export_in_mlflow_format": True},
+            }
+        ]
+
     with open("phi2_optimize.json", "w") as f:
         json.dump(template_json, f, indent=4)
 

olive/engine/engine.py

@@ -372,14 +372,22 @@ def run_accelerator(
         logger.debug("run_accelerator done")
         return output_footprint
 
+    def get_host_device(self):
+        if self.host_config.config.accelerators:
+            # for host device, we will always use the first accelerator device
+            return self.host_config.config.accelerators[0].device
+        else:
+            return None
+
     def setup_passes(self, accelerator_spec: "AcceleratorSpec"):
+        host_device = self.get_host_device()
         # clean the passes
         self.passes.clear()
         for name, config in self.pass_config.items():
             pass_cls: Type["Pass"] = config["type"]
             pass_cfg = config["config"]
             pass_cfg = pass_cls.generate_search_space(accelerator_spec, pass_cfg, config["disable_search"])
-            p = pass_cls(accelerator_spec, pass_cfg, config["disable_search"])
+            p = pass_cls(accelerator_spec, pass_cfg, config["disable_search"], host_device)
             self.register_pass(
                 p,
                 name=name,

olive/hardware/accelerator.py

@@ -173,7 +173,7 @@ def normalize_accelerators(system_config: "SystemConfig", skip_supported_eps_che
 
     * the accelerators is not specified, infer the device/ep based on the installed ORT in case of local/python system.
     * only device is specified, infer the execution providers based on the installed ORT in case of local/python system.
-    * only EP is specified, infer the device based on the installed ORT in case of local/python sytemm.
+    * only EP is specified, infer the device based on the installed ORT in case of local/python system.
     * For AzureML and Docker system, the accelerators and execution providers must be specified.
     """
     from olive.systems.common import SystemType