From 841bacce21f6f7e4d0e184a536fd1e432bb4fa6d Mon Sep 17 00:00:00 2001 From: HSYZhang <2908064130@qq.com> Date: Tue, 10 Nov 2026 12:06:33 +0800 Subject: [PATCH 1/2] Fix controller fetch blocking data updates --- README.md | 4 + scripts/test_controller_fetch_progress.sh | 7 ++ tests/test_controller_fetch_progress.py | 93 +++++++++++++++++++++++ transfer_queue/client.py | 53 ++++++++----- transfer_queue/controller.py | 35 ++++++--- 5 files changed, 162 insertions(+), 30 deletions(-) create mode 100755 scripts/test_controller_fetch_progress.sh create mode 100644 tests/test_controller_fetch_progress.py diff --git a/README.md b/README.md index d7785108..62543733 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,10 @@ See [tutorial/06_streaming_dataloader.py](tutorial/06_streaming_dataloader.py) f #### Low-Level Native API The native interfaces of TransferQueue are implemented in `TransferQueueClient`. It offers maximum flexibility through native, atomic operations. +For `get_meta(mode="fetch")`, the client retries until `TQ_CONTROLLER_GET_METADATA_TIMEOUT` +when data is insufficient and `polling_mode` is false. With `polling_mode` true, it +returns an empty `BatchMeta` immediately. Waiting happens in the client, so the +controller can continue processing producers' data-update notifications. Developers can leverage `TransferQueueClient` directly to implement advanced features that require fine-grained control and fully streamed data scheduling, as illustrated in the following tutorials: diff --git a/scripts/test_controller_fetch_progress.sh b/scripts/test_controller_fetch_progress.sh new file mode 100755 index 00000000..2599849d --- /dev/null +++ b/scripts/test_controller_fetch_progress.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +set -euo pipefail + +cd "$(dirname "$0")/.." +python_bin="${PYTHON:-python3}" +"$python_bin" -m compileall -q transfer_queue tutorial tests +"$python_bin" -m pytest -q tests/test_controller_fetch_progress.py diff --git a/tests/test_controller_fetch_progress.py b/tests/test_controller_fetch_progress.py new file mode 100644 index 00000000..8dd684bd --- /dev/null +++ b/tests/test_controller_fetch_progress.py @@ -0,0 +1,93 @@ +import asyncio +import time +from uuid import uuid4 + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from tensordict import TensorDict + +import transfer_queue as tq +import transfer_queue.client as client_module + + +@pytest.mark.asyncio +async def test_waiting_fetch_allows_put_and_notification(monkeypatch): + namespace = f"tq-fetch-progress-{uuid4().hex}" + monkeypatch.setenv("TQ_RAY_NAMESPACE", namespace) + ray.init(namespace=namespace, num_cpus=3, include_dashboard=False) + try: + tq.init( + OmegaConf.create( + { + "controller": {"polling_mode": False}, + "backend": {"SimpleStorage": {"num_data_storage_units": 1}}, + } + ) + ) + client = tq.get_client() + client.controller_rpc_pool._maxsize = 1 + monkeypatch.setattr(client_module, "TQ_CONTROLLER_GET_METADATA_TIMEOUT", 4) + + reader = asyncio.create_task( + asyncio.to_thread( + client.get_meta, + data_fields=["input_ids"], + batch_size=1, + partition_id="waiting_fetch", + task_name="trainer", + ) + ) + await asyncio.sleep(0.2) + data = TensorDict({"input_ids": torch.tensor([[1, 2]])}, batch_size=[1]) + await asyncio.wait_for(asyncio.to_thread(client.put, data, partition_id="waiting_fetch"), timeout=3) + metadata = await asyncio.wait_for(reader, timeout=5) + result = await asyncio.to_thread(client.get_data, metadata) + assert metadata.is_ready + assert torch.equal(result["input_ids"][0], data["input_ids"][0]) + + monkeypatch.setattr(client_module, "TQ_CONTROLLER_GET_METADATA_TIMEOUT", 0.25) + started = time.monotonic() + with pytest.raises(TimeoutError, match="Timeout while waiting for sufficient data"): + await asyncio.to_thread( + client.get_meta, + data_fields=["input_ids"], + batch_size=1, + partition_id="never_produced", + task_name="trainer", + ) + assert time.monotonic() - started < 2 + finally: + tq.close() + ray.shutdown() + + +@pytest.mark.asyncio +async def test_polling_mode_returns_empty_immediately(monkeypatch): + namespace = f"tq-polling-{uuid4().hex}" + monkeypatch.setenv("TQ_RAY_NAMESPACE", namespace) + ray.init(namespace=namespace, num_cpus=3, include_dashboard=False) + try: + tq.init( + OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": {"SimpleStorage": {"num_data_storage_units": 1}}, + } + ) + ) + metadata = await asyncio.wait_for( + asyncio.to_thread( + tq.get_client().get_meta, + data_fields=["input_ids"], + batch_size=1, + partition_id="empty", + task_name="trainer", + ), + timeout=2, + ) + assert metadata.size == 0 + finally: + tq.close() + ray.shutdown() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index d44a0544..66ba23f3 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -16,6 +16,7 @@ import asyncio import os import threading +import time import weakref from typing import Any, Callable @@ -24,7 +25,11 @@ import zmq.asyncio from tensordict import TensorDict -from transfer_queue.controller import PUBLIC_PARTITION +from transfer_queue.controller import ( + PUBLIC_PARTITION, + TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL, + TQ_CONTROLLER_GET_METADATA_TIMEOUT, +) from transfer_queue.metadata import BatchMeta from transfer_queue.storage import StorageManagerFactory from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager @@ -239,7 +244,6 @@ async def _request_controller( return response_msg # ==================== Basic API ==================== - @with_controller_socket async def async_get_meta( self, data_fields: list[str], @@ -248,7 +252,6 @@ async def async_get_meta( mode: str = "fetch", task_name: str | None = None, sampling_config: dict[str, Any] | None = None, - socket: zmq.asyncio.Socket | None = None, ) -> BatchMeta: """Asynchronously fetch data metadata from the controller via ZMQ. @@ -260,9 +263,10 @@ async def async_get_meta( - 'fetch': Get ready data only - 'force_fetch': Get data regardless of readiness (may return unready samples) - 'insert': Internal usage - should not be used by users + In fetch mode, insufficient data returns an empty batch immediately when + polling_mode is enabled; otherwise this client retries until the configured timeout. task_name: Optional task name associated with the request sampling_config: Optional sampling configuration for custom samplers. - socket: ZMQ async socket for message transmission (injected by decorator) Returns: BatchMeta: Metadata object containing data structure, sample information, and readiness status @@ -301,20 +305,33 @@ async def async_get_meta( ... ) ... print(batch_meta.is_ready) # May be False if some samples not ready """ - response_msg = await self._request_controller( - socket=socket, - request_type=ZMQRequestType.GET_META, - response_type=ZMQRequestType.GET_META_RESPONSE, - body={ - "data_fields": data_fields, - "batch_size": batch_size, - "partition_id": partition_id, - "mode": mode, - "task_name": task_name, - "sampling_config": sampling_config, - }, - ) - return response_msg.body["metadata"] + request = { + "data_fields": data_fields, + "batch_size": batch_size, + "partition_id": partition_id, + "mode": mode, + "task_name": task_name, + "sampling_config": sampling_config, + } + deadline = time.monotonic() + TQ_CONTROLLER_GET_METADATA_TIMEOUT + retry_interval = min(TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL, TQ_CONTROLLER_GET_METADATA_TIMEOUT / 4) + while True: + async with self.controller_rpc_pool.alease(self._controller) as socket: + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.GET_META, + response_type=ZMQRequestType.GET_META_RESPONSE, + body=request, + ) + if not response_msg.body.get("retry", False): + return response_msg.body["metadata"] + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"Timeout while waiting for sufficient data for task {task_name}. " + f"Required: {batch_size}, partition: {partition_id}, fields: {data_fields}" + ) + await asyncio.sleep(min(retry_interval, remaining)) @with_controller_socket async def async_set_custom_meta( diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index d500c027..053bf1ec 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1398,6 +1398,7 @@ def get_metadata( batch_size: int | None = None, sampling_config: dict[str, Any] | None = None, *args, + wait_for_data: bool = True, **kwargs, ) -> BatchMeta: """ @@ -1413,11 +1414,14 @@ def get_metadata( (excludes already consumed samples) task_name: Name of the consumer task (required for fetch modes) batch_size: Number of samples to retrieve + wait_for_data: Whether to wait for enough samples in fetch mode. The ZMQ + request loop disables this so it can process data updates. *args: Additional positional arguments **kwargs: Additional keyword arguments Returns: - BatchMeta object containing the requested metadata + BatchMeta object containing the requested metadata, or an empty batch when + wait_for_data is disabled and insufficient data is ready Raises: TimeoutError: If waiting for sufficient data times out in fetch mode @@ -1458,7 +1462,7 @@ def get_metadata( if batch_size is None: raise ValueError("must provide batch_size in fetch mode") - start_time = time.time() + deadline = time.monotonic() + TQ_CONTROLLER_GET_METADATA_TIMEOUT while True: # ready_for_consume_indexes: samples where all required fields are produced # (production status is ready) and not yet consumed @@ -1477,19 +1481,22 @@ def get_metadata( f" Returning None due to polling mode." ) return BatchMeta.empty() - else: - logger.warning( - f"[{self.controller_id}]: Insufficient data for task {task_name}. Required: {batch_size} " - f"samples with fields {data_fields} in partition {partition_id}, but only have " - f"{len(ready_for_consume_indexes)} samples meeting the criteria. " - f"Retrying in {TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL}s..." - ) - time.sleep(TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL) - if time.time() - start_time > TQ_CONTROLLER_GET_METADATA_TIMEOUT: + if not wait_for_data: + return BatchMeta.empty() + remaining = deadline - time.monotonic() + if remaining <= 0: raise TimeoutError( f"Timeout while waiting for sufficient data for task {task_name}. " f"Required: {batch_size}, Available: {len(ready_for_consume_indexes)}" ) + delay = min(TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL, remaining) + logger.warning( + f"[{self.controller_id}]: Insufficient data for task {task_name}. Required: {batch_size} " + f"samples with fields {data_fields} in partition {partition_id}, but only have " + f"{len(ready_for_consume_indexes)} samples meeting the criteria. " + f"Retrying in {delay:.2f}s..." + ) + time.sleep(delay) else: break @@ -2541,8 +2548,12 @@ def _handle_get_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage: mode=params.get("mode", "fetch"), task_name=params.get("task_name"), sampling_config=params.get("sampling_config", {}), + wait_for_data=False, + ) + retry = params.get("mode", "fetch") == "fetch" and not self.polling_mode and metadata.size == 0 + return self._make_response( + request_msg, ZMQRequestType.GET_META_RESPONSE, {"metadata": metadata, "retry": retry} ) - return self._make_response(request_msg, ZMQRequestType.GET_META_RESPONSE, {"metadata": metadata}) def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMessage: params = request_msg.body From 5223769dfa1fa091e22a4f77e982a10d6195b60c Mon Sep 17 00:00:00 2001 From: "zhanglong.443" Date: Tue, 10 Nov 2026 20:21:22 +0800 Subject: [PATCH 2/2] Respect metadata fetch retry interval and deadline --- README.md | 4 ---- tests/test_controller_fetch_progress.py | 12 ++++++++++++ transfer_queue/client.py | 14 +++++++------- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 62543733..d7785108 100644 --- a/README.md +++ b/README.md @@ -148,10 +148,6 @@ See [tutorial/06_streaming_dataloader.py](tutorial/06_streaming_dataloader.py) f #### Low-Level Native API The native interfaces of TransferQueue are implemented in `TransferQueueClient`. It offers maximum flexibility through native, atomic operations. -For `get_meta(mode="fetch")`, the client retries until `TQ_CONTROLLER_GET_METADATA_TIMEOUT` -when data is insufficient and `polling_mode` is false. With `polling_mode` true, it -returns an empty `BatchMeta` immediately. Waiting happens in the client, so the -controller can continue processing producers' data-update notifications. Developers can leverage `TransferQueueClient` directly to implement advanced features that require fine-grained control and fully streamed data scheduling, as illustrated in the following tutorials: diff --git a/tests/test_controller_fetch_progress.py b/tests/test_controller_fetch_progress.py index 8dd684bd..00b2f4d0 100644 --- a/tests/test_controller_fetch_progress.py +++ b/tests/test_controller_fetch_progress.py @@ -29,6 +29,7 @@ async def test_waiting_fetch_allows_put_and_notification(monkeypatch): client = tq.get_client() client.controller_rpc_pool._maxsize = 1 monkeypatch.setattr(client_module, "TQ_CONTROLLER_GET_METADATA_TIMEOUT", 4) + monkeypatch.setattr(client_module, "TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 1) reader = asyncio.create_task( asyncio.to_thread( @@ -48,6 +49,16 @@ async def test_waiting_fetch_allows_put_and_notification(monkeypatch): assert torch.equal(result["input_ids"][0], data["input_ids"][0]) monkeypatch.setattr(client_module, "TQ_CONTROLLER_GET_METADATA_TIMEOUT", 0.25) + monkeypatch.setattr(client_module, "TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 5) + request_count = 0 + original_request = client._request_controller + + async def count_requests(*args, **kwargs): + nonlocal request_count + request_count += 1 + return await original_request(*args, **kwargs) + + monkeypatch.setattr(client, "_request_controller", count_requests) started = time.monotonic() with pytest.raises(TimeoutError, match="Timeout while waiting for sufficient data"): await asyncio.to_thread( @@ -58,6 +69,7 @@ async def test_waiting_fetch_allows_put_and_notification(monkeypatch): task_name="trainer", ) assert time.monotonic() - started < 2 + assert request_count == 1 finally: tq.close() ray.shutdown() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 66ba23f3..22a5c21d 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -314,8 +314,12 @@ async def async_get_meta( "sampling_config": sampling_config, } deadline = time.monotonic() + TQ_CONTROLLER_GET_METADATA_TIMEOUT - retry_interval = min(TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL, TQ_CONTROLLER_GET_METADATA_TIMEOUT / 4) while True: + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timeout while waiting for sufficient data for task {task_name}. " + f"Required: {batch_size}, partition: {partition_id}, fields: {data_fields}" + ) async with self.controller_rpc_pool.alease(self._controller) as socket: response_msg = await self._request_controller( socket=socket, @@ -326,12 +330,8 @@ async def async_get_meta( if not response_msg.body.get("retry", False): return response_msg.body["metadata"] remaining = deadline - time.monotonic() - if remaining <= 0: - raise TimeoutError( - f"Timeout while waiting for sufficient data for task {task_name}. " - f"Required: {batch_size}, partition: {partition_id}, fields: {data_fields}" - ) - await asyncio.sleep(min(retry_interval, remaining)) + if remaining > 0: + await asyncio.sleep(min(TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL, remaining)) @with_controller_socket async def async_set_custom_meta(