Unify OS constants (microsoft#1183)

## Describe your changes Unify OS constants ## Checklist before requesting a review - [ ] Add unit tests for this change. - [ ] Make sure all tests can pass. - [ ] Update documents if necessary. - [ ] Lint and apply fixes to your code by running `lintrunner -a` - [ ] Is this a user-facing change? If yes, give a description of this change to be included in the release notes. - [ ] Is this PR including examples changes? If yes, please remember to update [example documentation](https://github.com/microsoft/Olive/blob/main/docs/source/examples.md) in a follow-up PR. ## (Optional) Issue link
yuwenzho · Jun 3, 2024 · 3d3277b · 3d3277b
1 parent 8dbf3e3
commit 3d3277b
Show file tree

Hide file tree

Showing 40 changed files with 133 additions and 82 deletions.
diff --git a/examples/inception/prepare_config.py b/examples/inception/prepare_config.py
@@ -6,6 +6,8 @@
 import platform
 from pathlib import Path
 
+from olive.common.constants import OS
+
 
 def resolve_windows_config():
 
@@ -18,5 +20,5 @@ def resolve_windows_config():
 
 
 if __name__ == "__main__":
-    if platform.system() == "Windows":
+    if platform.system() == OS.WINDOWS:
         resolve_windows_config()
diff --git a/examples/mobilenet/prepare_config.py b/examples/mobilenet/prepare_config.py
@@ -7,6 +7,8 @@
 import platform
 from pathlib import Path
 
+from olive.common.constants import OS
+
 
 def raw_qnn_config():
     # pylint: disable=redefined-outer-name
@@ -16,20 +18,20 @@ def raw_qnn_config():
 
     sys_platform = platform.system()
 
-    if sys_platform == "Linux":
+    if sys_platform == OS.LINUX:
         raw_qnn_config["passes"]["qnn_context_binary"] = {
             "type": "QNNContextBinaryGenerator",
             "config": {"backend": "libQnnHtp.so"},
         }
         raw_qnn_config["pass_flows"].append(["converter", "build_model_lib", "qnn_context_binary"])
         raw_qnn_config["passes"]["build_model_lib"]["config"]["lib_targets"] = "x86_64-linux-clang"
-    elif sys_platform == "Windows":
+    elif sys_platform == OS.WINDOWS:
         raw_qnn_config["passes"]["build_model_lib"]["config"]["lib_targets"] = "x86_64-windows-msvc"
 
     for metric_config in raw_qnn_config["evaluators"]["common_evaluator"]["metrics"]:
-        if sys_platform == "Windows":
+        if sys_platform == OS.WINDOWS:
             metric_config["user_config"]["inference_settings"]["qnn"]["backend"] = "QnnCpu"
-        elif sys_platform == "Linux":
+        elif sys_platform == OS.LINUX:
             metric_config["user_config"]["inference_settings"]["qnn"]["backend"] = "libQnnCpu"
 
     with Path("raw_qnn_sdk_config.json").open("w") as f:

diff --git a/examples/phi2/phi2.py b/examples/phi2/phi2.py
@@ -11,6 +11,7 @@
 from onnxruntime import __version__ as OrtVersion
 from packaging import version
 
+from olive.common.constants import OS
 from olive.workflows import run as olive_run
 
 # flake8: noqa: T201
@@ -30,25 +31,25 @@
         # -1 to use CPU
         "device_id": -1,
         "use_fp16": False,
-        "use_step": platform.system() == "Linux",
+        "use_step": platform.system() == OS.LINUX,
     },
     "cpu_int4": {
         "use_buffer_share": False,
         "device_id": -1,
         "use_fp16": False,
-        "use_step": platform.system() == "Linux",
+        "use_step": platform.system() == OS.LINUX,
     },
     "cuda_fp16": {
         "use_buffer_share": False,
         "device_id": 0,
         "use_fp16": True,
-        "use_step": platform.system() == "Linux",
+        "use_step": platform.system() == OS.LINUX,
     },
     "cuda_int4": {
         "use_buffer_share": False,
         "device_id": 0,
         "use_fp16": True,
-        "use_step": platform.system() == "Linux",
+        "use_step": platform.system() == OS.LINUX,
     },
 }
 
@@ -180,7 +181,7 @@ def main(raw_args=None):
         with open(json_file_template) as f:
             template_json = json.load(f)
 
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             legacy_optimization_setting(template_json)
 
         # add pass flows

diff --git a/examples/test/test_bert_ptq_cpu_docker.py b/examples/test/test_bert_ptq_cpu_docker.py
@@ -9,6 +9,8 @@
 import pytest
 from utils import check_output, patch_config
 
+from olive.common.constants import OS
+
 
 @pytest.fixture(scope="module", autouse=True)
 def setup():
@@ -25,7 +27,7 @@ def setup():
 @pytest.mark.parametrize("system", ["docker_system"])
 @pytest.mark.parametrize("olive_json", ["bert_ptq_cpu.json"])
 def test_bert(search_algorithm, execution_order, system, olive_json):
-    if system == "docker_system" and platform.system() == "Windows":
+    if system == "docker_system" and platform.system() == OS.WINDOWS:
         pytest.skip("Skip Linux containers on Windows host test case.")
 
     from olive.workflows import run as olive_run

diff --git a/examples/test/test_qnn_tooklit.py b/examples/test/test_qnn_tooklit.py
@@ -9,6 +9,7 @@
 import pytest
 from utils import check_output, download_azure_blob
 
+from olive.common.constants import OS
 from olive.common.utils import retry_func, run_subprocess
 from olive.logging import set_verbosity_debug
 
@@ -20,9 +21,9 @@ class TestQnnToolkit:
     def setup(self, tmp_path):
         """Download the qnn sdk."""
         blob, download_path = "", ""
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             blob, download_path = "qnn_sdk_windows.zip", "qnn_sdk_windows.zip"
-        elif platform.system() == "Linux":
+        elif platform.system() == OS.LINUX:
             blob, download_path = "qnn_sdk_linux.zip", "qnn_sdk_linux.zip"
 
         download_azure_blob(
@@ -32,10 +33,10 @@ def setup(self, tmp_path):
         )
         target_path = tmp_path / "qnn_sdk"
         target_path.mkdir(parents=True, exist_ok=True)
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             cmd = f"powershell Expand-Archive -Path {download_path} -DestinationPath {str(target_path)}"
             run_subprocess(cmd=cmd, check=True)
-        elif platform.system() == "Linux":
+        elif platform.system() == OS.LINUX:
             run_subprocess(cmd=f"unzip {download_path} -d {str(target_path)}", check=True)
 
         os.environ["QNN_SDK_ROOT"] = str(target_path / "opt" / "qcom" / "aistack")
@@ -55,9 +56,9 @@ def _setup_resource(self, use_olive_env):
             )
             # install dependencies
             python_cmd = ""
-            if platform.system() == "Windows":
+            if platform.system() == OS.WINDOWS:
                 python_cmd = str(Path(os.environ["QNN_SDK_ROOT"]) / "olive-pyenv" / "python.exe")
-            elif platform.system() == "Linux":
+            elif platform.system() == OS.LINUX:
                 python_cmd = str(Path(os.environ["QNN_SDK_ROOT"]) / "olive-pyenv" / "bin" / "python")
             install_cmd = [
                 python_cmd,
@@ -69,7 +70,7 @@ def _setup_resource(self, use_olive_env):
             packages = ["tensorflow==2.10.1", "numpy==1.23.5"]
             retry_func(run_subprocess, kwargs={"cmd": f"python -m pip install {' '.join(packages)}", "check": True})
             os.environ["PYTHONPATH"] = str(Path(os.environ["QNN_SDK_ROOT"]) / "lib" / "python")
-            if platform.system() == "Linux":
+            if platform.system() == OS.LINUX:
                 os.environ["PATH"] = (
                     str(Path(os.environ["QNN_SDK_ROOT"]) / "bin" / "x86_64-linux-clang")
                     + os.path.pathsep

diff --git a/examples/test/test_snpe_toolkit.py b/examples/test/test_snpe_toolkit.py
@@ -9,6 +9,7 @@
 import pytest
 from utils import check_output, download_azure_blob
 
+from olive.common.constants import OS
 from olive.common.utils import retry_func, run_subprocess
 from olive.logging import set_verbosity_debug
 
@@ -20,9 +21,9 @@ class TestSnpeToolkit:
     def setup(self, tmp_path):
         """Download the snpe sdk."""
         blob, download_path = "", ""
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             blob, download_path = "snpe_sdk_windows.zip", "snpe_sdk_windows.zip"
-        elif platform.system() == "Linux":
+        elif platform.system() == OS.LINUX:
             blob, download_path = "snpe_sdk_linux.zip", "snpe_sdk_linux.zip"
 
         download_azure_blob(
@@ -32,10 +33,10 @@ def setup(self, tmp_path):
         )
         target_path = tmp_path / "snpe_sdk"
         target_path.mkdir(parents=True, exist_ok=True)
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             cmd = f"powershell Expand-Archive -Path {download_path} -DestinationPath {str(target_path)}"
             run_subprocess(cmd=cmd, check=True)
-        elif platform.system() == "Linux":
+        elif platform.system() == OS.LINUX:
             run_subprocess(cmd=f"unzip {download_path} -d {str(target_path)}", check=True)
         os.environ["SNPE_ROOT"] = str(target_path)
 
@@ -54,9 +55,9 @@ def _setup_resource(self, use_olive_env):
             )
             # install dependencies
             python_cmd = ""
-            if platform.system() == "Windows":
+            if platform.system() == OS.WINDOWS:
                 python_cmd = str(Path(os.environ["SNPE_ROOT"]) / "olive-pyenv" / "python.exe")
-            elif platform.system() == "Linux":
+            elif platform.system() == OS.LINUX:
                 python_cmd = str(Path(os.environ["SNPE_ROOT"]) / "olive-pyenv" / "bin" / "python")
             install_cmd = [
                 python_cmd,
@@ -68,7 +69,7 @@ def _setup_resource(self, use_olive_env):
             packages = ["tensorflow==2.10.1", "numpy==1.23.5"]
             retry_func(run_subprocess, kwargs={"cmd": f"python -m pip install {' '.join(packages)}", "check": True})
             os.environ["PYTHONPATH"] = str(Path(os.environ["SNPE_ROOT"]) / "lib" / "python")
-            if platform.system() == "Linux":
+            if platform.system() == OS.LINUX:
                 os.environ["PATH"] = (
                     str(Path(os.environ["SNPE_ROOT"]) / "bin" / "x86_64-linux-clang")
                     + os.path.pathsep

diff --git a/examples/test/test_whisper.py b/examples/test/test_whisper.py
@@ -11,6 +11,8 @@
 import pytest
 from utils import check_output
 
+from olive.common.constants import OS
+
 
 @pytest.fixture(scope="module", autouse=True)
 def setup():
@@ -34,7 +36,7 @@ def setup():
 def test_whisper(device_precision):
     from olive.workflows import run as olive_run
 
-    if platform.system() == "Windows" and device_precision[1].startswith("inc_int8"):
+    if platform.system() == OS.WINDOWS and device_precision[1].startswith("inc_int8"):
         pytest.skip("Skip test on Windows. neural-compressor import is hanging on Windows.")
 
     device, precision = device_precision

diff --git a/examples/vgg/prepare_config.py b/examples/vgg/prepare_config.py
@@ -6,6 +6,8 @@
 import platform
 from pathlib import Path
 
+from olive.common.constants import OS
+
 
 def resolve_windows_config():
 
@@ -18,5 +20,5 @@ def resolve_windows_config():
 
 
 if __name__ == "__main__":
-    if platform.system() == "Windows":
+    if platform.system() == OS.WINDOWS:
         resolve_windows_config()
diff --git a/olive/common/constants.py b/olive/common/constants.py
@@ -2,5 +2,11 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
+from enum import Enum
 
 DEFAULT_WORKFLOW_ID = "default_workflow"
+
+
+class OS(str, Enum):
+    WINDOWS = "Windows"
+    LINUX = "Linux"
diff --git a/olive/common/utils.py b/olive/common/utils.py
@@ -18,14 +18,16 @@
 from pathlib import Path
 from typing import Dict, List, Tuple, Union
 
+from olive.common.constants import OS
+
 logger = logging.getLogger(__name__)
 
 
 def run_subprocess(cmd, env=None, cwd=None, check=False):
     logger.debug("Running command: %s", cmd)
 
     assert isinstance(cmd, (str, list)), f"cmd must be a string or a list, got {type(cmd)}."
-    windows = platform.system() == "Windows"
+    windows = platform.system() == OS.WINDOWS
     if isinstance(cmd, str):
         # In posix model, the cmd string will be handled with specific posix rules.
         # https://docs.python.org/3.8/library/shlex.html#parsing-rules

diff --git a/olive/engine/packaging/packaging_generator.py b/olive/engine/packaging/packaging_generator.py
@@ -16,6 +16,7 @@
 
 import pkg_resources
 
+from olive.common.constants import OS
 from olive.common.utils import copy_dir, retry_func, run_subprocess
 from olive.engine.packaging.packaging_config import (
     AzureMLDeploymentPackagingConfig,
@@ -707,14 +708,14 @@ def _download_ort_extensions_package(use_ort_extensions: bool, download_path: st
         # Hardcode the nightly version number for now until we have a better way to identify nightly version
         if version.startswith("0.8.0."):
             system = platform.system()
-            if system == "Windows":
+            if system == OS.WINDOWS:
                 download_command = (
                     f"{sys.executable} -m pip download -i "
                     "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/ "
                     f"onnxruntime-extensions=={version} --no-deps -d {download_path}"
                 )
                 run_subprocess(download_command)
-            elif system == "Linux":
+            elif system == OS.LINUX:
                 logger.warning(
                     "ONNXRuntime-Extensions nightly package is not available for Linux. "
                     "Skip packaging ONNXRuntime-Extensions package. Please manually install ONNXRuntime-Extensions."

diff --git a/olive/model/handler/qnn.py b/olive/model/handler/qnn.py
@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
+from olive.common.constants import OS
 from olive.constants import Framework, ModelFileFormat
 from olive.hardware.accelerator import Device
 from olive.model.config import IoConfig
@@ -55,10 +56,10 @@ def model_path(self):
                 logger.debug(
                     "QNNModelHandler: lib_targets is not provided, using default lib_targets x86_64-linux-clang"
                 )
-                if platform.system() == "Linux":
+                if platform.system() == OS.LINUX:
                     lib_targets = "x86_64-linux-clang"
                     model_lib_suffix = ".so"
-                elif platform.system() == "Windows":
+                elif platform.system() == OS.WINDOWS:
                     # might be different for arm devices
                     lib_targets = "x64"
                     model_lib_suffix = ".dll"

diff --git a/olive/passes/qnn/context_binary_generator.py b/olive/passes/qnn/context_binary_generator.py
@@ -8,6 +8,7 @@
 from pathlib import Path
 from typing import Any, Dict
 
+from olive.common.constants import OS
 from olive.constants import ModelFileFormat
 from olive.hardware import AcceleratorSpec
 from olive.model import QNNModelHandler
@@ -26,7 +27,7 @@ class QNNContextBinaryGenerator(Pass):
 
     @classmethod
     def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassConfigParam]:
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             raise NotImplementedError("QNNContextBinaryGenerator is not supported on Windows.")
 
         return {

diff --git a/olive/passes/qnn/conversion.py b/olive/passes/qnn/conversion.py
@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Union
 
+from olive.common.constants import OS
 from olive.constants import ModelFileFormat
 from olive.hardware import AcceleratorSpec
 from olive.model import ONNXModelHandler, PyTorchModelHandler, QNNModelHandler, TensorFlowModelHandler
@@ -82,7 +83,7 @@ def _run_for_config(
         converter_program = [f"qnn-{converter_platform}-converter"]
 
         runner = QNNSDKRunner(use_dev_tools=True)
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             converter_program = [
                 "python",
                 str(Path(runner.sdk_env.sdk_root_path) / "bin" / runner.sdk_env.target_arch / converter_program[0]),

diff --git a/olive/passes/qnn/model_lib_generator.py b/olive/passes/qnn/model_lib_generator.py
@@ -8,6 +8,7 @@
 from pathlib import Path
 from typing import Any, Dict
 
+from olive.common.constants import OS
 from olive.constants import ModelFileFormat
 from olive.hardware import AcceleratorSpec
 from olive.model import QNNModelHandler
@@ -53,7 +54,7 @@ def _run_for_config(
     ) -> QNNModelHandler:
         main_cmd = "qnn-model-lib-generator"
         runner = QNNSDKRunner(use_dev_tools=True)
-        if platform.system() == "Windows":
+        if platform.system() == OS.WINDOWS:
             main_cmd = "python " + str(
                 Path(runner.sdk_env.sdk_root_path) / "bin" / runner.sdk_env.target_arch / main_cmd
             )