[train] enable train + tune tests and examples (ray-project#39021)

Signed-off-by: Matthew Deng <matt@anyscale.com> Signed-off-by: Jim Thompson <jimthompson5802@gmail.com>
jimthompson5802 · Sep 12, 2023 · 4bd62dc · 4bd62dc
1 parent e09cdf0
commit 4bd62dc
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 25 deletions.
diff --git a/.buildkite/pipeline.ml.yml b/.buildkite/pipeline.ml.yml
@@ -47,16 +47,14 @@
       --test_tag_filters=-gpu_only,-gpu,-minimal,-tune,-doctest
       python/ray/train/...
 
-
-# Currently empty test suite
-#- label: ":steam_locomotive: :octopus: Train + Tune tests and examples"
-#  conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_TRAIN_AFFECTED"]
-#  instance_size: medium
-#  commands:
-#    - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
-#    - TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
-#    - ./ci/env/env_info.sh
-#    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=tune,-gpu_only,-ray_air,-gpu,-doctest,-new_storage python/ray/train/...
+- label: ":steam_locomotive: :octopus: Train + Tune tests and examples"
+  conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_TRAIN_AFFECTED"]
+  instance_size: medium
+  commands:
+    - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
+    - TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
+    - ./ci/env/env_info.sh
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=tune,-gpu_only,-ray_air,-gpu,-doctest python/ray/train/...
 
 
 - label: ":brain: RLlib: Benchmarks (Torch 2.x)"
@@ -361,19 +359,6 @@
 
 ##### STORAGE REFACTOR
 
-# TODO(krfricke): Add new test for this suite
-# - label: ":steam_locomotive: :octopus: :floppy_disk: New persistence mode: Train + Tune tests and examples"
-#   conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_TRAIN_AFFECTED"]
-#   instance_size: medium
-#   commands:
-#     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
-#     - TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
-#     - ./ci/env/env_info.sh
-#     - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only 
-#       --test_tag_filters=tune,-gpu_only,-ray_air,-gpu,-doctest,-no_new_storage
-#       --test_env=RAY_AIR_NEW_PERSISTENCE_MODE=1
-#       python/ray/train/...
-
 
 - label: ":octopus: :floppy_disk: New persistence mode: Tune tests and examples (small)"
   conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_TUNE_AFFECTED"]

diff --git a/python/ray/train/__init__.py b/python/ray/train/__init__.py
@@ -34,6 +34,7 @@
 
 usage_lib.record_library_usage("train")
 
+
 __all__ = [
     "get_checkpoint",
     "get_context",

diff --git a/python/ray/train/tests/test_tune.py b/python/ray/train/tests/test_tune.py
@@ -219,15 +219,20 @@ def train_func():
     assert len(df[TRAINING_ITERATION]) == 4
 
 
-def test_restore_with_new_trainer(ray_start_4_cpus, tmpdir, propagate_logs, caplog):
+def test_restore_with_new_trainer(
+    ray_start_4_cpus, tmpdir, propagate_logs, caplog, monkeypatch
+):
+
+    monkeypatch.setenv("RAY_AIR_LOCAL_CACHE_DIR", str(tmpdir))
+
     def train_func(config):
         raise RuntimeError("failing!")
 
     trainer = DataParallelTrainer(
         train_func,
         backend_config=TestConfig(),
         scaling_config=ScalingConfig(num_workers=1),
-        run_config=RunConfig(local_dir=str(tmpdir), name="restore_new_trainer"),
+        run_config=RunConfig(name="restore_new_trainer"),
         datasets={"train": ray.data.from_items([{"a": i} for i in range(10)])},
     )
     results = Tuner(trainer).fit()