[tune] Experiment stopping API (#6886)

python/ray/tune/tests/test_api.py

@@ -10,7 +10,7 @@ import ray
 from ray.rllib import _register_all
 
 from ray import tune
-from ray.tune import DurableTrainable, Trainable, TuneError
+from ray.tune import DurableTrainable, Trainable, TuneError, Stopper
 from ray.tune import register_env, register_trainable, run_experiments
 from ray.tune.schedulers import TrialScheduler, FIFOScheduler
 from ray.tune.trial import Trial
@@ -452,28 +452,52 @@ class TrainableFunctionApiTest(unittest.TestCase):
             for i in range(10):
                 reporter(test=i)
 
-        class Stopper:
+        class Stopclass:
             def stop(self, trial_id, result):
                 return result["test"] > 6
 
-        [trial] = tune.run(train, stop=Stopper().stop).trials
+        [trial] = tune.run(train, stop=Stopclass().stop).trials
         self.assertEqual(trial.last_result["training_iteration"], 8)
 
+    def testStopper(self):
+        def train(config, reporter):
+            for i in range(10):
+                reporter(test=i)
+
+        class CustomStopper(Stopper):
+            def __init__(self):
+                self._count = 0
+
+            def __call__(self, trial_id, result):
+                print("called")
+                self._count += 1
+                return result["test"] > 6
+
+            def stop_all(self):
+                return self._count > 5
+
+        trials = tune.run(train, num_samples=5, stop=CustomStopper()).trials
+        self.assertTrue(all(t.status == Trial.TERMINATED for t in trials))
+        self.assertTrue(
+            any(
+                t.last_result.get("training_iteration") is None
+                for t in trials))
+
     def testBadStoppingFunction(self):
         def train(config, reporter):
             for i in range(10):
                 reporter(test=i)
 
-        class Stopper:
+        class CustomStopper:
             def stop(self, result):
                 return result["test"] > 6
 
         def stop(result):
             return result["test"] > 6
 
-        with self.assertRaises(ValueError):
-            tune.run(train, stop=Stopper().stop)
-        with self.assertRaises(ValueError):
+        with self.assertRaises(TuneError):
+            tune.run(train, stop=CustomStopper().stop)
+        with self.assertRaises(TuneError):
             tune.run(train, stop=stop)
 
     def testEarlyReturn(self):