diff --git a/README.md b/README.md index 0e026e80..af8e2248 100644 --- a/README.md +++ b/README.md @@ -191,6 +191,48 @@ pip install TransferQueue ```bash pip install . ``` + +### Pin TransferQueue actors to Ray custom-resource nodes + +TransferQueue can optionally hard-pin the controller and its built-in +`SimpleStorageUnit` actors to alive Ray nodes that advertise a positive custom +resource. Configure either actor type independently: + +```yaml +controller: + required_node_resource: control_pool +backend: + SimpleStorage: + required_node_resource: storage_pool +``` + +When configured, actor placement uses hard node affinity (`soft=False`) and +initialization fails immediately if no alive node provides the requested +resource. Leaving either option unset or `null` preserves its existing Ray +scheduling behavior. The resource names are user-defined Ray resource labels; +TransferQueue does not reserve or consume their capacity. + +Placement is round-robin within each node's total one-CPU actor capacity. +Storage placement subtracts the running controller's CPU on its actual Ray +node, including when controller affinity is unset. An impossible layout raises +`ValueError` rather than assigning more persistent actors than a node can hold. +Storage affinity is checked before controller creation and checked again after +the controller starts. An affinity failure at that second check rolls back the +controller created by this attempt, so initialization can be retried. + +These checks are not atomic CPU reservations: other workloads may still occupy +the selected nodes and Ray may wait for their CPUs. If an affinity-bound actor +dies or becomes unschedulable during startup, the failed initialization cleans +up its owned actors; attaching processes discard a failed controller handle so +an explicit retry can discover a replacement. This is not automatic recovery +of an already initialized TransferQueue deployment. + +Set `TQ_LOGGING_LEVEL=INFO` before starting the process that initializes +TransferQueue to log each configured actor's affinity selection. The +`Applying node affinity:` message includes the actor name, +`required_node_resource`, selected `node_id`, and `soft=false`. This records +the scheduling strategy being submitted, not confirmation that the actor is +running on that node; verify the actor's actual placement through Ray state. ### Build wheel package from source code @@ -345,4 +387,4 @@ Please kindly cite our paper if you find this repo is useful: journal={arXiv preprint arXiv:2507.01663}, year={2025} } -``` \ No newline at end of file +``` diff --git a/tests/test_resource_affinity.py b/tests/test_resource_affinity.py new file mode 100644 index 00000000..78fad7a2 --- /dev/null +++ b/tests/test_resource_affinity.py @@ -0,0 +1,507 @@ +# 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. + +import logging +from unittest.mock import MagicMock, call + +import pytest +from omegaconf import OmegaConf +from ray.exceptions import ActorUnschedulableError, RayActorError + +from transfer_queue import interface +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 +_UNSET = object() + + +def _node(node_id: str, *, alive: bool = True, resources: dict[str, float] | None = None) -> dict: + return {"NodeID": node_id, "Alive": alive, "Resources": {"CPU": 8, **(resources or {})}} + + +def _simple_storage_conf(required_node_resource=_UNSET): + simple_storage = { + "num_data_storage_units": 2, + "total_storage_size": None, + } + if required_node_resource is not _UNSET: + simple_storage["required_node_resource"] = required_node_resource + return OmegaConf.create( + { + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": simple_storage, + } + } + ) + + +def _mock_storage_initialization(monkeypatch): + storage_unit = MagicMock() + storage_unit.options.return_value.remote.side_effect = [MagicMock(), MagicMock()] + monkeypatch.setattr(simple_storage_bootstrap, "SimpleStorageUnit", storage_unit) + monkeypatch.setattr(simple_storage_bootstrap, "process_zmq_server_info", lambda _: {}) + return storage_unit + + +def _mock_controller_initialization(monkeypatch): + controller = MagicMock() + controller_handle = MagicMock() + controller.options.return_value.remote.return_value = controller_handle + monkeypatch.setattr(interface, "TransferQueueController", controller) + monkeypatch.setattr(interface, "_init_from_existing", lambda: False) + monkeypatch.setattr(interface, "_maybe_create_tq_storage", lambda conf, **kwargs: conf) + monkeypatch.setattr(interface, "_maybe_create_tq_client", MagicMock()) + monkeypatch.setattr(interface, "process_zmq_server_info", lambda _: {}) + monkeypatch.setattr(interface.ray, "get", lambda value: value) + interface._TQ_CONTROLLER = None + interface._TQ_STORAGE = None + interface._TQ_CLIENT = None + return controller + + +@pytest.fixture(autouse=True) +def _reset_interface_globals(): + yield + interface._TQ_CONTROLLER = None + interface._TQ_STORAGE = None + interface._TQ_CLIENT = None + + +def test_affinity_filters_dead_zero_and_missing_resources(monkeypatch): + monkeypatch.setattr( + common.ray, + "nodes", + lambda: [ + _node(_NODE_C, resources={"storage_pool": 2}), + _node(_NODE_B, resources={"storage_pool": 0}), + _node(_NODE_A, resources={"storage_pool": 1}), + _node(_NODE_D, alive=False, resources={"storage_pool": 1}), + _node("05" * 28, resources={"compute_pool": 1}), + ], + ) + + strategies = common.get_node_round_robin_scheduling_strategies(5, required_node_resource="storage_pool") + + assert [strategy.node_id for strategy in strategies] == [ + _NODE_A, + _NODE_C, + _NODE_A, + _NODE_C, + _NODE_A, + ] + assert all(strategy.soft is False for strategy in strategies) + + +def test_affinity_fails_fast_when_no_alive_node_matches(monkeypatch): + monkeypatch.setattr( + common.ray, + "nodes", + lambda: [ + _node(_NODE_A, resources={"control_pool": 0}), + _node(_NODE_B, alive=False, resources={"control_pool": 1}), + ], + ) + + with pytest.raises(ValueError, match="No alive Ray nodes provide custom resource 'control_pool'"): + common.get_node_round_robin_scheduling_strategies(1, required_node_resource="control_pool") + + +@pytest.mark.parametrize("required_node_resource", [_UNSET, None], ids=["missing", "null"]) +def test_unconfigured_simple_storage_preserves_placement_group(monkeypatch, caplog, required_node_resource): + caplog.set_level(logging.INFO, logger=simple_storage_bootstrap.logger.name) + storage_unit = _mock_storage_initialization(monkeypatch) + placement_group = MagicMock() + get_placement_group = MagicMock(return_value=placement_group) + get_strategies = MagicMock() + monkeypatch.setattr(simple_storage_bootstrap, "get_placement_group", get_placement_group) + monkeypatch.setattr( + simple_storage_bootstrap, + "get_node_round_robin_scheduling_strategies", + get_strategies, + ) + + simple_storage_bootstrap.initialize_simple_storage(_simple_storage_conf(required_node_resource)) + + get_placement_group.assert_called_once_with(2, num_cpus_per_actor=1) + get_strategies.assert_not_called() + assert "Applying node affinity:" not in caplog.text + assert storage_unit.options.call_args_list == [ + call( + name="TransferQueueStorageUnit#0", + placement_group=placement_group, + placement_group_bundle_index=0, + ), + call( + name="TransferQueueStorageUnit#1", + placement_group=placement_group, + placement_group_bundle_index=1, + ), + ] + + +def test_simple_storage_uses_hard_affinity_when_configured(monkeypatch, caplog): + caplog.set_level(logging.INFO, logger=simple_storage_bootstrap.logger.name) + storage_unit = _mock_storage_initialization(monkeypatch) + get_placement_group = MagicMock() + monkeypatch.setattr(simple_storage_bootstrap, "get_placement_group", get_placement_group) + monkeypatch.setattr( + common.ray, + "nodes", + lambda: [_node(_NODE_A, resources={"storage_pool": 1})], + ) + + simple_storage_bootstrap.initialize_simple_storage(_simple_storage_conf("storage_pool")) + + get_placement_group.assert_not_called() + strategies = [options.kwargs["scheduling_strategy"] for options in storage_unit.options.call_args_list] + assert [strategy.node_id for strategy in strategies] == [_NODE_A, _NODE_A] + assert all(strategy.soft is False for strategy in strategies) + affinity_logs = [record for record in caplog.records if "Applying node affinity:" in record.getMessage()] + assert len(affinity_logs) == 2 + assert "has been created on node" not in caplog.text + for rank, record in enumerate(affinity_logs): + assert record.levelno == logging.INFO + assert record.getMessage() == ( + f"Applying node affinity: actor=TransferQueueStorageUnit#{rank} " + f"required_node_resource=storage_pool node_id={_NODE_A} soft=false" + ) + + +def test_simple_storage_fails_before_actor_creation_when_no_node_matches(monkeypatch): + storage_unit = _mock_storage_initialization(monkeypatch) + monkeypatch.setattr(common.ray, "nodes", lambda: []) + + with pytest.raises(ValueError, match="No alive Ray nodes provide custom resource 'storage_pool'"): + simple_storage_bootstrap.initialize_simple_storage(_simple_storage_conf("storage_pool")) + + storage_unit.options.assert_not_called() + + +@pytest.mark.parametrize("controller_conf", [None, {"required_node_resource": None}]) +def test_unconfigured_controller_preserves_default_ray_scheduling(monkeypatch, caplog, controller_conf): + caplog.set_level(logging.INFO, logger=interface.logger.name) + controller = _mock_controller_initialization(monkeypatch) + conf = None if controller_conf is None else OmegaConf.create({"controller": controller_conf}) + + interface.init(conf) + + controller.options.assert_called_once_with( + name="TransferQueueController", + namespace="transfer_queue", + ) + assert "Applying node affinity:" not in caplog.text + + +def test_controller_uses_hard_affinity_when_configured(monkeypatch, caplog): + caplog.set_level(logging.INFO, logger=interface.logger.name) + controller = _mock_controller_initialization(monkeypatch) + monkeypatch.setattr( + common.ray, + "nodes", + lambda: [_node(_NODE_A, resources={"control_pool": 1})], + ) + + interface.init(OmegaConf.create({"controller": {"required_node_resource": "control_pool"}})) + + options = controller.options.call_args.kwargs + assert options["name"] == "TransferQueueController" + assert options["namespace"] == "transfer_queue" + assert options["scheduling_strategy"].node_id == _NODE_A + assert options["scheduling_strategy"].soft is False + affinity_logs = [record for record in caplog.records if "Applying node affinity:" in record.getMessage()] + assert len(affinity_logs) == 1 + assert affinity_logs[0].levelno == logging.INFO + assert affinity_logs[0].getMessage() == ( + f"Applying node affinity: actor=TransferQueueController " + f"required_node_resource=control_pool node_id={_NODE_A} soft=false" + ) + + +def test_controller_fails_before_actor_creation_when_no_node_matches(monkeypatch): + controller = _mock_controller_initialization(monkeypatch) + monkeypatch.setattr(common.ray, "nodes", lambda: []) + + with pytest.raises(ValueError, match="No alive Ray nodes provide custom resource 'control_pool'"): + interface.init(OmegaConf.create({"controller": {"required_node_resource": "control_pool"}})) + + controller.options.assert_not_called() + + +def test_affinity_respects_heterogeneous_cpu_capacity(monkeypatch): + monkeypatch.setattr( + common.ray, + "nodes", + lambda: [ + _node(_NODE_A, resources={"storage_pool": 1, "CPU": 1}), + _node(_NODE_B, resources={"storage_pool": 1, "CPU": 3}), + ], + ) + strategies = common.get_node_round_robin_scheduling_strategies(4, "storage_pool") + assert [strategy.node_id for strategy in strategies] == [_NODE_A, _NODE_B, _NODE_B, _NODE_B] + + +def test_affinity_rejects_insufficient_cpu_capacity(monkeypatch): + monkeypatch.setattr( + common.ray, + "nodes", + lambda: [ + _node(_NODE_A, resources={"storage_pool": 1, "CPU": 1}), + ], + ) + with pytest.raises(ValueError, match="CPU"): + common.get_node_round_robin_scheduling_strategies(2, "storage_pool") + + +def _mock_full_initialization(monkeypatch): + """Mock Ray transport, keeping init, attach and storage bootstrap real.""" + controller = MagicMock() + handle = controller.options.return_value.remote.return_value + handle.get_config.remote.return_value = None + handle.get_node_id.remote.return_value = _NODE_A + registry = {} + + def create_controller(**kwargs): + registry["controller"] = handle + return handle + + def get_actor(*args, **kwargs): + if "controller" not in registry: + raise ValueError("no controller") + return registry["controller"] + + controller.options.return_value.remote.side_effect = create_controller + + def kill_controller(actor_handle, *, no_restart): + assert no_restart is True + if actor_handle is handle: + registry.pop("controller", None) + + monkeypatch.setattr(interface.ray, "kill", MagicMock(side_effect=kill_controller)) + handle.store_config.remote.side_effect = lambda conf: setattr(handle.get_config.remote, "return_value", conf) + monkeypatch.setattr(interface, "TransferQueueController", controller) + monkeypatch.setattr(interface.ray, "get_actor", get_actor) + monkeypatch.setattr(interface.ray, "get", lambda value: value) + monkeypatch.setattr(interface, "_maybe_create_tq_client", MagicMock()) + monkeypatch.setattr(interface, "process_zmq_server_info", lambda _: {}) + monkeypatch.setattr(interface.time, "sleep", lambda _: pytest.fail("initialization waits for unpublished config")) + storage = _mock_storage_initialization(monkeypatch) + return controller, storage + + +def test_missing_storage_resource_leaves_no_controller_and_can_retry(monkeypatch): + controller, storage = _mock_full_initialization(monkeypatch) + nodes = [_node(_NODE_A, resources={"CPU": 4, "control_pool": 1})] + monkeypatch.setattr(common.ray, "nodes", lambda: nodes) + conf = _simple_storage_conf("storage_pool") + with pytest.raises(ValueError, match="storage_pool"): + interface.init(conf) + controller.options.assert_not_called() + storage.options.assert_not_called() + assert interface._TQ_CONTROLLER is None + assert interface._TQ_STORAGE is None + nodes[0]["Resources"]["storage_pool"] = 1 + interface.init(conf) + assert controller.options.return_value.remote.call_count == 1 + assert storage.options.call_count == 2 + + +@pytest.mark.parametrize("controller_resource", [_UNSET, None, "storage_pool"]) +def test_storage_accounts_for_actual_controller_cpu(monkeypatch, controller_resource): + controller, storage = _mock_full_initialization(monkeypatch) + nodes = [ + {**_node(_NODE_A, resources={"CPU": 1, "storage_pool": 1}), "NodeManagerAddress": "10.0.0.1"}, + {**_node(_NODE_B, resources={"CPU": 2, "storage_pool": 1}), "NodeManagerAddress": "10.0.0.2"}, + ] + monkeypatch.setattr(common.ray, "nodes", lambda: nodes) + conf = _simple_storage_conf("storage_pool") + if controller_resource is not _UNSET: + conf.controller = {"required_node_resource": controller_resource} + interface.init(conf) + if controller_resource in (_UNSET, None): + controller.options.assert_called_once_with(name="TransferQueueController", namespace="transfer_queue") + assert [c.kwargs["scheduling_strategy"].node_id for c in storage.options.call_args_list] == [_NODE_B, _NODE_B] + + +@pytest.mark.parametrize("after_preflight", ["missing", "controller_occupancy"]) +def test_storage_recheck_failure_cleans_owned_controller_and_can_retry(monkeypatch, after_preflight): + controller, storage = _mock_full_initialization(monkeypatch) + nodes = [_node(_NODE_A, resources={"CPU": 2, "storage_pool": 1})] + snapshots = MagicMock(side_effect=[nodes, [] if after_preflight == "missing" else nodes]) + monkeypatch.setattr(common.ray, "nodes", snapshots) + conf = _simple_storage_conf("storage_pool") + with pytest.raises(ValueError): + interface.init(conf) + handle = controller.options.return_value.remote.return_value + interface.ray.kill.assert_called_once_with(handle, no_restart=True) + handle.store_config.remote.assert_not_called() + storage.options.assert_not_called() + assert interface._TQ_CONTROLLER is None + assert interface._TQ_STORAGE is None + nodes[0]["Resources"]["CPU"] = 3 + monkeypatch.setattr(common.ray, "nodes", lambda: nodes) + interface.init(conf) + assert controller.options.return_value.remote.call_count == 2 + assert storage.options.call_count == 2 + assert handle.store_config.remote.call_count == 1 + + +def test_existing_controller_reuse_does_not_validate_or_kill(monkeypatch): + controller, storage = _mock_full_initialization(monkeypatch) + handle = controller.options.return_value.remote.return_value + interface._TQ_CONTROLLER = handle + handle.get_config.remote.return_value = OmegaConf.create({}) + nodes = MagicMock(side_effect=AssertionError("reuse must not replan")) + monkeypatch.setattr(common.ray, "nodes", nodes) + interface.init(_simple_storage_conf("missing_pool")) + controller.options.assert_not_called() + storage.options.assert_not_called() + interface.ray.kill.assert_not_called() + + +@pytest.mark.parametrize("cpu", [0, 0.5]) +def test_affinity_rejects_nodes_without_one_cpu_slot(monkeypatch, cpu): + monkeypatch.setattr(common.ray, "nodes", lambda: [_node(_NODE_A, resources={"CPU": cpu, "control_pool": 1})]) + with pytest.raises(ValueError, match="CPU"): + common.get_node_round_robin_scheduling_strategies(1, "control_pool") + + +@pytest.mark.parametrize("error_type", [ActorUnschedulableError, RayActorError]) +def test_failed_controller_start_can_retry(monkeypatch, error_type): + controller, storage = _mock_full_initialization(monkeypatch) + monkeypatch.setattr(common.ray, "nodes", lambda: [_node(_NODE_A, resources={"control_pool": 1})]) + server_info = MagicMock(side_effect=error_type("node disappeared")) + monkeypatch.setattr(interface, "process_zmq_server_info", server_info) + conf = OmegaConf.create({"controller": {"required_node_resource": "control_pool"}}) + with pytest.raises(error_type): + interface.init(conf) + assert interface._TQ_CONTROLLER is None + assert interface._TQ_STORAGE is None + storage.options.assert_not_called() + interface.ray.kill.assert_called_once_with(controller.options.return_value.remote.return_value, no_restart=True) + server_info.side_effect = None + server_info.return_value = {} + monkeypatch.setattr(simple_storage_bootstrap, "get_placement_group", lambda *args, **kwargs: MagicMock()) + interface.init(conf) + assert controller.options.return_value.remote.call_count == 2 + assert storage.options.call_count == 2 + + +@pytest.mark.parametrize("error_type", [ActorUnschedulableError, RayActorError]) +def test_attacher_discards_dead_controller_without_killing_it(monkeypatch, error_type): + controller, storage = _mock_full_initialization(monkeypatch) + dead = MagicMock() + dead.get_config.remote.side_effect = error_type("initializer rolled back") + interface._TQ_CONTROLLER = dead + with pytest.raises(error_type): + interface.init() + assert interface._TQ_CONTROLLER is None + interface.ray.kill.assert_not_called() + replacement = MagicMock() + replacement.get_config.remote.return_value = OmegaConf.create({}) + lookup = MagicMock(return_value=replacement) + monkeypatch.setattr(interface.ray, "get_actor", lookup) + interface.init() + lookup.assert_called_once_with("TransferQueueController", namespace="transfer_queue") + assert interface._TQ_CONTROLLER is replacement + controller.options.assert_not_called() + storage.options.assert_not_called() + + +@pytest.mark.parametrize("failure_stage", ["storage_start", "store_config"]) +def test_initialization_actor_failure_removes_only_owned_actors(monkeypatch, failure_stage): + controller, storage = _mock_full_initialization(monkeypatch) + monkeypatch.setattr(common.ray, "nodes", lambda: [_node(_NODE_A, resources={"storage_pool": 1})]) + storage_handles = [MagicMock(), MagicMock()] + storage.options.return_value.remote.side_effect = storage_handles + controller_handle = controller.options.return_value.remote.return_value + if failure_stage == "storage_start": + monkeypatch.setattr( + simple_storage_bootstrap, + "process_zmq_server_info", + MagicMock(side_effect=ActorUnschedulableError("node disappeared")), + ) + else: + controller_handle.store_config.remote.side_effect = RayActorError("controller died") + with pytest.raises((ActorUnschedulableError, RayActorError)): + interface.init(_simple_storage_conf("storage_pool")) + assert interface.ray.kill.call_args_list == [ + call(storage_handles[1], no_restart=True), + call(storage_handles[0], no_restart=True), + call(controller_handle, no_restart=True), + ] + assert interface._TQ_CONTROLLER is None + assert interface._TQ_STORAGE is None + + +def test_racing_controller_creator_does_not_kill_winner(monkeypatch): + controller, storage = _mock_full_initialization(monkeypatch) + existing = MagicMock() + existing.get_config.remote.return_value = OmegaConf.create({}) + monkeypatch.setattr(interface.ray, "get_actor", MagicMock(side_effect=[ValueError("not found"), existing])) + controller.options.return_value.remote.side_effect = ValueError("name already exists") + monkeypatch.setattr(common.ray, "nodes", lambda: [_node(_NODE_A, resources={"storage_pool": 1})]) + interface.init(_simple_storage_conf("storage_pool")) + assert interface._TQ_CONTROLLER is existing + storage.options.assert_not_called() + interface.ray.kill.assert_not_called() + + +@pytest.mark.parametrize("storage_resource", [_UNSET, None, "storage_pool"]) +@pytest.mark.parametrize("failure_stage", ["second_actor", "storage_ready", "store_config"]) +def test_controller_affinity_rolls_back_storage_and_placement_group(monkeypatch, storage_resource, failure_stage): + controller, storage = _mock_full_initialization(monkeypatch) + monkeypatch.setattr(common.ray, "nodes", lambda: [_node(_NODE_A, resources={"control_pool": 1, "storage_pool": 1})]) + placement_group = MagicMock() + monkeypatch.setattr(simple_storage_bootstrap, "get_placement_group", lambda *args, **kwargs: placement_group) + remove_pg = MagicMock() + monkeypatch.setattr(interface.ray.util, "remove_placement_group", remove_pg) + handles = [MagicMock(), MagicMock()] + controller_handle = controller.options.return_value.remote.return_value + storage.options.return_value.remote.side_effect = handles + if failure_stage == "second_actor": + storage.options.return_value.remote.side_effect = [handles[0], ValueError("name belongs to another actor")] + elif failure_stage == "storage_ready": + monkeypatch.setattr( + simple_storage_bootstrap, "process_zmq_server_info", MagicMock(side_effect=RayActorError("node lost")) + ) + else: + controller_handle.store_config.remote.side_effect = RayActorError("controller lost") + conf = _simple_storage_conf(storage_resource) + conf.controller = {"required_node_resource": "control_pool"} + with pytest.raises((ValueError, RayActorError)): + interface.init(conf) + created_handles = handles[:1] if failure_stage == "second_actor" else handles + assert interface.ray.kill.call_args_list == [ + *[call(handle, no_restart=True) for handle in reversed(created_handles)], + call(controller_handle, no_restart=True), + ] + if storage_resource in (_UNSET, None): + remove_pg.assert_called_once_with(placement_group) + else: + remove_pg.assert_not_called() + assert interface._TQ_CONTROLLER is None + assert interface._TQ_STORAGE is None + + storage.options.return_value.remote.side_effect = [MagicMock(), MagicMock()] + monkeypatch.setattr(simple_storage_bootstrap, "process_zmq_server_info", lambda _: {}) + controller_handle.store_config.remote.side_effect = None + interface.init(conf) + assert controller.options.return_value.remote.call_count == 2 + # Successful initialization must not execute the rollback callbacks. + assert interface.ray.kill.call_count == len(created_handles) + 1 diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 43ba1ee9..dfa6b437 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -2460,6 +2460,10 @@ def get_zmq_server_info(self) -> ZMQServerInfo: """Get ZMQ server connection information.""" return self.zmq_server_info + def get_node_id(self) -> str: + """Return the Ray node holding this controller's CPU allocation.""" + return ray.get_runtime_context().get_node_id() + def store_config(self, conf: DictConfig) -> None: """Store the global config of TransferQueue.""" self.tq_config = conf diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index fd129560..9e6f20cd 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -18,6 +18,7 @@ import shutil import subprocess import time +from contextlib import ExitStack from importlib import resources from pathlib import Path from typing import Any, Callable @@ -25,6 +26,7 @@ import ray import torch from omegaconf import DictConfig, OmegaConf +from ray.exceptions import ActorUnschedulableError, RayActorError from tensordict import TensorDict from tensordict.tensorclass import NonTensorStack @@ -35,6 +37,7 @@ from transfer_queue.sampler import BaseSampler from transfer_queue.storage.bootstrap import StorageBootstrapProvider from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager +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.yuanrong_utils import cleanup_yuanrong_resources from transfer_queue.utils.zmq_utils import process_zmq_server_info @@ -68,7 +71,9 @@ def _maybe_create_tq_client(conf: DictConfig | None = None) -> TransferQueueClie return _TQ_CLIENT -def _maybe_create_tq_storage(conf: DictConfig) -> DictConfig: +def _maybe_create_tq_storage( + conf: DictConfig, reserved_cpus: dict[str, float] | None = None, rollback: ExitStack | None = None +) -> DictConfig: global _TQ_STORAGE if _TQ_STORAGE is None: @@ -76,7 +81,10 @@ def _maybe_create_tq_storage(conf: DictConfig) -> DictConfig: backend_name = conf.backend.storage_backend provider_fn = StorageBootstrapProvider.get_provider(backend_name) if provider_fn is not None: - backend_resources = provider_fn(conf) + if backend_name == "SimpleStorage" and rollback is not None: + backend_resources = provider_fn(conf, reserved_cpus=reserved_cpus, rollback=rollback) + else: + backend_resources = provider_fn(conf) if backend_resources is not None: _TQ_STORAGE[backend_name] = backend_resources else: @@ -107,7 +115,13 @@ def _init_from_existing() -> bool: conf = None while conf is None: - conf = ray.get(_TQ_CONTROLLER.get_config.remote()) + try: + conf = ray.get(_TQ_CONTROLLER.get_config.remote()) + except (ActorUnschedulableError, RayActorError): + # An initializer may have rolled back this unpublished controller. + # Fail this call, but let an explicit retry discover a replacement. + _TQ_CONTROLLER = None + raise if conf is not None: _maybe_create_tq_client(conf) @@ -149,6 +163,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: >>> metadata = tq.get_meta(...) >>> data = tq.get_data(metadata) """ + global _TQ_CONTROLLER, _TQ_STORAGE if _init_from_existing(): return conf @@ -176,23 +191,67 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: except KeyError: raise ValueError(f"Could not find sampler {final_conf.controller.sampler}") from None + # Reject invalid storage affinity before publishing a named controller whose + # config other ranks would wait for. Bootstrap rechecks the current layout + # after the controller starts, accounting for its actual CPU placement. + storage_resource = None + if final_conf.backend.storage_backend == "SimpleStorage": + storage_conf = final_conf.backend.SimpleStorage + storage_resource = storage_conf.get("required_node_resource", None) + if storage_resource is not None: + get_node_round_robin_scheduling_strategies( + storage_conf.num_data_storage_units, required_node_resource=storage_resource + ) + + controller_options: dict[str, Any] = { + "name": "TransferQueueController", + "namespace": "transfer_queue", + } + required_node_resource = final_conf.controller.get("required_node_resource", None) + if required_node_resource is not None: + strategy = get_node_round_robin_scheduling_strategies(1, required_node_resource=required_node_resource)[0] + controller_options["scheduling_strategy"] = strategy + logger.info( + f"Applying node affinity: actor={controller_options['name']} " + f"required_node_resource={required_node_resource} node_id={strategy.node_id} " + f"soft={str(strategy.soft).lower()}" + ) + try: - global _TQ_CONTROLLER - _TQ_CONTROLLER = TransferQueueController.options( # type: ignore[attr-defined] - name="TransferQueueController", namespace="transfer_queue" - ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) + _TQ_CONTROLLER = TransferQueueController.options(**controller_options).remote( # type: ignore[attr-defined] + sampler=sampler, polling_mode=final_conf.controller.polling_mode + ) logger.info("TransferQueueController has been created.") except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() return final_conf - controller_zmq_info = process_zmq_server_info(_TQ_CONTROLLER) - final_conf.controller.zmq_info = controller_zmq_info - - final_conf = _maybe_create_tq_storage(final_conf) + affinity_enabled = required_node_resource is not None or storage_resource is not None + previous_storage = _TQ_STORAGE + # Register only objects created by this attempt. Reuse/racing-creator paths + # returned above; callbacks are discarded once the config is published. + with ExitStack() as rollback: + if affinity_enabled: + rollback.callback(ray.kill, _TQ_CONTROLLER, no_restart=True) + try: + final_conf.controller.zmq_info = process_zmq_server_info(_TQ_CONTROLLER) + if affinity_enabled: + reserved_cpus = None + if storage_resource is not None: + controller_node_id = ray.get(_TQ_CONTROLLER.get_node_id.remote()) + reserved_cpus = {controller_node_id: 1.0} + final_conf = _maybe_create_tq_storage(final_conf, reserved_cpus=reserved_cpus, rollback=rollback) + else: + final_conf = _maybe_create_tq_storage(final_conf) + ray.get(_TQ_CONTROLLER.store_config.remote(final_conf)) + except Exception: + if affinity_enabled: + _TQ_CONTROLLER = None + _TQ_STORAGE = previous_storage + raise + rollback.pop_all() - ray.get(_TQ_CONTROLLER.store_config.remote(final_conf)) logger.info(f"TransferQueue config: {final_conf}") # start Prometheus metrics exporter if enabled diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index adbd539e..d57a32d9 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -14,13 +14,15 @@ # limitations under the License. import math +from contextlib import ExitStack from typing import Any +import ray from omegaconf import DictConfig 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, get_placement_group from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import process_zmq_server_info @@ -28,13 +30,27 @@ @StorageBootstrapProvider.register_provider("SimpleStorage") -def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: - """Initialize Simple storage with metastore mode.""" +def initialize_simple_storage( + conf: DictConfig, reserved_cpus: dict[str, float] | None = None, rollback: ExitStack | None = None +) -> dict[str, Any]: + """Initialize storage, accounting for CPU allocations and tracking owned resources for rollback.""" 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) + required_node_resource = conf.backend.SimpleStorage.get("required_node_resource", None) + if required_node_resource is None: + storage_placement_group = get_placement_group(num_data_storage_units, num_cpus_per_actor=1) + if rollback is not None: + rollback.callback(ray.util.remove_placement_group, storage_placement_group) + scheduling_strategies = None + else: + storage_placement_group = None + scheduling_strategies = get_node_round_robin_scheduling_strategies( + num_data_storage_units, + required_node_resource=required_node_resource, + reserved_cpus=reserved_cpus, + ) # Compute per-unit capacity: None means unlimited storage_unit_size = ( @@ -42,13 +58,26 @@ 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, - name=f"TransferQueueStorageUnit#{storage_unit_rank}", - ).remote( + actor_options: dict[str, Any] = {"name": f"TransferQueueStorageUnit#{storage_unit_rank}"} + if scheduling_strategies is None: + actor_options.update( + placement_group=storage_placement_group, + placement_group_bundle_index=storage_unit_rank, + ) + else: + strategy = scheduling_strategies[storage_unit_rank] + actor_options["scheduling_strategy"] = strategy + logger.info( + f"Applying node affinity: actor={actor_options['name']} " + f"required_node_resource={required_node_resource} node_id={strategy.node_id} " + f"soft={str(strategy.soft).lower()}" + ) + + storage_node = SimpleStorageUnit.options(**actor_options).remote( # type: ignore[attr-defined] storage_unit_size=storage_unit_size, ) + if rollback is not None: + rollback.callback(ray.kill, storage_node, no_restart=True) simple_storage_handles[f"TransferQueueStorageUnit#{storage_unit_rank}"] = storage_node logger.info(f"TransferQueueStorageUnit#{storage_unit_rank} has been created.") diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py index 291d2b30..fce41017 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,67 @@ 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, + required_node_resource: str, + *, + reserved_cpus: dict[str, float] | None = None, +) -> list[NodeAffinitySchedulingStrategy]: + """Create hard-affinity strategies across nodes providing a resource. + + Eligible nodes must be alive and advertise a positive capacity for + ``required_node_resource``. Actors are assigned to eligible nodes in + deterministic round-robin order, skipping nodes whose CPU slots are full. + Each actor requires one CPU, as declared by Controller and SimpleStorage. + This is a total-capacity check, not a reservation of currently free CPUs. + + Args: + num_actors: Number of Ray actors to schedule. + required_node_resource: Ray custom resource required on eligible nodes. + reserved_cpus: CPUs already assigned to persistent actors, keyed by node ID. + + Returns: + One hard node-affinity scheduling strategy per actor. + + Raises: + ValueError: If no eligible nodes exist or CPU capacity is insufficient. + """ + eligible_nodes = sorted( + ( + node + for node in ray.nodes() + if node.get("Alive", False) and node.get("Resources", {}).get(required_node_resource, 0) > 0 + ), + key=lambda node: node["NodeID"], + ) + if not eligible_nodes: + raise ValueError( + f"No alive Ray nodes provide custom resource {required_node_resource!r}. " + "Start an eligible node with a positive resource capacity or unset " + "the corresponding required_node_resource option." + ) + + reservations = reserved_cpus or {} + cpu_slots = { + node["NodeID"]: max(0, int(node.get("Resources", {}).get("CPU", 0) - reservations.get(node["NodeID"], 0))) + for node in eligible_nodes + } + if sum(cpu_slots.values()) < num_actors: + raise ValueError( + f"Insufficient CPU capacity for {num_actors} actors requiring resource {required_node_resource!r}: " + f"{sum(cpu_slots.values())} one-CPU slots remain after accounting for existing actors." + ) + strategies: list[NodeAffinitySchedulingStrategy] = [] + while len(strategies) < num_actors: + for node_id, slots in cpu_slots.items(): + if slots > 0: + strategies.append(NodeAffinitySchedulingStrategy(node_id=node_id, soft=False)) + cpu_slots[node_id] -= 1 + if len(strategies) == num_actors: + break + return strategies + + @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."""