From 5e55c90d77e5a26f643bf0e9f80597860f3d2e70 Mon Sep 17 00:00:00 2001
From: xudonghui <xudonghui@xiaohongshu.com>
Date: Fri, 4 Sep 2026 18:24:51 +0800
Subject: [PATCH 1/6] feat: add configurable Ray actor resource affinity

Allow the TransferQueue controller and SimpleStorage units to opt into hard node affinity based on a user-defined Ray custom resource. Preserve existing scheduling when unset and fail fast when no alive node matches.
---
 README.md                                     |  22 +-
 scripts/performance_test/perftest_config.yaml |   8 +-
 tests/test_resource_affinity.py               | 219 ++++++++++++++++++
 transfer_queue/config.yaml                    |   8 +-
 transfer_queue/interface.py                   |  17 +-
 .../bootstrap/simple_storage_bootstrap.py     |  35 ++-
 transfer_queue/utils/common.py                |  38 +++
 7 files changed, 333 insertions(+), 14 deletions(-)
 create mode 100644 tests/test_resource_affinity.py

diff --git a/README.md b/README.md
index 0e026e80..de20ced7 100644
--- a/README.md
+++ b/README.md
@@ -191,6 +191,26 @@ 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.
    
 ### Build wheel package from source code
 
@@ -345,4 +365,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/scripts/performance_test/perftest_config.yaml b/scripts/performance_test/perftest_config.yaml
index 5b8b98dd..e4c6bf54 100644
--- a/scripts/performance_test/perftest_config.yaml
+++ b/scripts/performance_test/perftest_config.yaml
@@ -12,6 +12,9 @@ controller:
   sampler: SequentialSampler
   # Whether return an empty BatchMeta to prevent request blocking when no enough data is available
   polling_mode: False
+  # Optional Ray custom resource required on the controller node.
+  # null preserves Ray's default scheduling behavior.
+  required_node_resource: null
   # ZMQ Server IP & Ports (automatically generated during init)
   zmq_info: null
 
@@ -28,6 +31,9 @@ backend:
     # Number of distributed storage units.
     # Recommended: >= 2 x number of nodes for load balancing.
     num_data_storage_units: 16
+    # Optional Ray custom resource required on SimpleStorageUnit nodes.
+    # null preserves the existing SPREAD placement-group behavior.
+    required_node_resource: null
     # ZMQ Server IP & Ports (automatically generated during init)
     zmq_info: null
 
@@ -119,4 +125,4 @@ backend:
     worker_args: "--shared_memory_size_mb 65536 --remote_h2d_device_ids 0 --enable_huge_tlb true"
 
   # For RayStore:
-  RayStore:
\ No newline at end of file
+  RayStore:
diff --git a/tests/test_resource_affinity.py b/tests/test_resource_affinity.py
new file mode 100644
index 00000000..d473930b
--- /dev/null
+++ b/tests/test_resource_affinity.py
@@ -0,0 +1,219 @@
+# 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, call
+
+import pytest
+from omegaconf import OmegaConf
+
+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": 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: 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, required_node_resource):
+    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 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):
+    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)
+
+
+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, controller_conf):
+    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",
+    )
+
+
+def test_controller_uses_hard_affinity_when_configured(monkeypatch):
+    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
+
+
+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()
diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml
index eebfc178..de59e0f3 100644
--- a/transfer_queue/config.yaml
+++ b/transfer_queue/config.yaml
@@ -12,6 +12,9 @@ controller:
   sampler: SequentialSampler
   # Whether return an empty BatchMeta to prevent request blocking when no enough data is available
   polling_mode: False
+  # Optional Ray custom resource required on the controller node.
+  # null preserves Ray's default scheduling behavior.
+  required_node_resource: null
   # ZMQ Server IP & Ports (automatically generated during init)
   zmq_info: null
 
@@ -29,6 +32,9 @@ backend:
     # Number of distributed storage units.
     # Recommended: >= 2 x number of nodes for load balancing.
     num_data_storage_units: 2
+    # Optional Ray custom resource required on SimpleStorageUnit nodes.
+    # null preserves the existing SPREAD placement-group behavior.
+    required_node_resource: null
     # ZMQ Server IP & Ports (automatically generated during init)
     zmq_info: null
 
@@ -134,4 +140,4 @@ backend:
     worker_args: "--shared_memory_size_mb 8192"
 
   # For RayStore:
-  RayStore:
\ No newline at end of file
+  RayStore:
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index fd129560..55dc320a 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -35,6 +35,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
@@ -176,11 +177,21 @@ def init(conf: DictConfig | None = None) -> DictConfig | None:
     except KeyError:
         raise ValueError(f"Could not find sampler {final_conf.controller.sampler}") from None
 
+    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:
+        controller_options["scheduling_strategy"] = get_node_round_robin_scheduling_strategies(
+            1, required_node_resource=required_node_resource
+        )[0]
+
     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.")
diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
index adbd539e..fd394859 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, get_placement_group
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import process_zmq_server_info
 
@@ -34,7 +34,15 @@ 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)
+    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)
+        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
+        )
 
     # Compute per-unit capacity: None means unlimited
     storage_unit_size = (
@@ -42,15 +50,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:
+            actor_options["scheduling_strategy"] = scheduling_strategies[storage_unit_rank]
+
+        storage_node = SimpleStorageUnit.options(**actor_options).remote(  # type: ignore[attr-defined]
             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.")
+        if scheduling_strategies is None:
+            logger.info(f"TransferQueueStorageUnit#{storage_unit_rank} has been created.")
+        else:
+            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..c1f80119 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,43 @@ 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
+) -> 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.
+
+    Args:
+        num_actors: Number of Ray actors to schedule.
+        required_node_resource: Ray custom resource required on eligible nodes.
+
+    Returns:
+        One hard node-affinity scheduling strategy per actor.
+
+    Raises:
+        ValueError: If no alive Ray node provides the required resource.
+    """
+    eligible_node_ids = sorted(
+        node["NodeID"]
+        for node in ray.nodes()
+        if node.get("Alive", False) and node.get("Resources", {}).get(required_node_resource, 0) > 0
+    )
+    if not eligible_node_ids:
+        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."
+        )
+
+    return [
+        NodeAffinitySchedulingStrategy(node_id=eligible_node_ids[i % len(eligible_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."""

From a6da53c983e25cb91c16a80b86fc4514387951b3 Mon Sep 17 00:00:00 2001
From: xudonghui <xudonghui@xiaohongshu.com>
Date: Mon, 7 Sep 2026 17:59:14 +0800
Subject: [PATCH 2/6] chore: keep resource affinity configuration optional

---
 scripts/performance_test/perftest_config.yaml | 8 +-------
 transfer_queue/config.yaml                    | 6 ------
 2 files changed, 1 insertion(+), 13 deletions(-)

diff --git a/scripts/performance_test/perftest_config.yaml b/scripts/performance_test/perftest_config.yaml
index e4c6bf54..5b8b98dd 100644
--- a/scripts/performance_test/perftest_config.yaml
+++ b/scripts/performance_test/perftest_config.yaml
@@ -12,9 +12,6 @@ controller:
   sampler: SequentialSampler
   # Whether return an empty BatchMeta to prevent request blocking when no enough data is available
   polling_mode: False
-  # Optional Ray custom resource required on the controller node.
-  # null preserves Ray's default scheduling behavior.
-  required_node_resource: null
   # ZMQ Server IP & Ports (automatically generated during init)
   zmq_info: null
 
@@ -31,9 +28,6 @@ backend:
     # Number of distributed storage units.
     # Recommended: >= 2 x number of nodes for load balancing.
     num_data_storage_units: 16
-    # Optional Ray custom resource required on SimpleStorageUnit nodes.
-    # null preserves the existing SPREAD placement-group behavior.
-    required_node_resource: null
     # ZMQ Server IP & Ports (automatically generated during init)
     zmq_info: null
 
@@ -125,4 +119,4 @@ backend:
     worker_args: "--shared_memory_size_mb 65536 --remote_h2d_device_ids 0 --enable_huge_tlb true"
 
   # For RayStore:
-  RayStore:
+  RayStore:
\ No newline at end of file
diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml
index de59e0f3..d97d7f89 100644
--- a/transfer_queue/config.yaml
+++ b/transfer_queue/config.yaml
@@ -12,9 +12,6 @@ controller:
   sampler: SequentialSampler
   # Whether return an empty BatchMeta to prevent request blocking when no enough data is available
   polling_mode: False
-  # Optional Ray custom resource required on the controller node.
-  # null preserves Ray's default scheduling behavior.
-  required_node_resource: null
   # ZMQ Server IP & Ports (automatically generated during init)
   zmq_info: null
 
@@ -32,9 +29,6 @@ backend:
     # Number of distributed storage units.
     # Recommended: >= 2 x number of nodes for load balancing.
     num_data_storage_units: 2
-    # Optional Ray custom resource required on SimpleStorageUnit nodes.
-    # null preserves the existing SPREAD placement-group behavior.
-    required_node_resource: null
     # ZMQ Server IP & Ports (automatically generated during init)
     zmq_info: null
 

From 8733160b1698d219746dae99c048b111de31c00a Mon Sep 17 00:00:00 2001
From: xudonghui <xudonghui@xiaohongshu.com>
Date: Mon, 7 Sep 2026 17:59:35 +0800
Subject: [PATCH 3/6] chore: avoid default config churn

---
 transfer_queue/config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml
index d97d7f89..eebfc178 100644
--- a/transfer_queue/config.yaml
+++ b/transfer_queue/config.yaml
@@ -134,4 +134,4 @@ backend:
     worker_args: "--shared_memory_size_mb 8192"
 
   # For RayStore:
-  RayStore:
+  RayStore:
\ No newline at end of file

From 975ec6fc760e5d325a0be5ae7b4d9e840e632de0 Mon Sep 17 00:00:00 2001
From: xudonghui <xudonghui@xiaohongshu.com>
Date: Tue, 8 Sep 2026 15:11:09 +0800
Subject: [PATCH 4/6] feat: log explicit actor resource affinity selection

---
 README.md                                     |  7 +++++
 tests/test_resource_affinity.py               | 30 ++++++++++++++++---
 transfer_queue/interface.py                   |  6 ++++
 .../bootstrap/simple_storage_bootstrap.py     |  6 ++++
 4 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index de20ced7..7053f703 100644
--- a/README.md
+++ b/README.md
@@ -211,6 +211,13 @@ 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.
+
+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
 
diff --git a/tests/test_resource_affinity.py b/tests/test_resource_affinity.py
index d473930b..ef5e78af 100644
--- a/tests/test_resource_affinity.py
+++ b/tests/test_resource_affinity.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from unittest.mock import MagicMock, call
 
 import pytest
@@ -121,7 +122,8 @@ def test_affinity_fails_fast_when_no_alive_node_matches(monkeypatch):
 
 
 @pytest.mark.parametrize("required_node_resource", [_UNSET, None], ids=["missing", "null"])
-def test_unconfigured_simple_storage_preserves_placement_group(monkeypatch, required_node_resource):
+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)
@@ -137,6 +139,7 @@ def test_unconfigured_simple_storage_preserves_placement_group(monkeypatch, requ
 
     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",
@@ -151,7 +154,8 @@ def test_unconfigured_simple_storage_preserves_placement_group(monkeypatch, requ
     ]
 
 
-def test_simple_storage_uses_hard_affinity_when_configured(monkeypatch):
+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)
@@ -167,6 +171,14 @@ def test_simple_storage_uses_hard_affinity_when_configured(monkeypatch):
     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
+    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):
@@ -180,7 +192,8 @@ def test_simple_storage_fails_before_actor_creation_when_no_node_matches(monkeyp
 
 
 @pytest.mark.parametrize("controller_conf", [None, {"required_node_resource": None}])
-def test_unconfigured_controller_preserves_default_ray_scheduling(monkeypatch, controller_conf):
+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})
 
@@ -190,9 +203,11 @@ def test_unconfigured_controller_preserves_default_ray_scheduling(monkeypatch, c
         name="TransferQueueController",
         namespace="transfer_queue",
     )
+    assert "Applying node affinity:" not in caplog.text
 
 
-def test_controller_uses_hard_affinity_when_configured(monkeypatch):
+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,
@@ -207,6 +222,13 @@ def test_controller_uses_hard_affinity_when_configured(monkeypatch):
     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):
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 55dc320a..fcec3e21 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -186,6 +186,12 @@ def init(conf: DictConfig | None = None) -> DictConfig | None:
         controller_options["scheduling_strategy"] = get_node_round_robin_scheduling_strategies(
             1, required_node_resource=required_node_resource
         )[0]
+        strategy = controller_options["scheduling_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
diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
index fd394859..d76cad46 100644
--- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
+++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
@@ -58,6 +58,12 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
             )
         else:
             actor_options["scheduling_strategy"] = scheduling_strategies[storage_unit_rank]
+            strategy = scheduling_strategies[storage_unit_rank]
+            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,

From e42f370f27d2f410d68301fc4c7d3fc534c2969f Mon Sep 17 00:00:00 2001
From: xudonghui <xudonghui@xiaohongshu.com>
Date: Tue, 8 Sep 2026 15:44:43 +0800
Subject: [PATCH 5/6] refactor: simplify actor affinity logging

---
 tests/test_resource_affinity.py                        |  1 +
 transfer_queue/interface.py                            |  6 ++----
 .../storage/bootstrap/simple_storage_bootstrap.py      | 10 ++--------
 3 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/tests/test_resource_affinity.py b/tests/test_resource_affinity.py
index ef5e78af..33a17d70 100644
--- a/tests/test_resource_affinity.py
+++ b/tests/test_resource_affinity.py
@@ -173,6 +173,7 @@ def test_simple_storage_uses_hard_affinity_when_configured(monkeypatch, caplog):
     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() == (
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index fcec3e21..b208afa3 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -183,10 +183,8 @@ def init(conf: DictConfig | None = None) -> DictConfig | None:
     }
     required_node_resource = final_conf.controller.get("required_node_resource", None)
     if required_node_resource is not None:
-        controller_options["scheduling_strategy"] = get_node_round_robin_scheduling_strategies(
-            1, required_node_resource=required_node_resource
-        )[0]
-        strategy = controller_options["scheduling_strategy"]
+        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} "
diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
index d76cad46..c875a56a 100644
--- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
+++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
@@ -57,8 +57,8 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
                 placement_group_bundle_index=storage_unit_rank,
             )
         else:
-            actor_options["scheduling_strategy"] = scheduling_strategies[storage_unit_rank]
             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} "
@@ -69,13 +69,7 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
             storage_unit_size=storage_unit_size,
         )
         simple_storage_handles[f"TransferQueueStorageUnit#{storage_unit_rank}"] = storage_node
-        if scheduling_strategies is None:
-            logger.info(f"TransferQueueStorageUnit#{storage_unit_rank} has been created.")
-        else:
-            logger.info(
-                f"TransferQueueStorageUnit#{storage_unit_rank} has been created "
-                f"on node {scheduling_strategies[storage_unit_rank].node_id}."
-            )
+        logger.info(f"TransferQueueStorageUnit#{storage_unit_rank} has been created.")
 
     storage_zmq_info = process_zmq_server_info(simple_storage_handles)
     backend_name = conf.backend.storage_backend

From b36d9ad7c3c8a3798335822f893a2be70a8d8625 Mon Sep 17 00:00:00 2001
From: xudonghui <xudonghui@xiaohongshu.com>
Date: Tue, 8 Sep 2026 17:42:48 +0800
Subject: [PATCH 6/6] fix: make actor affinity initialization capacity-safe

Preflight storage affinity and account for per-node CPU capacity, including
Controller CPU on its actual Ray node. Roll back only resources owned by
failed affinity initialization and invalidate failed attachment handles.

Add lifecycle, race, retry and heterogeneous-capacity regression tests.
Validation: 154 tests, pre-commit checks, and isolated multi-node Ray init/retry.
---
 README.md                                     |  15 +
 tests/test_resource_affinity.py               | 269 +++++++++++++++++-
 transfer_queue/controller.py                  |   4 +
 transfer_queue/interface.py                   |  62 +++-
 .../bootstrap/simple_storage_bootstrap.py     |  16 +-
 transfer_queue/utils/common.py                |  48 +++-
 6 files changed, 388 insertions(+), 26 deletions(-)

diff --git a/README.md b/README.md
index 7053f703..af8e2248 100644
--- a/README.md
+++ b/README.md
@@ -212,6 +212,21 @@ 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,
diff --git a/tests/test_resource_affinity.py b/tests/test_resource_affinity.py
index 33a17d70..78fad7a2 100644
--- a/tests/test_resource_affinity.py
+++ b/tests/test_resource_affinity.py
@@ -17,6 +17,7 @@
 
 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
@@ -30,7 +31,7 @@
 
 
 def _node(node_id: str, *, alive: bool = True, resources: dict[str, float] | None = None) -> dict:
-    return {"NodeID": node_id, "Alive": alive, "Resources": resources or {}}
+    return {"NodeID": node_id, "Alive": alive, "Resources": {"CPU": 8, **(resources or {})}}
 
 
 def _simple_storage_conf(required_node_resource=_UNSET):
@@ -64,7 +65,7 @@ def _mock_controller_initialization(monkeypatch):
     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: conf)
+    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)
@@ -240,3 +241,267 @@ def test_controller_fails_before_actor_creation_when_no_node_matches(monkeypatch
         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 b208afa3..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
 
@@ -69,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:
@@ -77,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:
@@ -108,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)
 
@@ -150,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
 
@@ -177,6 +191,18 @@ 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",
@@ -192,7 +218,6 @@ def init(conf: DictConfig | None = None) -> DictConfig | None:
         )
 
     try:
-        global _TQ_CONTROLLER
         _TQ_CONTROLLER = TransferQueueController.options(**controller_options).remote(  # type: ignore[attr-defined]
             sampler=sampler, polling_mode=final_conf.controller.polling_mode
         )
@@ -202,12 +227,31 @@ def init(conf: DictConfig | None = None) -> DictConfig | None:
         _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 c875a56a..d57a32d9 100644
--- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
+++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py
@@ -14,8 +14,10 @@
 # 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
@@ -28,8 +30,10 @@
 
 
 @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
@@ -37,11 +41,15 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
     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
+            num_data_storage_units,
+            required_node_resource=required_node_resource,
+            reserved_cpus=reserved_cpus,
         )
 
     # Compute per-unit capacity: None means unlimited
@@ -68,6 +76,8 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
         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 c1f80119..fce41017 100644
--- a/transfer_queue/utils/common.py
+++ b/transfer_queue/utils/common.py
@@ -46,40 +46,64 @@ def get_placement_group(num_ray_actors: int, num_cpus_per_actor: int = 1):
 
 
 def get_node_round_robin_scheduling_strategies(
-    num_actors: int, required_node_resource: str
+    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.
+    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 alive Ray node provides the required resource.
+        ValueError: If no eligible nodes exist or CPU capacity is insufficient.
     """
-    eligible_node_ids = sorted(
-        node["NodeID"]
-        for node in ray.nodes()
-        if node.get("Alive", False) and node.get("Resources", {}).get(required_node_resource, 0) > 0
+    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_node_ids:
+    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."
         )
 
-    return [
-        NodeAffinitySchedulingStrategy(node_id=eligible_node_ids[i % len(eligible_node_ids)], soft=False)
-        for i in range(num_actors)
-    ]
+    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