ray.air.session.get_world_rank

ray.air.session.get_world_rank() → int

Get the world rank of this worker.

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

def train_loop_per_worker():
    for iter in range(100):
        time.sleep(1)
        if session.get_world_rank() == 0:
            print("Worker 0")

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()