[train] Port over remaining tests to Train v2 (ray-project#57722)

Ports over the remaining unit tests that were marked as TODOs from this
series of PRs: ray-project#57534, ray-project#57256, ray-project#56868, ray-project#56820, ray-project#56816.

Notably:
* `test_new_dataset_config -> test_data_integration`
* `test_backend -> test_torch_trainer, test_worker_group`
* `test_gpu -> test_torch_gpu`

This PR also finishes migrating the Tune LightGBM/Keras examples which
were unblocked by ray-project#57042 and
ray-project#57121.

---------

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

Author: justinvyu

doc/source/train/examples/pytorch/tune_cifar_torch_pbt_example.rst

@@ -20,8 +20,6 @@ py_test_run_all_notebooks(
     exclude = [
         "pbt_ppo_example.ipynb",
         "tune-xgboost.ipynb",
-        "lightgbm_example.ipynb",  # TODO: [V2] Remove after fixing Tune lightgbm callback.
-        "tune_mnist_keras.ipynb",  # TODO: [V2] Remove after fixing Tune keras callback.
         "pbt_transformers.ipynb",  # Transformers uses legacy Tune APIs.
         "tune-aim.ipynb",  # CI does not have `aim`
         "bohb_example.ipynb",  # CI does not have bohb requirements

doc/source/tune/examples/BUILD.bazel

@@ -44,7 +44,6 @@ Ray Tune integrates with many popular machine learning frameworks. Here you find
   * - :doc:`Tuning XGBoost parameters with Tune <tune-xgboost>`
   * - :doc:`Tuning LightGBM parameters with Tune <lightgbm_example>`
   * - :doc:`Tuning Hugging Face Transformers with Tune <pbt_transformers>`
-  * - :doc:`End-to-end example for tuning a PyTorch model with PBT <../../train/examples/pytorch/tune_cifar_torch_pbt_example>`  
 
 .. _experiment-tracking-tools:
 

doc/source/tune/examples/index.rst

@@ -36,7 +36,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "19e3c389",
    "metadata": {
     "tags": [
@@ -150,7 +150,7 @@
     "\n",
     "from ray import tune\n",
     "from ray.tune.schedulers import AsyncHyperBandScheduler\n",
-    "from ray.air.integrations.keras import ReportCheckpointCallback\n",
+    "from ray.tune.integration.keras import TuneReportCheckpointCallback\n",
     "\n",
     "\n",
     "def train_mnist(config):\n",
@@ -186,7 +186,7 @@
     "        epochs=epochs,\n",
     "        verbose=0,\n",
     "        validation_data=(x_test, y_test),\n",
-    "        callbacks=[ReportCheckpointCallback(metrics={\"accuracy\": \"accuracy\"})],\n",
+    "        callbacks=[TuneReportCheckpointCallback(metrics={\"accuracy\": \"accuracy\"})],\n",
     "    )\n",
     "\n",
     "\n",

doc/source/tune/examples/tune_mnist_keras.ipynb

@@ -43,7 +43,7 @@ py_test(
     name = "test_new_dataset_config",
     size = "large",
     srcs = ["tests/test_new_dataset_config.py"],
-    # TODO: Move over
+    # NOTE: Relevant tests moved to train/v2/tests/test_data_integration.py
     env = {"RAY_TRAIN_V2_ENABLED": "0"},
     tags = [
         "exclusive",

python/ray/air/BUILD.bazel

@@ -181,23 +181,6 @@ py_test(
     deps = [":train_lib"],
 )
 
-# TODO: Delete this together with the long running PBT release test.
-py_test(
-    name = "tune_cifar_torch_pbt_example",
-    size = "medium",
-    srcs = ["examples/pytorch/tune_cifar_torch_pbt_example.py"],
-    args = ["--smoke-test"],
-    env = {"RAY_TRAIN_V2_ENABLED": "0"},
-    main = "examples/pytorch/tune_cifar_torch_pbt_example.py",
-    tags = [
-        "exclusive",
-        "pytorch",
-        "team:ml",
-        "tune",
-    ],
-    deps = [":train_lib"],
-)
-
 # Formerly AIR examples
 
 py_test(
@@ -362,7 +345,9 @@ py_test(
     name = "test_backend",
     size = "large",
     srcs = ["tests/test_backend.py"],
-    # TODO: should port over
+    # NOTE: Relevant tests have been migrated to
+    # test_torch_trainer.py and test_worker_group.py
+    # TODO: [V2] There are still some accelerator integration tests left.
     env = {"RAY_TRAIN_V2_ENABLED": "0"},
     tags = [
         "exclusive",
@@ -466,7 +451,7 @@ py_test(
     name = "test_gpu",
     size = "large",
     srcs = ["tests/test_gpu.py"],
-    # TODO: move over, rename to test_torch_gpu_utils
+    # NOTE: Migrated relevant tests to v2/tests/test_torch_gpu.py
     env = {"RAY_TRAIN_V2_ENABLED": "0"},
     tags = [
         "exclusive",
@@ -602,13 +587,12 @@ py_test(
     name = "test_torch_lightning_train",
     size = "large",
     srcs = ["tests/test_torch_lightning_train.py"],
-    # TODO: move over
-    env = {"RAY_TRAIN_V2_ENABLED": "0"},
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
     tags = [
         "exclusive",
-        "gpu",
         "ptl_v2",
         "team:ml",
+        "train_v2_gpu",
     ],
     deps = [":train_lib"],
 )
@@ -683,7 +667,8 @@ py_test(
     name = "test_minimal",
     size = "small",
     srcs = ["tests/test_minimal.py"],
-    # TODO: move over
+    # TODO: [V2] The minimal test needs to install pydantic,
+    # which is a Train V2 dependency.
     env = {"RAY_TRAIN_V2_ENABLED": "0"},
     tags = [
         "exclusive",