wehub-resource-sync
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import logging
|
|
from typing import Any
|
|
|
|
import ray
|
|
import ray.cloudpickle as pickle
|
|
from ray.train.v2._internal.execution.context import get_train_context
|
|
|
|
# For reference, {1:1} is 19 bytes, {"1":"1"} is 21 bytes,
|
|
# and {"12345": "12345"} is 25 bytes.
|
|
_MAX_BROADCAST_SIZE_BYTES = 1000
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def barrier() -> None:
|
|
"""
|
|
Create a barrier across all training workers.
|
|
"""
|
|
train_context = get_train_context()
|
|
sync_actor = train_context.get_synchronization_actor()
|
|
return ray.get(
|
|
sync_actor.broadcast_from_rank_zero.remote(
|
|
world_rank=train_context.get_world_rank(),
|
|
world_size=train_context.get_world_size(),
|
|
data=None,
|
|
caller_method_name="ray.train.collective.barrier",
|
|
)
|
|
)
|
|
|
|
|
|
def broadcast_from_rank_zero(data: Any) -> Any:
|
|
"""Broadcast data from the rank 0 worker to all other workers.
|
|
|
|
This method is used by the public API function :func:`ray.train.collective.broadcast_from_rank_zero`.
|
|
Users should typically call ``ray.train.collective.broadcast_from_rank_zero()`` instead of calling this method directly.
|
|
"""
|
|
# Validate data.
|
|
if data is not None:
|
|
data_bytes = len(pickle.dumps(data))
|
|
if data_bytes > _MAX_BROADCAST_SIZE_BYTES:
|
|
logger.warning(
|
|
f"Data size {data_bytes} bytes exceeds the maximum broadcast "
|
|
f"size of {_MAX_BROADCAST_SIZE_BYTES} bytes"
|
|
)
|
|
|
|
train_context = get_train_context()
|
|
sync_actor = train_context.get_synchronization_actor()
|
|
return ray.get(
|
|
sync_actor.broadcast_from_rank_zero.remote(
|
|
world_rank=train_context.get_world_rank(),
|
|
world_size=train_context.get_world_size(),
|
|
data=data,
|
|
caller_method_name="ray.train.collective.broadcast_from_rank_zero",
|
|
)
|
|
)
|