Merge branch 'add_rei' of https://github.com/SaiAakash/botorch into a…

…dd_rei
pytorch · Jan 31, 2025 · 281cba7 · 281cba7
2 parents 3133f4f + 01c813d
commit 281cba7
Show file tree

Hide file tree

Showing 10 changed files with 95 additions and 51 deletions.
diff --git a/.github/workflows/deploy_on_release.yml b/.github/workflows/deploy_on_release.yml
@@ -20,9 +20,14 @@ jobs:
     runs-on: ubuntu-latest
     permissions:
       id-token: write # This is required for PyPI OIDC authentication.
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     needs: tests-and-coverage
     steps:
     - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
     - name: Fetch all history for all tags and branches
       run: git fetch --prune --unshallow
     - name: Set up Python
@@ -31,8 +36,8 @@ jobs:
         python-version: "3.10"
     - name: Install dependencies
       run: |
-        pip install .[test]
-        pip install --upgrade build setuptools setuptools_scm wheel
+        uv pip install .[test]
+        uv pip install --upgrade build setuptools setuptools_scm wheel
     - name: Build packages (wheel and source distribution)
       run: |
         python -m build --sdist --wheel

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -15,8 +15,13 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     steps:
     - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
@@ -25,10 +30,10 @@ jobs:
       env:
         ALLOW_LATEST_GPYTORCH_LINOP: true
       run: |
-        pip install git+https://github.com/cornellius-gp/linear_operator.git
-        pip install git+https://github.com/cornellius-gp/gpytorch.git
-        pip install .[dev]
-        pip install beautifulsoup4 ipython nbconvert jinja2
+        uv pip install git+https://github.com/cornellius-gp/linear_operator.git
+        uv pip install git+https://github.com/cornellius-gp/gpytorch.git
+        uv pip install .[dev]
+        uv pip install beautifulsoup4 ipython nbconvert jinja2
     - name: Validate Sphinx
       run: |
         python scripts/validate_sphinx.py -p "$(pwd)"

diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -12,16 +12,20 @@ jobs:
 
   lint:
     runs-on: ubuntu-latest
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     steps:
       - uses: actions/checkout@v4
-
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
           python-version: "3.10"
 
       - name: Install dependencies
-        run: pip install pre-commit
+        run: uv pip install pre-commit
 
       - name: Run pre-commit
         run: pre-commit run --all-files --show-diff-on-failure
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
@@ -20,10 +20,15 @@ jobs:
     runs-on: ubuntu-latest
     permissions:
       id-token: write # This is required for PyPI OIDC authentication.
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     steps:
     - uses: actions/checkout@v4
     - name: Fetch all history for all tags and branches
       run: git fetch --prune --unshallow
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
@@ -32,11 +37,11 @@ jobs:
       env:
         ALLOW_LATEST_GPYTORCH_LINOP: true
       run: |
-        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
-        pip install git+https://github.com/cornellius-gp/linear_operator.git
-        pip install git+https://github.com/cornellius-gp/gpytorch.git
-        pip install .[test]
-        pip install --upgrade build setuptools setuptools_scm wheel
+        uv pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+        uv pip install git+https://github.com/cornellius-gp/linear_operator.git
+        uv pip install git+https://github.com/cornellius-gp/gpytorch.git
+        uv pip install .[test]
+        uv pip install --upgrade build setuptools setuptools_scm wheel
     - name: Extract reduced version and save to env var
       # strip the commit hash from the version to enable upload to pypi
       # env var will persist for subsequent steps

diff --git a/.github/workflows/reusable_test.yml b/.github/workflows/reusable_test.yml
@@ -21,25 +21,30 @@ jobs:
       matrix:
         os: ["ubuntu-latest", "macos-14", "windows-latest"]
         python-version: ["3.10", "3.12"]
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     steps:
     - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install (auto-install dependencies)
       if: ${{ !inputs.use_latest_pytorch_gpytorch }}
       run: |
-        pip install .[test]
+        uv pip install .[test]
     - name: Install dependencies with latest PyTorch & GPyTorch
       if: ${{ inputs.use_latest_pytorch_gpytorch }}
       env:
         ALLOW_LATEST_GPYTORCH_LINOP: true
       run: |
-        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
-        pip install git+https://github.com/cornellius-gp/linear_operator.git
-        pip install git+https://github.com/cornellius-gp/gpytorch.git
-        pip install .[test]
+        uv pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+        uv pip install git+https://github.com/cornellius-gp/linear_operator.git
+        uv pip install git+https://github.com/cornellius-gp/gpytorch.git
+        uv pip install .[test]
     - name: Unit tests and coverage -- BoTorch
       run: |
         pytest -ra test/ --cov botorch/ --cov-report term-missing --cov-report xml:botorch_cov.xml

diff --git a/.github/workflows/reusable_tutorials.yml b/.github/workflows/reusable_tutorials.yml
@@ -31,8 +31,13 @@ jobs:
   tutorials:
     name: Run tutorials
     runs-on: ubuntu-latest
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     steps:
     - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
@@ -43,9 +48,9 @@ jobs:
     - if: ${{ !inputs.use_stable_pytorch_gpytorch }}
       name: Install latest PyTorch & GPyTorch
       run: |
-        pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
-        pip install git+https://github.com/cornellius-gp/linear_operator.git
-        pip install git+https://github.com/cornellius-gp/gpytorch.git
+        uv pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
+        uv pip install git+https://github.com/cornellius-gp/linear_operator.git
+        uv pip install git+https://github.com/cornellius-gp/gpytorch.git
     - if: ${{ inputs.use_stable_pytorch_gpytorch }}
       name: Install min required PyTorch, GPyTorch, and linear_operator
       run: |
@@ -54,26 +59,26 @@ jobs:
         min_torch_version=$(grep '\btorch[>=]=' ${req_txt} | sed 's/[^0-9.]//g')
         min_gpytorch_version=$(grep '\bgpytorch[>=]=' ${req_txt} | sed 's/[^0-9.]//g')
         min_linear_operator_version=$(grep '\blinear_operator[>=]=' ${req_txt} | sed 's/[^0-9.]//g')
-        pip install "numpy<2"  # Numpy >2.0 is not compatible with PyTorch <2.2.
-        pip install "torch==${min_torch_version}" "gpytorch==${min_gpytorch_version}" "linear_operator==${min_linear_operator_version}" torchvision
+        uv pip install "numpy<2"  # Numpy >2.0 is not compatible with PyTorch <2.2.
+        uv pip install "torch==${min_torch_version}" "gpytorch==${min_gpytorch_version}" "linear_operator==${min_linear_operator_version}" torchvision
     - name: Install BoTorch with tutorials dependencies
       env:
         ALLOW_LATEST_GPYTORCH_LINOP: true
       run: |
-        pip install .[tutorials]
+        uv pip install .[tutorials]
     - if: ${{ !inputs.use_stable_ax }}
       name: Install latest Ax
       env:
         # This is so Ax's setup doesn't install a pinned BoTorch version.
         ALLOW_BOTORCH_LATEST: true
       run: |
-        pip install git+https://github.com/facebook/Ax.git
+        uv pip install git+https://github.com/facebook/Ax.git
     - if: ${{ inputs.use_stable_ax }}
       name: Install stable Ax
       env:
         ALLOW_BOTORCH_LATEST: true
       run: |
-        pip install ax-platform --no-binary ax-platform
+        uv pip install ax-platform --no-binary ax-platform
     - if: ${{ inputs.smoke_test }}
       name: Run tutorials with smoke test
       run: |

diff --git a/.github/workflows/reusable_website.yml b/.github/workflows/reusable_website.yml
@@ -24,8 +24,13 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: true
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     steps:
     - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
@@ -35,17 +40,17 @@ jobs:
     - if: ${{ !inputs.publish_versioned_website }}
       name: Install latest GPyTorch and Linear Operator
       run: |
-        pip install git+https://github.com/cornellius-gp/linear_operator.git
-        pip install git+https://github.com/cornellius-gp/gpytorch.git
+        uv pip install git+https://github.com/cornellius-gp/linear_operator.git
+        uv pip install git+https://github.com/cornellius-gp/gpytorch.git
     - name: Install common dependencies
       env:
         ALLOW_LATEST_GPYTORCH_LINOP: true
         ALLOW_BOTORCH_LATEST: true  # Allow Ax to install w/ new BoTorch release.
       run: |
-        pip install ."[dev, tutorials]"
-        # There may not be a compatible Ax pip version, so we use the development version.
-        pip install git+https://github.com/facebook/Ax.git
-        pip install beautifulsoup4 ipython nbconvert jinja2
+        uv pip install ."[dev, tutorials]"
+        # There may not be a compatible Ax uv pip version, so we use the development version.
+        uv pip install git+https://github.com/facebook/Ax.git
+        uv pip install beautifulsoup4 ipython nbconvert jinja2
     - if: ${{ inputs.publish_versioned_website }}
       name: Publish versioned website
       env:

diff --git a/.github/workflows/test_stable.yml b/.github/workflows/test_stable.yml
@@ -21,25 +21,30 @@ jobs:
       matrix:
         os: ["ubuntu-latest", "macos-14"]
         python-version: ["3.10", "3.12"]
+    env:
+      # `uv pip ...` requires venv by default. This skips that requirement.
+      UV_SYSTEM_PYTHON: 1
     steps:
     - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install setuptools  # Needed for next line on Python 3.12.
+        uv pip install setuptools  # Needed for next line on Python 3.12.
         python setup.py egg_info
         req_txt="botorch.egg-info/requires.txt"
         min_torch_version=$(grep '\btorch[>=]=' ${req_txt} | sed 's/[^0-9.]//g')
         # The earliest PyTorch version on Python 3.12 is 2.2.0.
         min_torch_version="$(if ${{ matrix.python-version == '3.12' }}; then echo "2.2.0"; else echo "${min_torch_version}"; fi)"
         min_gpytorch_version=$(grep '\bgpytorch[>=]=' ${req_txt} | sed 's/[^0-9.]//g')
         min_linear_operator_version=$(grep '\blinear_operator[>=]=' ${req_txt} | sed 's/[^0-9.]//g')
-        pip install "numpy<2"  # Numpy >2.0 is not compatible with PyTorch <2.2.
-        pip install "torch==${min_torch_version}" "gpytorch==${min_gpytorch_version}" "linear_operator==${min_linear_operator_version}"
-        pip install .[test]
+        uv pip install "numpy<2"  # Numpy >2.0 is not compatible with PyTorch <2.2.
+        uv pip install "torch==${min_torch_version}" "gpytorch==${min_gpytorch_version}" "linear_operator==${min_linear_operator_version}"
+        uv pip install .[test]
     - name: Unit tests and coverage -- BoTorch
       run: |
         pytest -ra test/ --cov botorch/ --cov-report term-missing --cov-report xml:botorch_cov.xml

diff --git a/botorch/optim/utils/timeout.py b/botorch/optim/utils/timeout.py
@@ -14,6 +14,7 @@
 import numpy.typing as npt
 from botorch.exceptions.errors import OptimizationTimeoutError
 from scipy import optimize
+from threadpoolctl import threadpool_limits
 
 
 def minimize_with_timeout(
@@ -79,20 +80,23 @@ def wrapped_callback(xk: npt.NDArray) -> None:
 
     try:
         warnings.filterwarnings("error", message="Method .* cannot handle")
-        return optimize.minimize(
-            fun=fun,
-            x0=x0,
-            args=args,
-            method=method,
-            jac=jac,
-            hess=hess,
-            hessp=hessp,
-            bounds=bounds,
-            constraints=constraints,
-            tol=tol,
-            callback=wrapped_callback,
-            options=options,
-        )
+        # To prevent slowdowns after scipy 1.15.
+        # See https://github.com/scipy/scipy/issues/22438.
+        with threadpool_limits(limits=1, user_api="blas"):
+            return optimize.minimize(
+                fun=fun,
+                x0=x0,
+                args=args,
+                method=method,
+                jac=jac,
+                hess=hess,
+                hessp=hessp,
+                bounds=bounds,
+                constraints=constraints,
+                tol=tol,
+                callback=wrapped_callback,
+                options=options,
+            )
     except OptimizationTimeoutError as e:
         msg = f"Optimization timed out after {e.runtime} seconds."
         current_fun, *_ = fun(e.current_x, *args)

diff --git a/requirements.txt b/requirements.txt
@@ -4,5 +4,6 @@ gpytorch==1.14
 linear_operator==0.6
 torch>=2.0.1
 pyro-ppl>=1.8.4
-scipy<1.15
+scipy
 multipledispatch
+threadpoolctl