[fix] Let concurrent jobs keep their TransferQueue actors apart with TQ_RAY_NAMESPACE

recipe/simple_use_case/relax_demo.py

@@ -40,6 +40,7 @@
 import transfer_queue as tq
 from transfer_queue import RankAwareSampler, StreamingDataLoader, StreamingDataset
 from transfer_queue.metadata import BatchMeta
+from transfer_queue.utils.common import ray_namespace
 from transfer_queue.utils.logging_utils import get_logger
 
 logger = get_logger("MAIN", default_level="INFO")
@@ -131,7 +132,7 @@ class BaseStageWorker:
     def __init__(self, tq_config, tracker, worker_id: int, config: "DemoConfig"):
         tq.init(tq_config)
         self.tq_client = tq.get_client()
-        controller = ray.get_actor("TransferQueueController", namespace="transfer_queue")
+        controller = ray.get_actor("TransferQueueController", namespace=ray_namespace())
         self.cfg = ray.get(controller.get_config.remote())
         self.tracker = tracker
         self.worker_id = worker_id

transfer_queue/interface.py

@@ -54,6 +54,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 ray_namespace
 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
@@ -123,7 +124,7 @@ def _init_from_existing() -> bool:
     global _TQ_CONTROLLER
     try:
         if _TQ_CONTROLLER is None:
-            _TQ_CONTROLLER = ray.get_actor("TransferQueueController", namespace="transfer_queue")
+            _TQ_CONTROLLER = ray.get_actor("TransferQueueController", namespace=ray_namespace())
 
     except ValueError:
         logger.info("Called _init_from_existing() but TransferQueueController has not been initialized yet.")
@@ -207,7 +208,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None:
     try:
         global _TQ_CONTROLLER, _TQ_IS_OWNER
         _TQ_CONTROLLER = TransferQueueController.options(  # type: ignore[attr-defined]
-            name="TransferQueueController", namespace="transfer_queue"
+            name="TransferQueueController", namespace=ray_namespace()
         ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode)
         _TQ_IS_OWNER = True
         logger.info("TransferQueueController has been created.")

transfer_queue/kv_lock.py

@@ -42,7 +42,7 @@
 from ray.exceptions import GetTimeoutError, RayActorError
 
 from transfer_queue.lock_manager import _RESERVED_SLOTS, TransferQueueLockManager
-from transfer_queue.utils.common import get_env_bool, get_node_round_robin_scheduling_strategies
+from transfer_queue.utils.common import get_env_bool, get_node_round_robin_scheduling_strategies, ray_namespace
 
 # Not asyncio.Lock: it is bound to one event loop and is not thread-safe, while a process
 # may run several loops plus plain threads. A key is held exactly while it is in _table,
@@ -292,7 +292,7 @@ def _manager_name(shard: int) -> str:
 def _lock_manager(shard: int):
     if shard not in _managers:
         try:
-            _managers[shard] = ray.get_actor(_manager_name(shard), namespace="transfer_queue")
+            _managers[shard] = ray.get_actor(_manager_name(shard), namespace=ray_namespace())
         except ValueError:
             raise RuntimeError(_NOT_INITIALIZED) from None
     return _managers[shard]
@@ -556,7 +556,7 @@ def _start_lock_managers(settings: tuple[int, int] | None) -> None:
         _owned = [
             TransferQueueLockManager.options(  # type: ignore[attr-defined]
                 name=_manager_name(shard),
-                namespace="transfer_queue",
+                namespace=ray_namespace(),
                 get_if_exists=True,
                 max_concurrency=max_waiting_acquires + _RESERVED_SLOTS,
                 scheduling_strategy=strategy,

transfer_queue/storage/clients/ray_storage_client.py

@@ -20,6 +20,7 @@
 import torch
 
 from transfer_queue.storage.clients.base import StorageClientFactory, StorageKVClient
+from transfer_queue.utils.common import ray_namespace
 
 
 @ray.remote(max_concurrency=8)
@@ -57,10 +58,10 @@ def __init__(self, config=None):
 
         # initialize actor
         try:
-            self.storage_actor = ray.get_actor("RayObjectRefStorage", namespace="transfer_queue")
+            self.storage_actor = ray.get_actor("RayObjectRefStorage", namespace=ray_namespace())
         except ValueError:
             self.storage_actor = RayObjectRefStorage.options(
-                name="RayObjectRefStorage", namespace="transfer_queue", get_if_exists=False
+                name="RayObjectRefStorage", namespace=ray_namespace(), get_if_exists=False
             ).remote()
 
     def put(self, keys: list[str], values: list[Any]) -> list[Any] | None:

transfer_queue/utils/common.py

@@ -122,6 +122,15 @@ def limit_pytorch_auto_parallel_threads(target_num_threads: int | None = None, i
         )
 
 
+def ray_namespace() -> str:
+    """Ray namespace of the TransferQueue named actors.
+
+    Defaults to ``transfer_queue``, shared by every job on the Ray cluster (see #115). Set
+    ``TQ_RAY_NAMESPACE`` in every process of a job to keep concurrent jobs apart.
+    """
+    return os.environ.get("TQ_RAY_NAMESPACE") or "transfer_queue"
+
+
 def get_env_bool(env_key: str, default: bool = False) -> bool:
     """Robustly get a boolean from an environment variable."""
     env_value = os.getenv(env_key)

tutorial/06_streaming_dataloader.py

@@ -70,6 +70,7 @@
     StreamingDataLoader,
     StreamingDataset,
 )
+from transfer_queue.utils.common import ray_namespace  # noqa: E402
 
 
 def setup_transfer_queue():
@@ -184,7 +185,7 @@ def update_worker(
     # Step 1: Create StreamingDataset
     # This dataset integrates with TransferQueue and handles batch retrieval
 
-    controller = ray.get_actor("TransferQueueController", namespace="transfer_queue")
+    controller = ray.get_actor("TransferQueueController", namespace=ray_namespace())
     config = ray.get(controller.get_config.remote())
 
     dataset = StreamingDataset(