[2/n] Lightweight Ray AIR API refactor

doc/external/pytorch_tutorials_hyperparameter_tuning_tutorial.py

@@ -49,8 +49,8 @@
 from torch.utils.data import random_split
 import torchvision
 import torchvision.transforms as transforms
-from ray import tune
-from ray.air import Checkpoint, session
+from ray import train, tune
+from ray.train import Checkpoint
 from ray.tune.schedulers import ASHAScheduler
 
 ######################################################################
@@ -125,7 +125,7 @@ def forward(self, x):
 #
 #     net = Net(config["l1"], config["l2"])
 #
-#     checkpoint = session.get_checkpoint()
+#     checkpoint = ray.train.get_checkpoint()
 #
 #     if checkpoint:
 #         checkpoint_state = checkpoint.to_dict()
@@ -189,7 +189,7 @@ def forward(self, x):
 #     }
 #     checkpoint = Checkpoint.from_dict(checkpoint_data)
 #
-#     session.report(
+#     ray.train.report(
 #         {"loss": val_loss / val_steps, "accuracy": correct / total},
 #         checkpoint=checkpoint,
 #     )
@@ -226,7 +226,7 @@ def train_cifar(config, data_dir=None):
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(net.parameters(), lr=config["lr"], momentum=0.9)
 
-    checkpoint = session.get_checkpoint()
+    checkpoint = train.get_checkpoint()
 
     if checkpoint:
         checkpoint_state = checkpoint.to_dict()
@@ -303,7 +303,7 @@ def train_cifar(config, data_dir=None):
         }
         checkpoint = Checkpoint.from_dict(checkpoint_data)
 
-        session.report(
+        train.report(
             {"loss": val_loss / val_steps, "accuracy": correct / total},
             checkpoint=checkpoint,
         )

doc/source/data/doc_code/preprocessors.py

@@ -50,7 +50,7 @@
 
 from ray.data.preprocessors import MinMaxScaler
 from ray.train.xgboost import XGBoostTrainer
-from ray.air.config import ScalingConfig
+from ray.train import ScalingConfig
 
 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)])

doc/source/data/working-with-pytorch.rst

@@ -47,7 +47,8 @@ Ray Data integrates with :ref:`Ray Train <train-docs>` for easy data ingest for
     import torch
     from torch import nn
     import ray
-    from ray.air import session, ScalingConfig
+    from ray import train
+    from ray.train import ScalingConfig
     from ray.train.torch import TorchTrainer
 
     def train_func(config):
@@ -56,7 +57,7 @@ Ray Data integrates with :ref:`Ray Train <train-docs>` for easy data ingest for
         optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
 
         # Datasets can be accessed in your train_func via ``get_dataset_shard``.
-        train_data_shard = session.get_dataset_shard("train")
+        train_data_shard = train.get_dataset_shard("train")
 
         for epoch_idx in range(2):
             for batch in train_data_shard.iter_torch_batches(batch_size=128, dtypes=torch.float32):

doc/source/index.md

@@ -132,7 +132,7 @@ predictions.show(limit=1)
           </div>
           <div class="tab-pane fade" id="v-pills-training" role="tabpanel" aria-labelledby="v-pills-training-tab" style="user-select:none;">
             <pre style="margin:0;"><code class="language-python not-selectable">
-from ray.air.config import ScalingConfig
+from ray.train import ScalingConfig
 from ray.train.torch import TorchTrainer
 
 # Step 1: setup PyTorch model training as you normally would
@@ -166,7 +166,7 @@ result = trainer.fit()
           <div class="tab-pane fade" id="v-pills-tuning" role="tabpanel" aria-labelledby="v-pills-tuning-tab" style="user-select:none;" style="user-select:none;">
             <pre style="margin:0;"><code class="language-python not-selectable">
 from ray import tune
-from ray.air.config import ScalingConfig
+from ray.train import ScalingConfig
 from ray.train.lightgbm import LightGBMTrainer
 
 train_dataset, eval_dataset = ...

doc/source/ray-air/computer-vision.rst

@@ -207,13 +207,13 @@ Training vision models
 Creating checkpoints
 --------------------
 
-:class:`Checkpoints <ray.air.checkpoint.Checkpoint>` are required for batch inference and model
+:class:`Checkpoints <ray.train.Checkpoint>` are required for batch inference and model
 serving. They contain model state and optionally a preprocessor.
 
 If you're going from training to prediction, don't create a new checkpoint.
 :meth:`Trainer.fit() <ray.train.trainer.BaseTrainer.fit>` returns a
-:class:`~ray.air.result.Result` object. Use
-:attr:`Result.checkpoint <ray.air.result.Result.checkpoint>` instead.
+:class:`~ray.train.Result` object. Use
+:attr:`Result.checkpoint <ray.train.Result.checkpoint>` instead.
 
 .. tab-set::
 

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

@@ -186,12 +186,11 @@ def train_torch_model(dataset, preprocessor, per_epoch_preprocessor):
     from torchvision import models
 
     from ray import train
-    from ray.air import session
-    from ray.air.config import ScalingConfig
+    from ray.train import ScalingConfig
     from ray.train.torch import TorchCheckpoint, TorchTrainer
 
     def train_one_epoch(model, *, criterion, optimizer, batch_size, epoch):
-        dataset_shard = session.get_dataset_shard("train")
+        dataset_shard = train.get_dataset_shard("train")
 
         running_loss = 0
         for i, batch in enumerate(
@@ -210,7 +209,7 @@ def train_one_epoch(model, *, criterion, optimizer, batch_size, epoch):
 
             running_loss += loss.item()
             if i % 2000 == 1999:
-                session.report(
+                train.report(
                     metrics={
                         "epoch": epoch,
                         "batch": i,
@@ -254,13 +253,13 @@ def train_tensorflow_model(dataset, preprocessor, per_epoch_preprocessor):
     # __tensorflow_training_loop_start__
     import tensorflow as tf
 
-    from ray.air import session
+    from ray import train
     from ray.air.integrations.keras import ReportCheckpointCallback
 
     def train_loop_per_worker(config):
         strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
 
-        train_shard = session.get_dataset_shard("train")
+        train_shard = train.get_dataset_shard("train")
         train_dataset = train_shard.to_tf(
             "image",
             "label",
@@ -286,7 +285,7 @@ def train_loop_per_worker(config):
     # __tensorflow_training_loop_stop__
 
     # __tensorflow_trainer_start__
-    from ray.air import ScalingConfig
+    from ray.train import ScalingConfig
     from ray.train.tensorflow import TensorflowTrainer
 
     # The following transform operation is lazy.

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

@@ -1,12 +1,11 @@
+import horovod.torch as hvd
 import ray
-import ray.train as train
+from ray import train
+from ray.train import Checkpoint, ScalingConfig
 import ray.train.torch  # Need this to use `train.torch.get_device()`
-import horovod.torch as hvd
+from ray.train.horovod import HorovodTrainer
 import torch
 import torch.nn as nn
-from ray.air import session, Checkpoint
-from ray.train.horovod import HorovodTrainer
-from ray.air.config import ScalingConfig
 
 # If using GPUs, set this to True.
 use_gpu = False
@@ -31,7 +30,7 @@ def forward(self, input):
 
 def train_loop_per_worker():
     hvd.init()
-    dataset_shard = session.get_dataset_shard("train")
+    dataset_shard = train.get_dataset_shard("train")
     model = NeuralNetwork()
     device = train.torch.get_device()
     model.to(device)
@@ -56,7 +55,7 @@ def train_loop_per_worker():
             loss.backward()
             optimizer.step()
             print(f"epoch: {epoch}, loss: {loss.item()}")
-        session.report(
+        train.report(
             {},
             checkpoint=Checkpoint.from_dict(dict(model=model.state_dict())),
         )

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

@@ -4,9 +4,8 @@
 # __air_session_start__
 
 import tensorflow as tf
-from ray.air import session
-from ray.air.checkpoint import Checkpoint
-from ray.air.config import ScalingConfig
+from ray import train
+from ray.train import Checkpoint, ScalingConfig
 from ray.train.tensorflow import TensorflowTrainer
 
 
@@ -22,7 +21,7 @@ def build_model() -> tf.keras.Model:
 
 
 def train_func():
-    ckpt = session.get_checkpoint()
+    ckpt = train.get_checkpoint()
     if ckpt:
         with ckpt.as_directory() as loaded_checkpoint_dir:
             import tensorflow as tf
@@ -32,9 +31,7 @@ def train_func():
         model = build_model()
 
     model.save("my_model", overwrite=True)
-    session.report(
-        metrics={"iter": 1}, checkpoint=Checkpoint.from_directory("my_model")
-    )
+    train.report(metrics={"iter": 1}, checkpoint=Checkpoint.from_directory("my_model"))
 
 
 scaling_config = ScalingConfig(num_workers=2)

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

@@ -5,10 +5,10 @@
 import ray
 import tensorflow as tf
 
-from ray.air import session
+from ray import train
+from ray.train import ScalingConfig
 from ray.air.integrations.keras import ReportCheckpointCallback
 from ray.train.tensorflow import TensorflowTrainer
-from ray.air.config import ScalingConfig
 
 
 # If using GPUs, set this to True.
@@ -46,7 +46,7 @@ def train_func(config: dict):
             metrics=[tf.keras.metrics.mean_squared_error],
         )
 
-    dataset = session.get_dataset_shard("train")
+    dataset = train.get_dataset_shard("train")
 
     results = []
     for _ in range(epochs):

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

@@ -3,9 +3,8 @@
 
 import ray
 from ray import train
-from ray.air import session, Checkpoint
+from ray.train import Checkpoint, ScalingConfig
 from ray.train.torch import TorchTrainer
-from ray.air.config import ScalingConfig
 
 
 # If using GPUs, set this to True.
@@ -30,7 +29,7 @@ def forward(self, input):
 
 
 def train_loop_per_worker():
-    dataset_shard = session.get_dataset_shard("train")
+    dataset_shard = train.get_dataset_shard("train")
     model = NeuralNetwork()
     loss_fn = nn.MSELoss()
     optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
@@ -49,7 +48,7 @@ def train_loop_per_worker():
             optimizer.step()
             print(f"epoch: {epoch}, loss: {loss.item()}")
 
-        session.report(
+        train.report(
             {},
             checkpoint=Checkpoint.from_dict(
                 dict(epoch=epoch, model=model.state_dict())