[Train] Clean up ray.train package

.buildkite/pipeline.ml.yml

@@ -319,12 +319,13 @@
     - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only python/ray/util/horovod/...
     - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only python/ray/util/ray_lightning/...
 
-- label: ":octopus: Ludwig tests and examples. Python 3.7"
-  conditions: ["RAY_CI_TUNE_AFFECTED"]
-  commands:
-    - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
-    - PYTHON=3.7 INSTALL_LUDWIG=1 INSTALL_HOROVOD=1 ./ci/env/install-dependencies.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only python/ray/tests/ludwig/...
+# TODO(amogkam): Re-enable Ludwig tests after Ludwig supports Ray 2.0
+#- label: ":octopus: Ludwig tests and examples. Python 3.7"
+#  conditions: ["RAY_CI_TUNE_AFFECTED"]
+#  commands:
+#    - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
+#    - PYTHON=3.7 INSTALL_LUDWIG=1 INSTALL_HOROVOD=1 ./ci/env/install-dependencies.sh
+#    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only python/ray/tests/ludwig/...
 
 - label: ":tropical_fish: ML Libraries w/ Ray Client Examples (Python 3.7)."
   conditions: ["RAY_CI_TUNE_AFFECTED"]

ci/lint/check_api_annotations.py

@@ -88,8 +88,8 @@ def verify(symbol, scanned, ok, output, prefix=None):
 
     verify(ray.rllib, set(), ok, output)
     verify(ray.air, set(), ok, output)
+    verify(ray.train, set(), ok, output)
     # TODO(ekl) enable it for all modules.
-    #    verify(ray.train, set(), ok, output)
     #    verify(ray.serve, set(), ok, output)
     #    verify(ray.tune, set(), ok, output)
     #    verify(ray, set(), ok, output)

doc/source/train/api.rst

@@ -109,32 +109,6 @@ TorchTensorboardProfilerCallback
 
 .. autoclass:: ray.train.callbacks.TorchTensorboardProfilerCallback
 
-ResultsPreprocessors
-~~~~~~~~~~~~~~~~~~~~
-
-.. _train-api-results-preprocessor:
-
-ResultsPreprocessor
-+++++++++++++++++++
-
-.. autoclass:: ray.train.callbacks.results_preprocessors.ResultsPreprocessor
-    :members:
-
-SequentialResultsPreprocessor
-+++++++++++++++++++++++++++++++
-
-.. autoclass:: ray.train.callbacks.results_preprocessors.SequentialResultsPreprocessor
-
-IndexedResultsPreprocessor
-+++++++++++++++++++++++++++++++
-
-.. autoclass:: ray.train.callbacks.results_preprocessors.IndexedResultsPreprocessor
-
-ExcludedKeysResultsPreprocessor
-+++++++++++++++++++++++++++++++
-
-.. autoclass:: ray.train.callbacks.results_preprocessors.ExcludedKeysResultsPreprocessor
-
 Checkpointing
 -------------
 

python/ray/air/train/data_parallel_trainer.py

@@ -26,9 +26,9 @@
 from ray.air.checkpoint import Checkpoint
 from ray.air.train.data_parallel_ingest import _DataParallelIngestSpec
 from ray.train import BackendConfig, TrainingIterator
-from ray.train.backend import BackendExecutor
-from ray.train._checkpoint import TuneCheckpointManager
-from ray.train.utils import construct_train_func
+from ray.train._internal.backend_executor import BackendExecutor
+from ray.train._internal.checkpoint import TuneCheckpointManager
+from ray.train._internal.utils import construct_train_func
 from ray.util.annotations import DeveloperAPI
 from ray.util.ml_utils.checkpoint_manager import CheckpointStrategy, _TrackedCheckpoint
 

python/ray/train/__init__.py

@@ -1,6 +1,6 @@
 from ray.train.backend import BackendConfig
 from ray.train.callbacks import TrainingCallback
-from ray.train.session import (
+from ray.train.train_loop_utils import (
     get_dataset_shard,
     local_rank,
     load_checkpoint,