From 3faa55b5c9c1a790c32dcf4a06d167e1f499733e Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 16 Jan 2024 15:08:39 -0800 Subject: [PATCH] move arch/os types to enums --- smartsim/_core/_install/builder.py | 85 ++++++++++++++++++++---------- 1 file changed, 57 insertions(+), 28 deletions(-) diff --git a/smartsim/_core/_install/builder.py b/smartsim/_core/_install/builder.py index bf96135b1..41291ad08 100644 --- a/smartsim/_core/_install/builder.py +++ b/smartsim/_core/_install/builder.py @@ -25,6 +25,7 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import concurrent.futures +import enum import os import platform import re @@ -52,8 +53,6 @@ TRedisAIBackendStr = t.Literal["tensorflow", "torch", "onnxruntime", "tflite"] TDeviceStr = t.Literal["cpu", "gpu"] -TArchitectureStr = t.Literal["x64"] -TOperatingSystemStr = t.Literal["linux", "darwin"] _T = t.TypeVar("_T") _U = t.TypeVar("_U") @@ -83,6 +82,32 @@ class BuildError(Exception): pass +class Architecture(enum.Enum): + X64 = ("x86_64", "amd64") + ARM64 = ("arm64", "aarch64") + + @classmethod + def from_str(cls, string: str, /) -> "Architecture": + string = string.lower() + for type_ in cls: + if string in type_.value: + return type_ + raise BuildError(f"Unrecognized or unsupported architecture: {string}") + + +class OperatingSystem(enum.Enum): + LINUX = ("linux", "linux2") + DARWIN = ("darwin",) + + @classmethod + def from_str(cls, string: str, /) -> "OperatingSystem": + string = string.lower() + for type_ in cls: + if string in type_.value: + return type_ + raise BuildError(f"Unrecognized or unsupported operating system: {string}") + + class Builder: """Base class for building third-party libraries""" @@ -362,24 +387,16 @@ def __init__( # TODO: It might be worth making these constructor args so that users # of this class can configure exactly _what_ they are building. - os_ = platform.system().lower() - if "darwin" in os_: - self._os: TOperatingSystemStr = "darwin" - elif "linux" in os_: - self._os = "linux" - else: - raise BuildError(f"Unrecognized or unsupported operating system '{os_}'") - - machine = platform.machine().lower() - if machine not in ("amd64", "x86_64") and any( + self._os = OperatingSystem.from_str(platform.system()) + self._architecture = Architecture.from_str(platform.machine()) + if self._architecture == Architecture.ARM64 and any( (self.fetch_tf, self.fetch_torch, self.fetch_onnx) ): raise BuildError( - "SmartSim currently only supports building ML backends for " + "SmartSim currently only supports fetching ML backends for " "the 'x64' architecture; found unrecognized or unsupported " - f"architecture '{machine}'" + f"architecture '{platform.system()}'" ) - self._architecture: TArchitectureStr = "x64" @property def rai_build_path(self) -> Path: @@ -416,8 +433,20 @@ def fetch_onnx(self) -> bool: return self.build_onnx def get_deps_dir_path_for(self, device: TDeviceStr) -> Path: - os_ = "macos" if self._os == "darwin" else self._os - return self.rai_build_path / f"deps/{os_}-{self._architecture}-{device}" + def fail_to_format(reason: str) -> BuildError: # pragma: no cover + return BuildError(f"Failed to format RedisAI dependency path: {reason}") + + if self._os == OperatingSystem.DARWIN: + os_ = "macos" + elif self._os == OperatingSystem.LINUX: + os_ = "linux" + else: # pragma: no cover + raise fail_to_format(f"Unknown operating system: {self._os}") + if self._architecture == Architecture.X64: + arch = "x64" + else: # pragma: no cover + raise fail_to_format(f"Unknown architecture: {self._architecture}") + return self.rai_build_path / f"deps/{os_}-{arch}-{device}" def _get_deps_to_fetch_for( self, device: TDeviceStr @@ -738,20 +767,20 @@ def _extract_download( @t.final @dataclass(frozen=True) class _PTArchive(_WebZip, _RAIBuildDependency): - os_: TOperatingSystemStr + os_: OperatingSystem device: TDeviceStr version: str @property def url(self) -> str: - if "linux" in self.os_: + if self.os_ == OperatingSystem.LINUX: if self.device == "gpu": pt_build = "cu117" else: pt_build = "cpu" # pylint: disable-next=line-too-long libtorch_arch = f"libtorch-cxx11-abi-shared-without-deps-{self.version}%2B{pt_build}.zip" - elif "darwin" in self.os_: + elif self.os_ == OperatingSystem.DARWIN: if self.device == "gpu": raise BuildError("RedisAI does not currently support GPU on Macos") pt_build = "cpu" @@ -775,24 +804,24 @@ def __place_for_rai__(self, target: t.Union[str, "os.PathLike[str]"]) -> Path: @t.final @dataclass(frozen=True) class _TFArchive(_WebTGZ, _RAIBuildDependency): - os_: TOperatingSystemStr - architecture: TArchitectureStr + os_: OperatingSystem + architecture: Architecture device: TDeviceStr version: str @property def url(self) -> str: - if self.architecture == "x64": + if self.architecture == Architecture.X64: tf_arch = "x86_64" else: raise BuildError( "Unexpected Architecture for TF Archive: {self.architecture}" ) - if "linux" in self.os_: + if self.os_ == OperatingSystem.LINUX: tf_os = "linux" tf_device = self.device - elif "darwin" in self.os_: + elif self.os_ == OperatingSystem.DARWIN: tf_os = "darwin" if self.device == "gpu": raise BuildError("RedisAI does not currently support GPU on Macos") @@ -818,7 +847,7 @@ def __place_for_rai__(self, target: t.Union[str, "os.PathLike[str]"]) -> Path: @t.final @dataclass(frozen=True) class _ORTArchive(_WebTGZ, _RAIBuildDependency): - os_: TOperatingSystemStr + os_: OperatingSystem device: TDeviceStr version: str @@ -828,11 +857,11 @@ def url(self) -> str: "https://github.com/microsoft/onnxruntime/releases/" f"download/v{self.version}" ) - if "linux" in self.os_: + if self.os_ == OperatingSystem.LINUX: ort_os = "linux" ort_arch = "x64" ort_build = "-gpu" if self.device == "gpu" else "" - elif "darwin" in self.os_: + elif self.os_ == OperatingSystem.DARWIN: ort_os = "osx" ort_arch = "x86_64" ort_build = ""