[StreamingDataLoader, 4/N] feat: Introduce sample pre-allocation for dynamic streaming (#16)

## Background

In PR #9, we introduced initial support for the `StreamingDataLoader`
interface. Currently, the system assumes prompts are pre-loaded into the
TransferQueue. However, a critical use case involves generation workers
put both prompts and responses into `TransferQueue` on the run (e.g.,
`rollout_buffer` mechanism in
[Slime](https://github.com/THUDM/slime/blob/main/slime_plugins/rollout_buffer/README.md)).

Since TransferQueue supports dynamic expansion, if the producer has not
yet pushed any data to the TransferQueue, the TransferQueue appears
empty. Consequently, the consumer's `check_consumption_status` API
incorrectly assumes no data is available and prematurely terminates the
data retrieval iteration.

## Solution

This PR introduces a new environment variable,
`TQ_PRE_ALLOC_SAMPLE_NUM`, to handle sample pre-allocation in
TransferQueue.

- **Mechanism**: When set (typically to `global_batch_size`), the
controller pre-allocates a fixed number of global indexes before data
production begins.
- **Effect**: The `check_consumption_status` API now accounts for these
pre-allocated slots. This ensures the `StreamingDataLoader` waits for
the pending data instead of exiting immediately when the TransferQueue
is temporarily empty.

## Other Changes
Deprecate `TQ_INIT_SAMPLE_NUM`, `TQ_INIT_FIELD_NUM`,
`TQ_SAMPLE_MIN_EXPANSION_SIZE` and `TQ_SAMPLE_MIN_EXPANSION_SIZE` for
simplicity.

---

CC: @NINGBENZHE

---------

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

tests/test_controller_data_partitions.py

@@ -653,3 +653,172 @@ def test_get_consumption_status_parameter():
     print("✓ get_consumption_status with mask works")
 
     print("Consumption status mask parameter tests passed!\n")
+
+
+def test_pre_allocated_indexes_basic():
+    """Test basic pre-allocated indexes functionality in DataPartitionStatus."""
+    from transfer_queue.controller import DataPartitionStatus
+
+    print("Testing pre-allocated indexes basic functionality...")
+
+    partition = DataPartitionStatus(partition_id="prealloc_test")
+
+    # Initially, pre_allocated_global_indexes should be empty
+    assert len(partition.pre_allocated_global_indexes) == 0
+    assert partition.total_samples_num == 0
+
+    print("✓ Initial state correct")
+
+    # Register pre-allocated indexes
+    pre_allocated = [0, 1, 2, 3, 4]
+    partition.register_pre_allocated_indexes(pre_allocated)
+
+    assert partition.pre_allocated_global_indexes == set(pre_allocated)
+    # global_indexes should still be empty until retrieved
+    assert partition.total_samples_num == 0
+
+    print("✓ Pre-allocated indexes registered")
+
+    # activate pre-allocated indexes
+    retrieved = partition.activate_pre_allocated_indexes(3)
+
+    assert len(retrieved) == 3
+    assert set(retrieved) == {0, 1, 2}
+    assert partition.global_indexes == {0, 1, 2}
+    assert partition.pre_allocated_global_indexes == {3, 4}
+    assert partition.total_samples_num == 3
+
+    print("✓ Pre-allocated indexes activate & retrieved correctly")
+
+    # Activate remaining indexes
+    retrieved = partition.activate_pre_allocated_indexes(5)
+
+    assert len(retrieved) == 2  # Only 2 remaining
+    assert set(retrieved) == {3, 4}
+    assert partition.global_indexes == {0, 1, 2, 3, 4}
+    assert partition.pre_allocated_global_indexes == set()
+    assert partition.total_samples_num == 5
+
+    print("✓ All pre-allocated indexes retrieved")
+
+    print("Pre-allocated indexes basic tests passed!\n")
+
+
+def test_pre_allocated_indexes_consumption_status():
+    """Test that pre-allocated indexes are included in consumption status."""
+    import torch
+
+    from transfer_queue.controller import DataPartitionStatus
+
+    print("Testing pre-allocated indexes in consumption status...")
+
+    partition = DataPartitionStatus(partition_id="consumption_test")
+
+    # Register pre-allocated indexes
+    partition.register_pre_allocated_indexes([0, 1, 2, 3, 4])
+
+    # Get consumption status - should include pre-allocated indexes
+    global_index, consumption_status = partition.get_consumption_status("test_task", mask=True)
+
+    # global_index should include all pre-allocated indexes
+    assert torch.equal(global_index, torch.tensor([0, 1, 2, 3, 4], dtype=torch.long))
+    # All consumption statuses should be 0 (not consumed yet)
+    assert torch.all(consumption_status == 0)
+
+    print("✓ Consumption status includes pre-allocated indexes")
+
+    # Mark some samples as consumed
+    partition.mark_consumed("test_task", [0, 2, 4])
+
+    # Get consumption status again
+    global_index, consumption_status = partition.get_consumption_status("test_task", mask=True)
+
+    assert consumption_status[0].item() == 1  # consumed
+    assert consumption_status[1].item() == 0  # not consumed
+    assert consumption_status[2].item() == 1  # consumed
+    assert consumption_status[3].item() == 0  # not consumed
+    assert consumption_status[4].item() == 1  # consumed
+
+    print("✓ Marked consumed works with pre-allocated indexes")
+
+    print("Pre-allocated indexes consumption status tests passed!\n")
+
+
+def test_pre_allocated_indexes_in_scan_data_status():
+    """Test that pre-allocated indexes affect scan_data_status behavior."""
+    from transfer_queue.controller import DataPartitionStatus
+
+    print("Testing pre-allocated indexes in scan_data_status...")
+
+    partition = DataPartitionStatus(partition_id="scan_test")
+
+    # Register pre-allocated indexes (5 samples)
+    partition.register_pre_allocated_indexes([0, 1, 2, 3, 4])
+
+    # Before any production, scan should return empty (no samples produced yet)
+    ready = partition.scan_data_status(field_names=["input_ids"], task_name="test_task")
+    assert ready == []
+
+    print("✓ Scan returns empty before production")
+
+    # Now produce some samples (0, 2, 4)
+    partition.update_production_status(
+        global_indices=[0, 2, 4],
+        field_names=["input_ids"],
+        dtypes={i: {"input_ids": "torch.int32"} for i in [0, 2, 4]},
+        shapes={i: {"input_ids": (32,)} for i in [0, 2, 4]},
+    )
+
+    # Scan should return produced and unconsumed samples
+    ready = partition.scan_data_status(field_names=["input_ids"], task_name="test_task")
+    assert set(ready) == {0, 2, 4}
+
+    print("✓ Scan returns produced samples correctly")
+
+    # Mark sample 2 as consumed
+    partition.mark_consumed("test_task", [2])
+
+    # Scan should now return only 0 and 4
+    ready = partition.scan_data_status(field_names=["input_ids"], task_name="test_task")
+    assert set(ready) == {0, 4}
+
+    print("✓ Scan respects consumption status")
+
+    print("Pre-allocated indexes scan_data_status tests passed!\n")
+
+
+def test_pre_allocated_indexes_mixed_with_dynamic():
+    """Test mixing pre-allocated indexes with dynamically allocated ones."""
+    from transfer_queue.controller import DataPartitionStatus
+
+    print("Testing mixed pre-allocated and dynamic indexes...")
+
+    partition = DataPartitionStatus(partition_id="mixed_test")
+
+    # Register 3 pre-allocated indexes
+    partition.register_pre_allocated_indexes([0, 1, 2])
+
+    # Simulate adding more samples (indexes 5, 6, 7)
+    # This would happen when producer calls update_production_status
+    partition.update_production_status(
+        global_indices=[5, 6, 7],
+        field_names=["input_ids"],
+        dtypes={i: {"input_ids": "torch.int32"} for i in [5, 6, 7]},
+        shapes={i: {"input_ids": (32,)} for i in [5, 6, 7]},
+    )
+
+    # Now global_indexes should only contain dynamically generated in (5,6,7)
+    assert partition.global_indexes == {5, 6, 7}
+    assert partition.total_samples_num == 3
+
+    # all pre-allocated
+    retrieved = partition.activate_pre_allocated_indexes(3)
+    assert set(retrieved) == {0, 1, 2}
+
+    # Now global_indexes should have both pre-allocated (0,1,2) and dynamic (5,6,7)
+    assert partition.global_indexes == {0, 1, 2, 5, 6, 7}
+    assert partition.total_samples_num == 6
+
+    print("✓ Mixed pre-allocated and dynamic indexes work correctly")
+
+    print("Mixed indexes tests passed!\n")