[AIR] Replace train. with session.

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

@@ -86,15 +86,15 @@
 
 # __config_4__
 import ray
-from ray import train
+from ray.air import session
 from ray.data import Dataset
 from ray.train.torch import TorchTrainer
 from ray.air.config import DatasetConfig
 
 
 def train_loop_per_worker():
     # By default, bulk loading is used and returns a Dataset object.
-    data_shard: Dataset = train.get_dataset_shard("train")
+    data_shard: Dataset = session.get_dataset_shard("train")
 
     # Manually iterate over the data 10 times (10 epochs).
     for _ in range(10):
@@ -117,15 +117,15 @@ def train_loop_per_worker():
 
 # __config_5__
 import ray
-from ray import train
+from ray.air import session
 from ray.data import DatasetPipeline
 from ray.train.torch import TorchTrainer
 from ray.air.config import DatasetConfig
 
 
 def train_loop_per_worker():
     # A DatasetPipeline object is returned when `use_stream_api` is set.
-    data_shard: DatasetPipeline = train.get_dataset_shard("train")
+    data_shard: DatasetPipeline = session.get_dataset_shard("train")
 
     # Use iter_epochs(10) to iterate over 10 epochs of data.
     for epoch in data_shard.iter_epochs(10):

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

@@ -29,6 +29,7 @@
 from torch import nn
 from torch.utils.data import DataLoader
 import ray.train as train
+from ray.air import session
 from ray.train.torch import TorchTrainer
 
 # Define model
@@ -52,7 +53,7 @@ def forward(self, x):
 
 
 def train_epoch(dataloader, model, loss_fn, optimizer):
-    size = len(dataloader.dataset) // train.world_size()
+    size = len(dataloader.dataset) // session.get_world_size()
     model.train()
     for batch, (X, y) in enumerate(dataloader):
         # Compute prediction error
@@ -70,7 +71,7 @@ def train_epoch(dataloader, model, loss_fn, optimizer):
 
 
 def validate_epoch(dataloader, model, loss_fn):
-    size = len(dataloader.dataset) // train.world_size()
+    size = len(dataloader.dataset) // session.get_world_size()
     num_batches = len(dataloader)
     model.eval()
     test_loss, correct = 0, 0
@@ -94,7 +95,7 @@ def train_func(config):
     lr = config["lr"]
     epochs = config["epochs"]
 
-    worker_batch_size = batch_size // train.world_size()
+    worker_batch_size = batch_size // session.get_world_size()
 
     # Create data loaders.
     train_dataloader = DataLoader(training_data, batch_size=worker_batch_size)
@@ -113,7 +114,7 @@ def train_func(config):
     for _ in range(epochs):
         train_epoch(train_dataloader, model, loss_fn, optimizer)
         loss = validate_epoch(test_dataloader, model, loss_fn)
-        train.report(loss=loss)
+        session.report(dict(loss=loss))
 
 
 num_workers = 2

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

@@ -15,9 +15,9 @@
 
 # __air_tf_train_start__
 import tensorflow as tf
-from tensorflow.keras.callbacks import Callback
 
-import ray.train as train
+from ray.air import session
+from ray.air.callbacks.keras import Callback
 from ray.train.tensorflow import prepare_dataset_shard
 from ray.train.tensorflow import TensorflowTrainer
 
@@ -33,12 +33,6 @@ def build_model() -> tf.keras.Model:
     return model
 
 
-class TrainCheckpointReportCallback(Callback):
-    def on_epoch_end(self, epoch, logs=None):
-        train.save_checkpoint(**{"model": self.model.get_weights()})
-        train.report(**logs)
-
-
 def train_func(config: dict):
     batch_size = config.get("batch_size", 64)
     epochs = config.get("epochs", 3)
@@ -53,7 +47,7 @@ def train_func(config: dict):
             metrics=[tf.keras.metrics.mean_squared_error],
         )
 
-    dataset = train.get_dataset_shard("train")
+    dataset = session.get_dataset_shard("train")
 
     results = []
     for _ in range(epochs):
@@ -67,9 +61,7 @@ def train_func(config: dict):
                 batch_size=batch_size,
             )
         )
-        history = multi_worker_model.fit(
-            tf_dataset, callbacks=[TrainCheckpointReportCallback()]
-        )
+        history = multi_worker_model.fit(tf_dataset, callbacks=[Callback()])
         results.append(history.history)
     return results
 

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

@@ -674,10 +674,11 @@
     "\n",
     "To facilitate this, we only need a few changes to the code:\n",
     "\n",
-    "1. We import Ray Train:\n",
+    "1. We import Ray Train and Ray AIR Session:\n",
     "\n",
     "```python\n",
     "import ray.train as train\n",
+    "from ray.air import session\n",
     "```\n",
     "\n",
     "\n",
@@ -693,7 +694,7 @@
     "3. We dynamically adjust the worker batch size according to the number of workers:\n",
     "\n",
     "```python\n",
-    "    batch_size_per_worker = batch_size // train.world_size()\n",
+    "    batch_size_per_worker = batch_size // session.get_world_size()\n",
     "```\n",
     "\n",
     "4. We prepare the data loader for distributed data sharding:\n",
@@ -716,13 +717,13 @@
     "\n",
     "```python\n",
     "        test_loss = test(test_dataloader, model, loss_fn)\n",
-    "        train.report(loss=test_loss)\n",
+    "        session.report(dict(loss=test_loss))\n",
     "```\n",
     "\n",
     "7. In the `train_epoch()` and `test_epoch()` functions we divide the `size` by the world size:\n",
     "\n",
     "```python\n",
-    "    size = len(dataloader.dataset) // train.world_size()  # Divide by word size\n",
+    "    size = len(dataloader.dataset) // session.get_world_size()  # Divide by word size\n",
     "```\n",
     "\n",
     "8. In the `train_epoch()` function we can get rid of the device mapping. Ray Train does this for us:\n",
@@ -745,7 +746,7 @@
    "outputs": [],
    "source": [
     "def train_epoch(dataloader, model, loss_fn, optimizer):\n",
-    "    size = len(dataloader.dataset) // train.world_size()  # Divide by word size\n",
+    "    size = len(dataloader.dataset) // session.get_world_size()  # Divide by word size\n",
     "    model.train()\n",
     "    for batch, (X, y) in enumerate(dataloader):\n",
     "        # We don't need this anymore! Ray Train does this automatically:\n",
@@ -781,7 +782,7 @@
    "outputs": [],
    "source": [
     "def test_epoch(dataloader, model, loss_fn):\n",
-    "    size = len(dataloader.dataset) // train.world_size()  # Divide by word size\n",
+    "    size = len(dataloader.dataset) // session.get_world_size()  # Divide by word size\n",
     "    num_batches = len(dataloader)\n",
     "    model.eval()\n",
     "    test_loss, correct = 0, 0\n",
@@ -821,14 +822,14 @@
    ],
    "source": [
     "import ray.train as train\n",
-    "\n",
+    "from ray.air import session\n",
     "\n",
     "def train_func(config: dict):\n",
     "    batch_size = config[\"batch_size\"]\n",
     "    lr = config[\"lr\"]\n",
     "    epochs = config[\"epochs\"]\n",
     "    \n",
-    "    batch_size_per_worker = batch_size // train.world_size()\n",
+    "    batch_size_per_worker = batch_size // session.get_world_size()\n",
     "    \n",
     "    # Create data loaders.\n",
     "    train_dataloader = DataLoader(training_data, batch_size=batch_size_per_worker)\n",
@@ -846,7 +847,7 @@
     "    for t in range(epochs):\n",
     "        train_epoch(train_dataloader, model, loss_fn, optimizer)\n",
     "        test_loss = test_epoch(test_dataloader, model, loss_fn)\n",
-    "        train.report(loss=test_loss)\n",
+    "        session.report(dict(loss=test_loss))\n",
     "\n",
     "    print(\"Done!\")"
    ]
@@ -1062,10 +1063,15 @@
    "metadata": {},
    "source": [
     "### Enabling checkpointing to retrieve the model\n",
-    "Enabling checkpointing is pretty easy - we just need to call the `train.save_checkpoint()` API and pass the model state to it:\n",
+    "Enabling checkpointing is pretty easy - we just need to pass a `Checkpoint` object with the model state to the `session.report()` API.\n",
     "\n",
     "```python\n",
-    "    train.save_checkpoint(epoch=t, model=model.module.state_dict())\n",
+    "    from ray.air import Checkpoint\n",
+    "\n",
+    "    checkpoint = Checkpoint.from_dict(\n",
+    "        dict(epoch=t, model=model.module.state_dict())\n",
+    "    )\n",
+    "    session.report(dict(loss=test_loss), checkpoint=checkpoint)\n",
     "```\n",
     "\n",
     "Note that the `model.module` part is needed because the model gets wrapped in `torch.nn.DistributedDataParallel` by `train.torch.prepare_model`.\n",
@@ -1086,6 +1092,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from ray.air import Checkpoint\n",
+    "\n",
     "def load_data():\n",
     "    # Download training data from open datasets.\n",
     "    training_data = datasets.FashionMNIST(\n",
@@ -1110,7 +1118,7 @@
     "    lr = config[\"lr\"]\n",
     "    epochs = config[\"epochs\"]\n",
     "    \n",
-    "    batch_size_per_worker = batch_size // train.world_size()\n",
+    "    batch_size_per_worker = batch_size // session.get_world_size()\n",
     "    \n",
     "    training_data, test_data = load_data()  # <- this is new!\n",
     "    \n",
@@ -1130,8 +1138,10 @@
     "    for t in range(epochs):\n",
     "        train_epoch(train_dataloader, model, loss_fn, optimizer)\n",
     "        test_loss = test_epoch(test_dataloader, model, loss_fn)\n",
-    "        train.save_checkpoint(epoch=t, model=model.module.state_dict())  # <- this is new!\n",
-    "        train.report(loss=test_loss)\n",
+    "        checkpoint = Checkpoint.from_dict(\n",
+    "            dict(epoch=t, model=model.module.state_dict())\n",
+    "        )\n",
+    "        session.report(dict(loss=test_loss), checkpoint=checkpoint)\n",
     "\n",
     "    print(\"Done!\")"
    ]

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

@@ -619,12 +619,11 @@
    },
    "outputs": [],
    "source": [
-    "from ray import train\n",
+    "from ray.air import session, Checkpoint\n",
     "from ray.train.tensorflow import prepare_dataset_shard\n",
-    "from ray.tune.integration.keras import TuneReportCallback\n",
     "\n",
     "def train_loop_per_worker():\n",
-    "    dataset_shard = train.get_dataset_shard(\"train\")\n",
+    "    dataset_shard = session.get_dataset_shard(\"train\")\n",
     "\n",
     "    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()\n",
     "    with strategy.scope():\n",
@@ -653,7 +652,12 @@
     "\n",
     "        model.fit(tf_dataset, verbose=0)\n",
     "        # This saves checkpoint in a way that can be used by Ray Serve coherently.\n",
-    "        train.save_checkpoint(epoch=epoch, model=model.get_weights())"
+    "        session.report(\n",
+    "            {},\n",
+    "            checkpoint=Checkpoint.from_dict(\n",
+    "                dict(epoch=epoch, model=model.get_weights())\n",
+    "            ),\n",
+    "        )"
    ]
   },
   {

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

@@ -253,8 +253,8 @@
     "\n",
     "`train_loop_per_worker` contains regular PyTorch code with a few notable exceptions:\n",
     "* We wrap our model with {py:func}`train.torch.prepare_model <ray.train.torch.prepare_model>`.\n",
-    "* We call {py:func}`train.get_dataset_shard <ray.train.get_dataset_shard>` and {py:meth}`Dataset.to_torch <ray.data.Dataset.to_torch>` to convert a subset of our training data to a Torch dataset.\n",
-    "* We save model state using {py:func}`train.save_checkpoint <ray.train.save_checkpoint>`."
+    "* We call {py:func}`session.get_dataset_shard <ray.air.session.get_dataset_shard>` and {py:meth}`Dataset.to_torch <ray.data.Dataset.to_torch>` to convert a subset of our training data to a Torch dataset.\n",
+    "* We save model state using {py:func}`session.report <ray.air.session.report>`."
    ]
   },
   {
@@ -265,6 +265,7 @@
    "outputs": [],
    "source": [
     "from ray import train\n",
+    "from ray.air import session, Checkpoint\n",
     "import torch.optim as optim\n",
     "\n",
     "\n",
@@ -274,7 +275,7 @@
     "    criterion = nn.CrossEntropyLoss()\n",
     "    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
     "\n",
-    "    train_dataset_shard: torch.utils.data.Dataset = train.get_dataset_shard(\"train\").to_torch(\n",
+    "    train_dataset_shard: torch.utils.data.Dataset = session.get_dataset_shard(\"train\").to_torch(\n",
     "        feature_columns=[\"image\"],\n",
     "        label_column=\"label\",\n",
     "        batch_size=config[\"batch_size\"],\n",
@@ -303,7 +304,10 @@
     "                print(f\"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}\")\n",
     "                running_loss = 0.0\n",
     "\n",
-    "        train.save_checkpoint(model=model.module.state_dict())"
+    "        session.report(\n",
+    "            dict(running_loss=running_loss),\n",
+    "            checkpoint=Checkpoint.from_dict(dict(model=model.module.state_dict())),\n",
+    "        )"
    ]
   },
   {