[fix] Guarantee even node distribution for simple storage units

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:
+  RayStore:

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

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."""