fix

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

Author: 0oshowero0

tests/test_controller.py

@@ -89,6 +89,7 @@ def test_controller_with_single_partition(self, ray_setup):
                 field_names=metadata.field_names,
                 dtypes=dtypes,
                 shapes=shapes,
+                custom_meta=None,
             )
         )
         assert success
@@ -97,13 +98,18 @@ def test_controller_with_single_partition(self, ray_setup):
         assert partition.production_status.size(0) == gbs * num_n_samples
 
         # Test for get production status
-        production_status = ray.get(
+        global_index, production_status = ray.get(
             tq_controller.get_production_status.remote(
                 partition_id=partition_id,
                 data_fields=data_fields,
             )
         )
-        assert production_status
+        # Verify global_index contains all expected indexes
+        assert torch.equal(global_index, torch.tensor(range(gbs * num_n_samples), dtype=torch.long))
+        # Verify all samples are produced for all fields (status should be 1)
+        expected_production_status = torch.ones(gbs * num_n_samples, len(metadata.field_names), dtype=torch.int8)
+        assert torch.equal(production_status, expected_production_status)
+        print("✓ Get production status returns correct global_index and production_status")
 
         # Total fields should match the number of fields we added
         assert partition.total_fields_num == len(data_fields)
@@ -126,14 +132,19 @@ def test_controller_with_single_partition(self, ray_setup):
 
         print(f"✓ Updated production status for partition {partition_id}")
 
-        # Test for get consumption status
-        consumption_status = ray.get(
+        # Test for get consumption status BEFORE consumption
+        global_index, consumption_status = ray.get(
             tq_controller.get_consumption_status.remote(
                 partition_id=partition_id,
                 task_name="generate_sequences",
             )
         )
-        assert torch.equal(consumption_status, torch.zeros(gbs * num_n_samples))
+        # Verify global_index
+        assert torch.equal(global_index, torch.tensor(range(gbs * num_n_samples), dtype=torch.long))
+        # Verify all samples are NOT consumed yet (status should be 0)
+        expected_consumption_status_before = torch.zeros(gbs * num_n_samples, dtype=torch.int8)
+        assert torch.equal(consumption_status, expected_consumption_status_before)
+        print("✓ Get consumption status returns correct global_index and status (before consumption)")
 
         # Test get metadate in fetch mode
         gen_meta = ray.get(
@@ -153,14 +164,19 @@ def test_controller_with_single_partition(self, ray_setup):
         assert torch.equal(partition.consumption_status["generate_sequences"], torch.ones(gbs * num_n_samples))
         print("✓ Get metadata in fetch mode correct")
 
-        # Test for get consumption status
-        consumption_status = ray.get(
+        # Test for get consumption status AFTER consumption
+        global_index, consumption_status = ray.get(
             tq_controller.get_consumption_status.remote(
                 partition_id=partition_id,
                 task_name="generate_sequences",
             )
         )
-        assert torch.equal(consumption_status, torch.ones(gbs * num_n_samples))
+        # Verify global_index
+        assert torch.equal(global_index, torch.tensor(range(gbs * num_n_samples), dtype=torch.long))
+        # Verify all samples are consumed (status should be 1)
+        expected_consumption_status_after = torch.ones(gbs * num_n_samples, dtype=torch.int8)
+        assert torch.equal(consumption_status, expected_consumption_status_after)
+        print("✓ Get consumption status returns correct global_index and status (after consumption)")
 
         # Test get clear meta
         clear_meta = ray.get(
@@ -222,6 +238,19 @@ def test_controller_with_multi_partitions(self, ray_setup):
         )
         assert success
 
+        # Verify get production status returns correct data
+        global_index_1, production_status_1 = ray.get(
+            tq_controller.get_production_status.remote(
+                partition_id=partition_id_1,
+                data_fields=data_fields,
+            )
+        )
+        expected_global_index_1 = torch.tensor(range(gbs_1 * num_n_samples_1), dtype=torch.long)
+        assert torch.equal(global_index_1, expected_global_index_1)
+        expected_production_status_1 = torch.ones(gbs_1 * num_n_samples_1, len(data_fields), dtype=torch.int8)
+        assert torch.equal(production_status_1, expected_production_status_1)
+        print("✓ Get production status for partition_1 returns correct global_index and status")
+
         # Test get metadate in fetch mode
         gen_meta = ray.get(
             tq_controller.get_metadata.remote(
@@ -234,6 +263,18 @@ def test_controller_with_multi_partitions(self, ray_setup):
         )
         assert gen_meta
 
+        # Verify get consumption status after fetch (samples should be consumed)
+        global_index_1_consumed, consumption_status_1 = ray.get(
+            tq_controller.get_consumption_status.remote(
+                partition_id=partition_id_1,
+                task_name="generate_sequences",
+            )
+        )
+        assert torch.equal(global_index_1_consumed, expected_global_index_1)
+        expected_consumption_status_1 = torch.ones(gbs_1 * num_n_samples_1, dtype=torch.int8)
+        assert torch.equal(consumption_status_1, expected_consumption_status_1)
+        print("✓ Get consumption status for partition_1 returns correct global_index and status (after fetch)")
+
         # Test get clear meta
         clear_meta = ray.get(
             tq_controller.get_metadata.remote(
@@ -282,6 +323,31 @@ def test_controller_with_multi_partitions(self, ray_setup):
         )
         assert success
 
+        # Verify get production status for partition_2
+        global_index_2, production_status_2 = ray.get(
+            tq_controller.get_production_status.remote(
+                partition_id=partition_id_2,
+                data_fields=data_fields,
+            )
+        )
+        expected_global_index_2 = torch.tensor(range(part1_index_range, part2_index_range + part1_index_range), dtype=torch.long)
+        assert torch.equal(global_index_2, expected_global_index_2)
+        expected_production_status_2 = torch.ones(part2_index_range, len(data_fields), dtype=torch.int8)
+        assert torch.equal(production_status_2, expected_production_status_2)
+        print("✓ Get production status for partition_2 returns correct global_index and status")
+
+        # Verify get consumption status for partition_2 (before consumption - should be all zeros)
+        global_index_2_consumed, consumption_status_2 = ray.get(
+            tq_controller.get_consumption_status.remote(
+                partition_id=partition_id_2,
+                task_name="generate_sequences",
+            )
+        )
+        assert torch.equal(global_index_2_consumed, expected_global_index_2)
+        expected_consumption_status_2 = torch.zeros(part2_index_range, dtype=torch.int8)
+        assert torch.equal(consumption_status_2, expected_consumption_status_2)
+        print("✓ Get consumption status for partition_2 returns correct global_index and status (before consumption)")
+
         # Clear partition 1
         partition_index_range_1 = ray.get(tq_controller.get_partition_index_range.remote(partition_id_1))
         assert partition_index_range_1

tests/test_controller_data_partitions.py

@@ -86,7 +86,7 @@ def test_data_partition_status():
     print("✓ Field metadata retrieval works")
 
     # Test consumption status
-    consumption_tensor = partition.get_consumption_status("test_task")
+    global_index, consumption_tensor = partition.get_consumption_status("test_task", )
     assert consumption_tensor is not None
     assert consumption_tensor.shape[0] == partition.total_samples_num
 
@@ -240,7 +240,7 @@ def test_data_partition_status_advanced():
 
     # Initial consumption tracking
     partition.mark_consumed(task_name, [0, 1])
-    initial_consumption = partition.get_consumption_status(task_name)
+    global_index, initial_consumption = partition.get_consumption_status(task_name)
     assert initial_consumption[0] == 1
     assert initial_consumption[1] == 1
 
@@ -258,7 +258,7 @@ def test_data_partition_status_advanced():
         12: {"field_d": (32,)},
     }
     partition.update_production_status([10, 11, 12], ["field_d"], dtypes, shapes, None)  # Triggers sample expansion
-    expanded_consumption = partition.get_consumption_status(task_name)
+    global_index, expanded_consumption = partition.get_consumption_status(task_name)
     assert expanded_consumption[0] == 1  # Preserved
     assert expanded_consumption[1] == 1  # Preserved
     assert expanded_consumption.shape[0] >= 13  # Expanded to accommodate new samples
@@ -358,7 +358,7 @@ def test_edge_cases_and_error_handling():
     # Test 3: Consumption status edge cases
     # Test consumption status creation before production status
     task_name = "early_task"
-    consumption_tensor = partition.get_consumption_status(task_name)
+    _, consumption_tensor = partition.get_consumption_status(task_name)
     assert consumption_tensor is not None
     assert consumption_tensor.shape[0] == partition.allocated_samples_num
 

transfer_queue/controller.py

@@ -367,7 +367,7 @@ def update_production_status(
             # Update production status
             if self.production_status is not None and global_indices and field_names:
                 field_indices = [self.field_name_mapping.get(field) for field in field_names]
-                self.production_status[Tensor(global_indices)[:, None], Tensor(field_indices)] = 1
+                self.production_status[torch.tensor(global_indices)[:, None], torch.tensor(field_indices)] = 1
 
             # Update field metadata
             self._update_field_metadata(global_indices, dtypes, shapes, custom_meta)
@@ -487,12 +487,12 @@ def get_consumption_status(self, task_name: str, mask:bool=False) -> tuple[Tenso
             else:
                 self.consumption_status[task_name] = torch.zeros(0, dtype=torch.int8)
 
+        # Get consumption status for requested task
         consumption_status = self.consumption_status[task_name]
 
+        partition_global_index = torch.tensor(sorted(self.global_indexes), dtype=torch.long)
+
         if mask:
-            # Get mask for target partition
-            partition_global_index = Tensor(sorted(self.global_indexes), dtype=torch.long)
-            # Get consumption status for requested task
             consumption_status = consumption_status[partition_global_index]
 
         return partition_global_index, consumption_status
@@ -527,10 +527,9 @@ def get_production_status_for_fields(self, field_names: list[str], mask:bool=Fal
 
         production_status = self.production_status[:, col_mask]
 
+        partition_global_index = torch.tensor(sorted(self.global_indexes), dtype=torch.long)
+
         if mask:
-            # Get mask for target partition
-            partition_global_index = Tensor(sorted(self.global_indexes), dtype=torch.long)
-            # Get production status for requested fields
             production_status = production_status[partition_global_index]
 
         return partition_global_index, production_status
@@ -561,7 +560,7 @@ def scan_data_status(self, field_names: list[str], task_name: str) -> list[int]:
             row_mask = torch.ones(self.allocated_samples_num, dtype=torch.bool)
 
             # Apply consumption filter (exclude already consumed samples)
-            consumption_status = self.get_consumption_status(task_name, mask=False)
+            _, consumption_status = self.get_consumption_status(task_name, mask=False)
             if consumption_status is not None:
                 unconsumed_mask = consumption_status == 0
                 row_mask &= unconsumed_mask