From 2b8cd08f9b298b81d1cc8d2d6b8678d06b771ef9 Mon Sep 17 00:00:00 2001 From: Andrew Shao Date: Fri, 5 Jan 2024 14:51:23 -0600 Subject: [PATCH 1/4] Fix torch index when installing through smart build --- smartsim/_core/_cli/build.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/smartsim/_core/_cli/build.py b/smartsim/_core/_cli/build.py index 474d96c8a..8ffa27dc7 100644 --- a/smartsim/_core/_cli/build.py +++ b/smartsim/_core/_cli/build.py @@ -260,10 +260,15 @@ def check_py_torch_version(versions: Versioner, device: _TDeviceStr = "cpu") -> "Torch version not found in python environment. " "Attempting to install via `pip`" ) + if device.lower() == "cpu": + index_url = "https://download.pytorch.org/whl/cpu" + else: + index_url = f"https://download.pytorch.org/whl/{device_suffix}" + pip( "install", - "-f", - "https://download.pytorch.org/whl/torch_stable.html", + "--extra-index-url", + index_url, *(f"{package}=={version}" for package, version in torch_deps.items()), ) elif missing or conflicts: From 7b9e0b92b7ec28906650b1e3c1a0da8d5d536412 Mon Sep 17 00:00:00 2001 From: Andrew Shao Date: Fri, 5 Jan 2024 16:00:45 -0600 Subject: [PATCH 2/4] Make the change more compact --- smartsim/_core/_cli/build.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/smartsim/_core/_cli/build.py b/smartsim/_core/_cli/build.py index 8ffa27dc7..285cb35dd 100644 --- a/smartsim/_core/_cli/build.py +++ b/smartsim/_core/_cli/build.py @@ -229,6 +229,7 @@ def build_redis_ai( def check_py_torch_version(versions: Versioner, device: _TDeviceStr = "cpu") -> None: """Check Python environment for TensorFlow installation""" + device = device.lower() if BuildEnv.is_macos(): if device == "gpu": raise BuildError("SmartSim does not support GPU on MacOS") @@ -260,15 +261,11 @@ def check_py_torch_version(versions: Versioner, device: _TDeviceStr = "cpu") -> "Torch version not found in python environment. " "Attempting to install via `pip`" ) - if device.lower() == "cpu": - index_url = "https://download.pytorch.org/whl/cpu" - else: - index_url = f"https://download.pytorch.org/whl/{device_suffix}" - + wheel_device = device if device == "cpu" else device_suffix pip( "install", "--extra-index-url", - index_url, + f"https://download.pytorch.org/whl/{wheel_device}" *(f"{package}=={version}" for package, version in torch_deps.items()), ) elif missing or conflicts: From b7fe003c93d424ce4614936955110accae20e567 Mon Sep 17 00:00:00 2001 From: Andrew Shao Date: Fri, 5 Jan 2024 17:49:45 -0600 Subject: [PATCH 3/4] Fix missing comma --- smartsim/_core/_cli/build.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/smartsim/_core/_cli/build.py b/smartsim/_core/_cli/build.py index 285cb35dd..82b3c7786 100644 --- a/smartsim/_core/_cli/build.py +++ b/smartsim/_core/_cli/build.py @@ -261,11 +261,11 @@ def check_py_torch_version(versions: Versioner, device: _TDeviceStr = "cpu") -> "Torch version not found in python environment. " "Attempting to install via `pip`" ) - wheel_device = device if device == "cpu" else device_suffix + wheel_device = device if device == "cpu" else device_suffix.replace("+","") pip( "install", "--extra-index-url", - f"https://download.pytorch.org/whl/{wheel_device}" + f"https://download.pytorch.org/whl/{wheel_device}", *(f"{package}=={version}" for package, version in torch_deps.items()), ) elif missing or conflicts: From 73846b90f2e864ef702b9cdc6eb8c77097755bc6 Mon Sep 17 00:00:00 2001 From: Andrew Shao Date: Fri, 5 Jan 2024 18:00:31 -0600 Subject: [PATCH 4/4] Avoid lowering an input variable --- smartsim/_core/_cli/build.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/smartsim/_core/_cli/build.py b/smartsim/_core/_cli/build.py index 82b3c7786..87fbff5fb 100644 --- a/smartsim/_core/_cli/build.py +++ b/smartsim/_core/_cli/build.py @@ -226,10 +226,10 @@ def build_redis_ai( logger.info("ML Backends and RedisAI build complete!") -def check_py_torch_version(versions: Versioner, device: _TDeviceStr = "cpu") -> None: +def check_py_torch_version(versions: Versioner, device_in: _TDeviceStr = "cpu") -> None: """Check Python environment for TensorFlow installation""" - device = device.lower() + device = device_in.lower() if BuildEnv.is_macos(): if device == "gpu": raise BuildError("SmartSim does not support GPU on MacOS")