Add optuna integration dependency

coiled · Aug 2, 2024 · 8b24fdc · 8b24fdc
1 parent 859f2b8
commit 8b24fdc
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 1 deletion.
diff --git a/AB_environments/AB_baseline.conda.yaml b/AB_environments/AB_baseline.conda.yaml
@@ -37,6 +37,7 @@ dependencies:
   - zict ==3.0.0
   - xgboost ==2.1.1
   - optuna ==3.6.1
+  - optuna-integration ==3.6.0
   - scipy ==1.14.0
   - sqlalchemy ==2.0.31
   - pynvml ==11.5.3

diff --git a/AB_environments/AB_sample.conda.yaml b/AB_environments/AB_sample.conda.yaml
@@ -43,6 +43,7 @@ dependencies:
   - zict ==3.0.0
   - xgboost ==2.1.1
   - optuna ==3.6.1
+  - optuna-integration ==3.6.0
   - scipy ==1.14.0
   - sqlalchemy ==2.0.31
   - pynvml ==11.5.3

diff --git a/ci/environment.yml b/ci/environment.yml
@@ -39,6 +39,7 @@ dependencies:
   - zict ==3.0.0
   - xgboost ==2.1.1
   - optuna ==3.6.1
+  - optuna-integration ==3.6.0
   - scipy ==1.14.0
   - sqlalchemy ==2.0.31
   - pynvml ==11.5.3

diff --git a/tests/workflows/test_xgboost_optuna.py b/tests/workflows/test_xgboost_optuna.py
@@ -7,8 +7,8 @@
 xgb = pytest.importorskip("xgboost")
 pytest.importorskip("sklearn")
 
-from optuna.integration.dask import DaskStorage  # noqa: E402
 from optuna.samplers import RandomSampler  # noqa: E402
+from optuna_integration import DaskStorage  # noqa: E402
 from sklearn.datasets import fetch_covtype  # noqa: E402
 from sklearn.model_selection import KFold, cross_val_score  # noqa: E402
 from sklearn.preprocessing import LabelEncoder  # noqa: E402