[air] Move python/ray/ml to python/ray/air (#25449)

The package "ml" should be renamed to "air".

Main question: Keep a `ml.py` with `from ray.air import *` for some level of backwards compatibility?
I'd go for no to force people to use the new structure.

.buildkite/pipeline.gpu.large.yml

@@ -32,13 +32,13 @@
       --test_tag_filters=multi_gpu --test_env=RAY_USE_MULTIPROCESSING_CPU_COUNT=1 rllib/...
 
 
-- label: ":tv: :airplane: ML GPU tests (ray/ml)"
+- label: ":tv: :airplane: ML GPU tests (ray/air)"
   conditions: ["RAY_CI_ML_AFFECTED"]
   commands:
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - DATA_PROCESSING_TESTING=1 TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
     - pip install -Ur ./python/requirements_ml_docker.txt
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=gpu python/ray/ml/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=gpu python/ray/air/...
 
 
 - label: ":tv: :book: Doc GPU tests and examples"

.buildkite/pipeline.ml.yml

@@ -1,9 +1,9 @@
-- label: ":airplane: ML tests (ray/ml)"
+- label: ":airplane: ML tests (ray/air)"
   conditions: ["RAY_CI_ML_AFFECTED"]
   commands:
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - DATA_PROCESSING_TESTING=1 INSTALL_HOROVOD=1 ./ci/env/install-dependencies.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=-gpu,-needs_credentials python/ray/ml/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=-gpu,-needs_credentials python/ray/air/...
 
 - label: ":brain: RLlib: Learning discr. actions TF2-static-graph"
   conditions: ["RAY_CI_RLLIB_AFFECTED"]

.github/dependabot.yml

@@ -4,7 +4,7 @@ updates:
   - package-ecosystem: "pip"
     # If we want to add more requirements here (Core, Serve, etc.), then we should
     # make additional subdirectories for each one.
-    directory: "/python/requirements/ml"
+    directory: "/python/requirements/air"
     schedule:
       # Automatic upgrade checks Saturday at 12 AM.
       # Dependabot updates can still be manually triggered via Github at any time.

ci/env/cleanup_test_state.py

@@ -9,7 +9,7 @@ import sys
 def clear_wandb_project():
     import wandb
 
-    # This is hardcoded in the `ray/ml/examples/upload_to_wandb.py` example
+    # This is hardcoded in the `ray/air/examples/upload_to_wandb.py` example
     wandb_project = "ray_air_example"
 
     api = wandb.Api()
@@ -20,7 +20,7 @@ def clear_wandb_project():
 def clear_comet_ml_project():
     import comet_ml
 
-    # This is hardcoded in the `ray/ml/examples/upload_to_comet_ml.py` example
+    # This is hardcoded in the `ray/air/examples/upload_to_comet_ml.py` example
     comet_ml_project = "ray-air-example"
 
     api = comet_ml.API()

ci/lint/check_api_annotations.py

@@ -88,7 +88,7 @@ if __name__ == "__main__":
 
     verify(ray.rllib, set(), ok, output)
     # TODO(ekl) enable it for all modules.
-    #    verify(ray.ml, set(), ok, output)
+    #    verify(ray.air, set(), ok, output)
     #    verify(ray.train, set(), ok, output)
     #    verify(ray.serve, set(), ok, output)
     #    verify(ray.tune, set(), ok, output)

ci/pipeline/determine_tests_to_run.py

@@ -125,7 +125,7 @@ if __name__ == "__main__":
         skip_prefix_list = ["doc/", "examples/", "dev/", "kubernetes/", "site/"]
 
         for changed_file in files:
-            if changed_file.startswith("python/ray/ml"):
+            if changed_file.startswith("python/ray/air"):
                 RAY_CI_ML_AFFECTED = 1
                 RAY_CI_TRAIN_AFFECTED = 1
                 RAY_CI_TUNE_AFFECTED = 1

doc/source/ray-air/check-ingest.rst

@@ -110,7 +110,7 @@ Debugging Ingest with the ``DummyTrainer``
 ------------------------------------------
 
 Data ingest problems can be challenging to debug when combined in a full training pipeline. To isolate data
-ingest issues from other possible training problems, we provide the ``ray.ml.util.check_ingest.DummyTrainer``
+ingest issues from other possible training problems, we provide the ``ray.air.util.check_ingest.DummyTrainer``
 utility class that can be used to debug ingest problems. Let's walk through using DummyTrainer to understand
 and resolve an ingest misconfiguration.
 

doc/source/ray-air/doc_code/air_ingest.py

@@ -2,7 +2,7 @@
 
 # __shared_dataset_start__
 import ray
-from ray.ml.utils.check_ingest import DummyTrainer
+from ray.air.utils.check_ingest import DummyTrainer
 from ray.tune.tuner import Tuner, TuneConfig
 
 ray.init(num_cpus=5)
@@ -26,7 +26,7 @@ ray.shutdown()
 # __indep_dataset_start__
 import ray
 from ray import tune
-from ray.ml.utils.check_ingest import DummyTrainer
+from ray.air.utils.check_ingest import DummyTrainer
 from ray.tune.tuner import Tuner, TuneConfig
 
 ray.init(num_cpus=5)
@@ -60,8 +60,8 @@ tuner.fit()
 
 # __check_ingest_1__
 import ray
-from ray.ml.preprocessors import Chain, BatchMapper
-from ray.ml.utils.check_ingest import DummyTrainer
+from ray.air.preprocessors import Chain, BatchMapper
+from ray.air.utils.check_ingest import DummyTrainer
 
 # Generate a synthetic dataset of ~10GiB of float64 data. The dataset is sharded
 # into 100 blocks (parallelism=100).

doc/source/ray-air/doc_code/air_key_concepts.py

@@ -6,7 +6,7 @@ import pandas as pd
 from sklearn.datasets import load_breast_cancer
 from sklearn.model_selection import train_test_split
 
-from ray.ml.preprocessors import *
+from ray.air.preprocessors import *
 
 data_raw = load_breast_cancer()
 dataset_df = pd.DataFrame(data_raw["data"], columns=data_raw["feature_names"])
@@ -21,7 +21,7 @@ preprocessor = StandardScaler(columns=columns_to_scale)
 # __air_preprocessors_end__
 
 # __air_trainer_start__
-from ray.ml.train.integrations.xgboost import XGBoostTrainer
+from ray.air.train.integrations.xgboost import XGBoostTrainer
 
 num_workers = 2
 use_gpu = False
@@ -68,8 +68,8 @@ print(best_result)
 # __air_tuner_end__
 
 # __air_batch_predictor_start__
-from ray.ml.batch_predictor import BatchPredictor
-from ray.ml.predictors.integrations.xgboost import XGBoostPredictor
+from ray.air.batch_predictor import BatchPredictor
+from ray.air.predictors.integrations.xgboost import XGBoostPredictor
 
 batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, XGBoostPredictor)
 

doc/source/ray-air/doc_code/preprocessors.py

@@ -4,7 +4,7 @@
 # __preprocessor_setup_start__
 import pandas as pd
 import ray
-from ray.ml.preprocessors import MinMaxScaler
+from ray.air.preprocessors import MinMaxScaler
 
 # Generate two simple datasets.
 dataset = ray.data.range_table(8)
@@ -47,8 +47,8 @@ print(batch_transformed)
 # __trainer_start__
 import ray
 
-from ray.ml.train.integrations.xgboost import XGBoostTrainer
-from ray.ml.preprocessors import MinMaxScaler
+from ray.air.train.integrations.xgboost import XGBoostTrainer
+from ray.air.preprocessors import MinMaxScaler
 
 train_dataset = ray.data.from_items([{"x": x, "y": 2 * x} for x in range(0, 32, 3)])
 valid_dataset = ray.data.from_items([{"x": x, "y": 2 * x} for x in range(1, 32, 3)])
@@ -67,7 +67,7 @@ result = trainer.fit()
 
 
 # __checkpoint_start__
-from ray.ml.utils.checkpointing import load_preprocessor_from_dir
+from ray.air.utils.checkpointing import load_preprocessor_from_dir
 
 checkpoint = result.checkpoint
 with checkpoint.as_directory() as checkpoint_path:
@@ -78,8 +78,8 @@ with checkpoint.as_directory() as checkpoint_path:
 
 
 # __predictor_start__
-from ray.ml.batch_predictor import BatchPredictor
-from ray.ml.predictors.integrations.xgboost import XGBoostPredictor
+from ray.air.batch_predictor import BatchPredictor
+from ray.air.predictors.integrations.xgboost import XGBoostPredictor
 
 test_dataset = ray.data.from_items([{"x": x} for x in range(2, 32, 3)])
 
@@ -102,7 +102,7 @@ print(predicted_labels.to_pandas())
 
 # __chain_start__
 import ray
-from ray.ml.preprocessors import Chain, MinMaxScaler, SimpleImputer
+from ray.air.preprocessors import Chain, MinMaxScaler, SimpleImputer
 
 # Generate one simple dataset.
 dataset = ray.data.from_items(
@@ -121,7 +121,7 @@ print(dataset_transformed.take())
 
 # __custom_stateless_start__
 import ray
-from ray.ml.preprocessors import BatchMapper
+from ray.air.preprocessors import BatchMapper
 
 # Generate a simple dataset.
 dataset = ray.data.range_table(4)

doc/source/ray-air/doc_code/pytorch_starter.py

@@ -28,7 +28,7 @@ import torch
 from torch import nn
 from torch.utils.data import DataLoader
 import ray.train as train
-from ray.ml.train.integrations.torch import TorchTrainer
+from ray.air.train.integrations.torch import TorchTrainer
 
 # Define model
 class NeuralNetwork(nn.Module):
@@ -130,8 +130,8 @@ print(f"Last result: {result.metrics}")
 
 # # __air_pytorch_batchpred_start__
 # import random
-# from ray.ml.batch_predictor import BatchPredictor
-# from ray.ml.predictors.integrations.torch import TorchPredictor
+# from ray.air.batch_predictor import BatchPredictor
+# from ray.air.predictors.integrations.torch import TorchPredictor
 
 # batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, TorchPredictor)
 

doc/source/ray-air/doc_code/tf_starter.py

@@ -18,7 +18,7 @@ from tensorflow.keras.callbacks import Callback
 
 import ray.train as train
 from ray.train.tensorflow import prepare_dataset_shard
-from ray.ml.train.integrations.tensorflow import TensorflowTrainer
+from ray.air.train.integrations.tensorflow import TensorflowTrainer
 
 
 def build_model() -> tf.keras.Model:
@@ -91,8 +91,8 @@ print(result.metrics)
 # __air_tf_batchpred_start__
 import numpy as np
 
-from ray.ml.batch_predictor import BatchPredictor
-from ray.ml.predictors.integrations.tensorflow import TensorflowPredictor
+from ray.air.batch_predictor import BatchPredictor
+from ray.air.predictors.integrations.tensorflow import TensorflowPredictor
 
 
 batch_predictor = BatchPredictor.from_checkpoint(

doc/source/ray-air/doc_code/xgboost_starter.py

@@ -2,7 +2,7 @@
 
 # __air_xgb_preprocess_start__
 import ray
-from ray.ml.preprocessors import StandardScaler
+from ray.air.preprocessors import StandardScaler
 
 import pandas as pd
 
@@ -25,7 +25,7 @@ preprocessor = StandardScaler(columns=columns_to_scale)
 
 
 # __air_xgb_train_start__
-from ray.ml.train.integrations.xgboost import XGBoostTrainer
+from ray.air.train.integrations.xgboost import XGBoostTrainer
 
 # XGBoost specific params
 params = {
@@ -53,8 +53,8 @@ print(result.metrics)
 # __air_xgb_train_end__
 
 # __air_xgb_batchpred_start__
-from ray.ml.batch_predictor import BatchPredictor
-from ray.ml.predictors.integrations.xgboost import XGBoostPredictor
+from ray.air.batch_predictor import BatchPredictor
+from ray.air.predictors.integrations.xgboost import XGBoostPredictor
 
 batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, XGBoostPredictor)
 

doc/source/ray-air/examples/analyze_tuning_results.ipynb

@@ -71,8 +71,8 @@
     "\n",
     "import ray\n",
     "from ray import tune\n",
-    "from ray.ml import RunConfig\n",
-    "from ray.ml.train.integrations.xgboost import XGBoostTrainer\n",
+    "from ray.air import RunConfig\n",
+    "from ray.air.train.integrations.xgboost import XGBoostTrainer\n",
     "from ray.tune.tune_config import TuneConfig\n",
     "from ray.tune.tuner import Tuner"
    ]

doc/source/ray-air/examples/huggingface_text_classification.ipynb

@@ -797,8 +797,8 @@
       },
       "outputs": [],
       "source": [
-        "from ray.ml.train.integrations.huggingface import HuggingFaceTrainer\n",
-        "from ray.ml import RunConfig\n",
+        "from ray.air.train.integrations.huggingface import HuggingFaceTrainer\n",
+        "from ray.air import RunConfig\n",
         "from ray.tune.integration.mlflow import MLflowLoggerCallback\n",
         "\n",
         "trainer = HuggingFaceTrainer(\n",
@@ -1171,7 +1171,7 @@
         {
           "data": {
             "text/plain": [
-              "Result(metrics={'loss': 0.1575, 'learning_rate': 1.308411214953271e-06, 'epoch': 5.0, 'step': 2675, 'eval_loss': 0.8616615533828735, 'eval_matthews_correlation': 0.5420036503219092, 'eval_runtime': 1.2577, 'eval_samples_per_second': 829.302, 'eval_steps_per_second': 52.477, 'train_runtime': 187.8585, 'train_samples_per_second': 227.592, 'train_steps_per_second': 14.239, 'train_loss': 0.30010223103460865, '_timestamp': 1652380513, '_time_this_iter_s': 39.63672137260437, '_training_iteration': 5, 'time_this_iter_s': 39.64510202407837, 'should_checkpoint': True, 'done': True, 'timesteps_total': None, 'episodes_total': None, 'training_iteration': 5, 'trial_id': 'bb9dd_00000', 'experiment_id': 'db0c5ea784a44980819bf5e1bfb72c04', 'date': '2022-05-12_18-35-13', 'timestamp': 1652380513, 'time_total_s': 222.39091277122498, 'pid': 419, 'hostname': 'e618da00601e', 'node_ip': '172.28.0.2', 'config': {}, 'time_since_restore': 222.39091277122498, 'timesteps_since_restore': 0, 'iterations_since_restore': 5, 'warmup_time': 0.004034996032714844, 'experiment_tag': '0'}, checkpoint=<ray.ml.checkpoint.Checkpoint object at 0x7f9ffd9d9c90>, error=None)"
+              "Result(metrics={'loss': 0.1575, 'learning_rate': 1.308411214953271e-06, 'epoch': 5.0, 'step': 2675, 'eval_loss': 0.8616615533828735, 'eval_matthews_correlation': 0.5420036503219092, 'eval_runtime': 1.2577, 'eval_samples_per_second': 829.302, 'eval_steps_per_second': 52.477, 'train_runtime': 187.8585, 'train_samples_per_second': 227.592, 'train_steps_per_second': 14.239, 'train_loss': 0.30010223103460865, '_timestamp': 1652380513, '_time_this_iter_s': 39.63672137260437, '_training_iteration': 5, 'time_this_iter_s': 39.64510202407837, 'should_checkpoint': True, 'done': True, 'timesteps_total': None, 'episodes_total': None, 'training_iteration': 5, 'trial_id': 'bb9dd_00000', 'experiment_id': 'db0c5ea784a44980819bf5e1bfb72c04', 'date': '2022-05-12_18-35-13', 'timestamp': 1652380513, 'time_total_s': 222.39091277122498, 'pid': 419, 'hostname': 'e618da00601e', 'node_ip': '172.28.0.2', 'config': {}, 'time_since_restore': 222.39091277122498, 'timesteps_since_restore': 0, 'iterations_since_restore': 5, 'warmup_time': 0.004034996032714844, 'experiment_tag': '0'}, checkpoint=<ray.air.checkpoint.Checkpoint object at 0x7f9ffd9d9c90>, error=None)"
             ]
           },
           "execution_count": 17,
@@ -1368,8 +1368,8 @@
         }
       ],
       "source": [
-        "from ray.ml.predictors.integrations.huggingface import HuggingFacePredictor\n",
-        "from ray.ml.batch_predictor import BatchPredictor\n",
+        "from ray.air.predictors.integrations.huggingface import HuggingFacePredictor\n",
+        "from ray.air.batch_predictor import BatchPredictor\n",
         "import pandas as pd\n",
         "\n",
         "sentences = ['Bill whistled past the house.',\n",
@@ -1458,7 +1458,7 @@
       },
       "outputs": [],
       "source": [
-        "from ray.ml.train.integrations.huggingface import load_checkpoint\n",
+        "from ray.air.train.integrations.huggingface import load_checkpoint\n",
         "\n",
         "hf_trainer = load_checkpoint(\n",
         "    checkpoint=result.checkpoint,\n",

doc/source/ray-air/examples/lightgbm_example.ipynb

@@ -49,14 +49,14 @@
     "import pandas as pd\n",
     "\n",
     "import ray\n",
-    "from ray.ml.batch_predictor import BatchPredictor\n",
-    "from ray.ml.predictors.integrations.lightgbm import LightGBMPredictor\n",
-    "from ray.ml.preprocessors.chain import Chain\n",
-    "from ray.ml.preprocessors.encoder import Categorizer\n",
-    "from ray.ml.train.integrations.lightgbm import LightGBMTrainer\n",
+    "from ray.air.batch_predictor import BatchPredictor\n",
+    "from ray.air.predictors.integrations.lightgbm import LightGBMPredictor\n",
+    "from ray.air.preprocessors.chain import Chain\n",
+    "from ray.air.preprocessors.encoder import Categorizer\n",
+    "from ray.air.train.integrations.lightgbm import LightGBMTrainer\n",
     "from ray.data.dataset import Dataset\n",
-    "from ray.ml.result import Result\n",
-    "from ray.ml.preprocessors import StandardScaler\n",
+    "from ray.air.result import Result\n",
+    "from ray.air.preprocessors import StandardScaler\n",
     "from sklearn.datasets import load_breast_cancer\n",
     "from sklearn.model_selection import train_test_split"
    ]

doc/source/ray-air/examples/rl_offline_example.ipynb

@@ -59,11 +59,11 @@
     "\n",
     "import numpy as np\n",
     "import ray\n",
-    "from ray.ml import Checkpoint\n",
-    "from ray.ml.config import RunConfig\n",
-    "from ray.ml.predictors.integrations.rl.rl_predictor import RLPredictor\n",
-    "from ray.ml.train.integrations.rl.rl_trainer import RLTrainer\n",
-    "from ray.ml.result import Result\n",
+    "from ray.air import Checkpoint\n",
+    "from ray.air.config import RunConfig\n",
+    "from ray.air.predictors.integrations.rl.rl_predictor import RLPredictor\n",
+    "from ray.air.train.integrations.rl.rl_trainer import RLTrainer\n",
+    "from ray.air.result import Result\n",
     "from ray.rllib.agents.marwil import BCTrainer\n",
     "from ray.tune.tuner import Tuner"
    ]

doc/source/ray-air/examples/rl_online_example.ipynb

@@ -59,11 +59,11 @@
     "\n",
     "import numpy as np\n",
     "import ray\n",
-    "from ray.ml import Checkpoint\n",
-    "from ray.ml.config import RunConfig\n",
-    "from ray.ml.predictors.integrations.rl.rl_predictor import RLPredictor\n",
-    "from ray.ml.train.integrations.rl.rl_trainer import RLTrainer\n",
-    "from ray.ml.result import Result\n",
+    "from ray.air import Checkpoint\n",
+    "from ray.air.config import RunConfig\n",
+    "from ray.air.predictors.integrations.rl.rl_predictor import RLPredictor\n",
+    "from ray.air.train.integrations.rl.rl_trainer import RLTrainer\n",
+    "from ray.air.result import Result\n",
     "from ray.rllib.agents.marwil import BCTrainer\n",
     "from ray.tune.tuner import Tuner"
    ]

doc/source/ray-air/examples/rl_serving_example.ipynb

@@ -50,11 +50,11 @@
     "import numpy as np\n",
     "import requests\n",
     "\n",
-    "from ray.ml.checkpoint import Checkpoint\n",
-    "from ray.ml.config import RunConfig\n",
-    "from ray.ml.train.integrations.rl.rl_trainer import RLTrainer\n",
-    "from ray.ml.predictors.integrations.rl.rl_predictor import RLPredictor\n",
-    "from ray.ml.result import Result\n",
+    "from ray.air.checkpoint import Checkpoint\n",
+    "from ray.air.config import RunConfig\n",
+    "from ray.air.train.integrations.rl.rl_trainer import RLTrainer\n",
+    "from ray.air.predictors.integrations.rl.rl_predictor import RLPredictor\n",
+    "from ray.air.result import Result\n",
     "from ray.serve.model_wrappers import ModelWrapperDeployment\n",
     "from ray import serve\n",
     "from ray.tune.tuner import Tuner"

doc/source/ray-air/examples/sklearn_example.ipynb

@@ -54,11 +54,11 @@
     "\n",
     "import ray\n",
     "from ray.data.dataset import Dataset\n",
-    "from ray.ml.batch_predictor import BatchPredictor\n",
-    "from ray.ml.predictors.integrations.sklearn import SklearnPredictor\n",
-    "from ray.ml.preprocessors import Chain, OrdinalEncoder, StandardScaler\n",
-    "from ray.ml.result import Result\n",
-    "from ray.ml.train.integrations.sklearn import SklearnTrainer\n",
+    "from ray.air.batch_predictor import BatchPredictor\n",
+    "from ray.air.predictors.integrations.sklearn import SklearnPredictor\n",
+    "from ray.air.preprocessors import Chain, OrdinalEncoder, StandardScaler\n",
+    "from ray.air.result import Result\n",
+    "from ray.air.train.integrations.sklearn import SklearnTrainer\n",
     "\n",
     "\n",
     "from sklearn.datasets import load_breast_cancer\n",

doc/source/ray-air/examples/tfx_tabular_train_to_serve.ipynb

@@ -470,7 +470,7 @@
    },
    "outputs": [],
    "source": [
-    "from ray.ml.preprocessors import (\n",
+    "from ray.air.preprocessors import (\n",
     "    BatchMapper,\n",
     "    Chain,\n",
     "    OneHotEncoder,\n",
@@ -746,7 +746,7 @@
     }
    ],
    "source": [
-    "from ray.ml.train.integrations.tensorflow import TensorflowTrainer\n",
+    "from ray.air.train.integrations.tensorflow import TensorflowTrainer\n",
     "\n",
     "trainer = TensorflowTrainer(\n",
     "    train_loop_per_worker=train_loop_per_worker,\n",
@@ -774,7 +774,7 @@
    "source": [
     "We will use Ray Serve to serve the trained model. A core concept of Ray Serve is [Deployment](https://docs.ray.io/en/latest/serve/core-apis.html). It allows you to define and update your business logic or models that will handle incoming requests as well as how this is exposed over HTTP or in Python.\n",
     "\n",
-    "In the case of serving model, `ray.serve.model_wrappers.ModelWrapper` and `ray.serve.model_wrappers.ModelWrapperDeployment` wrap a `ray.ml.checkpoint.Checkpoint` into a Ray Serve deployment that can readily serve HTTP requests.\n",
+    "In the case of serving model, `ray.serve.model_wrappers.ModelWrapper` and `ray.serve.model_wrappers.ModelWrapperDeployment` wrap a `ray.air.checkpoint.Checkpoint` into a Ray Serve deployment that can readily serve HTTP requests.\n",
     "Note, ``Checkpoint`` captures both model and preprocessing steps in a way compatible with Ray Serve and ensures that ml workload can transition seamlessly between training and\n",
     "serving.\n",
     "\n",
@@ -818,8 +818,8 @@
    "outputs": [],
    "source": [
     "from ray import serve\n",
-    "from ray.ml.checkpoint import Checkpoint\n",
-    "from ray.ml.predictors.integrations.tensorflow import TensorflowPredictor\n",
+    "from ray.air.checkpoint import Checkpoint\n",
+    "from ray.air.predictors.integrations.tensorflow import TensorflowPredictor\n",
     "from ray.serve.model_wrappers import ModelWrapperDeployment\n",
     "\n",
     "\n",

doc/source/ray-air/examples/torch_image_example.ipynb

@@ -410,7 +410,7 @@
     }
    ],
    "source": [
-    "from ray.ml.train.integrations.torch import TorchTrainer\n",
+    "from ray.air.train.integrations.torch import TorchTrainer\n",
     "\n",
     "trainer = TorchTrainer(\n",
     "    train_loop_per_worker=train_loop_per_worker,\n",
@@ -437,9 +437,9 @@
     "\n",
     "Let's see how our model performs.\n",
     "\n",
-    "To classify images in the test dataset, we'll need to create a {py:class}`Predictor <ray.ml.predictor.Predictor>`.\n",
+    "To classify images in the test dataset, we'll need to create a {py:class}`Predictor <ray.air.predictor.Predictor>`.\n",
     "\n",
-    "{py:class}`Predictors <ray.ml.predictor.Predictor>` load data from checkpoints and efficiently perform inference. In contrast to {py:class}`TorchPredictor <ray.ml.predictors.integrations.torch.TorchPredictor>`, which performs inference on a single batch, {py:class}`BatchPredictor <ray.ml.batch_predictor.BatchPredictor>` performs inference on an entire dataset. Because we want to classify all of the images in the test dataset, we'll use a {py:class}`BatchPredictor <ray.ml.batch_predictor.BatchPredictor>`."
+    "{py:class}`Predictors <ray.air.predictor.Predictor>` load data from checkpoints and efficiently perform inference. In contrast to {py:class}`TorchPredictor <ray.air.predictors.integrations.torch.TorchPredictor>`, which performs inference on a single batch, {py:class}`BatchPredictor <ray.air.batch_predictor.BatchPredictor>` performs inference on an entire dataset. Because we want to classify all of the images in the test dataset, we'll use a {py:class}`BatchPredictor <ray.air.batch_predictor.BatchPredictor>`."
    ]
   },
   {
@@ -458,8 +458,8 @@
     }
    ],
    "source": [
-    "from ray.ml.predictors.integrations.torch import TorchPredictor\n",
-    "from ray.ml.batch_predictor import BatchPredictor\n",
+    "from ray.air.predictors.integrations.torch import TorchPredictor\n",
+    "from ray.air.batch_predictor import BatchPredictor\n",
     "\n",
     "batch_predictor = BatchPredictor.from_checkpoint(\n",
     "    checkpoint=latest_checkpoint,\n",

doc/source/ray-air/examples/torch_incremental_learning.ipynb

@@ -590,7 +590,7 @@
    },
    "outputs": [],
    "source": [
-    "from ray.ml.preprocessors import BatchMapper\n",
+    "from ray.air.preprocessors import BatchMapper\n",
     "\n",
     "from torchvision import transforms\n",
     "\n",
@@ -617,7 +617,7 @@
     "\n",
     "After training on each task, we want to use our trained model to do batch (i.e. offline) inference on a test dataset. \n",
     "\n",
-    "To do this, we leverage the built-in `ray.ml.BatchPredictor`. We define a `batch_predict` function that will take in a Checkpoint and a Test Dataset and outputs the accuracy our model achieves on the test dataset."
+    "To do this, we leverage the built-in `ray.air.BatchPredictor`. We define a `batch_predict` function that will take in a Checkpoint and a Test Dataset and outputs the accuracy our model achieves on the test dataset."
    ]
   },
   {
@@ -628,10 +628,10 @@
    },
    "outputs": [],
    "source": [
-    "from ray.ml.batch_predictor import BatchPredictor\n",
-    "from ray.ml.predictors.integrations.torch import TorchPredictor\n",
+    "from ray.air.batch_predictor import BatchPredictor\n",
+    "from ray.air.predictors.integrations.torch import TorchPredictor\n",
     "\n",
-    "def batch_predict(checkpoint: ray.ml.Checkpoint, test_dataset: ray.data.Dataset) -> float:\n",
+    "def batch_predict(checkpoint: ray.air.Checkpoint, test_dataset: ray.data.Dataset) -> float:\n",
     "  \"\"\"Perform batch prediction on the provided test dataset, and return accuracy results.\"\"\"\n",
     "\n",
     "  batch_predictor = BatchPredictor.from_checkpoint(checkpoint, predictor_cls=TorchPredictor, model=SimpleMLP(num_classes=10))\n",
@@ -691,7 +691,7 @@
     "from ray.serve.http_adapters import NdArray\n",
     "\n",
     "\n",
-    "def deploy_model(checkpoint: ray.ml.Checkpoint) -> str:\n",
+    "def deploy_model(checkpoint: ray.air.Checkpoint) -> str:\n",
     "  \"\"\"Deploys the model from the provided Checkpoint and returns the URL for the endpoint of the model deployment.\"\"\"\n",
     "  def json_to_pandas(payload: NdArray) -> pd.DataFrame:\n",
     "      \"\"\"Accepts an NdArray JSON from an HTTP body and converts it to a Pandas dataframe.\"\"\"\n",
@@ -1235,9 +1235,9 @@
     }
    ],
    "source": [
-    "from ray.ml.train.integrations.torch import TorchTrainer\n",
-    "from ray.ml.predictors.integrations.torch import TorchPredictor\n",
-    "from ray.ml import Checkpoint\n",
+    "from ray.air.train.integrations.torch import TorchTrainer\n",
+    "from ray.air.predictors.integrations.torch import TorchPredictor\n",
+    "from ray.air import Checkpoint\n",
     "from ray import serve\n",
     "from ray.serve.model_wrappers import ModelWrapperDeployment\n",
     "from ray.serve.http_adapters import json_to_ndarray\n",

doc/source/ray-air/examples/upload_to_comet_ml.ipynb

@@ -46,9 +46,9 @@
    "source": [
     "import ray\n",
     "\n",
-    "from ray.ml import RunConfig\n",
-    "from ray.ml.result import Result\n",
-    "from ray.ml.train.integrations.xgboost import XGBoostTrainer\n",
+    "from ray.air import RunConfig\n",
+    "from ray.air.result import Result\n",
+    "from ray.air.train.integrations.xgboost import XGBoostTrainer\n",
     "from ray.tune.integration.comet import CometLoggerCallback\n",
     "from sklearn.datasets import load_breast_cancer"
    ]

doc/source/ray-air/examples/upload_to_wandb.ipynb

@@ -46,9 +46,9 @@
    "source": [
     "import ray\n",
     "\n",
-    "from ray.ml import RunConfig\n",
-    "from ray.ml.result import Result\n",
-    "from ray.ml.train.integrations.xgboost import XGBoostTrainer\n",
+    "from ray.air import RunConfig\n",
+    "from ray.air.result import Result\n",
+    "from ray.air.train.integrations.xgboost import XGBoostTrainer\n",
     "from ray.tune.integration.wandb import WandbLoggerCallback\n",
     "from sklearn.datasets import load_breast_cancer"
    ]

doc/source/ray-air/examples/xgboost_example.ipynb

@@ -64,12 +64,12 @@
     "import pandas as pd\n",
     "\n",
     "import ray\n",
-    "from ray.ml.batch_predictor import BatchPredictor\n",
-    "from ray.ml.predictors.integrations.xgboost import XGBoostPredictor\n",
-    "from ray.ml.train.integrations.xgboost import XGBoostTrainer\n",
+    "from ray.air.batch_predictor import BatchPredictor\n",
+    "from ray.air.predictors.integrations.xgboost import XGBoostPredictor\n",
+    "from ray.air.train.integrations.xgboost import XGBoostTrainer\n",
     "from ray.data.dataset import Dataset\n",
-    "from ray.ml.result import Result\n",
-    "from ray.ml.preprocessors import StandardScaler\n",
+    "from ray.air.result import Result\n",
+    "from ray.air.preprocessors import StandardScaler\n",
     "from sklearn.datasets import load_breast_cancer\n",
     "from sklearn.model_selection import train_test_split"
    ]

doc/source/ray-air/getting-started.rst

@@ -126,7 +126,7 @@ Next Steps
 ----------
 
 - :ref:`air-key-concepts`
-- `Examples <https://github.com/ray-project/ray/tree/master/python/ray/ml/examples>`__
+- `Examples <https://github.com/ray-project/ray/tree/master/python/ray/air/examples>`__
 - :ref:`Deployment Guide <air-deployment>`
 - :ref:`API reference <air-api-ref>`
 

doc/source/ray-air/package-ref.rst

@@ -12,14 +12,14 @@ Components
 Preprocessors
 ~~~~~~~~~~~~~
 
-.. autoclass:: ray.ml.preprocessor.Preprocessor
+.. autoclass:: ray.air.preprocessor.Preprocessor
     :members:
 
-.. automodule:: ray.ml.preprocessors
+.. automodule:: ray.air.preprocessors
     :members:
     :show-inheritance:
 
-.. autofunction:: ray.ml.train_test_split
+.. autofunction:: ray.air.train_test_split
 
 
 .. _air-trainer-ref:
@@ -27,38 +27,38 @@ Preprocessors
 Trainer
 ~~~~~~~
 
-.. autoclass:: ray.ml.trainer.Trainer
+.. autoclass:: ray.air.trainer.Trainer
     :members:
 
-.. automodule:: ray.ml.train.integrations.xgboost
+.. automodule:: ray.air.train.integrations.xgboost
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.train.integrations.lightgbm
+.. automodule:: ray.air.train.integrations.lightgbm
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.train.integrations.tensorflow
+.. automodule:: ray.air.train.integrations.tensorflow
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.train.integrations.torch
+.. automodule:: ray.air.train.integrations.torch
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.train.integrations.huggingface
+.. automodule:: ray.air.train.integrations.huggingface
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.train.integrations.sklearn
+.. automodule:: ray.air.train.integrations.sklearn
     :members:
     :show-inheritance:
 
-.. autoclass:: ray.ml.train.data_parallel_trainer.DataParallelTrainer
+.. autoclass:: ray.air.train.data_parallel_trainer.DataParallelTrainer
     :members:
     :show-inheritance:
 
-.. autoclass:: ray.ml.train.gbdt_trainer.GBDTTrainer
+.. autoclass:: ray.air.train.gbdt_trainer.GBDTTrainer
     :members:
     :show-inheritance:
 
@@ -78,35 +78,35 @@ Tuner
 Predictors
 ~~~~~~~~~~
 
-.. autoclass:: ray.ml.predictor.Predictor
+.. autoclass:: ray.air.predictor.Predictor
     :members:
 
-.. autoclass:: ray.ml.predictor.DataBatchType
+.. autoclass:: ray.air.predictor.DataBatchType
 
-.. autoclass:: ray.ml.batch_predictor.BatchPredictor
+.. autoclass:: ray.air.batch_predictor.BatchPredictor
     :members:
 
-.. automodule:: ray.ml.predictors.integrations.xgboost
+.. automodule:: ray.air.predictors.integrations.xgboost
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.predictors.integrations.lightgbm
+.. automodule:: ray.air.predictors.integrations.lightgbm
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.predictors.integrations.tensorflow
+.. automodule:: ray.air.predictors.integrations.tensorflow
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.predictors.integrations.torch
+.. automodule:: ray.air.predictors.integrations.torch
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.predictors.integrations.sklearn
+.. automodule:: ray.air.predictors.integrations.sklearn
     :members:
     :show-inheritance:
 
-.. automodule:: ray.ml.predictors.integrations.huggingface
+.. automodule:: ray.air.predictors.integrations.huggingface
     :members:
     :show-inheritance:
 
@@ -124,17 +124,17 @@ Serving
 Outputs
 ~~~~~~~
 
-.. automodule:: ray.ml.checkpoint
+.. automodule:: ray.air.checkpoint
     :members:
 
 
-.. automodule:: ray.ml.result
+.. automodule:: ray.air.result
     :members:
 
 
 Configs
 ~~~~~~~
 
-.. automodule:: ray.ml.config
+.. automodule:: ray.air.config
     :members:
 

doc/source/ray-air/preprocessors.rst

@@ -20,7 +20,7 @@ Ray AIR exposes a ``Preprocessor`` class for preprocessing. The ``Preprocessor``
 #. ``transform()``: Apply a transformation to a ``Dataset``.
    If the ``Preprocessor`` is stateful, then ``fit()`` must be called first.
    *This is typically called on the training, validation, test datasets.*
-#. ``transform_batch()``: Apply a transformation to a single :class:`batch <ray.ml.predictor.DataBatchType>` of data.
+#. ``transform_batch()``: Apply a transformation to a single :class:`batch <ray.air.predictor.DataBatchType>` of data.
    *This is typically called on online or offline inference data.*
 #. ``fit_transform()``: Syntactic sugar for calling both ``fit()`` and ``transform()`` on a ``Dataset``.
 
@@ -63,7 +63,7 @@ The same logic is applicable to other integrations as well.
 Trainer
 ~~~~~~~
 
-The journey of the ``Preprocessor`` starts with the :class:`Trainer <ray.ml.trainer.Trainer>`.
+The journey of the ``Preprocessor`` starts with the :class:`Trainer <ray.air.trainer.Trainer>`.
 If the ``Trainer`` is instantiated with a ``Preprocessor``, then the following logic will be executed when ``Trainer.fit()`` is called:
 
 #. If a ``"train"`` ``Dataset`` is passed in, then the ``Preprocessor`` will call ``fit()`` on it.
@@ -127,31 +127,31 @@ Ray AIR provides a handful of ``Preprocessor``\s that you can use out of the box
 
 .. tabbed:: Common APIs
 
-    #. :class:`Preprocessor <ray.ml.preprocessor.Preprocessor>`
-    #. :class:`BatchMapper <ray.ml.preprocessors.BatchMapper>`
-    #. :class:`Chain <ray.ml.preprocessors.Chain>`
+    #. :class:`Preprocessor <ray.air.preprocessor.Preprocessor>`
+    #. :class:`BatchMapper <ray.air.preprocessors.BatchMapper>`
+    #. :class:`Chain <ray.air.preprocessors.Chain>`
 
 .. tabbed:: Tabular
 
-    #. :class:`Categorizer <ray.ml.preprocessors.Categorizer>`
-    #. :class:`FeatureHasher <ray.ml.preprocessors.FeatureHasher>`
-    #. :class:`LabelEncoder <ray.ml.preprocessors.LabelEncoder>`
-    #. :class:`MaxAbsScaler <ray.ml.preprocessors.MaxAbsScaler>`
-    #. :class:`MinMaxScaler <ray.ml.preprocessors.MinMaxScaler>`
-    #. :class:`Normalizer <ray.ml.preprocessors.Normalizer>`
-    #. :class:`OneHotEncoder <ray.ml.preprocessors.OneHotEncoder>`
-    #. :class:`OrdinalEncoder <ray.ml.preprocessors.OrdinalEncoder>`
-    #. :class:`PowerTransformer <ray.ml.preprocessors.PowerTransformer>`
-    #. :class:`RobustScaler <ray.ml.preprocessors.RobustScaler>`
-    #. :class:`SimpleImputer <ray.ml.preprocessors.SimpleImputer>`
-    #. :class:`StandardScaler <ray.ml.preprocessors.StandardScaler>`
-    #. :class:`SimpleImputer <ray.ml.preprocessors.SimpleImputer>`
+    #. :class:`Categorizer <ray.air.preprocessors.Categorizer>`
+    #. :class:`FeatureHasher <ray.air.preprocessors.FeatureHasher>`
+    #. :class:`LabelEncoder <ray.air.preprocessors.LabelEncoder>`
+    #. :class:`MaxAbsScaler <ray.air.preprocessors.MaxAbsScaler>`
+    #. :class:`MinMaxScaler <ray.air.preprocessors.MinMaxScaler>`
+    #. :class:`Normalizer <ray.air.preprocessors.Normalizer>`
+    #. :class:`OneHotEncoder <ray.air.preprocessors.OneHotEncoder>`
+    #. :class:`OrdinalEncoder <ray.air.preprocessors.OrdinalEncoder>`
+    #. :class:`PowerTransformer <ray.air.preprocessors.PowerTransformer>`
+    #. :class:`RobustScaler <ray.air.preprocessors.RobustScaler>`
+    #. :class:`SimpleImputer <ray.air.preprocessors.SimpleImputer>`
+    #. :class:`StandardScaler <ray.air.preprocessors.StandardScaler>`
+    #. :class:`SimpleImputer <ray.air.preprocessors.SimpleImputer>`
 
 .. tabbed:: Text
 
-    #. :class:`CountVectorizer <ray.ml.preprocessors.CountVectorizer>`
-    #. :class:`HashingVectorizer <ray.ml.preprocessors.HashingVectorizer>`
-    #. :class:`Tokenizer <ray.ml.preprocessors.Tokenizer>`
+    #. :class:`CountVectorizer <ray.air.preprocessors.CountVectorizer>`
+    #. :class:`HashingVectorizer <ray.air.preprocessors.HashingVectorizer>`
+    #. :class:`Tokenizer <ray.air.preprocessors.Tokenizer>`
 
 .. tabbed:: Image
 
@@ -159,7 +159,7 @@ Ray AIR provides a handful of ``Preprocessor``\s that you can use out of the box
 
 .. tabbed:: Utilities
 
-    #. :func:`train_test_split <ray.ml.train_test_split>`
+    #. :func:`train_test_split <ray.air.train_test_split>`
 
 Chaining Preprocessors
 ~~~~~~~~~~~~~~~~~~~~~~

doc/source/tune/examples/horovod_simple.ipynb

@@ -42,7 +42,7 @@
     "import ray\n",
     "from ray import train\n",
     "from ray import tune\n",
-    "from ray.ml.train.integrations.horovod import HorovodTrainer\n",
+    "from ray.air.train.integrations.horovod import HorovodTrainer\n",
     "from ray.tune.tune_config import TuneConfig\n",
     "from ray.tune.tuner import Tuner\n",
     "\n",

python/ray/air/BUILD

@@ -1,5 +1,5 @@
 # --------------------------------------------------------------------
-# Tests from the python/ray/ml/examples directory.
+# Tests from the python/ray/air/examples directory.
 # Please keep these sorted alphabetically.
 # --------------------------------------------------------------------
 py_test(
@@ -148,7 +148,7 @@ py_test(
 )
 
 # --------------------------------------------------------------------
-# Tests from the python/ray/ml/tests directory.
+# Tests from the python/ray/air/tests directory.
 # Covers all tests starting with `test_`.
 # Please keep these sorted alphabetically.
 # --------------------------------------------------------------------

python/ray/air/__init__.py

@@ -0,0 +1,12 @@
+from ray.air.checkpoint import Checkpoint
+from ray.air.config import RunConfig, ScalingConfig
+from ray.air.preprocessor import Preprocessor
+from ray.air.utils.datasets import train_test_split
+
+__all__ = [
+    "Checkpoint",
+    "Preprocessor",
+    "RunConfig",
+    "ScalingConfig",
+    "train_test_split",
+]

python/ray/air/batch_predictor.py

@@ -1,8 +1,8 @@
 from typing import Type, Optional, Dict, Any
 
 import ray
-from ray.ml import Checkpoint
-from ray.ml.predictor import Predictor
+from ray.air import Checkpoint
+from ray.air.predictor import Predictor
 
 
 class BatchPredictor(Predictor):

python/ray/air/checkpoint.py

@@ -12,7 +12,7 @@ from typing import Any, Dict, Iterator, Optional, Tuple, Union
 
 import ray
 from ray import cloudpickle as pickle
-from ray.ml.utils.remote_storage import (
+from ray.air.utils.remote_storage import (
     download_from_uri,
     fs_hint,
     is_non_local_path_uri,
@@ -53,7 +53,7 @@ class Checkpoint:
 
         .. code-block:: python
 
-            from ray.ml.checkpoint import Checkpoint
+            from ray.air.checkpoint import Checkpoint
 
             # Create checkpoint data dict
             checkpoint_data = {"data": 123}
@@ -79,7 +79,7 @@ class Checkpoint:
 
         .. code-block:: python
 
-            from ray.ml.checkpoint import Checkpoint
+            from ray.air.checkpoint import Checkpoint
             from sklearn.ensemble import RandomForestClassifier
             import mlflow.sklearn
 

python/ray/air/examples/custom_trainer.py

@@ -5,7 +5,7 @@
 # __custom_trainer_begin__
 import torch
 
-from ray.ml.trainer import Trainer
+from ray.air.trainer import Trainer
 from ray import tune
 
 

python/ray/air/examples/horovod/horovod_pytorch_example.py

@@ -10,7 +10,7 @@ from torchvision import datasets, transforms
 
 import ray
 from ray import train
-from ray.ml.train.integrations.horovod import HorovodTrainer
+from ray.air.train.integrations.horovod import HorovodTrainer
 
 
 def metric_average(val, name):

python/ray/air/examples/horovod/horovod_tune_example.py

@@ -5,7 +5,7 @@ import torch
 import ray
 from ray import train
 from ray import tune
-from ray.ml.train.integrations.horovod import HorovodTrainer
+from ray.air.train.integrations.horovod import HorovodTrainer
 from ray.tune.tune_config import TuneConfig
 from ray.tune.tuner import Tuner
 

python/ray/air/examples/huggingface/huggingface_basic_language_modeling_example.py

@@ -17,9 +17,9 @@ import torch
 
 import ray
 import ray.data
-from ray.ml.train.integrations.huggingface import HuggingFaceTrainer
-from ray.ml.predictors.integrations.huggingface import HuggingFacePredictor
-from ray.ml.batch_predictor import BatchPredictor
+from ray.air.train.integrations.huggingface import HuggingFaceTrainer
+from ray.air.predictors.integrations.huggingface import HuggingFacePredictor
+from ray.air.batch_predictor import BatchPredictor
 
 
 def main(

python/ray/air/examples/pytorch/torch_fashion_mnist_example.py

@@ -8,7 +8,7 @@ from torchvision import datasets
 from torchvision.transforms import ToTensor
 
 import ray.train as train
-from ray.ml.train.integrations.torch import TorchTrainer
+from ray.air.train.integrations.torch import TorchTrainer
 
 # Download training data from open datasets.
 training_data = datasets.FashionMNIST(

python/ray/air/examples/pytorch/torch_linear_dataset_example.py

@@ -8,11 +8,11 @@ import torch.nn as nn
 import ray
 import ray.train as train
 from ray.data import Dataset
-from ray.ml import train_test_split
-from ray.ml.batch_predictor import BatchPredictor
-from ray.ml.predictors.integrations.torch import TorchPredictor
-from ray.ml.result import Result
-from ray.ml.train.integrations.torch import TorchTrainer
+from ray.air import train_test_split
+from ray.air.batch_predictor import BatchPredictor
+from ray.air.predictors.integrations.torch import TorchPredictor
+from ray.air.result import Result
+from ray.air.train.integrations.torch import TorchTrainer
 
 
 def get_datasets(a=5, b=10, size=1000, split=0.8) -> Tuple[Dataset]:

python/ray/air/examples/pytorch/torch_linear_example.py

@@ -4,7 +4,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import ray.train as train
-from ray.ml.train.integrations.torch import TorchTrainer
+from ray.air.train.integrations.torch import TorchTrainer
 
 
 class LinearDataset(torch.utils.data.Dataset):

python/ray/air/examples/pytorch/tune_torch_linear_dataset_example.py

@@ -2,7 +2,7 @@ import argparse
 
 import ray
 from ray import tune
-from ray.ml.train.integrations.torch import TorchTrainer
+from ray.air.train.integrations.torch import TorchTrainer
 from ray.tune.tune_config import TuneConfig
 from ray.tune.tuner import Tuner
 

python/ray/air/examples/pytorch_geometric/distributed_sage_example.py

@@ -13,7 +13,7 @@ from torch_geometric.loader import NeighborSampler
 from torch_geometric.nn import SAGEConv
 
 from ray import train
-from ray.ml.train.integrations.torch import TorchTrainer
+from ray.air.train.integrations.torch import TorchTrainer
 from torch_geometric.transforms import RandomNodeSplit
 
 

python/ray/air/examples/tf/tensorflow_linear_dataset_example.py

@@ -3,16 +3,16 @@ import numpy as np
 
 
 import tensorflow as tf
-from ray.ml.batch_predictor import BatchPredictor
+from ray.air.batch_predictor import BatchPredictor
 from tensorflow.keras.callbacks import Callback
 
 import ray
 import ray.train as train
 from ray.data import Dataset
 from ray.train.tensorflow import prepare_dataset_shard
-from ray.ml.train.integrations.tensorflow import TensorflowTrainer
-from ray.ml.predictors.integrations.tensorflow import TensorflowPredictor
-from ray.ml.result import Result
+from ray.air.train.integrations.tensorflow import TensorflowTrainer
+from ray.air.predictors.integrations.tensorflow import TensorflowPredictor
+from ray.air.result import Result
 
 
 class TrainCheckpointReportCallback(Callback):

python/ray/air/examples/tf/tensorflow_mnist_example.py

@@ -6,12 +6,12 @@ import json
 import os
 
 import numpy as np
-from ray.ml.result import Result
+from ray.air.result import Result
 import tensorflow as tf
 from tensorflow.keras.callbacks import Callback
 
 import ray.train as train
-from ray.ml.train.integrations.tensorflow import TensorflowTrainer
+from ray.air.train.integrations.tensorflow import TensorflowTrainer
 
 
 class TrainCheckpointReportCallback(Callback):

python/ray/air/examples/tf/tune_tensorflow_mnist_example.py

@@ -2,9 +2,9 @@ import argparse
 
 import ray
 from ray import tune
-from ray.ml.train.integrations.tensorflow import TensorflowTrainer
+from ray.air.train.integrations.tensorflow import TensorflowTrainer
 
-from ray.ml.examples.tf.tensorflow_mnist_example import train_func
+from ray.air.examples.tf.tensorflow_mnist_example import train_func
 from ray.tune.tune_config import TuneConfig
 from ray.tune.tuner import Tuner
 

python/ray/air/examples/wandb/latest-run

@@ -0,0 +1 @@
+run-20220506_100917-33977_00000

python/ray/air/predictor.py

@@ -1,7 +1,7 @@
 import abc
 from typing import Union, TYPE_CHECKING
 
-from ray.ml.checkpoint import Checkpoint
+from ray.air.checkpoint import Checkpoint
 
 if TYPE_CHECKING:
     import numpy as np

python/ray/air/predictors/integrations/huggingface/__init__.py

@@ -0,0 +1,5 @@
+from ray.air.predictors.integrations.huggingface.huggingface_predictor import (
+    HuggingFacePredictor,
+)
+
+__all__ = ["HuggingFacePredictor"]

python/ray/air/predictors/integrations/huggingface/huggingface_predictor.py

@@ -8,10 +8,10 @@ from transformers.pipelines.table_question_answering import (
     TableQuestionAnsweringPipeline,
 )
 
-from ray.ml.predictor import DataBatchType, Predictor
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.checkpoint import Checkpoint
-from ray.ml.utils.checkpointing import load_preprocessor_from_dir
+from ray.air.predictor import DataBatchType, Predictor
+from ray.air.preprocessor import Preprocessor
+from ray.air.checkpoint import Checkpoint
+from ray.air.utils.checkpointing import load_preprocessor_from_dir
 
 
 class HuggingFacePredictor(Predictor):
@@ -127,7 +127,7 @@ class HuggingFacePredictor(Predictor):
             import pandas as pd
             from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
             from transformers.pipelines import pipeline
-            from ray.ml.predictors.integrations.huggingface import HuggingFacePredictor
+            from ray.air.predictors.integrations.huggingface import HuggingFacePredictor
 
             model_checkpoint = "gpt2"
             tokenizer_checkpoint = "sgugger/gpt2-like-tokenizer"

python/ray/air/predictors/integrations/lightgbm/__init__.py

@@ -0,0 +1,5 @@
+from ray.air.predictors.integrations.lightgbm.lightgbm_predictor import (
+    LightGBMPredictor,
+)
+
+__all__ = ["LightGBMPredictor"]

python/ray/air/predictors/integrations/lightgbm/lightgbm_predictor.py

@@ -4,10 +4,10 @@ import pandas as pd
 
 import lightgbm
 
-from ray.ml.checkpoint import Checkpoint
-from ray.ml.predictor import Predictor, DataBatchType
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.train.integrations.lightgbm import load_checkpoint
+from ray.air.checkpoint import Checkpoint
+from ray.air.predictor import Predictor, DataBatchType
+from ray.air.preprocessor import Preprocessor
+from ray.air.train.integrations.lightgbm import load_checkpoint
 
 
 class LightGBMPredictor(Predictor):
@@ -63,7 +63,7 @@ class LightGBMPredictor(Predictor):
 
             import numpy as np
             import lightgbm as lgbm
-            from ray.ml.predictors.lightgbm import LightGBMPredictor
+            from ray.air.predictors.lightgbm import LightGBMPredictor
 
             train_X = np.array([[1, 2], [3, 4]])
             train_y = np.array([0, 1])
@@ -82,7 +82,7 @@ class LightGBMPredictor(Predictor):
 
             import pandas as pd
             import lightgbm as lgbm
-            from ray.ml.predictors.lightgbm import LightGBMPredictor
+            from ray.air.predictors.lightgbm import LightGBMPredictor
 
             train_X = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
             train_y = pd.Series([0, 1])

python/ray/air/predictors/integrations/rl/__init__.py

@@ -0,0 +1,3 @@
+from ray.air.predictors.integrations.rl.rl_predictor import RLPredictor
+
+__all__ = ["RLPredictor"]

python/ray/air/predictors/integrations/rl/rl_predictor.py

@@ -3,9 +3,9 @@ from typing import Optional
 import numpy
 import numpy as np
 import pandas as pd
-from ray.ml import Preprocessor, Checkpoint
-from ray.ml.predictor import Predictor, DataBatchType
-from ray.ml.train.integrations.rl import load_checkpoint
+from ray.air import Preprocessor, Checkpoint
+from ray.air.predictor import Predictor, DataBatchType
+from ray.air.train.integrations.rl import load_checkpoint
 from ray.rllib.policy.policy import Policy
 from ray.rllib.utils.typing import EnvType
 

python/ray/air/predictors/integrations/sklearn/__init__.py

@@ -0,0 +1,5 @@
+from ray.air.predictors.integrations.sklearn.sklearn_predictor import (
+    SklearnPredictor,
+)
+
+__all__ = ["SklearnPredictor"]

python/ray/air/predictors/integrations/sklearn/sklearn_predictor.py

@@ -4,11 +4,11 @@ import pandas as pd
 import numpy as np
 from joblib import parallel_backend
 
-from ray.ml.checkpoint import Checkpoint
-from ray.ml.predictor import Predictor, DataBatchType
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.train.integrations.sklearn import load_checkpoint
-from ray.ml.utils.sklearn_utils import set_cpu_params
+from ray.air.checkpoint import Checkpoint
+from ray.air.predictor import Predictor, DataBatchType
+from ray.air.preprocessor import Preprocessor
+from ray.air.train.integrations.sklearn import load_checkpoint
+from ray.air.utils.sklearn_utils import set_cpu_params
 from ray.util.joblib import register_ray
 
 from sklearn.base import BaseEstimator
@@ -71,7 +71,7 @@ class SklearnPredictor(Predictor):
 
             import numpy as np
             from sklearn.ensemble import RandomForestClassifier
-            from ray.ml.predictors.sklearn import SklearnPredictor
+            from ray.air.predictors.sklearn import SklearnPredictor
 
             train_X = np.array([[1, 2], [3, 4]])
             train_y = np.array([0, 1])
@@ -90,7 +90,7 @@ class SklearnPredictor(Predictor):
 
             import pandas as pd
             from sklearn.ensemble import RandomForestClassifier
-            from ray.ml.predictors.sklearn import SklearnPredictor
+            from ray.air.predictors.sklearn import SklearnPredictor
 
             train_X = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
             train_y = pd.Series([0, 1])

python/ray/air/predictors/integrations/tensorflow/__init__.py

@@ -0,0 +1,5 @@
+from ray.air.predictors.integrations.tensorflow.tensorflow_predictor import (
+    TensorflowPredictor,
+)
+
+__all__ = ["TensorflowPredictor"]

python/ray/air/predictors/integrations/tensorflow/tensorflow_predictor.py

@@ -3,10 +3,10 @@ from typing import Callable, Optional, Union, List, Type
 import pandas as pd
 import tensorflow as tf
 
-from ray.ml.predictor import Predictor, DataBatchType
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.checkpoint import Checkpoint
-from ray.ml.train.data_parallel_trainer import _load_checkpoint
+from ray.air.predictor import Predictor, DataBatchType
+from ray.air.preprocessor import Preprocessor
+from ray.air.checkpoint import Checkpoint
+from ray.air.train.data_parallel_trainer import _load_checkpoint
 
 
 class TensorflowPredictor(Predictor):
@@ -82,7 +82,7 @@ class TensorflowPredictor(Predictor):
 
             import numpy as np
             import tensorflow as tf
-            from ray.ml.predictors.tensorflow import TensorflowPredictor
+            from ray.air.predictors.tensorflow import TensorflowPredictor
 
             def build_model(self):
                 return tf.keras.Sequential(
@@ -101,7 +101,7 @@ class TensorflowPredictor(Predictor):
 
             import pandas as pd
             import tensorflow as tf
-            from ray.ml.predictors.tensorflow import TensorflowPredictor
+            from ray.air.predictors.tensorflow import TensorflowPredictor
 
             def build_model(self):
                 return tf.keras.Sequential(

python/ray/air/predictors/integrations/torch/__init__.py

@@ -0,0 +1,3 @@
+from ray.air.predictors.integrations.torch.torch_predictor import TorchPredictor
+
+__all__ = ["TorchPredictor"]

python/ray/air/predictors/integrations/torch/torch_predictor.py

@@ -4,11 +4,11 @@ import numpy as np
 import pandas as pd
 import torch
 
-from ray.ml.predictor import Predictor, DataBatchType
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.checkpoint import Checkpoint
-from ray.ml.train.integrations.torch import load_checkpoint
-from ray.ml.utils.torch_utils import convert_pandas_to_torch_tensor
+from ray.air.predictor import Predictor, DataBatchType
+from ray.air.preprocessor import Preprocessor
+from ray.air.checkpoint import Checkpoint
+from ray.air.train.integrations.torch import load_checkpoint
+from ray.air.utils.torch_utils import convert_pandas_to_torch_tensor
 
 
 class TorchPredictor(Predictor):
@@ -127,7 +127,7 @@ class TorchPredictor(Predictor):
 
             import numpy as np
             import torch
-            from ray.ml.predictors.integrations.torch import TorchPredictor
+            from ray.air.predictors.integrations.torch import TorchPredictor
 
             model = torch.nn.Linear(2, 1)
             predictor = TorchPredictor(model=model)
@@ -139,7 +139,7 @@ class TorchPredictor(Predictor):
 
             import pandas as pd
             import torch
-            from ray.ml.predictors.integrations.torch import TorchPredictor
+            from ray.air.predictors.integrations.torch import TorchPredictor
 
             model = torch.nn.Linear(1, 1)
             predictor = TorchPredictor(model=model)

python/ray/air/predictors/integrations/xgboost/__init__.py

@@ -0,0 +1,3 @@
+from ray.air.predictors.integrations.xgboost.xgboost_predictor import XGBoostPredictor
+
+__all__ = ["XGBoostPredictor"]

python/ray/air/predictors/integrations/xgboost/xgboost_predictor.py

@@ -4,10 +4,10 @@ import pandas as pd
 
 import xgboost
 
-from ray.ml.checkpoint import Checkpoint
-from ray.ml.predictor import Predictor, DataBatchType
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.train.integrations.xgboost import load_checkpoint
+from ray.air.checkpoint import Checkpoint
+from ray.air.predictor import Predictor, DataBatchType
+from ray.air.preprocessor import Preprocessor
+from ray.air.train.integrations.xgboost import load_checkpoint
 
 
 class XGBoostPredictor(Predictor):
@@ -67,7 +67,7 @@ class XGBoostPredictor(Predictor):
 
             import numpy as np
             import xgboost as xgb
-            from ray.ml.predictors.xgboost import XGBoostPredictor
+            from ray.air.predictors.xgboost import XGBoostPredictor
 
             train_X = np.array([[1, 2], [3, 4]])
             train_y = np.array([0, 1])
@@ -86,7 +86,7 @@ class XGBoostPredictor(Predictor):
 
             import pandas as pd
             import xgboost as xgb
-            from ray.ml.predictors.xgboost import XGBoostPredictor
+            from ray.air.predictors.xgboost import XGBoostPredictor
 
             train_X = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
             train_y = pd.Series([0, 1])

python/ray/air/preprocessor.py

@@ -7,7 +7,7 @@ if TYPE_CHECKING:
     import pandas as pd
 
 from ray.data import Dataset
-from ray.ml.predictor import DataBatchType
+from ray.air.predictor import DataBatchType
 
 
 class PreprocessorNotFittedException(RuntimeError):

python/ray/air/preprocessors/__init__.py

@@ -1,24 +1,24 @@
-from ray.ml.preprocessors.batch_mapper import BatchMapper
-from ray.ml.preprocessors.chain import Chain
-from ray.ml.preprocessors.encoder import (
+from ray.air.preprocessors.batch_mapper import BatchMapper
+from ray.air.preprocessors.chain import Chain
+from ray.air.preprocessors.encoder import (
     Categorizer,
     LabelEncoder,
     MultiHotEncoder,
     OneHotEncoder,
     OrdinalEncoder,
 )
-from ray.ml.preprocessors.hasher import FeatureHasher
-from ray.ml.preprocessors.imputer import SimpleImputer
-from ray.ml.preprocessors.normalizer import Normalizer
-from ray.ml.preprocessors.scaler import (
+from ray.air.preprocessors.hasher import FeatureHasher
+from ray.air.preprocessors.imputer import SimpleImputer
+from ray.air.preprocessors.normalizer import Normalizer
+from ray.air.preprocessors.scaler import (
     StandardScaler,
     MinMaxScaler,
     MaxAbsScaler,
     RobustScaler,
 )
-from ray.ml.preprocessors.tokenizer import Tokenizer
-from ray.ml.preprocessors.transformer import PowerTransformer
-from ray.ml.preprocessors.vectorizer import CountVectorizer, HashingVectorizer
+from ray.air.preprocessors.tokenizer import Tokenizer
+from ray.air.preprocessors.transformer import PowerTransformer
+from ray.air.preprocessors.vectorizer import CountVectorizer, HashingVectorizer
 
 __all__ = [
     "BatchMapper",

python/ray/air/preprocessors/batch_mapper.py

@@ -1,6 +1,6 @@
 from typing import Callable, TYPE_CHECKING
 
-from ray.ml.preprocessor import Preprocessor
+from ray.air.preprocessor import Preprocessor
 
 if TYPE_CHECKING:
     import pandas

python/ray/air/preprocessors/chain.py

@@ -1,5 +1,5 @@
 from ray.data import Dataset
-from ray.ml.preprocessor import Preprocessor, DataBatchType
+from ray.air.preprocessor import Preprocessor, DataBatchType
 
 
 class Chain(Preprocessor):

python/ray/air/preprocessors/encoder.py

@@ -6,7 +6,7 @@ import pandas as pd
 import pandas.api.types
 
 from ray.data import Dataset
-from ray.ml.preprocessor import Preprocessor
+from ray.air.preprocessor import Preprocessor
 
 
 class OrdinalEncoder(Preprocessor):
@@ -25,7 +25,7 @@ class OrdinalEncoder(Preprocessor):
     .. code-block:: python
 
         import ray.data
-        from ray.ml.preprocessors import OrdinalEncoder
+        from ray.air.preprocessors import OrdinalEncoder
         import pandas as pd
         batch = pd.DataFrame(
             {
@@ -202,7 +202,7 @@ class MultiHotEncoder(Preprocessor):
     .. code-block:: python
 
         import ray.data
-        from ray.ml.preprocessors import MultiHotEncoder
+        from ray.air.preprocessors import MultiHotEncoder
         import pandas as pd
         mhe = MultiHotEncoder(columns=["A", "B"])
         batch = pd.DataFrame(

python/ray/air/preprocessors/hasher.py

@@ -3,9 +3,9 @@ from typing import List
 
 import pandas as pd
 
-from ray.ml.preprocessor import Preprocessor
+from ray.air.preprocessor import Preprocessor
 
-from ray.ml.preprocessors.utils import simple_hash
+from ray.air.preprocessors.utils import simple_hash
 
 
 class FeatureHasher(Preprocessor):

python/ray/air/preprocessors/imputer.py

@@ -6,7 +6,7 @@ import pandas as pd
 
 from ray.data import Dataset
 from ray.data.aggregate import Mean
-from ray.ml.preprocessor import Preprocessor
+from ray.air.preprocessor import Preprocessor
 
 
 class SimpleImputer(Preprocessor):

python/ray/air/preprocessors/normalizer.py

@@ -3,7 +3,7 @@ from typing import List
 import numpy as np
 import pandas as pd
 
-from ray.ml.preprocessor import Preprocessor
+from ray.air.preprocessor import Preprocessor
 
 
 class Normalizer(Preprocessor):

python/ray/air/preprocessors/scaler.py

@@ -5,7 +5,7 @@ import pandas as pd
 
 from ray.data import Dataset
 from ray.data.aggregate import Mean, Std, Min, Max, AbsMax
-from ray.ml.preprocessor import Preprocessor
+from ray.air.preprocessor import Preprocessor
 
 
 class StandardScaler(Preprocessor):

python/ray/air/preprocessors/tokenizer.py

@@ -2,8 +2,8 @@ from typing import List, Callable, Optional
 
 import pandas as pd
 
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.preprocessors.utils import simple_split_tokenizer
+from ray.air.preprocessor import Preprocessor
+from ray.air.preprocessors.utils import simple_split_tokenizer
 
 
 class Tokenizer(Preprocessor):

python/ray/air/preprocessors/transformer.py

@@ -3,7 +3,7 @@ from typing import List
 import numpy as np
 import pandas as pd
 
-from ray.ml.preprocessor import Preprocessor
+from ray.air.preprocessor import Preprocessor
 
 
 class PowerTransformer(Preprocessor):

python/ray/air/preprocessors/vectorizer.py

@@ -4,8 +4,8 @@ from typing import List, Callable, Optional
 import pandas as pd
 
 from ray.data import Dataset
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.preprocessors.utils import simple_split_tokenizer, simple_hash
+from ray.air.preprocessor import Preprocessor
+from ray.air.preprocessors.utils import simple_split_tokenizer, simple_hash
 
 
 class HashingVectorizer(Preprocessor):

python/ray/air/result.py

@@ -1,7 +1,7 @@
 from typing import Any, Dict, Optional
 from dataclasses import dataclass
 
-from ray.ml.checkpoint import Checkpoint
+from ray.air.checkpoint import Checkpoint
 
 
 @dataclass

python/ray/air/tests/test_api.py

@@ -1,11 +1,11 @@
 import pytest
 
 import ray
-from ray.ml import Checkpoint
-from ray.ml.config import ScalingConfigDataClass
-from ray.ml.trainer import Trainer
-from ray.ml.preprocessor import Preprocessor
-from ray.ml.utils.config import (
+from ray.air import Checkpoint
+from ray.air.config import ScalingConfigDataClass
+from ray.air.trainer import Trainer
+from ray.air.preprocessor import Preprocessor
+from ray.air.utils.config import (
     ensure_only_allowed_dataclass_keys_updated,
     ensure_only_allowed_dict_keys_set,
 )
@@ -38,7 +38,7 @@ def test_run_config():
     DummyTrainer(run_config=None)
 
     # Succeed
-    DummyTrainer(run_config=ray.ml.RunConfig())
+    DummyTrainer(run_config=ray.air.RunConfig())
 
 
 def test_scaling_config():

python/ray/air/tests/test_batchpredictor_runtime.py

@@ -8,9 +8,9 @@ from pathlib import Path
 from typing import Type
 
 import ray
-from ray.ml.checkpoint import Checkpoint
-from ray.ml.batch_predictor import BatchPredictor
-from ray.ml.predictor import Predictor
+from ray.air.checkpoint import Checkpoint
+from ray.air.batch_predictor import BatchPredictor
+from ray.air.predictor import Predictor
 
 NUM_REPEATS = 3