GeezerVison

Walkthrough

This pull request introduces a pluggable sampling system to the TransferQueue architecture. It replaces the boolean get_n_samples parameter with a sampling_config dictionary parameter across client and controller APIs. Three sampler implementations are added: a BaseSampler abstract base class, SequentialSampler, and GRPOGroupNSampler. The controller is updated to accept and utilize samplers during metadata retrieval operations.

Changes

Sequence Diagram

sequenceDiagram
    actor Client
    participant AsyncClient as AsyncTransferQueueClient
    participant Controller as TransferQueueController
    participant Sampler as BaseSampler

    Client->>AsyncClient: async_get_meta(data_fields, batch_size, partition_id, mode="fetch", sampling_config={...})
    AsyncClient->>Controller: GET_META request with sampling_config payload
    
    alt mode == "fetch"
        Controller->>Sampler: sample(ready_indexes, batch_size, **sampling_config)
        Sampler-->>Controller: (sampled_indexes, consumed_indexes)
        Controller->>Controller: generate_batch_meta(batch_global_indexes, consumed_indexes, ...)
    else mode == "insert"
        Controller->>Controller: generate_batch_meta(batch_global_indexes, consumed_indexes, ...)
    end
    
    Controller-->>AsyncClient: BatchMeta response
    AsyncClient-->>Client: BatchMeta

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

• Areas requiring extra attention:
  • transfer_queue/controller.py: Substantial refactor with new sampler initialization, sampling_config routing logic, and five new validation helpers; requires verification of state handling and error paths
  • transfer_queue/sampler/grpo_group_n_sampler.py: PyTorch tensor reshaping and grouping logic for n_samples_per_prompt; verify batch_size divisibility enforcement and index flattening correctness
  • transfer_queue/client.py: Signature changes across two public methods (async and sync); ensure backward compatibility implications are understood
  • Cross-file consistency: Verify that batch_global_indices → batch_global_indexes renaming is applied consistently throughout the controller and consumed_indexes is threaded correctly through all paths

Possibly related PRs

• [Misc] Refactor: a better storage folder structure #93: Overlapping edits in transfer_queue/controller.__init__ and transfer_queue/client.py for metadata path modifications.
• Merge dev 1029 #94: Both PRs modify client/controller metadata APIs by changing get_meta signatures and introduce sampler-related exports.

Poem

🐰 A sampler's quest through indexes neat,
From SequentialHops to GRPO's beat,
Config replaces flags of old,
Pluggable strategies, fresh and bold,
Metadata flows where samplers lead! 🎯

Pre-merge checks and finishing touches

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

@coderabbitai review