[BREAKING][Feat] Support managing multiple data partitions for Train/Val/Test in controller by LLLLxmmm · Pull Request #98 · TransferQueue/TransferQueue · GitHub

No description provided.

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

…oller

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

✨ Finishing touches 🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

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

Why not make get_metadata a private method?

Why not make scan_data_status a private method?

Why not make generate_batch_meta a private method?

Maybe we need to support put to & get from multiple partitions.
We can simply point out that currently we support only one partition and may consider supporting it in the future

Suggested change

      
                        partition_id: Current data partition id (required if metadata is not provided)
          
                        partition_id: Target data partition id (required if metadata is not provided)

better to raise an error. assert might be ignored

Suggested change

      
                async def check_current_step_consumption(self, task_name: str, partition_id: str):
          
                async def check_data_consumption_status(self, task_name: str, partition_id: str):

Suggested change

      
                async def check_current_step_production(self, data_fields: list[str], partition_id: str):
          
                async def check_data_production_status(self, data_fields: list[str], partition_id: str):

Suggested change

      
                        partition_id: Current data partition id (required if metadata is not provided)
          
                        partition_id: Target data partition id (required if metadata is not provided)

Suggested change

      
                        partition_id: Current data partition id
          
                        partition_id: Target data partition

Put this near TQ_INIT_FIELD_NUM ~

why allocated_indexes is tracking reusable_indexes?

maybe we can just release them during release_indexes

need to tackle count <=0 by rasing an error

can be optimized. using while loop and append list is slow

why?

If not used, we can simply delete this func

sample_idx -> global_idx？

if not used we can delete this func

Suggested change

      
                def total_fields(self) -> int:
          
                def total_fields_num(self) -> int:

Suggested change

      
                def total_samples(self) -> int:
          
                def total_samples_num(self) -> int:

better do this during init process

we don't need a return value. same for field expansion

add TODO

Suggested change

      
                    sample_indices: list[int],
          
                    global_indices: list[int],

the d type and shape will be always required

because many storage backends need these info to speed up transport

add TODO

extend consumption status should not be here (in get_consumption_status)? When we update data status, we should extend the consumption status as well

Just raise an error. same for other functions

add TODO

When merge with Sampler design, maybe we can just pass the production_status and consumption_status to the Sampler.sample function.

need more thinking here

暂时不用改

if it is not necessary, we can delete it

need to optimize. Now it will be very slow

add TODO

0oshowero0 and others added 2 commits October 28, 2025 20:20

provide draft for controller refactor

f9db83e

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

Support managing multiple data partitions for Train/Val/Test in contr…

c388140

…oller

liuximeng added 6 commits October 31, 2025 14:57

modify UT

ef19735

fix partition UT

5174fdb

fix async simple storage manager UT

626b631

fix codecheck

6987bb3

Complete the Controller's UT and improve the docstrings

79e47e0

fix codecheck

56496bc

jianjunzhong reviewed Nov 3, 2025

View reviewed changes

liuximeng added 3 commits November 3, 2025 17:32

fix CLEAR_META_RESPONSE bug

f68d1be

Enable all demo and UT

d6fa381

Remove invalid TODOs

8a3c327

0oshowero0 reviewed Nov 3, 2025

View reviewed changes

0oshowero0 reviewed Nov 4, 2025

View reviewed changes

liuximeng added 3 commits November 4, 2025 21:11

fix comments

22218df

fix codecheck

f86354d

delete unused code

36088dc

0oshowero0 changed the title Support managing multiple data partitions for Train/Val/Test in controller [BREAKING] Support managing multiple data partitions for Train/Val/Test in controller Nov 4, 2025

0oshowero0 changed the title [BREAKING] Support managing multiple data partitions for Train/Val/Test in controller [BREAKING][Feat] Support managing multiple data partitions for Train/Val/Test in controller Nov 4, 2025

0oshowero0 merged commit 445edb4 into TransferQueue:dev Nov 4, 2025
3 checks passed

0oshowero0 mentioned this pull request Nov 10, 2025

[Feature] Support managing multiple data partitions for Train/Val/Test in controller #81

Closed

2 tasks

0oshowero0 mentioned this pull request Nov 18, 2025

[recipe, data] feat: TransferQueue - Support managing multiple data partitions for Train/Val/Test in controller verl-project/verl#4175

Merged

7 tasks

0oshowero0 mentioned this pull request Dec 11, 2025

[Feature] Support data initializing in multiple chunks #80

Closed

0oshowero0 mentioned this pull request Dec 18, 2025

[BREAKING][Fix] Fix controller concurrency error when expanding production/consumption status #145

Merged

                           metadata: Records the metadata of a batch of data samples, containing index and
                                     storage unit information. If None, metadata will be auto-generated.
-                          global_step: Current processing step (required if metadata is not provided)
+                          partition_id: Current data partition id (required if metadata is not provided)

                   @dynamic_socket(socket_name="request_handle_socket")
-                  async def check_current_step_production(self, data_fields: list[str], global_step: int):
-                      """Check if all samples for current step are ready for consumption.
+                  async def check_current_step_production(self, data_fields: list[str], partition_id: str):

                           data: Data to write as TensorDict
                           metadata: Optional metadata containing index and storage unit information
-                          global_step: Current processing step (required if metadata is not provided)
+                          partition_id: Current data partition id (required if metadata is not provided)

                           data_fields: List of data field names to retrieve metadata for
                           batch_size: Number of samples to request in the batch
-                          global_step: Current training/processing step
+                          partition_id: Current data partition id

+              TQ_FIELD_MIN_EXPANSION_SIZE = int(
+                  os.environ.get("TQ_FIELD_MIN_EXPANSION_SIZE", 5)
+              )  # Minimum expansion size for fields (columns)
+              TQ_INIT_SAMPLE_NUM = int(os.environ.get("TQ_INIT_SAMPLE_NUM", 10))  # Initial number of samples

+                      return partition.generate_data_status_mask(field_names, task_name, sample_filter)
-                  def _get_metadata(
+                  def get_metadata(

-                  def _scan_data_status(
-                      self, data_fields: list[str], global_step: int, task_name: str, get_n_samples: bool
+                  def scan_data_status(

-                  def _generate_batch_meta(
-                      self, global_step: int, global_indexes: list[int], data_fields: list[str], mode: str
+                  def generate_batch_meta(

                               per_field_shapes[global_idx][field] = data_item.shape if hasattr(data_item, "shape") else None
+                      # Get current data partition id
+                      partition_id = metadata.samples[0].partition_id

	partition_id: Current data partition id (required if metadata is not provided)
	partition_id: Target data partition id (required if metadata is not provided)

                       if metadata is None:
-                          assert global_step is not None, "global_steps must be provided if metadata is not given"
+                          assert partition_id is not None, "partition_id must be provided if metadata is not given"

	async def check_current_step_consumption(self, task_name: str, partition_id: str):
	async def check_data_consumption_status(self, task_name: str, partition_id: str):

	async def check_current_step_production(self, data_fields: list[str], partition_id: str):
	async def check_data_production_status(self, data_fields: list[str], partition_id: str):

	partition_id: Current data partition id
	partition_id: Target data partition

+                      indexes = []
+                      # Get indexes from reusable pool
+                      if self.reusable_indexes and count > 0:

		for idx in indexes:
		self.allocated_indexes.discard(idx)

		# Add released indexes to allocated_indexes set
		self.allocated_indexes.update(indexes)


		# ==================== Data Scanning and Query Methods ====================

		def scan_data_status(


		return ready_sample_indices

		def generate_data_status_mask(

+                          while len(new_indexes) < needed:
+                              # Check if current counter points to an already used index
+                              if self.global_index_counter not in self.allocated_indexes:
+                                  new_indexes.append(self.global_index_counter)
+                                  self.allocated_indexes.add(self.global_index_counter)
+                                  self.global_index_counter += 1
+                              else:
+                                  # If already used, increment counter until finding available index
+                                  self.global_index_counter += 1
+                          indexes.extend(new_indexes)

+                      """
+                      return self.partition_to_indexes.get(partition_id, set()).copy()
+                  def get_allocated_indexes(self):

+                  # Field metadata
+                  field_name_mapping: dict[str, int] = field(default_factory=dict)  # field_name -> column_index
+                  field_dtypes: dict[int, dict[str, Any]] = field(default_factory=dict)  # sample_idx -> {field: dtype}

+                      return len(self.field_name_mapping)
+                  @property
+                  def allocated_fields(self) -> int:

	def total_fields(self) -> int:
	def total_fields_num(self) -> int:

	def total_samples(self) -> int:
	def total_samples_num(self) -> int:

+                      Returns:
+                          True if expansion was successful or not needed, False on error
+                      """
+                      if self.production_status is None:

-                      num_n_samples: int = 1,
-                  ) -> None:
-                      """Initialize the TransferQueueController.
+                      sample_indices: list[int],

+                      # Ensure consumption tensor has same number of rows as production tensor
+                      consumption_tensor = self.consumption_status[task_name]
+                      if self.production_status is not None and consumption_tensor.shape[0] < self.total_samples:
+                          expanded_consumption = torch.zeros(self.total_samples, dtype=torch.int8)
+                          expanded_consumption[: consumption_tensor.shape[0]] = consumption_tensor
+                          self.consumption_status[task_name] = expanded_consumption

+                          sample_indices: List of sample indices to mark as consumed
+                      Returns:
+                          True if successful, False on error

[BREAKING][Feat] Support managing multiple data partitions for Train/Val/Test in controller by LLLLxmmm · Pull Request #98 · TransferQueue/TransferQueue · GitHub

Uh oh!

[BREAKING][Feat] Support managing multiple data partitions for Train/Val/Test in controller - #98

[BREAKING][Feat] Support managing multiple data partitions for Train/Val/Test in controller#980oshowero0 merged 14 commits intoTransferQueue:devTransferQueue/TransferQueue:devfrom LLLLxmmm:lxm_controller_refactorLLLLxmmm/TransferQueue:lxm_controller_refactorCopy head branch name to clipboard

Conversation

LLLLxmmm commented Oct 31, 2025

Uh oh!

coderabbitai Bot commented Oct 31, 2025

Review skipped

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

0oshowero0 Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

0oshowero0 Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

[BREAKING][Feat] Support managing multiple data partitions for Train/Val/Test in controller#98
0oshowero0 merged 14 commits into
TransferQueue:devfrom
LLLLxmmm:lxm_controller_refactor

0oshowero0 Nov 3, 2025 •
edited
Loading

0oshowero0 Nov 3, 2025 •
edited
Loading

+                              return []
+                      # Create row mask
+                      if sample_filter is not None:

+                          logger.error(f"Error updating production status for partition {self.partition_id}: {e}")
+                          return False
+                  def _update_field_metadata(