Merge pull request #435 from datamol-io/simple-ipu-install

Simple ipu install

README.md

@@ -54,30 +54,16 @@ pip install --no-deps -e .
 ```
 
 ### For IPU developers
-
 ```bash
-mkdir ~/.venv                               # Create the folder for the environment
-python3 -m venv ~/.venv/graphium_ipu        # Create the environment
-source ~/.venv/graphium_ipu/bin/activate    # Activate the environment
-
-# Update pip to the latest version
-python3 -m pip install --upgrade pip
-
-# Install the PopTorch wheel
-# Make sure this is the 3.3 SDK
-# Change the link according to your operating system and the `PATH_TO_SDK`
-pip install PATH_TO_SDK/poptorch-3.3.0+113432_960e9c294b_ubuntu_20_04-cp38-cp38-linux_x86_64.whl
-
-# Enable Poplar SDK (including Poplar and PopART)
-source PATH_TO_SDK/enable
+# Install Graphcore's SDK and Graphium dependencies in a new environment called `.graphium_ipu`
+./install_ipu.sh .graphium_ipu
+```
 
-# Install the IPU specific and graphium requirements
-pip install -r requirements_ipu.txt
+The above step needs to be done once. After that, enable the SDK and the environment as follows:
 
-# Install Graphium in dev mode
-pip install --no-deps -e .
+```bash
+source enable_ipu.sh .graphium_ipu
 ```
-If you are new to Graphcore IPUs, you can find more details in the section below: `First Time Running On IPUs`. 
 
 ## Training a model
 
@@ -112,42 +98,6 @@ graphium-train --config-path [PATH] --config-name [CONFIG]
 Thanks to the modular nature of `hydra` you can reuse many of our config settings for your own experiments with Graphium.
 
 
-## First Time Running on IPUs
-For new IPU developers this section helps provide some more explanation on how to set up an environment to use Graphcore IPUs with Graphium. 
-
-```bash
-# Set up a virtual environment as normal
-mkdir ~/.venv                               # Create the folder for the environment
-python3 -m venv ~/.venv/graphium_ipu        # Create the environment
-source ~/.venv/graphium_ipu/bin/activate    # Activate the environment
-
-python3 -m pip install --upgrade pip
-# We can download the Poplar SDK directly using `wget` - more details on the various Graphcore downloads can be found here `https://www.graphcore.ai/downloads`
-
-# NOTE: For simplicity this will download the SDK directly where you run this command, we recommend doing this outside the Graphium directory. 
-# Make sure to download the right file according to your operating system
-wget -q -O 'poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz' 'https://downloads.graphcore.ai/direct?package=poplar-poplar_sdk_ubuntu_20_04_3.3.0_208993bbb7-3.3.0&file=poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz'
-
-# Unzip the SDK file
-tar -xzf poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz
-# Then use pip to install the wheel 
-python3 -m pip install poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7/poptorch-3.3.0+113432_960e9c294b_ubuntu_20_04-cp38-cp38-linux_x86_64.whl
-# Enable Poplar SDK (including Poplar and PopART)
-source poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7/enable 
-
-# Then as a quick test make sure poptorch is correctly installed
-# If it is, this will not execute properly. 
-python3 -c "import poptorch;print('poptorch installed correctly')"
-
-# Install the IPU specific and graphium requirements
-pip install -r requirements_ipu.txt
-# Install Graphium in dev mode
-python -m pip install --no-deps -e .
-
-```
-
-
-
 ## License
 
 Under the Apache-2.0 license. See [LICENSE](LICENSE).

enable_ipu.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+# Default location for the virtual environment
+default_venv_name=".graphium_ipu"
+
+# Allow the user to specify the location of their virtual environment
+# If not specified, use the default location
+venv_name=${1:-$default_venv_name}
+
+# Constants
+sdk_path="${venv_name}/poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7"
+
+# Source the virtual environment
+source ${venv_name}/bin/activate
+source ${sdk_path}/enable

install_ipu.sh

@@ -0,0 +1,98 @@
+#!/bin/bash
+
+# Default location for the virtual environment
+default_venv_name=".graphium_ipu"
+
+# Allow the user to specify the location of their virtual environment
+# If not specified, use the default location
+venv_name=${1:-$default_venv_name}
+
+# Constants
+sdk_compressed_file="poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz"
+sdk_wheel_file="poptorch-3.3.0+113432_960e9c294b_ubuntu_20_04-cp38-cp38-linux_x86_64.whl"
+sdk_url="https://downloads.graphcore.ai/direct?package=poplar-poplar_sdk_ubuntu_20_04_3.3.0_208993bbb7-3.3.0&file=${sdk_compressed_file}"
+sdk_path="${venv_name}/poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7"
+
+# Check for Python3 and pip
+if ! command -v python3 &>/dev/null; then
+    echo "Python3 is required but it's not installed. Exiting."
+    exit 1
+fi
+
+if ! command -v pip3 &>/dev/null; then
+    echo "pip3 is required but it's not installed. Exiting."
+    exit 1
+fi
+
+# Remove existing venv directory if it exists
+if [[ -d $venv_name ]]; then
+    echo "Removing existing virtual environment directory..."
+    rm -rf $venv_name
+fi
+
+# Create the virtual environment
+echo "Creating virtual environment..."
+mkdir -p $venv_name
+python3 -m venv $venv_name
+source $venv_name/bin/activate
+
+# Update pip to the latest version
+echo "Upgrading pip..."
+python3 -m pip install --upgrade pip
+
+# Download the Poplar SDK
+echo "Downloading Poplar SDK..."
+wget -q -O "${venv_name}/${sdk_compressed_file}" "$sdk_url"
+
+# Check the wget exit status
+if [ $? -ne 0 ]; then
+    echo "Failed to download Poplar SDK. Exiting."
+    exit 1
+fi
+
+# Unzip the SDK file
+echo "Extracting Poplar SDK..."
+tar -xzf "$venv_name/$sdk_compressed_file" -C $venv_name
+
+# Install the PopTorch wheel
+echo "Installing PopTorch..."
+python3 -m pip install "${sdk_path}/${sdk_wheel_file}"
+
+# Enable Poplar SDK (including Poplar and PopART)
+echo "Enabling Poplar SDK..."
+source ${sdk_path}/enable
+
+# Install the IPU specific and Graphium requirements
+echo "Installing IPU specific and Graphium requirements..."
+python3 -m pip install -r requirements_ipu.txt
+
+# Install Graphium in dev mode
+echo "Installing Graphium in dev mode..."
+python3 -m pip install --no-deps -e .
+
+# This is a quick test make sure poptorch is correctly installed
+if python3 -c "import poptorch;print('poptorch installed correctly')" &> /dev/null; then
+    echo "Installation completed successfully."
+else
+    echo "Installation was not successful. Please check the logs and try again."
+    exit 1  # Exit with status code 1 to indicate failure
+fi
+
+# Download the datafiles (Total ~ 10Mb - nothing compared to the libraries)
+echo "Downloading the sub-datasets consisting on the ToyMix dataset"
+toymix_dir=expts/data/neurips2023/small-dataset/
+mkdir -p $toymix_dir
+
+base_url="https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/"
+files=("ZINC12k.csv.gz" "Tox21-7k-12-labels.csv.gz" "qm9.csv.gz" "qm9_random_splits.pt" "Tox21_random_splits.pt" "ZINC12k_random_splits.pt")
+
+for file in "${files[@]}"; do
+    if [ ! -f "${toymix_dir}${file}" ]; then
+        echo "Downloading ${file}..."
+        wget -P "${toymix_dir}" "${base_url}${file}"
+    else
+        echo "${file} already exists. Skipping..."
+    fi
+done
+
+echo "Data has been successfully downloaded."