diff --git a/recipe/simple_use_case/relax_demo.py b/recipe/simple_use_case/relax_demo.py index 7e3efa83..6cb1766e 100644 --- a/recipe/simple_use_case/relax_demo.py +++ b/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 diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 21e5a116..5aeadc42 100644 --- a/transfer_queue/interface.py +++ b/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.") diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 3586d2af..fc68839f 100644 --- a/transfer_queue/kv_lock.py +++ b/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, diff --git a/transfer_queue/storage/clients/ray_storage_client.py b/transfer_queue/storage/clients/ray_storage_client.py index add04974..81bad177 100644 --- a/transfer_queue/storage/clients/ray_storage_client.py +++ b/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: diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py index 1278dfd4..40574413 100644 --- a/transfer_queue/utils/common.py +++ b/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) diff --git a/tutorial/06_streaming_dataloader.py b/tutorial/06_streaming_dataloader.py index 0571e108..5b5caee5 100644 --- a/tutorial/06_streaming_dataloader.py +++ b/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(