[fix] metadata fetch blocking controller data-update notifications

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

tests/test_controller_fetch_progress.py

@@ -0,0 +1,105 @@
+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)
+        monkeypatch.setattr(client_module, "TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 1)
+
+        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)
+        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(
+                client.get_meta,
+                data_fields=["input_ids"],
+                batch_size=1,
+                partition_id="never_produced",
+                task_name="trainer",
+            )
+        assert time.monotonic() - started < 2
+        assert request_count == 1
+    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()

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
+        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,
+                    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:
+                await asyncio.sleep(min(TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL, remaining))
 
     @with_controller_socket
     async def async_set_custom_meta(

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