[train] New persistence mode: Tune tests + examples (small) (batch 1) (ray-project#38807)

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: e428265 <arvind.chandramouli@lmco.com>

Author: justinvyu

python/ray/train/tests/test_training_iterator.py

@@ -1,6 +1,5 @@
 import functools
 import time
-import tempfile
 from unittest.mock import patch
 import pytest
 from ray.train._internal.worker_group import WorkerGroup
@@ -12,7 +11,6 @@
 from ray.air._internal.util import StartTraceback
 from ray.train.backend import BackendConfig
 from ray.train._internal.session import init_session, get_session
-from ray.train._internal.storage import StorageContext
 from ray.train._internal.backend_executor import BackendExecutor
 from ray.train._internal.utils import construct_train_func
 from ray.train._internal.checkpoint import CheckpointManager
@@ -23,12 +21,13 @@
     train_func as linear_train_func,
 )
 
+from ray.train.tests.util import mock_storage_context
+
 MAX_RETRIES = 3
 
 
 @pytest.fixture(autouse=True, scope="module")
 def patch_tune_session():
-    tempdir = tempfile.mkdtemp()
     if not get_session():
         init_session(
             training_func=None,
@@ -37,11 +36,7 @@ def patch_tune_session():
             node_rank=None,
             local_world_size=None,
             world_size=None,
-            storage=StorageContext(
-                storage_path=tempdir,
-                experiment_dir_name="exp_name",
-                trial_dir_name="trial_name",
-            ),
+            storage=mock_storage_context(),
         )
     yield
 

python/ray/train/tests/util.py

@@ -26,8 +26,14 @@ def load_dict_checkpoint(checkpoint: Checkpoint) -> Dict[str, Any]:
 
 
 def mock_storage_context() -> StorageContext:
-    return StorageContext(
-        storage_path=tempfile.mkdtemp(),
-        experiment_dir_name="exp_name",
-        trial_dir_name="trial_name",
+    storage_path = tempfile.mkdtemp()
+    exp_name = "exp_name"
+    trial_name = "trial_name"
+    storage = StorageContext(
+        storage_path=storage_path,
+        experiment_dir_name=exp_name,
+        trial_dir_name=trial_name,
     )
+    storage.storage_local_path = storage_path
+    os.makedirs(os.path.join(storage_path, exp_name, trial_name), exist_ok=True)
+    return storage

python/ray/tune/BUILD

@@ -50,7 +50,7 @@ py_test(
     name = "test_actor_reuse",
     size = "large",
     srcs = ["tests/test_actor_reuse.py"],
-    tags = ["team:ml", "exclusive", "no_new_storage"],
+    tags = ["team:ml", "exclusive", "new_storage"],
     deps = [":tune_lib", ":conftest"],
 )
 
@@ -94,15 +94,6 @@ py_test(
     tags = ["team:ml", "exclusive", "rllib"],
 )
 
-py_test(
-    name = "test_cluster_searcher",
-    size = "large",
-    srcs = ["tests/test_cluster_searcher.py"],
-    data = ["tests/_test_cluster_interrupt_searcher.py"],
-    deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"],
-)
-
 py_test(
     name = "test_commands",
     size = "medium",
@@ -406,8 +397,8 @@ py_test(
     name = "test_tuner_restore",
     size = "large",
     srcs = ["tests/test_tuner_restore.py"],
-    deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"],
+    deps = [":tune_lib", ":conftest"],
+    tags = ["team:ml", "exclusive", "new_storage"],
 )
 
 py_test(
@@ -453,15 +444,15 @@ py_test(
     size = "small",
     srcs = ["tests/execution/test_actor_caching.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"]
+    tags = ["team:ml", "exclusive", "new_storage"]
 )
 
 py_test(
     name = "test_controller_callback_integration",
     size = "large",
     srcs = ["tests/execution/test_controller_callback_integration.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"]
+    tags = ["team:ml", "exclusive", "new_storage"]
 )
 
 py_test(
@@ -477,31 +468,31 @@ py_test(
     size = "large",
     srcs = ["tests/execution/test_controller_control_integration.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"]
+    tags = ["team:ml", "exclusive", "new_storage"]
 )
 
 py_test(
     name = "test_controller_errors_integration",
     size = "large",
     srcs = ["tests/execution/test_controller_errors_integration.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"]
+    tags = ["team:ml", "exclusive", "new_storage"]
 )
 
 py_test(
     name = "test_controller_resources_integration",
     size = "large",
     srcs = ["tests/execution/test_controller_resources_integration.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"]
+    tags = ["team:ml", "exclusive", "new_storage"]
 )
 
 py_test(
     name = "test_controller_search_alg_integration",
     size = "large",
     srcs = ["tests/execution/test_controller_search_alg_integration.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"]
+    tags = ["team:ml", "exclusive", "new_storage"]
 )
 
 # --------------------------------------------------------------------

python/ray/tune/tests/_test_cluster_interrupt_searcher.py

@@ -2,6 +2,7 @@
 import sys
 
 from ray.tune import PlacementGroupFactory
+
 from ray.tune.tests.execution.utils import create_execution_test_objects, TestingTrial
 
 

python/ray/tune/tests/execution/test_actor_caching.py

@@ -9,6 +9,8 @@
 from ray.tune.execution.tune_controller import TuneController
 from ray.tune.experiment import Trial
 
+from ray.train.tests.util import mock_storage_context
+
 
 @pytest.fixture(scope="function")
 def ray_start_4_cpus_2_gpus_extra():
@@ -43,21 +45,16 @@ def test_callback_save_restore(
 
     Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCallbackSaveRestore
     """
-    runner = TuneController(
-        callbacks=[StatefulCallback()],
-        experiment_path=str(tmpdir),
-    )
-    runner.add_trial(Trial("__fake", stub=True))
+    storage = mock_storage_context()
+    runner = TuneController(callbacks=[StatefulCallback()], storage=storage)
+    runner.add_trial(Trial("__fake", stub=True, storage=storage))
     for i in range(3):
         runner._callbacks.on_trial_result(
             iteration=i, trials=None, trial=None, result=None
         )
     runner.checkpoint(force=True)
     callback = StatefulCallback()
-    runner2 = TuneController(
-        callbacks=[callback],
-        experiment_path=str(tmpdir),
-    )
+    runner2 = TuneController(callbacks=[callback], storage=storage)
     assert callback.counter == 0
     runner2.resume()
     assert callback.counter == 3

python/ray/tune/tests/execution/test_controller_callback_integration.py

@@ -10,6 +10,12 @@
 from ray.tune.experiment import Trial
 
 
+from ray.train.tests.util import mock_storage_context
+
+
+STORAGE = mock_storage_context()
+
+
 @pytest.fixture(scope="function")
 def ray_start_4_cpus_2_gpus_extra():
     address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
@@ -26,12 +32,13 @@ def test_stop_trial(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
     Legacy test: test_trial_runner_3.py::TrialRunnerTest::testStopTrial
     """
     runner = TuneController(
-        resource_manager_factory=lambda: resource_manager_cls(),
+        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
     )
     kwargs = {
         "stopping_criterion": {"training_iteration": 10},
         "placement_group_factory": PlacementGroupFactory([{"CPU": 2, "GPU": 1}]),
         "config": {"sleep": 1},
+        "storage": STORAGE,
     }
     trials = [
         Trial("__fake", **kwargs),
@@ -107,7 +114,9 @@ def test_remove_actor_tracking(ray_start_4_cpus_2_gpus_extra, resource_manager_c
     in ``self._stopping_trials``.
     """
     runner = TuneController(
-        resource_manager_factory=lambda: resource_manager_cls(), reuse_actors=True
+        resource_manager_factory=lambda: resource_manager_cls(),
+        reuse_actors=True,
+        storage=STORAGE,
     )
 
     def train(config):
@@ -117,6 +126,7 @@ def train(config):
 
     kwargs = {
         "placement_group_factory": PlacementGroupFactory([{"CPU": 4, "GPU": 2}]),
+        "storage": STORAGE,
     }
     trials = [Trial("test_remove_actor_tracking", **kwargs) for i in range(4)]
     for t in trials:

python/ray/tune/tests/execution/test_controller_control_integration.py

@@ -13,9 +13,14 @@
 from ray.tune.registry import TRAINABLE_CLASS, _global_registry
 from ray.tune.schedulers import FIFOScheduler
 from ray.tune.search import BasicVariantGenerator
+
+from ray.train.tests.util import mock_storage_context
 from ray.tune.tests.execution.utils import BudgetResourceManager
 
 
+STORAGE = mock_storage_context()
+
+
 @pytest.fixture(scope="function")
 def ray_start_4_cpus_2_gpus_extra():
     address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
@@ -53,11 +58,12 @@ def test_invalid_trainable(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
     Legacy test: test_trial_runner_2.py::TrialRunnerTest::testErrorHandling
     """
     runner = TuneController(
-        resource_manager_factory=lambda: resource_manager_cls(),
+        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
     )
     kwargs = {
         "stopping_criterion": {"training_iteration": 1},
         "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
+        "storage": STORAGE,
     }
     _global_registry.register(TRAINABLE_CLASS, "asdf", None)
     trials = [Trial("asdf", **kwargs), Trial("__fake", **kwargs)]
@@ -78,6 +84,7 @@ def test_overstep(ray_start_4_cpus_2_gpus_extra):
     os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
     runner = TuneController(
         resource_manager_factory=lambda: BudgetResourceManager({"CPU": 4}),
+        storage=STORAGE,
     )
     runner.step()
     with pytest.raises(TuneError):
@@ -106,13 +113,15 @@ def test_failure_recovery(
         search_alg=searchalg,
         scheduler=scheduler,
         resource_manager_factory=lambda: resource_manager_cls(),
+        storage=STORAGE,
     )
     kwargs = {
         "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
         "stopping_criterion": {"training_iteration": 2},
         "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
         "max_failures": max_failures,
         "config": {"mock_error": True, "persistent_error": persistent_error},
+        "storage": STORAGE,
     }
     runner.add_trial(Trial("__fake", **kwargs))
     trials = runner.get_trials()
@@ -155,7 +164,9 @@ def test_fail_fast(ray_start_4_cpus_2_gpus_extra, resource_manager_cls, fail_fas
     """
 
     runner = TuneController(
-        resource_manager_factory=lambda: resource_manager_cls(), fail_fast=fail_fast
+        resource_manager_factory=lambda: resource_manager_cls(),
+        fail_fast=fail_fast,
+        storage=STORAGE,
     )
     kwargs = {
         "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
@@ -165,6 +176,7 @@ def test_fail_fast(ray_start_4_cpus_2_gpus_extra, resource_manager_cls, fail_fas
             "mock_error": True,
             "persistent_error": True,
         },
+        "storage": STORAGE,
     }
     runner.add_trial(Trial("__fake", **kwargs))
     runner.add_trial(Trial("__fake", **kwargs))