diff --git a/Makefile b/Makefile index fef69eab3..4f0baf6ef 100644 --- a/Makefile +++ b/Makefile @@ -66,6 +66,10 @@ clobber: clean # help: # help: Style # help: ------- +# help: check-all - Performs all the style-related checks +.PHONY: check-all +check-all: check-style check-format check-sort-imports check-lint check-mypy + $(info All style checks PASSED) # help: style - Sort imports and format with black .PHONY: style diff --git a/doc/changelog.rst b/doc/changelog.rst index cc16fa095..c287e73a5 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -26,6 +26,8 @@ Description - Updated SmartSim's machine learning backends - Added ONNX support for Python 3.10 - Added support for Python 3.11 +- Added support for SmartSim with Torch on Apple Silicon + Detailed Notes @@ -44,12 +46,16 @@ Detailed Notes there is now an available ONNX wheel for use with Python 3.10, and wheels for all of SmartSim's machine learning backends with Python 3.11. (SmartSim-PR451_) (SmartSim-PR461_) +- SmartSim can now be built and used on platforms using Apple Silicon + (ARM64). Currently, only the PyTorch backend is supported. Note that libtorch + will be downloaded from a CrayLabs github repo. (SmartSim-PR465_) .. _SmartSim-PR446: https://github.com/CrayLabs/SmartSim/pull/446 .. _SmartSim-PR448: https://github.com/CrayLabs/SmartSim/pull/448 .. _SmartSim-PR451: https://github.com/CrayLabs/SmartSim/pull/451 .. _SmartSim-PR453: https://github.com/CrayLabs/SmartSim/pull/453 .. _SmartSim-PR461: https://github.com/CrayLabs/SmartSim/pull/461 +.. _SmartSim-PR465: https://github.com/CrayLabs/SmartSim/pull/465 .. _SmartSim-PR472: https://github.com/CrayLabs/SmartSim/pull/472 diff --git a/smartsim/_core/_install/builder.py b/smartsim/_core/_install/builder.py index 27c11c791..fa8abf4fa 100644 --- a/smartsim/_core/_install/builder.py +++ b/smartsim/_core/_install/builder.py @@ -26,6 +26,7 @@ import concurrent.futures import enum +import itertools import os import platform import re @@ -44,12 +45,10 @@ from shutil import which from subprocess import SubprocessError -# NOTE: This will be imported by setup.py and hence no -# smartsim related items should be imported into -# this file. +# NOTE: This will be imported by setup.py and hence no smartsim related +# items should be imported into this file. -# TODO: -# - check cmake version and use system if possible to avoid conflicts +# TODO: check cmake version and use system if possible to avoid conflicts TRedisAIBackendStr = t.Literal["tensorflow", "torch", "onnxruntime", "tflite"] TDeviceStr = t.Literal["cpu", "gpu"] @@ -84,6 +83,7 @@ class BuildError(Exception): class Architecture(enum.Enum): X64 = ("x86_64", "amd64") + ARM64 = ("arm64",) @classmethod def from_str(cls, string: str, /) -> "Architecture": @@ -341,6 +341,10 @@ def __rai_dependency_name__(self) -> str: ... @abstractmethod def __place_for_rai__(self, target: t.Union[str, "os.PathLike[str]"]) -> Path: ... + @staticmethod + @abstractmethod + def supported_platforms() -> t.Sequence[t.Tuple[OperatingSystem, Architecture]]: ... + def _place_rai_dep_at( target: t.Union[str, "os.PathLike[str]"], verbose: bool @@ -366,6 +370,8 @@ class RedisAIBuilder(Builder): def __init__( self, + _os: OperatingSystem = OperatingSystem.from_str(platform.system()), + architecture: Architecture = Architecture.from_str(platform.machine()), build_env: t.Optional[t.Dict[str, t.Any]] = None, torch_dir: str = "", libtf_dir: str = "", @@ -376,7 +382,10 @@ def __init__( verbose: bool = False, ) -> None: super().__init__(build_env or {}, jobs=jobs, verbose=verbose) + self.rai_install_path: t.Optional[Path] = None + self._os = _os + self._architecture = architecture # convert to int for RAI build script self._torch = build_torch @@ -385,10 +394,25 @@ def __init__( self.libtf_dir = libtf_dir self.torch_dir = torch_dir - # TODO: It might be worth making these constructor args so that users - # of this class can configure exactly _what_ they are building. - self._os = OperatingSystem.from_str(platform.system()) - self._architecture = Architecture.from_str(platform.machine()) + # Sanity checks + self._validate_platform() + + def _validate_platform(self) -> None: + platform_ = (self._os, self._architecture) + unsupported = [] + if platform_ not in _DLPackRepository.supported_platforms(): + unsupported.append("DLPack") + if self.fetch_tf and (platform_ not in _TFArchive.supported_platforms()): + unsupported.append("Tensorflow") + if self.fetch_onnx and (platform_ not in _ORTArchive.supported_platforms()): + unsupported.append("ONNX") + if self.fetch_torch and (platform_ not in _PTArchive.supported_platforms()): + unsupported.append("PyTorch") + if unsupported: + raise BuildError( + f"The {', '.join(unsupported)} backend(s) are not " + f"supported on {self._os} with {self._architecture}" + ) @property def rai_build_path(self) -> Path: @@ -436,6 +460,8 @@ def fail_to_format(reason: str) -> BuildError: # pragma: no cover raise fail_to_format(f"Unknown operating system: {self._os}") if self._architecture == Architecture.X64: arch = "x64" + elif self._architecture == Architecture.ARM64: + arch = "arm64v8" else: # pragma: no cover raise fail_to_format(f"Unknown architecture: {self._architecture}") return self.rai_build_path / f"deps/{os_}-{arch}-{device}" @@ -450,13 +476,18 @@ def _get_deps_to_fetch_for( # dependency versions were declared in single location. # Unfortunately importing into this module is non-trivial as it # is used as script in the SmartSim `setup.py`. - fetchable_deps: t.Sequence[t.Tuple[bool, _RAIBuildDependency]] = ( - (True, _DLPackRepository("v0.5_RAI")), - (self.fetch_torch, _PTArchive(os_, device, "2.0.1")), - (self.fetch_tf, _TFArchive(os_, arch, device, "2.13.1")), - (self.fetch_onnx, _ORTArchive(os_, device, "1.16.3")), - ) - return tuple(dep for should_fetch, dep in fetchable_deps if should_fetch) + + # DLPack is always required + fetchable_deps: t.List[_RAIBuildDependency] = [_DLPackRepository("v0.5_RAI")] + if self.fetch_torch: + pt_dep = _choose_pt_variant(os_) + fetchable_deps.append(pt_dep(arch, device, "2.0.1")) + if self.fetch_tf: + fetchable_deps.append(_TFArchive(os_, arch, device, "2.13.1")) + if self.fetch_onnx: + fetchable_deps.append(_ORTArchive(os_, device, "1.16.3")) + + return tuple(fetchable_deps) def symlink_libtf(self, device: str) -> None: """Add symbolic link to available libtensorflow in RedisAI deps. @@ -698,6 +729,14 @@ def clone( class _DLPackRepository(_WebGitRepository, _RAIBuildDependency): version: str + @staticmethod + def supported_platforms() -> t.Sequence[t.Tuple[OperatingSystem, Architecture]]: + return ( + (OperatingSystem.LINUX, Architecture.X64), + (OperatingSystem.DARWIN, Architecture.X64), + (OperatingSystem.DARWIN, Architecture.ARM64), + ) + @property def url(self) -> str: return "https://github.com/RedisAI/dlpack.git" @@ -756,30 +795,20 @@ def _extract_download( zip_file.extractall(target) -@t.final @dataclass(frozen=True) class _PTArchive(_WebZip, _RAIBuildDependency): - os_: OperatingSystem + architecture: Architecture device: TDeviceStr version: str - @property - def url(self) -> str: - 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 self.os_ == OperatingSystem.DARWIN: - if self.device == "gpu": - raise BuildError("RedisAI does not currently support GPU on Macos") - pt_build = "cpu" - libtorch_arch = f"libtorch-macos-{self.version}.zip" - else: - raise BuildError(f"Unexpected OS for the PT Archive: {self.os_}") - return f"https://download.pytorch.org/libtorch/{pt_build}/{libtorch_arch}" + @staticmethod + def supported_platforms() -> t.Sequence[t.Tuple[OperatingSystem, Architecture]]: + # TODO: This will need to be revisited if the inheritance tree gets deeper + return tuple( + itertools.chain.from_iterable( + var.supported_platforms() for var in _PTArchive.__subclasses__() + ) + ) @property def __rai_dependency_name__(self) -> str: @@ -793,6 +822,66 @@ def __place_for_rai__(self, target: t.Union[str, "os.PathLike[str]"]) -> Path: return target +@t.final +class _PTArchiveLinux(_PTArchive): + @staticmethod + def supported_platforms() -> t.Sequence[t.Tuple[OperatingSystem, Architecture]]: + return ((OperatingSystem.LINUX, Architecture.X64),) + + @property + def url(self) -> str: + if self.device == "gpu": + pt_build = "cu117" + else: + pt_build = "cpu" + # pylint: disable-next=line-too-long + libtorch_archive = ( + f"libtorch-cxx11-abi-shared-without-deps-{self.version}%2B{pt_build}.zip" + ) + return f"https://download.pytorch.org/libtorch/{pt_build}/{libtorch_archive}" + + +@t.final +class _PTArchiveMacOSX(_PTArchive): + @staticmethod + def supported_platforms() -> t.Sequence[t.Tuple[OperatingSystem, Architecture]]: + return ( + (OperatingSystem.DARWIN, Architecture.ARM64), + (OperatingSystem.DARWIN, Architecture.X64), + ) + + @property + def url(self) -> str: + if self.device == "gpu": + raise BuildError("RedisAI does not currently support GPU on Mac OSX") + if self.architecture == Architecture.X64: + pt_build = "cpu" + libtorch_archive = f"libtorch-macos-{self.version}.zip" + root_url = "https://download.pytorch.org/libtorch" + return f"{root_url}/{pt_build}/{libtorch_archive}" + if self.architecture == Architecture.ARM64: + libtorch_archive = f"libtorch-macos-arm64-{self.version}.zip" + # pylint: disable-next=line-too-long + root_url = ( + "https://github.com/CrayLabs/ml_lib_builder/releases/download/v0.1/" + ) + return f"{root_url}/{libtorch_archive}" + + raise BuildError("Unsupported architecture for Pytorch: {self.architecture}") + + +def _choose_pt_variant( + os_: OperatingSystem, +) -> t.Union[t.Type[_PTArchiveLinux], t.Type[_PTArchiveMacOSX]]: + + if os_ == OperatingSystem.DARWIN: + return _PTArchiveMacOSX + if os_ == OperatingSystem.LINUX: + return _PTArchiveLinux + + raise BuildError(f"Unsupported OS for PyTorch: {os_}") + + @t.final @dataclass(frozen=True) class _TFArchive(_WebTGZ, _RAIBuildDependency): @@ -801,6 +890,13 @@ class _TFArchive(_WebTGZ, _RAIBuildDependency): device: TDeviceStr version: str + @staticmethod + def supported_platforms() -> t.Sequence[t.Tuple[OperatingSystem, Architecture]]: + return ( + (OperatingSystem.LINUX, Architecture.X64), + (OperatingSystem.DARWIN, Architecture.X64), + ) + @property def url(self) -> str: if self.architecture == Architecture.X64: @@ -843,6 +939,13 @@ class _ORTArchive(_WebTGZ, _RAIBuildDependency): device: TDeviceStr version: str + @staticmethod + def supported_platforms() -> t.Sequence[t.Tuple[OperatingSystem, Architecture]]: + return ( + (OperatingSystem.LINUX, Architecture.X64), + (OperatingSystem.DARWIN, Architecture.X64), + ) + @property def url(self) -> str: ort_url_base = ( diff --git a/tests/install/test_builder.py b/tests/install/test_builder.py index 446cdf776..f6ad93b71 100644 --- a/tests/install/test_builder.py +++ b/tests/install/test_builder.py @@ -34,10 +34,12 @@ import pytest import smartsim._core._install.builder as build +from smartsim._core._install.buildenv import RedisAIVersion # The tests in this file belong to the group_a group pytestmark = pytest.mark.group_a +RAI_VERSIONS = RedisAIVersion("1.2.7") for_each_device = pytest.mark.parametrize("device", ["cpu", "gpu"]) @@ -56,23 +58,21 @@ @pytest.mark.parametrize( "mock_os", [pytest.param(os_, id=f"os='{os_}'") for os_ in ("Windows", "Java", "")] ) -def test_rai_builder_raises_on_unsupported_op_sys(monkeypatch, mock_os): - monkeypatch.setattr(platform, "system", lambda: mock_os) +def test_os_enum_raises_on_unsupported(mock_os): with pytest.raises(build.BuildError, match="operating system") as err_info: - build.RedisAIBuilder() + build.OperatingSystem.from_str(mock_os) @pytest.mark.parametrize( "mock_arch", [ pytest.param(arch_, id=f"arch='{arch_}'") - for arch_ in ("i386", "i686", "i86pc", "aarch64", "arm64", "armv7l", "") + for arch_ in ("i386", "i686", "i86pc", "aarch64", "armv7l", "") ], ) -def test_rai_builder_raises_on_unsupported_architecture(monkeypatch, mock_arch): - monkeypatch.setattr(platform, "machine", lambda: mock_arch) +def test_arch_enum_raises_on_unsupported(mock_arch): with pytest.raises(build.BuildError, match="architecture"): - build.RedisAIBuilder() + build.Architecture.from_str(mock_arch) @pytest.fixture @@ -84,6 +84,7 @@ def p_test_dir(test_dir): def test_rai_builder_raises_if_attempting_to_place_deps_when_build_dir_dne( monkeypatch, p_test_dir, device ): + monkeypatch.setattr(build.RedisAIBuilder, "_validate_platform", lambda a: None) monkeypatch.setattr( build.RedisAIBuilder, "rai_build_path", @@ -99,6 +100,7 @@ def test_rai_builder_raises_if_attempting_to_place_deps_in_nonempty_dir( monkeypatch, p_test_dir, device ): (p_test_dir / "some_file.txt").touch() + monkeypatch.setattr(build.RedisAIBuilder, "_validate_platform", lambda a: None) monkeypatch.setattr( build.RedisAIBuilder, "rai_build_path", property(lambda self: p_test_dir) ) @@ -111,6 +113,27 @@ def test_rai_builder_raises_if_attempting_to_place_deps_in_nonempty_dir( rai_builder._fetch_deps_for(device) +invalid_build_arm64 = [ + dict(build_tf=True, build_onnx=True), + dict(build_tf=False, build_onnx=True), + dict(build_tf=True, build_onnx=False), +] +invalid_build_ids = [ + ",".join([f"{key}={value}" for key, value in d.items()]) + for d in invalid_build_arm64 +] + + +@pytest.mark.parametrize("build_options", invalid_build_arm64, ids=invalid_build_ids) +def test_rai_builder_raises_if_unsupported_deps_on_arm64(build_options): + with pytest.raises(build.BuildError, match=r"are not supported on.*ARM64"): + build.RedisAIBuilder( + _os=build.OperatingSystem.DARWIN, + architecture=build.Architecture.ARM64, + **build_options, + ) + + def _confirm_inst_presence(type_, should_be_present, seq): expected_num_occurrences = 1 if should_be_present else 0 occurrences = filter(lambda item: isinstance(item, type_), seq) @@ -133,8 +156,10 @@ def _confirm_inst_presence(type_, should_be_present, seq): @toggle_build_pt @toggle_build_ort def test_rai_builder_will_add_dep_if_backend_requested_wo_duplicates( - device, build_tf, build_pt, build_ort + monkeypatch, device, build_tf, build_pt, build_ort ): + monkeypatch.setattr(build.RedisAIBuilder, "_validate_platform", lambda a: None) + rai_builder = build.RedisAIBuilder( build_tf=build_tf, build_torch=build_pt, build_onnx=build_ort ) @@ -149,8 +174,9 @@ def test_rai_builder_will_add_dep_if_backend_requested_wo_duplicates( @toggle_build_tf @toggle_build_pt def test_rai_builder_will_not_add_dep_if_custom_dep_path_provided( - device, p_test_dir, build_tf, build_pt + monkeypatch, device, p_test_dir, build_tf, build_pt ): + monkeypatch.setattr(build.RedisAIBuilder, "_validate_platform", lambda a: None) mock_ml_lib = p_test_dir / "some/ml/lib" mock_ml_lib.mkdir(parents=True) rai_builder = build.RedisAIBuilder( @@ -171,6 +197,7 @@ def test_rai_builder_will_not_add_dep_if_custom_dep_path_provided( def test_rai_builder_raises_if_it_fetches_an_unexpected_number_of_ml_deps( monkeypatch, p_test_dir ): + monkeypatch.setattr(build.RedisAIBuilder, "_validate_platform", lambda a: None) monkeypatch.setattr( build.RedisAIBuilder, "rai_build_path", property(lambda self: p_test_dir) ) @@ -205,3 +232,58 @@ def _some_long_io_op(_): build._threaded_map(_some_long_io_op, []) end = time.time() assert end - start < sleep_duration + + +def test_correct_pt_variant_os(): + # Check that all Linux variants return Linux + for linux_variant in build.OperatingSystem.LINUX.value: + os_ = build.OperatingSystem.from_str(linux_variant) + assert build._choose_pt_variant(os_) == build._PTArchiveLinux + + # Check that ARM64 and X86_64 Mac OSX return the Mac variant + all_archs = (build.Architecture.ARM64, build.Architecture.X64) + for arch in all_archs: + os_ = build.OperatingSystem.DARWIN + assert build._choose_pt_variant(os_) == build._PTArchiveMacOSX + + +def test_PTArchiveMacOSX_url(): + arch = build.Architecture.X64 + pt_version = RAI_VERSIONS.torch + + pt_linux_cpu = build._PTArchiveLinux(build.Architecture.X64, "cpu", pt_version) + x64_prefix = "https://download.pytorch.org/libtorch/" + assert x64_prefix in pt_linux_cpu.url + + pt_macosx_cpu = build._PTArchiveMacOSX(build.Architecture.ARM64, "cpu", pt_version) + arm64_prefix = "https://github.com/CrayLabs/ml_lib_builder/releases/download/" + assert arm64_prefix in pt_macosx_cpu.url + + +def test_PTArchiveMacOSX_gpu_error(): + with pytest.raises(build.BuildError, match="support GPU on Mac OSX"): + build._PTArchiveMacOSX(build.Architecture.ARM64, "gpu", RAI_VERSIONS.torch).url + + +def test_valid_platforms(): + assert build.RedisAIBuilder( + _os=build.OperatingSystem.LINUX, + architecture=build.Architecture.X64, + build_tf=True, + build_torch=True, + build_onnx=True, + ) + assert build.RedisAIBuilder( + _os=build.OperatingSystem.DARWIN, + architecture=build.Architecture.X64, + build_tf=True, + build_torch=True, + build_onnx=False, + ) + assert build.RedisAIBuilder( + _os=build.OperatingSystem.DARWIN, + architecture=build.Architecture.X64, + build_tf=False, + build_torch=True, + build_onnx=False, + )