Changed tune.report import for ray>=2 (#1392)

* Changed tune.report import for ray>=2

* env: Changed pydantic restriction in env

* Reverted Pydantic install conditions

* Reverted Pydantic install conditions

* test: Check if GPU is available

* tests: uncommented a line

* tests: Better fix for Ray GPU checking

* tests: Added timeout to dataset loading

* tests: Deleted _test_hf_data()

* test: Reduce lrl2 dataset size

* bug: timeout error

* bug: timeout error

* fix: Added threading check for timout issue

* Undo old commits

* Timeout fix from #1406

---------

Co-authored-by: Daniel Grindrod <dannycg1996@gmail.com>

flaml/automl/model.py

@@ -2833,4 +2833,4 @@ class suppress_stdout_stderr:
         os.dup2(self.save_fds[1], 2)
         # Close the null files
         os.close(self.null_fds[0])
-        os.close(self.null_fds[1])
+        os.close(self.null_fds[1])

flaml/tune/tune.py

@@ -197,9 +197,16 @@ def report(_metric=None, **kwargs):
     global _training_iteration
     if _use_ray:
         try:
-            from ray import tune
+            from ray import __version__ as ray_version
 
-            return tune.report(_metric, **kwargs)
+            if ray_version.startswith("1."):
+                from ray import tune
+
+                return tune.report(_metric, **kwargs)
+            else:  # ray>=2
+                from ray.air import session
+
+                return session.report(metrics={"metric": _metric, **kwargs})
         except ImportError:
             # calling tune.report() outside tune.run()
             return

test/test_gpu.py

@@ -59,6 +59,17 @@ def _test_hf_data():
     except requests.exceptions.ConnectionError:
         return
 
+    # Tests will only run if there is a GPU available
+    try:
+        import ray
+
+        pg = ray.util.placement_group([{"CPU": 1, "GPU": 1}])
+
+        if not pg.wait(timeout_seconds=10):  # Wait 10 seconds for resources
+            raise RuntimeError("No available node types can fulfill resource request!")
+    except RuntimeError:
+        return
+
     custom_sent_keys = ["sentence1", "sentence2"]
     label_key = "label"