[train] Migrate/disable a few remaining tests for V2 (#57764)

There's a few more tests that haven't been migrated -- this PR migrates
them as preparation to flip the V2 switch.

---------

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

Author: justinvyu

ci/ray_ci/doc/cmd_check_api_discrepancy.py

@@ -63,6 +63,7 @@
             "ray.train.error.SessionMisuseError",
             "ray.train.TrainingFailedError",
             "ray.train.context.TrainContext",
+            "ray.train.context.get_context",
         },
     },
     "tune": {

doc/BUILD.bazel

@@ -416,6 +416,7 @@ py_test_run_all_subdirectory(
 
 py_test_run_all_subdirectory(
     size = "large",
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
     include = ["source/train/doc_code/*.py"],
     exclude = [
         "source/train/doc_code/hvd_trainer.py",  # CI do not have Horovod
@@ -485,17 +486,17 @@ py_test(
 # Test all doc/external code
 # --------------------------------------------------------------------
 
-py_test_run_all_subdirectory(
-    size = "enormous",
-    include = ["external/*.py"],
-    exclude = ["external/test_hashes.py"],
-    extra_srcs = [],
-    tags = [
-        "exclusive",
-        "external",
-        "team:ml",
-    ],
-)
+# py_test_run_all_subdirectory(
+#     size = "enormous",
+#     include = ["external/*.py"],
+#     exclude = ["external/test_hashes.py"],
+#     extra_srcs = [],
+#     tags = [
+#         "exclusive",
+#         "external",
+#         "team:ml",
+#     ],
+# )
 
 py_test(
     name = "test_external_hashes",
@@ -643,6 +644,8 @@ doctest(
 
 doctest(
     name = "doctest[train]",
+    # TODO: [V2] Migrate
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
     files = glob(
         include = [
             "source/train/**/*.md",
@@ -665,6 +668,8 @@ doctest(
         "source/train/user-guides/data-loading-preprocessing.rst",
         "source/train/user-guides/using-gpus.rst",
     ],
+    # TODO: [V2] Migrate
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
     tags = ["team:ml"],
     gpu = True,
 )

doc/source/rllib/getting-started.rst

@@ -210,7 +210,7 @@ one for each of the configured learning rates:
 
 .. testcode::
 
-    from ray import train, tune
+    from ray import tune
     from ray.rllib.algorithms.ppo import PPOConfig
 
     config = (
@@ -230,7 +230,7 @@ one for each of the configured learning rates:
         # pretty printed result metrics from the results returned previously by
         # ``.train()``. Also note that -1100 is not a good episode return for
         # Pendulum-v1, we are using it here to shorten the experiment time.
-        run_config=train.RunConfig(
+        run_config=tune.RunConfig(
             stop={"env_runners/episode_return_mean": -1100.0},
         ),
     )

doc/source/train/distributed-tensorflow-keras.rst

@@ -339,15 +339,6 @@ Load checkpoints
     result = trainer.fit()
     print(result.checkpoint)
 
-    # Start a new run from a loaded checkpoint
-    trainer = TensorflowTrainer(
-        train_func,
-        train_loop_config={"num_epochs": 5},
-        scaling_config=ScalingConfig(num_workers=2),
-        resume_from_checkpoint=result.checkpoint,
-    )
-    result = trainer.fit()
-
 
 Further reading
 ---------------

doc/source/train/doc_code/dl_guide.py

@@ -1,5 +1,10 @@
 # flake8: noqa
 
+# TODO: [V2] Deprecated doc code to delete.
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
+
 MOCK = True
 
 # __ft_initial_run_start__

doc/source/train/doc_code/hvd_trainer.py

@@ -1,4 +1,8 @@
+# TODO: [V2] Deprecated doc code to delete.
 import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
+
 import tempfile
 
 import horovod.torch as hvd

doc/source/train/doc_code/key_concepts.py

@@ -6,7 +6,7 @@
 
 from ray import train
 from ray.train import Checkpoint
-from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
 
 
 def train_fn(config):
@@ -34,8 +34,6 @@ def train_fn(config):
     # The experiment results will be saved to: storage_path/name
     storage_path=os.path.expanduser("~/ray_results"),
     # storage_path="s3://my_bucket/tune_results",
-    # Stopping criteria
-    stop={"training_iteration": 10},
 )
 # __run_config_end__
 

doc/source/train/doc_code/tuner.py

@@ -1,6 +1,11 @@
 # flake8: noqa
 # isort: skip_file
 
+# TODO: [V2] Deprecated doc code to delete.
+import os
+
+os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
+
 # __basic_start__
 import ray
 import ray.tune