tools: Autotuner: Add dependencies for Vizier

Signed-off-by: Eryk Szpotanski <[email protected]>

Author: eszpotanski

tools/AutoTuner/installer.sh

@@ -3,9 +3,20 @@
 # Get the directory where the script is located
 script_dir="$(dirname "${BASH_SOURCE[0]}")"
 
+dependencies=""
+if [[ "$#" -eq 0 ]]; then
+  echo "Installing dependencies for Ray Tune and Vizier"
+  dependencies="ray,vizier"
+elif [[ "$#" -ne 1 ]] || ([[ "$1" != "ray" ]] && [[ "$1" != "vizier" ]]); then
+  echo "Please specify whether 'ray' or 'vizier' dependencies should be installed" >&2
+  exit 1
+else
+  dependencies="$1"
+fi
+
 # Define the virtual environment name
 venv_name="autotuner_env"
 python3 -m venv "$script_dir/$venv_name"
 source "$script_dir/$venv_name/bin/activate"
-pip3 install -e "$script_dir"
+pip3 install -e "$script_dir[$dependencies]"
 deactivate

tools/AutoTuner/pyproject.toml

@@ -12,6 +12,8 @@ dynamic = ["dependencies", "optional-dependencies"]
 
 [tool.setuptools.dynamic]
 dependencies = { file = ["requirements.txt"] }
+optional-dependencies.ray = { file = ["requirements-ray.txt"] }
+optional-dependencies.vizier = { file = ["requirements-vizier.txt"] }
 optional-dependencies.dev = { file = ["requirements-dev.txt"] }
 
 [build-system]

tools/AutoTuner/requirements-ray.txt

@@ -0,0 +1,11 @@
+ray[tune]==2.9.3
+ax-platform>=0.3.3,<=0.3.7
+hyperopt==0.2.7
+optuna==3.6.0
+pandas>=2.0,<=2.2.1
+bayesian-optimization==1.4.0
+colorama==0.4.6
+tensorboard>=2.14.0,<=2.16.2
+protobuf==3.20.3
+SQLAlchemy==1.4.17
+urllib3<=1.26.15

tools/AutoTuner/requirements-vizier.txt

@@ -0,0 +1,3 @@
+jax<=0.4.33
+google-vizier[jax]
+tqdm

tools/AutoTuner/requirements.txt

@@ -1,11 +1 @@
-ray[default,tune]==2.9.3
-ax-platform>=0.3.3,<=0.3.7
-hyperopt==0.2.7
-optuna==3.6.0
-pandas>=2.0,<=2.2.1
-bayesian-optimization==1.4.0
-colorama==0.4.6
-tensorboard>=2.14.0,<=2.16.2
-protobuf==3.20.3
-SQLAlchemy==1.4.17
-urllib3<=1.26.15
+ray[default]==2.9.3