diff --git a/scripts/performance_test/perftest_config.yaml b/scripts/performance_test/perftest_config.yaml index 5b8b98dd..32d1e516 100644 --- a/scripts/performance_test/perftest_config.yaml +++ b/scripts/performance_test/perftest_config.yaml @@ -25,8 +25,9 @@ backend: SimpleStorage: # Maximum number of experience samples to hold across all storage units total_storage_size: 100000 - # Number of distributed storage units. - # Recommended: >= 2 x number of nodes for load balancing. + # Number of distributed storage units. Units are round-robin scheduled across all + # alive Ray nodes, guaranteeing an even split of memory/bandwidth usage per node. + # Recommended: >= 2 x number of nodes so each node hosts multiple units. num_data_storage_units: 16 # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index eebfc178..b7356b60 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -26,8 +26,9 @@ backend: # Maximum number of experience samples to hold across all storage units. # Set to null for unlimited capacity (no capacity check). total_storage_size: null - # Number of distributed storage units. - # Recommended: >= 2 x number of nodes for load balancing. + # Number of distributed storage units. Units are round-robin scheduled across all + # alive Ray nodes, guaranteeing an even split of memory/bandwidth usage per node. + # Recommended: >= 2 x number of nodes so each node hosts multiple units. num_data_storage_units: 2 # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null @@ -134,4 +135,4 @@ backend: worker_args: "--shared_memory_size_mb 8192" # For RayStore: - RayStore: \ No newline at end of file + RayStore: diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index adbd539e..3082120f 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -20,7 +20,7 @@ from transfer_queue.storage.bootstrap.provider import StorageBootstrapProvider from transfer_queue.storage.simple_storage import SimpleStorageUnit -from transfer_queue.utils.common import get_placement_group +from transfer_queue.utils.common import get_node_round_robin_scheduling_strategies from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import process_zmq_server_info @@ -34,7 +34,7 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: simple_storage_handles = {} num_data_storage_units = conf.backend.SimpleStorage.num_data_storage_units total_storage_size = conf.backend.SimpleStorage.get("total_storage_size", None) - storage_placement_group = get_placement_group(num_data_storage_units, num_cpus_per_actor=1) + scheduling_strategies = get_node_round_robin_scheduling_strategies(num_data_storage_units) # Compute per-unit capacity: None means unlimited storage_unit_size = ( @@ -43,14 +43,16 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: for storage_unit_rank in range(num_data_storage_units): storage_node = SimpleStorageUnit.options( # type: ignore[attr-defined] - placement_group=storage_placement_group, - placement_group_bundle_index=storage_unit_rank, + scheduling_strategy=scheduling_strategies[storage_unit_rank], name=f"TransferQueueStorageUnit#{storage_unit_rank}", ).remote( storage_unit_size=storage_unit_size, ) simple_storage_handles[f"TransferQueueStorageUnit#{storage_unit_rank}"] = storage_node - logger.info(f"TransferQueueStorageUnit#{storage_unit_rank} has been created.") + logger.info( + f"TransferQueueStorageUnit#{storage_unit_rank} has been created " + f"on node {scheduling_strategies[storage_unit_rank].node_id}." + ) storage_zmq_info = process_zmq_server_info(simple_storage_handles) backend_name = conf.backend.storage_backend diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py index 291d2b30..a5d98a7e 100644 --- a/transfer_queue/utils/common.py +++ b/transfer_queue/utils/common.py @@ -19,6 +19,7 @@ import psutil import ray import torch +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from transfer_queue.utils.logging_utils import get_logger @@ -44,6 +45,33 @@ def get_placement_group(num_ray_actors: int, num_cpus_per_actor: int = 1): return placement_group +def get_node_round_robin_scheduling_strategies(num_actors: int) -> list[NodeAffinitySchedulingStrategy]: + """ + Compute one scheduling strategy per actor that round-robins actors across all + currently alive Ray nodes, in order. + + Unlike a placement group with SPREAD (best-effort) or STRICT_SPREAD (fails when + num_actors > num_nodes), this guarantees each node is assigned floor(num_actors / + num_nodes) or ceil(num_actors / num_nodes) actors, regardless of how num_actors + compares to the number of nodes. + + Args: + num_actors (int): Number of Ray actors to schedule. + + Returns: + list[NodeAffinitySchedulingStrategy]: One scheduling strategy per actor. + """ + nodes = ray.nodes() + alive_node_ids = sorted(node["NodeID"] for node in nodes if node.get("Alive", False)) + if not alive_node_ids: + raise RuntimeError("No alive Ray nodes found. Is Ray initialized?") + + return [ + NodeAffinitySchedulingStrategy(node_id=alive_node_ids[i % len(alive_node_ids)], soft=False) + for i in range(num_actors) + ] + + @contextmanager def limit_pytorch_auto_parallel_threads(target_num_threads: int | None = None, info: str = ""): """Prevent PyTorch from overdoing the automatic parallelism during tensor aggregation operations."""