chore: Set the upper bound of sklearn to unbreak presubmit

yinghsienwu · copybara-github · commit 52ce05a92574 · 2024-12-10T10:44:29.000-08:00
PiperOrigin-RevId: 704765085
diff --git a/setup.py b/setup.py
@@ -131,7 +131,7 @@
     # ray train extras required for prediction tests
     "ray[train]",
     # Framework version constraints copied from testing_extra_require
-    "scikit-learn",
+    "scikit-learn<1.6.0",
     "tensorflow",
     "torch >= 2.0.0, < 2.1.0",
     "xgboost",
@@ -204,7 +204,8 @@
         "kfp >= 2.6.0, < 3.0.0",
         "pytest-asyncio",
         "pytest-xdist",
-        "scikit-learn",
+        "scikit-learn<1.6.0; python_version<='3.10'",
+        "scikit-learn; python_version>'3.10'",
         # Lazy import requires > 2.12.0
         "tensorflow == 2.13.0; python_version<='3.11'",
         "tensorflow == 2.16.1; python_version>'3.11'",
diff --git a/testing/constraints-ray-2.9.3.txt b/testing/constraints-ray-2.9.3.txt
@@ -11,3 +11,4 @@ mlflow==1.30.1 # Pinned to speed up installation
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 IPython # Added to test supernova rich html buttons
 pandas==2.1.4 # Pandas must be <2.2.0 to be compatible with ray 2.9.3
+scikit-learn<1.6.0  # Breaking Ray 2.9.3
diff --git a/tests/unit/vertexai/test_feature_monitor.py b/tests/unit/vertexai/test_feature_monitor.py
@@ -17,10 +17,8 @@
 
 import re
 from typing import Dict, List, Tuple
-from unittest import mock
-from unittest.mock import call, patch
+from unittest.mock import patch
 
-from google.api_core import operation as ga_operation
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
 
@@ -84,9 +82,7 @@ def create_feature_monitor_job_mock():
         FeatureRegistryServiceClient,
         "create_feature_monitor_job",
     ) as create_feature_monitor_job_mock:
-        create_feature_monitor_job_mock.return_value = (
-            _TEST_FG1_FMJ1
-        )
+        create_feature_monitor_job_mock.return_value = _TEST_FG1_FMJ1
         yield create_feature_monitor_job_mock
 
 
