From d2d9ca44546385baaf03b297c37655401b108c29 Mon Sep 17 00:00:00 2001 From: nataliekung <4054699+nataliekung@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:15:27 -0700 Subject: [PATCH 1/2] [feat] Filter SimpleStorage nodes by custom resource Allow SimpleStorage to restrict hard round-robin placement to alive Ray nodes with a positive custom resource while preserving the all-alive-node default when the option is unset. Closes #150 Signed-off-by: nataliekung <4054699+nataliekung@users.noreply.github.com> --- scripts/performance_test/perftest_config.yaml | 4 +- tests/test_simple_storage_scheduling.py | 117 ++++++++++++++++++ transfer_queue/config.yaml | 4 +- .../bootstrap/simple_storage_bootstrap.py | 5 +- transfer_queue/utils/common.py | 21 +++- 5 files changed, 144 insertions(+), 7 deletions(-) create mode 100644 tests/test_simple_storage_scheduling.py diff --git a/scripts/performance_test/perftest_config.yaml b/scripts/performance_test/perftest_config.yaml index 32d1e516..6ad01a9d 100644 --- a/scripts/performance_test/perftest_config.yaml +++ b/scripts/performance_test/perftest_config.yaml @@ -25,10 +25,12 @@ backend: SimpleStorage: # Maximum number of experience samples to hold across all storage units total_storage_size: 100000 - # Number of distributed storage units. Units are round-robin scheduled across all + # Number of distributed storage units. Units are round-robin scheduled across eligible # 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 + # Optional Ray custom resource required on storage nodes. null keeps all alive nodes eligible. + node_resource: null # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null diff --git a/tests/test_simple_storage_scheduling.py b/tests/test_simple_storage_scheduling.py new file mode 100644 index 00000000..91eb3913 --- /dev/null +++ b/tests/test_simple_storage_scheduling.py @@ -0,0 +1,117 @@ +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import pytest +from omegaconf import OmegaConf + +from transfer_queue.storage.bootstrap import simple_storage_bootstrap +from transfer_queue.utils import common + +_NODE_A = "01" * 28 +_NODE_B = "02" * 28 +_NODE_C = "03" * 28 +_NODE_D = "04" * 28 + + +def _node(node_id: str, *, alive: bool = True, resources: dict[str, float] | None = None) -> dict: + return {"NodeID": node_id, "Alive": alive, "Resources": resources or {}} + + +def _node_ids(strategies) -> list[str]: + return [strategy.node_id for strategy in strategies] + + +def test_round_robin_uses_all_alive_nodes_by_default(monkeypatch): + nodes = [_node(_NODE_B), _node(_NODE_C, alive=False), _node(_NODE_A)] + monkeypatch.setattr(common.ray, "nodes", lambda: nodes) + + strategies = common.get_node_round_robin_scheduling_strategies(5) + + assert _node_ids(strategies) == [_NODE_A, _NODE_B, _NODE_A, _NODE_B, _NODE_A] + + +def test_node_resource_filters_nodes_and_zero_capacity(monkeypatch): + nodes = [ + _node(_NODE_C, resources={"storage_pool": 2}), + _node(_NODE_B, resources={"storage_pool": 0}), + _node(_NODE_A, resources={"storage_pool": 1}), + _node(_NODE_D, resources={"compute_pool": 1}), + ] + monkeypatch.setattr(common.ray, "nodes", lambda: nodes) + + strategies = common.get_node_round_robin_scheduling_strategies(4, node_resource="storage_pool") + + assert _node_ids(strategies) == [_NODE_A, _NODE_C, _NODE_A, _NODE_C] + + +def test_node_resource_excludes_dead_nodes(monkeypatch): + nodes = [ + _node(_NODE_A, alive=False, resources={"storage_pool": 1}), + _node(_NODE_B, resources={"storage_pool": 1}), + ] + monkeypatch.setattr(common.ray, "nodes", lambda: nodes) + + strategies = common.get_node_round_robin_scheduling_strategies(2, node_resource="storage_pool") + + assert _node_ids(strategies) == [_NODE_B, _NODE_B] + + +def test_node_resource_raises_when_no_alive_node_matches(monkeypatch): + nodes = [ + _node(_NODE_A, resources={"storage_pool": 0}), + _node(_NODE_B, alive=False, resources={"storage_pool": 1}), + ] + monkeypatch.setattr(common.ray, "nodes", lambda: nodes) + + with pytest.raises(ValueError, match="No alive Ray nodes provide custom resource 'storage_pool'"): + common.get_node_round_robin_scheduling_strategies(1, node_resource="storage_pool") + + +def test_default_no_alive_node_error_is_unchanged(monkeypatch): + monkeypatch.setattr(common.ray, "nodes", lambda: []) + + with pytest.raises(RuntimeError, match="No alive Ray nodes found. Is Ray initialized?"): + common.get_node_round_robin_scheduling_strategies(1) + + +def test_simple_storage_initialization_forwards_node_resource(monkeypatch): + strategy = MagicMock(node_id=_NODE_A) + get_strategies = MagicMock(return_value=[strategy]) + storage_unit = MagicMock() + storage_handle = MagicMock() + storage_unit.options.return_value.remote.return_value = storage_handle + + monkeypatch.setattr(simple_storage_bootstrap, "get_node_round_robin_scheduling_strategies", get_strategies) + monkeypatch.setattr(simple_storage_bootstrap, "SimpleStorageUnit", storage_unit) + monkeypatch.setattr(simple_storage_bootstrap, "process_zmq_server_info", lambda _: {}) + + conf = OmegaConf.create( + { + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "num_data_storage_units": 1, + "total_storage_size": None, + "node_resource": "storage_pool", + }, + } + } + ) + + handles = simple_storage_bootstrap.initialize_simple_storage(conf) + + get_strategies.assert_called_once_with(1, node_resource="storage_pool") + assert handles == {"TransferQueueStorageUnit#0": storage_handle} diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index b7356b60..be60fafc 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -26,10 +26,12 @@ 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. Units are round-robin scheduled across all + # Number of distributed storage units. Units are round-robin scheduled across eligible # 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 + # Optional Ray custom resource required on storage nodes. null keeps all alive nodes eligible. + node_resource: null # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index 3082120f..538d322c 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -34,7 +34,10 @@ 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) - scheduling_strategies = get_node_round_robin_scheduling_strategies(num_data_storage_units) + node_resource = conf.backend.SimpleStorage.get("node_resource", None) + scheduling_strategies = get_node_round_robin_scheduling_strategies( + num_data_storage_units, node_resource=node_resource + ) # Compute per-unit capacity: None means unlimited storage_unit_size = ( diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py index a5d98a7e..7cc6d0aa 100644 --- a/transfer_queue/utils/common.py +++ b/transfer_queue/utils/common.py @@ -45,10 +45,12 @@ 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]: +def get_node_round_robin_scheduling_strategies( + num_actors: int, node_resource: str | None = None +) -> list[NodeAffinitySchedulingStrategy]: """ - Compute one scheduling strategy per actor that round-robins actors across all - currently alive Ray nodes, in order. + Compute one scheduling strategy per actor that round-robins actors across + eligible 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 / @@ -57,13 +59,24 @@ def get_node_round_robin_scheduling_strategies(num_actors: int) -> list[NodeAffi Args: num_actors (int): Number of Ray actors to schedule. + node_resource (str | None): Optional Ray custom resource required on eligible nodes. 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)) + alive_node_ids = sorted( + node["NodeID"] + for node in nodes + if node.get("Alive", False) and (node_resource is None or node.get("Resources", {}).get(node_resource, 0) > 0) + ) if not alive_node_ids: + if node_resource is not None: + raise ValueError( + f"No alive Ray nodes provide custom resource {node_resource!r}. " + "Start eligible nodes with a positive resource capacity or unset " + "backend.SimpleStorage.node_resource." + ) raise RuntimeError("No alive Ray nodes found. Is Ray initialized?") return [ From 18d3d17b8a5cfe715f559ea1a8abfb38ca78e142 Mon Sep 17 00:00:00 2001 From: nataliekung <4054699+nataliekung@users.noreply.github.com> Date: Wed, 12 Aug 2026 01:21:21 -0700 Subject: [PATCH 2/2] refactor: rename SimpleStorage node resource option Signed-off-by: nataliekung <4054699+nataliekung@users.noreply.github.com> --- scripts/performance_test/perftest_config.yaml | 2 +- tests/test_simple_storage_scheduling.py | 18 +++++++++--------- transfer_queue/config.yaml | 2 +- .../bootstrap/simple_storage_bootstrap.py | 4 ++-- transfer_queue/utils/common.py | 13 +++++++------ 5 files changed, 20 insertions(+), 19 deletions(-) diff --git a/scripts/performance_test/perftest_config.yaml b/scripts/performance_test/perftest_config.yaml index 6ad01a9d..a01a566c 100644 --- a/scripts/performance_test/perftest_config.yaml +++ b/scripts/performance_test/perftest_config.yaml @@ -30,7 +30,7 @@ backend: # Recommended: >= 2 x number of nodes so each node hosts multiple units. num_data_storage_units: 16 # Optional Ray custom resource required on storage nodes. null keeps all alive nodes eligible. - node_resource: null + required_node_resource: null # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null diff --git a/tests/test_simple_storage_scheduling.py b/tests/test_simple_storage_scheduling.py index 91eb3913..56665bc8 100644 --- a/tests/test_simple_storage_scheduling.py +++ b/tests/test_simple_storage_scheduling.py @@ -43,7 +43,7 @@ def test_round_robin_uses_all_alive_nodes_by_default(monkeypatch): assert _node_ids(strategies) == [_NODE_A, _NODE_B, _NODE_A, _NODE_B, _NODE_A] -def test_node_resource_filters_nodes_and_zero_capacity(monkeypatch): +def test_required_node_resource_filters_nodes_and_zero_capacity(monkeypatch): nodes = [ _node(_NODE_C, resources={"storage_pool": 2}), _node(_NODE_B, resources={"storage_pool": 0}), @@ -52,24 +52,24 @@ def test_node_resource_filters_nodes_and_zero_capacity(monkeypatch): ] monkeypatch.setattr(common.ray, "nodes", lambda: nodes) - strategies = common.get_node_round_robin_scheduling_strategies(4, node_resource="storage_pool") + strategies = common.get_node_round_robin_scheduling_strategies(4, required_node_resource="storage_pool") assert _node_ids(strategies) == [_NODE_A, _NODE_C, _NODE_A, _NODE_C] -def test_node_resource_excludes_dead_nodes(monkeypatch): +def test_required_node_resource_excludes_dead_nodes(monkeypatch): nodes = [ _node(_NODE_A, alive=False, resources={"storage_pool": 1}), _node(_NODE_B, resources={"storage_pool": 1}), ] monkeypatch.setattr(common.ray, "nodes", lambda: nodes) - strategies = common.get_node_round_robin_scheduling_strategies(2, node_resource="storage_pool") + strategies = common.get_node_round_robin_scheduling_strategies(2, required_node_resource="storage_pool") assert _node_ids(strategies) == [_NODE_B, _NODE_B] -def test_node_resource_raises_when_no_alive_node_matches(monkeypatch): +def test_required_node_resource_raises_when_no_alive_node_matches(monkeypatch): nodes = [ _node(_NODE_A, resources={"storage_pool": 0}), _node(_NODE_B, alive=False, resources={"storage_pool": 1}), @@ -77,7 +77,7 @@ def test_node_resource_raises_when_no_alive_node_matches(monkeypatch): monkeypatch.setattr(common.ray, "nodes", lambda: nodes) with pytest.raises(ValueError, match="No alive Ray nodes provide custom resource 'storage_pool'"): - common.get_node_round_robin_scheduling_strategies(1, node_resource="storage_pool") + common.get_node_round_robin_scheduling_strategies(1, required_node_resource="storage_pool") def test_default_no_alive_node_error_is_unchanged(monkeypatch): @@ -87,7 +87,7 @@ def test_default_no_alive_node_error_is_unchanged(monkeypatch): common.get_node_round_robin_scheduling_strategies(1) -def test_simple_storage_initialization_forwards_node_resource(monkeypatch): +def test_simple_storage_initialization_forwards_required_node_resource(monkeypatch): strategy = MagicMock(node_id=_NODE_A) get_strategies = MagicMock(return_value=[strategy]) storage_unit = MagicMock() @@ -105,7 +105,7 @@ def test_simple_storage_initialization_forwards_node_resource(monkeypatch): "SimpleStorage": { "num_data_storage_units": 1, "total_storage_size": None, - "node_resource": "storage_pool", + "required_node_resource": "storage_pool", }, } } @@ -113,5 +113,5 @@ def test_simple_storage_initialization_forwards_node_resource(monkeypatch): handles = simple_storage_bootstrap.initialize_simple_storage(conf) - get_strategies.assert_called_once_with(1, node_resource="storage_pool") + get_strategies.assert_called_once_with(1, required_node_resource="storage_pool") assert handles == {"TransferQueueStorageUnit#0": storage_handle} diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index be60fafc..66fb748c 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -31,7 +31,7 @@ backend: # Recommended: >= 2 x number of nodes so each node hosts multiple units. num_data_storage_units: 2 # Optional Ray custom resource required on storage nodes. null keeps all alive nodes eligible. - node_resource: null + required_node_resource: null # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index 538d322c..381341ec 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -34,9 +34,9 @@ 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) - node_resource = conf.backend.SimpleStorage.get("node_resource", None) + required_node_resource = conf.backend.SimpleStorage.get("required_node_resource", None) scheduling_strategies = get_node_round_robin_scheduling_strategies( - num_data_storage_units, node_resource=node_resource + num_data_storage_units, required_node_resource=required_node_resource ) # Compute per-unit capacity: None means unlimited diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py index 7cc6d0aa..999e8183 100644 --- a/transfer_queue/utils/common.py +++ b/transfer_queue/utils/common.py @@ -46,7 +46,7 @@ def get_placement_group(num_ray_actors: int, num_cpus_per_actor: int = 1): def get_node_round_robin_scheduling_strategies( - num_actors: int, node_resource: str | None = None + num_actors: int, required_node_resource: str | None = None ) -> list[NodeAffinitySchedulingStrategy]: """ Compute one scheduling strategy per actor that round-robins actors across @@ -59,7 +59,7 @@ def get_node_round_robin_scheduling_strategies( Args: num_actors (int): Number of Ray actors to schedule. - node_resource (str | None): Optional Ray custom resource required on eligible nodes. + required_node_resource (str | None): Optional Ray custom resource required on eligible nodes. Returns: list[NodeAffinitySchedulingStrategy]: One scheduling strategy per actor. @@ -68,14 +68,15 @@ def get_node_round_robin_scheduling_strategies( alive_node_ids = sorted( node["NodeID"] for node in nodes - if node.get("Alive", False) and (node_resource is None or node.get("Resources", {}).get(node_resource, 0) > 0) + if node.get("Alive", False) + and (required_node_resource is None or node.get("Resources", {}).get(required_node_resource, 0) > 0) ) if not alive_node_ids: - if node_resource is not None: + if required_node_resource is not None: raise ValueError( - f"No alive Ray nodes provide custom resource {node_resource!r}. " + f"No alive Ray nodes provide custom resource {required_node_resource!r}. " "Start eligible nodes with a positive resource capacity or unset " - "backend.SimpleStorage.node_resource." + "backend.SimpleStorage.required_node_resource." ) raise RuntimeError("No alive Ray nodes found. Is Ray initialized?")