wehub-resource-sync
72 lines
1.9 KiB
Python
72 lines
1.9 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
import ray.train.collective
|
|
from ray.train.v2._internal.execution import collective_impl
|
|
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
|
|
|
|
|
def test_barrier(ray_start_4_cpus):
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.num_reached_barrier = 0
|
|
|
|
def increment(self):
|
|
self.num_reached_barrier += 1
|
|
|
|
def get_num_reached_barrier(self):
|
|
return self.num_reached_barrier
|
|
|
|
counter = Counter.remote()
|
|
|
|
def train_fn():
|
|
counter.increment.remote()
|
|
ray.train.collective.barrier()
|
|
assert ray.get(counter.get_num_reached_barrier.remote()) == 2
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_fn,
|
|
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_broadcast_from_rank_zero(ray_start_4_cpus):
|
|
def train_fn():
|
|
rank = ray.train.get_context().get_world_rank()
|
|
value = ray.train.collective.broadcast_from_rank_zero({"key": rank})
|
|
assert value == {"key": 0}
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_fn,
|
|
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_broadcast_from_rank_zero_data_too_big(ray_start_4_cpus):
|
|
def train_fn():
|
|
collective_impl.logger = mock.create_autospec(
|
|
collective_impl.logger, instance=True
|
|
)
|
|
collective_impl._MAX_BROADCAST_SIZE_BYTES = 0
|
|
rank = ray.train.get_context().get_world_rank()
|
|
value = ray.train.collective.broadcast_from_rank_zero({"key": rank})
|
|
assert value == {"key": 0}
|
|
collective_impl.logger.warning.assert_called_once()
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_fn,
|
|
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|