chore: import upstream snapshot with attribution

python/ray/tune/tests/test_function_api.py

@@ -0,0 +1,277 @@
+import json
+import os
+import sys
+import tempfile
+import unittest
+
+import ray
+from ray import tune
+from ray.air.constants import TRAINING_ITERATION
+from ray.rllib import _register_all
+from ray.train.tests.util import mock_storage_context
+from ray.tune import Checkpoint, CheckpointConfig
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.result import DEFAULT_METRIC
+from ray.tune.schedulers import ResourceChangingScheduler
+from ray.tune.trainable import with_parameters, wrap_function
+
+
+class FunctionCheckpointingTest(unittest.TestCase):
+    def setUp(self):
+        self.tmpdir = tempfile.TemporaryDirectory()
+
+    def create_trainable(self, train_fn):
+        return wrap_function(train_fn)(storage=mock_storage_context())
+
+    def tearDown(self):
+        self.tmpdir.cleanup()
+
+    def testCheckpointReuse(self):
+        """Test that repeated save/restore never reuses same checkpoint dir."""
+
+        def train_fn(config):
+            checkpoint = ray.tune.get_checkpoint()
+            if checkpoint:
+                with checkpoint.as_directory() as checkpoint_dir:
+                    count = sum(
+                        "checkpoint-" in path for path in os.listdir(checkpoint_dir)
+                    )
+                    assert count == 1, os.listdir(checkpoint_dir)
+
+            for step in range(20):
+                with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                    path = os.path.join(
+                        temp_checkpoint_dir, "checkpoint-{}".format(step)
+                    )
+                    open(path, "a").close()
+                    ray.tune.report(
+                        dict(test=step),
+                        checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+                    )
+
+        checkpoint = None
+        for i in range(5):
+            new_trainable = self.create_trainable(train_fn)
+            if checkpoint:
+                new_trainable.restore(checkpoint)
+            for i in range(2):
+                result = new_trainable.train()
+            checkpoint = new_trainable.save()
+            new_trainable.stop()
+        assert result[TRAINING_ITERATION] == 10
+
+    def testFunctionRecurringSave(self):
+        """This tests that save and restore are commutative."""
+
+        def train_fn(config):
+            for step in range(10):
+                with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                    if step % 3 == 0:
+                        path = os.path.join(temp_checkpoint_dir, "checkpoint.json")
+                        with open(path, "w") as f:
+                            json.dump({"step": step}, f)
+                    ray.tune.report(
+                        dict(test=step),
+                        checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+                    )
+
+        new_trainable = self.create_trainable(train_fn)
+        new_trainable.train()
+        checkpoint_obj = new_trainable.save()
+        new_trainable.restore(checkpoint_obj)
+        checkpoint = new_trainable.save()
+
+        new_trainable.stop()
+
+        new_trainable2 = self.create_trainable(train_fn)
+        new_trainable2.restore(checkpoint)
+        new_trainable2.train()
+        new_trainable2.stop()
+
+
+class FunctionApiTest(unittest.TestCase):
+    def setUp(self):
+        ray.init(num_cpus=4, num_gpus=0, object_store_memory=150 * 1024 * 1024)
+
+    def tearDown(self):
+        ray.shutdown()
+        _register_all()  # re-register the evicted objects
+
+    def testCheckpointError(self):
+        def train_fn(config):
+            pass
+
+        with self.assertRaises(ValueError):
+            tune.run(
+                train_fn, checkpoint_config=CheckpointConfig(checkpoint_frequency=1)
+            )
+        with self.assertRaises(ValueError):
+            tune.run(
+                train_fn, checkpoint_config=CheckpointConfig(checkpoint_at_end=True)
+            )
+
+    def testWithParameters(self):
+        class Data:
+            def __init__(self):
+                self.data = [0] * 500_000
+
+        data = Data()
+        data.data[100] = 1
+
+        def train_fn(config, data=None):
+            data.data[101] = 2  # Changes are local
+            ray.tune.report(dict(metric=len(data.data), hundred=data.data[100]))
+
+        trial_1, trial_2 = tune.run(
+            with_parameters(train_fn, data=data), num_samples=2
+        ).trials
+
+        self.assertEqual(data.data[101], 0)
+        self.assertEqual(trial_1.last_result["metric"], 500_000)
+        self.assertEqual(trial_1.last_result["hundred"], 1)
+        self.assertEqual(trial_2.last_result["metric"], 500_000)
+        self.assertEqual(trial_2.last_result["hundred"], 1)
+        self.assertTrue(str(trial_1).startswith("train_"))
+
+        # With checkpoint dir parameter
+        def train_fn(config, data=None):
+            data.data[101] = 2  # Changes are local
+            ray.tune.report(dict(metric=len(data.data)))
+
+        trial_1, trial_2 = tune.run(
+            with_parameters(train_fn, data=data), num_samples=2
+        ).trials
+
+        self.assertEqual(data.data[101], 0)
+        self.assertEqual(trial_1.last_result["metric"], 500_000)
+        self.assertEqual(trial_2.last_result["metric"], 500_000)
+        self.assertTrue(str(trial_1).startswith("train_"))
+
+    def testWithParameters2(self):
+        class Data:
+            def __init__(self):
+                import numpy as np
+
+                self.data = np.random.rand((2 * 1024 * 1024))
+
+        def train_fn(config, data=None):
+            pass
+
+        trainable = tune.with_parameters(train_fn, data=Data())
+        # ray.cloudpickle will crash for some reason
+        import cloudpickle as cp
+
+        dumped = cp.dumps(trainable)
+        assert sys.getsizeof(dumped) < 100 * 1024
+
+    def testNewResources(self):
+        sched = ResourceChangingScheduler(
+            resources_allocation_function=(
+                lambda a, b, c, d: PlacementGroupFactory([{"CPU": 2}])
+            )
+        )
+
+        def train_fn(config):
+            ray.tune.report(
+                dict(metric=1, resources=ray.tune.get_context().get_trial_resources())
+            )
+
+        analysis = tune.run(
+            train_fn,
+            scheduler=sched,
+            stop={"training_iteration": 2},
+            resources_per_trial=PlacementGroupFactory([{"CPU": 1}]),
+            num_samples=1,
+        )
+
+        results_list = list(analysis.results.values())
+        assert results_list[0]["resources"].head_cpus == 2.0
+
+    def testWithParametersTwoRuns1(self):
+        # Makes sure two runs in the same script but different ray sessions
+        # pass (https://github.com/ray-project/ray/issues/16609)
+        def train_fn(config, extra=4):
+            ray.tune.report(dict(metric=extra))
+
+        trainable = tune.with_parameters(train_fn, extra=8)
+        out = tune.run(trainable, metric="metric", mode="max")
+        self.assertEqual(out.best_result["metric"], 8)
+
+        self.tearDown()
+        self.setUp()
+
+        def train_fn_2(config, extra=5):
+            ray.tune.report(dict(metric=extra))
+
+        trainable = tune.with_parameters(train_fn_2, extra=9)
+        out = tune.run(trainable, metric="metric", mode="max")
+        self.assertEqual(out.best_result["metric"], 9)
+
+    def testWithParametersTwoRuns2(self):
+        # Makes sure two runs in the same script
+        # pass (https://github.com/ray-project/ray/issues/16609)
+        def train_fn(config, extra=4):
+            ray.tune.report(dict(metric=extra))
+
+        def train_fn_2(config, extra=5):
+            ray.tune.report(dict(metric=extra))
+
+        trainable1 = tune.with_parameters(train_fn, extra=8)
+        trainable2 = tune.with_parameters(train_fn_2, extra=9)
+
+        out1 = tune.run(trainable1, metric="metric", mode="max")
+        out2 = tune.run(trainable2, metric="metric", mode="max")
+        self.assertEqual(out1.best_result["metric"], 8)
+        self.assertEqual(out2.best_result["metric"], 9)
+
+    def testReturnAnonymous(self):
+        def train_fn(config):
+            return config["a"]
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result[DEFAULT_METRIC], 4)
+        self.assertEqual(trial_2.last_result[DEFAULT_METRIC], 8)
+
+    def testReturnSpecific(self):
+        def train_fn(config):
+            return {"m": config["a"]}
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result["m"], 4)
+        self.assertEqual(trial_2.last_result["m"], 8)
+
+    def testYieldAnonymous(self):
+        def train_fn(config):
+            for i in range(10):
+                yield config["a"] + i
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result[DEFAULT_METRIC], 4 + 9)
+        self.assertEqual(trial_2.last_result[DEFAULT_METRIC], 8 + 9)
+
+    def testYieldSpecific(self):
+        def train_fn(config):
+            for i in range(10):
+                yield {"m": config["a"] + i}
+
+        trial_1, trial_2 = tune.run(
+            train_fn, config={"a": tune.grid_search([4, 8])}
+        ).trials
+
+        self.assertEqual(trial_1.last_result["m"], 4 + 9)
+        self.assertEqual(trial_2.last_result["m"], 8 + 9)
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))