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..00b2f4d0 --- /dev/null +++ b/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() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index d44a0544..22a5c21d 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 + 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( 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