ray.air.session.get_world_size

ray.air.session.get_world_size() → int

Get the current world size (i.e. total number of workers) for this run.

import time
from ray.air import session
from ray.air.config import ScalingConfig

def train_loop_per_worker(config):
    assert session.get_world_size() == 4

train_dataset = ray.data.from_items(
    [{"x": x, "y": x + 1} for x in range(32)])
trainer = TensorflowTrainer(train_loop_per_worker,
    scaling_config=ScalingConfig(num_workers=1),
    datasets={"train": train_dataset})
trainer.fit()