TransferQueue · 0oshowero0 · Dec 24, 2025 · Dec 24, 2025 · Dec 24, 2025 · Dec 24, 2025
diff --git a/recipe/simple_use_case/async_demo.py b/recipe/simple_use_case/async_demo.py
@@ -306,7 +306,7 @@ def fit(self):
 
                 # Client notifies controller to clear data status, controller returns metadata;
                 # Client then notifies the storage plane to clear based on metadata
-                asyncio.run(self.data_system_client.async_clear(partition_id=f"train_{step}"))
+                asyncio.run(self.data_system_client.async_clear_partition(partition_id=f"train_{step}"))
                 logger.info("clear ok! ")
         logger.info("demo done!")
 

diff --git a/recipe/simple_use_case/sync_demo.py b/recipe/simple_use_case/sync_demo.py
@@ -190,7 +190,7 @@ def fit(config, data_system_client):
             # Client then notifies the storage plane to clear based on metadata
             # Client selects one master controller to get metadata,
             # other controllers directly clear without returning metadata
-            data_system_client.clear(partition_id=f"train_{step}")
+            data_system_client.clear_partition(partition_id=f"train_{step}")
             logger.info("clear ok! ")
     logger.info("demo done!")
 

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -101,12 +101,16 @@ def _handle_requests(self):
                     if request_msg.request_type == ZMQRequestType.GET_META:
                         response_body = self._mock_batch_meta(request_msg.body)
                         response_type = ZMQRequestType.GET_META_RESPONSE
-                    elif request_msg.request_type == ZMQRequestType.GET_CLEAR_META:
-                        response_body = self._mock_batch_meta(request_msg.body)
-                        response_type = ZMQRequestType.GET_CLEAR_META_RESPONSE
                     elif request_msg.request_type == ZMQRequestType.CLEAR_META:
-                        response_body = {"message": "clear ok"}
+                        response_body = {"message": "clear meta ok"}
                         response_type = ZMQRequestType.CLEAR_META_RESPONSE
+                    elif request_msg.request_type == ZMQRequestType.CLEAR_PARTITION:
+                        response_body = {"message": "clear partition ok"}
+                        response_type = ZMQRequestType.CLEAR_PARTITION_RESPONSE
+                    elif request_msg.request_type == ZMQRequestType.GET_PARTITION_META:
+                        # Mock partition metadata response
+                        response_body = {"metadata": self._mock_batch_meta(request_msg.body)}
+                        response_type = ZMQRequestType.GET_PARTITION_META_RESPONSE
                     elif request_msg.request_type == ZMQRequestType.CHECK_CONSUMPTION:
                         # Mock consumption status check - all consumed
                         response_body = {
@@ -375,14 +379,6 @@ def test_get_meta(client_setup):
     assert len(metadata.global_indexes) == 10
 
 
-def test_clear_operation(client_setup):
-    """Test clear operation"""
-    client, _, _ = client_setup
-
-    # Test clear operation
-    client.clear(partition_id="0")
-
-
 # Test with single controller and multiple storage units
 def test_single_controller_multiple_storages():
     """Test client with single controller and multiple storage units"""
@@ -517,3 +513,72 @@ async def test_async_get_partition_list(client_setup):
     assert "partition_0" in partition_list
     assert "partition_1" in partition_list
     assert "test_partition" in partition_list
+
+
+# Test clear methods
+@pytest.mark.asyncio
+async def test_async_clear_partition(client_setup):
+    """Test async clear partition operation"""
+    client, _, _ = client_setup
+
+    # Test async_clear_partition
+    await client.async_clear_partition(partition_id="test_partition")
+
+    # If no exception is raised, the test passes
+    assert True
+
+
+@pytest.mark.asyncio
+async def test_async_clear_samples(client_setup):
+    """Test async clear samples operation"""
+    client, _, _ = client_setup
+
+    # First get metadata to create a BatchMeta object
+    metadata = await client.async_get_meta(data_fields=["tokens", "labels"], batch_size=2, partition_id="0")
+
+    # Test async_clear_samples
+    await client.async_clear_samples(metadata=metadata)
+
+    # If no exception is raised, the test passes
+    assert True
+
+
+def test_clear_partition(client_setup):
+    """Test synchronous clear partition operation"""
+    client, _, _ = client_setup
+
+    # Test synchronous clear_partition
+    client.clear_partition(partition_id="test_partition")
+
+    # If no exception is raised, the test passes
+    assert True
+
+
+def test_clear_samples(client_setup):
+    """Test synchronous clear samples operation"""
+    client, _, _ = client_setup
+
+    # First get metadata to create a BatchMeta object
+    metadata = client.get_meta(data_fields=["tokens", "labels"], batch_size=2, partition_id="0")
+
+    # Test synchronous clear_samples
+    client.clear_samples(metadata=metadata)
+
+    # If no exception is raised, the test passes
+    assert True
+
+
+@pytest.mark.asyncio
+async def test_async_clear_samples_with_empty_metadata(client_setup):
+    """Test async_clear_samples with empty BatchMeta"""
+    client, _, _ = client_setup
+
+    # Create empty BatchMeta
+    metadata = BatchMeta(samples=[])
+
+    # The clear operation should complete without raising an exception
+    # because the mock storage manager is configured to handle this
+    await client.async_clear_samples(metadata=metadata)
+
+    # If no exception is raised, the test passes
+    assert True
diff --git a/tests/test_controller.py b/tests/test_controller.py
@@ -173,13 +173,13 @@ def test_controller_with_single_partition(self, ray_setup):
         assert [sample.fields for sample in clear_meta.samples] == [{}] * (gbs * num_n_samples)
         print("✓ Clear metadata correct")
 
-        # Test clear
-        ray.get(tq_controller.clear.remote(partition_id))
+        # Test clear_partition
+        ray.get(tq_controller.clear_partition.remote(partition_id))
         partition = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
         partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id))
         assert partition_index_range == set()
         assert partition is None
-        print("✓ Clear correct")
+        print("✓ Clear partition correct")
 
     def test_controller_with_multi_partitions(self, ray_setup):
         gbs_1 = 8
@@ -284,7 +284,7 @@ def test_controller_with_multi_partitions(self, ray_setup):
         # Clear partition 1
         partition_index_range_1 = ray.get(tq_controller.get_partition_index_range.remote(partition_id_1))
         assert partition_index_range_1
-        ray.get(tq_controller.clear.remote(partition_id_1))
+        ray.get(tq_controller.clear_partition.remote(partition_id_1))
         partition_1_after_clear = ray.get(tq_controller.get_partition_snapshot.remote(partition_id_1))
         partition_index_range_1_after_clear = ray.get(tq_controller.get_partition_index_range.remote(partition_id_1))
 
@@ -320,3 +320,64 @@ def test_controller_with_multi_partitions(self, ray_setup):
         partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id_3))
         assert partition_index_range == set(list(range(32)) + list(range(48, 80)))
         print("✓ Correctly assign partition_3")
+
+    def test_controller_clear_meta(self, ray_setup):
+        """Test clear_meta functionality for individual samples"""
+        gbs = 4
+        num_n_samples = 2
+        partition_id = "test_clear_meta"
+
+        tq_controller = TransferQueueController.remote()
+
+        # Create metadata in insert mode
+        data_fields = ["prompt_ids", "attention_mask"]
+        metadata = ray.get(
+            tq_controller.get_metadata.remote(
+                data_fields=data_fields,
+                batch_size=gbs * num_n_samples,
+                partition_id=partition_id,
+                mode="insert",
+            )
+        )
+
+        assert metadata.global_indexes == list(range(gbs * num_n_samples))
+
+        # Update production status
+        dtypes = {k: {"prompt_ids": "torch.int64", "attention_mask": "torch.bool"} for k in metadata.global_indexes}
+        shapes = {k: {"prompt_ids": (32,), "attention_mask": (32,)} for k in metadata.global_indexes}
+        success = ray.get(
+            tq_controller.update_production_status.remote(
+                partition_id=partition_id,
+                global_indexes=metadata.global_indexes,
+                field_names=metadata.field_names,
+                dtypes=dtypes,
+                shapes=shapes,
+            )
+        )
+        assert success
+
+        # Get partition snapshot before clear
+        partition_before = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
+        assert partition_before is not None
+        assert len(partition_before.global_indexes) == gbs * num_n_samples
+        assert set(partition_before.global_indexes) == set(range(gbs * num_n_samples))
+
+        # Test clear_meta - clear first 4 samples (indexes 0-3)
+        global_indexes_to_clear = [0, 1, 2, 3, 6]
+        partition_ids_to_clear = [partition_id] * len(global_indexes_to_clear)
+
+        ray.get(
+            tq_controller.clear_meta.remote(
+                global_indexes=global_indexes_to_clear,
+                partition_ids=partition_ids_to_clear,
+            )
+        )
+
+        # Check that only the cleared samples are affected
+        partition_after = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
+        assert partition_after is not None
+
+        # Verify production status is cleared for the specified indexes
+        assert set(partition_after.global_indexes) == set([4, 5, 7])
+
+        print("✓ Clear meta correct")
diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py
@@ -314,8 +314,7 @@ def test_data_partition_status_advanced():
     initial_consumption_sum = sum(t.sum().item() for t in partition.consumption_status.values())
 
     # Clear only production data
-    success = partition.clear_data(list(range(4)), clear_consumption=False)
-    assert success
+    partition.clear_data(list(range(4)), clear_consumption=False)
     assert partition.production_status[:4, :].sum().item() == 0
 
     # Consumption data should remain