__init__.py 560 B

12345678910111213141516171819
  1. from ray.train.v2._internal.constants import is_v2_enabled
  2. if is_v2_enabled():
  3. from ray.train.collective.collectives import barrier, broadcast_from_rank_zero
  4. __all__ = [
  5. "broadcast_from_rank_zero",
  6. "barrier",
  7. ]
  8. broadcast_from_rank_zero.__module__ = "ray.train.collective"
  9. barrier.__module__ = "ray.train.collective"
  10. else:
  11. raise ImportError(
  12. "`ray.train.collective` is only available in Ray Train v2. "
  13. "To enable it, please set `RAY_TRAIN_V2_ENABLED=1`."
  14. )
  15. # DO NOT ADD ANYTHING AFTER THIS LINE.