[fix] Fail fast when SimpleStorage lacks Ray CPU capacity

tests/test_simple_storage_scheduling.py

@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import MagicMock
+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
 
@@ -96,7 +97,8 @@ def test_simple_storage_initialization_forwards_required_node_resource(monkeypat
 
     monkeypatch.setattr(simple_storage_bootstrap, "get_node_round_robin_scheduling_strategies", get_strategies)
     monkeypatch.setattr(simple_storage_bootstrap, "SimpleStorageUnit", storage_unit)
-    monkeypatch.setattr(simple_storage_bootstrap, "process_zmq_server_info", lambda _: {})
+    monkeypatch.setattr(simple_storage_bootstrap, "process_zmq_server_info", lambda _, timeout=None: {})
+    monkeypatch.setattr(simple_storage_bootstrap.ray, "available_resources", lambda: {"CPU": 1.0})
 
     conf = OmegaConf.create(
         {
@@ -115,3 +117,68 @@ def test_simple_storage_initialization_forwards_required_node_resource(monkeypat
 
     get_strategies.assert_called_once_with(1, required_node_resource="storage_pool")
     assert handles == {"TransferQueueStorageUnit#0": storage_handle}
+
+
+def test_simple_storage_start_timeout_kills_units_and_reports_cpu_requirement(monkeypatch):
+    strategies = [MagicMock(node_id=_NODE_A), MagicMock(node_id=_NODE_A)]
+    storage_handles = [MagicMock(), MagicMock()]
+    server_info_ref = MagicMock()
+    storage_handles[0].get_zmq_server_info.remote.return_value = server_info_ref
+    storage_unit = MagicMock()
+    storage_unit.options.return_value.remote.side_effect = storage_handles
+    get = MagicMock(side_effect=simple_storage_bootstrap.ray.exceptions.GetTimeoutError())
+    kill = MagicMock()
+
+    monkeypatch.setattr(
+        simple_storage_bootstrap, "get_node_round_robin_scheduling_strategies", lambda *_args, **_kwargs: strategies
+    )
+    monkeypatch.setattr(simple_storage_bootstrap, "SimpleStorageUnit", storage_unit)
+    monkeypatch.setattr(simple_storage_bootstrap.ray, "available_resources", lambda: {"CPU": 1.0})
+    monkeypatch.setattr(simple_storage_bootstrap.ray, "get", get)
+    monkeypatch.setattr(simple_storage_bootstrap.ray, "kill", kill)
+
+    conf = OmegaConf.create(
+        {
+            "backend": {
+                "storage_backend": "SimpleStorage",
+                "SimpleStorage": {"num_data_storage_units": 2, "total_storage_size": None},
+            }
+        }
+    )
+
+    with pytest.raises(RuntimeError) as exc_info:
+        simple_storage_bootstrap.initialize_simple_storage(conf)
+
+    assert str(exc_info.value) == (
+        "SimpleStorage startup timed out after 60 seconds. Each SimpleStorageUnit requires 1 Ray CPU; "
+        "backend.SimpleStorage.num_data_storage_units=2 therefore requires Ray CPU capacity of 2, but Ray reported "
+        "available CPU capacity of 1 before startup. Reduce backend.SimpleStorage.num_data_storage_units or make "
+        "more CPUs available on the eligible Ray nodes."
+    )
+    get.assert_called_once_with(server_info_ref, timeout=simple_storage_bootstrap.SIMPLE_STORAGE_START_TIMEOUT_SECONDS)
+    assert kill.call_args_list == [call(storage_handles[0]), call(storage_handles[1])]
+
+
+def test_init_rolls_back_controller_when_storage_initialization_fails(monkeypatch):
+    controller = MagicMock()
+    controller_class = MagicMock()
+    controller_class.options.return_value.remote.return_value = controller
+    storage_error = RuntimeError("storage initialization failed")
+    kill = MagicMock()
+
+    monkeypatch.setattr(interface, "_TQ_CLIENT", None)
+    monkeypatch.setattr(interface, "_TQ_STORAGE", None)
+    monkeypatch.setattr(interface, "_TQ_CONTROLLER", None)
+    monkeypatch.setattr(interface, "_init_from_existing", lambda: False)
+    monkeypatch.setattr(interface, "TransferQueueController", controller_class)
+    monkeypatch.setattr(interface, "process_zmq_server_info", lambda _: {})
+    monkeypatch.setattr(interface, "_maybe_create_tq_storage", MagicMock(side_effect=storage_error))
+    monkeypatch.setattr(interface.ray, "kill", kill)
+
+    with pytest.raises(RuntimeError) as exc_info:
+        interface.init()
+
+    assert exc_info.value is storage_error
+    kill.assert_called_once_with(controller)
+    assert interface._TQ_CONTROLLER is None
+    assert interface._TQ_STORAGE is None

transfer_queue/interface.py

@@ -190,7 +190,11 @@ def init(conf: DictConfig | None = None) -> DictConfig | None:
     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)
+    try:
+        final_conf = _maybe_create_tq_storage(final_conf)
+    except Exception:
+        close()
+        raise
 
     ray.get(_TQ_CONTROLLER.store_config.remote(final_conf))
     logger.info(f"TransferQueue config: {final_conf}")

transfer_queue/storage/bootstrap/simple_storage_bootstrap.py

@@ -16,6 +16,7 @@
 import math
 from typing import Any
 
+import ray
 from omegaconf import DictConfig
 
 from transfer_queue.storage.bootstrap.provider import StorageBootstrapProvider
@@ -26,6 +27,8 @@
 
 logger = get_logger(__name__)
 
+SIMPLE_STORAGE_START_TIMEOUT_SECONDS = 60
+
 
 @StorageBootstrapProvider.register_provider("SimpleStorage")
 def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
@@ -38,6 +41,7 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
     scheduling_strategies = get_node_round_robin_scheduling_strategies(
         num_data_storage_units, required_node_resource=required_node_resource
     )
+    available_cpus = ray.available_resources().get("CPU", 0.0)
 
     # Compute per-unit capacity: None means unlimited
     storage_unit_size = (
@@ -57,7 +61,23 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]:
             f"on node {scheduling_strategies[storage_unit_rank].node_id}."
         )
 
-    storage_zmq_info = process_zmq_server_info(simple_storage_handles)
+    try:
+        storage_zmq_info = process_zmq_server_info(simple_storage_handles, timeout=SIMPLE_STORAGE_START_TIMEOUT_SECONDS)
+    except ray.exceptions.GetTimeoutError as error:
+        for storage_node in simple_storage_handles.values():
+            try:
+                ray.kill(storage_node)
+            except Exception:
+                logger.exception("Failed to kill SimpleStorageUnit after startup timeout.")
+        raise RuntimeError(
+            f"SimpleStorage startup timed out after {SIMPLE_STORAGE_START_TIMEOUT_SECONDS} seconds. "
+            "Each SimpleStorageUnit requires 1 Ray CPU; "
+            f"backend.SimpleStorage.num_data_storage_units={num_data_storage_units} therefore requires "
+            f"Ray CPU capacity of {num_data_storage_units}, but Ray reported "
+            f"available CPU capacity of {available_cpus:g} "
+            "before startup. "
+            "Reduce backend.SimpleStorage.num_data_storage_units or make more CPUs available on the eligible Ray nodes."
+        ) from error
     backend_name = conf.backend.storage_backend
     conf.backend[backend_name].zmq_info = storage_zmq_info
 

transfer_queue/utils/zmq_utils.py

@@ -431,12 +431,13 @@ async def wrapper(self, *args, **kwargs):
     return decorator
 
 
-def process_zmq_server_info(handlers: dict[Any, Any] | Any):
+def process_zmq_server_info(handlers: dict[Any, Any] | Any, timeout: float | None = None):
     """Extract ZMQ server information from handler objects.
 
     Args:
         handlers: Dictionary of handler objects (controllers, storage managers or storage units),
                   or a single handler object
+        timeout: Maximum seconds to wait for each handler to return its server information.
 
     Returns:
         If handlers is a dictionary: Dictionary mapping handler names to their ZMQ server information
@@ -451,9 +452,11 @@ def process_zmq_server_info(handlers: dict[Any, Any] | Any):
         >>> handlers = {"storage_0": storage_0, "storage_1": storage_1}
         >>> info_dict = process_zmq_server_info(handlers)"""
     if not isinstance(handlers, dict):
-        return ray.get(handlers.get_zmq_server_info.remote())  # type: ignore[union-attr, attr-defined]
+        return ray.get(handlers.get_zmq_server_info.remote(), timeout=timeout)  # type: ignore[union-attr, attr-defined]
     else:
         server_info = {}
         for name, handler in handlers.items():
-            server_info[name] = ray.get(handler.get_zmq_server_info.remote())  # type: ignore[union-attr, attr-defined]
+            server_info[name] = ray.get(  # type: ignore[union-attr, attr-defined]
+                handler.get_zmq_server_info.remote(), timeout=timeout
+            )
         return server_info