TransferQueue 0.1.4.dev2__py3-none-any.whl → 0.1.5.dev3__py3-none-any.whl

recipe/simple_use_case/async_demo.py

@@ -306,7 +306,7 @@ class Trainer:
 
                 # 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!")
 

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!")
 

tests/test_client.py

@@ -101,11 +101,38 @@ class MockController:
                     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 = {
+                            "partition_id": request_msg.body.get("partition_id"),
+                            "consumed": True,
+                        }
+                        response_type = ZMQRequestType.CONSUMPTION_RESPONSE
+                    elif request_msg.request_type == ZMQRequestType.CHECK_PRODUCTION:
+                        # Mock production status check - all produced
+                        response_body = {
+                            "partition_id": request_msg.body.get("partition_id"),
+                            "produced": True,
+                        }
+                        response_type = ZMQRequestType.PRODUCTION_RESPONSE
+                    elif request_msg.request_type == ZMQRequestType.GET_LIST_PARTITIONS:
+                        # Mock partition list
+                        response_body = {
+                            "partition_ids": ["partition_0", "partition_1", "test_partition"],
+                        }
+                        response_type = ZMQRequestType.LIST_PARTITIONS_RESPONSE
+                    else:
+                        response_body = {"error": f"Unknown request type: {request_msg.request_type}"}
                         response_type = ZMQRequestType.CLEAR_META_RESPONSE
 
                     # Send response
@@ -352,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"""
@@ -426,3 +445,140 @@ def test_put_without_required_params(client_setup):
     # Test put without partition id (should fail)
     with pytest.raises(ValueError):
         client.put(data=test_data)
+
+
+# Test new status checking methods
+def test_check_consumption_status(client_setup):
+    """Test consumption status checking"""
+    client, _, _ = client_setup
+
+    # Test synchronous check_consumption_status
+    is_consumed = client.check_consumption_status(task_name="generate_sequences", partition_id="train_0")
+    assert is_consumed is True
+
+
+def test_check_production_status(client_setup):
+    """Test production status checking"""
+    client, _, _ = client_setup
+
+    # Test synchronous check_production_status
+    is_produced = client.check_production_status(data_fields=["prompt_ids", "attention_mask"], partition_id="train_0")
+    assert is_produced is True
+
+
+def test_get_partition_list(client_setup):
+    """Test partition list retrieval"""
+    client, _, _ = client_setup
+
+    # Test synchronous get_partition_list
+    partition_list = client.get_partition_list()
+    assert isinstance(partition_list, list)
+    assert len(partition_list) > 0
+    assert "partition_0" in partition_list
+    assert "partition_1" in partition_list
+    assert "test_partition" in partition_list
+
+
+@pytest.mark.asyncio
+async def test_async_check_consumption_status(client_setup):
+    """Test async consumption status checking"""
+    client, _, _ = client_setup
+
+    # Test async_check_consumption_status
+    is_consumed = await client.async_check_consumption_status(task_name="generate_sequences", partition_id="train_0")
+    assert is_consumed is True
+
+
+@pytest.mark.asyncio
+async def test_async_check_production_status(client_setup):
+    """Test async production status checking"""
+    client, _, _ = client_setup
+
+    # Test async_check_production_status
+    is_produced = await client.async_check_production_status(
+        data_fields=["prompt_ids", "attention_mask"], partition_id="train_0"
+    )
+    assert is_produced is True
+
+
+@pytest.mark.asyncio
+async def test_async_get_partition_list(client_setup):
+    """Test async partition list retrieval"""
+    client, _, _ = client_setup
+
+    # Test async_get_partition_list
+    partition_list = await client.async_get_partition_list()
+    assert isinstance(partition_list, list)
+    assert len(partition_list) > 0
+    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

tests/test_controller.py

@@ -28,7 +28,6 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 from transfer_queue import TransferQueueController  # noqa: E402
-from transfer_queue.controller import TQ_INIT_FIELD_NUM  # noqa: E402
 from transfer_queue.utils.utils import ProductionStatus  # noqa: E402
 
 
@@ -92,21 +91,49 @@ class TestTransferQueueController:
             )
         )
         assert success
-        partition = ray.get(tq_controller.get_partition.remote(partition_id))
+        partition = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
         assert partition.production_status is not None
         assert partition.production_status.size(0) == gbs * num_n_samples
-        assert partition.production_status.size(1) == TQ_INIT_FIELD_NUM
+
+        # Test for get production status
+        production_status = ray.get(
+            tq_controller.get_production_status.remote(
+                partition_id=partition_id,
+                data_fields=data_fields,
+            )
+        )
+        assert production_status
+
+        # Total fields should match the number of fields we added
+        assert partition.total_fields_num == len(data_fields)
+
+        # Allocated fields should be at least the number of actual fields
+        assert partition.allocated_fields_num >= partition.total_fields_num
+
+        # Check production status for the fields we added
         assert torch.equal(
             sum(partition.production_status[:, : len(data_fields)]),
             torch.Tensor([gbs * num_n_samples, gbs * num_n_samples]),
         )
-        assert torch.equal(
-            sum(partition.production_status[:, len(data_fields) :]),
-            torch.zeros(1 * (TQ_INIT_FIELD_NUM - len(data_fields))),
-        )
+
+        # Any additional allocated fields should be zero (unused)
+        if partition.allocated_fields_num > len(data_fields):
+            assert torch.equal(
+                sum(partition.production_status[:, len(data_fields) :]),
+                torch.zeros(1 * (partition.allocated_fields_num - len(data_fields))),
+            )
 
         print(f"✓ Updated production status for partition {partition_id}")
 
+        # Test for get consumption status
+        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))
+
         # Test get metadate in fetch mode
         gen_meta = ray.get(
             tq_controller.get_metadata.remote(
@@ -117,13 +144,23 @@ class TestTransferQueueController:
                 task_name="generate_sequences",
             )
         )
+
         assert gen_meta.global_indexes == list(range(gbs * num_n_samples))
         assert gen_meta.samples[0].partition_id == "train_0"
         assert gen_meta.field_names == ["prompt_ids"]
-        partition = ray.get(tq_controller.get_partition.remote(partition_id))
+        partition = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
         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(
+            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))
+
         # Test get clear meta
         clear_meta = ray.get(
             tq_controller.get_metadata.remote(
@@ -136,14 +173,13 @@ class TestTransferQueueController:
         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))
-        partition = ray.get(tq_controller.get_partition.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 torch.all(partition.production_status == 0)
-        assert torch.all(partition.consumption_status["generate_sequences"] == 0)
-        print("✓ Clear correct")
+        assert partition is None
+        print("✓ Clear partition correct")
 
     def test_controller_with_multi_partitions(self, ray_setup):
         gbs_1 = 8
@@ -248,15 +284,15 @@ class TestTransferQueueController:
         # 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))
-        partition_1_after_clear = ray.get(tq_controller.get_partition.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))
 
         assert not partition_index_range_1_after_clear
-        assert torch.all(partition_1_after_clear.production_status[list(partition_index_range_1), :] == 0)
-        assert torch.all(partition_1_after_clear.consumption_status["generate_sequences"] == 0)
+        assert partition_1_after_clear is None
+        assert partition_index_range_1_after_clear == set()
 
-        partition_2 = ray.get(tq_controller.get_partition.remote(partition_id_2))
+        partition_2 = ray.get(tq_controller.get_partition_snapshot.remote(partition_id_2))
         partition_index_range_2 = ray.get(tq_controller.get_partition_index_range.remote(partition_id_2))
         assert partition_index_range_2 == set([32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47])
         assert torch.all(
@@ -284,3 +320,64 @@ class TestTransferQueueController:
         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")

tests/test_controller_data_partitions.py

@@ -26,8 +26,8 @@ sys.path.append(str(parent_dir))
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-TQ_INIT_SAMPLE_NUM = int(os.environ.get("TQ_INIT_SAMPLE_NUM", 10))  # Initial number of samples
-TQ_INIT_FIELD_NUM = int(os.environ.get("TQ_INIT_FIELD_NUM", 10))
+TQ_INIT_SAMPLE_NUM = int(os.environ.get("TQ_INIT_SAMPLE_NUM", 1))  # Initial number of samples
+TQ_INIT_FIELD_NUM = int(os.environ.get("TQ_INIT_FIELD_NUM", 1))
 
 
 def test_data_partition_status():
@@ -40,7 +40,8 @@ def test_data_partition_status():
     partition = DataPartitionStatus(partition_id="test@partition_1")
 
     # Test initial state
-    assert partition.total_samples_num == TQ_INIT_SAMPLE_NUM
+    assert partition.total_samples_num == 0
+    assert partition.allocated_samples_num == TQ_INIT_SAMPLE_NUM
     assert partition.total_fields_num == 0
     assert partition.allocated_fields_num == TQ_INIT_FIELD_NUM
     assert partition.production_status is not None
@@ -127,7 +128,7 @@ def test_partition_interface():
 
     # Test that the class can be imported and has expected methods
     assert hasattr(TransferQueueController, "create_partition")
-    assert hasattr(TransferQueueController, "get_partition")
+    assert hasattr(TransferQueueController, "get_partition_snapshot")
     assert hasattr(TransferQueueController, "update_production_status")
     assert hasattr(TransferQueueController, "scan_data_status")
     assert hasattr(TransferQueueController, "generate_batch_meta")
@@ -171,8 +172,8 @@ def test_dynamic_expansion_scenarios():
             10: {"field_1": (32,)},
         },
     )
-    assert partition.total_samples_num >= 11  # Should accommodate index 10
-
+    assert partition.total_samples_num == 3
+    assert partition.allocated_samples_num >= 11  # Should accommodate index 10
     print("✓ Large index gaps handled correctly")
 
     # Scenario 2: Adding many fields dynamically
@@ -212,7 +213,8 @@ def test_data_partition_status_advanced():
     partition = DataPartitionStatus(partition_id="advanced_test")
 
     # Initially empty
-    assert partition.total_samples_num == TQ_INIT_SAMPLE_NUM
+    assert partition.total_samples_num == 0
+    assert partition.allocated_samples_num == TQ_INIT_SAMPLE_NUM
     assert partition.total_fields_num == 0
     assert partition.allocated_fields_num == TQ_INIT_FIELD_NUM
 
@@ -289,6 +291,7 @@ def test_data_partition_status_advanced():
         "created_at",
         "total_samples_num",
         "total_fields_num",
+        "allocated_samples_num",
         "allocated_fields_num",
         "registered_tasks",
         "produced_samples",
@@ -311,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
@@ -353,7 +355,7 @@ def test_edge_cases_and_error_handling():
     task_name = "early_task"
     consumption_tensor = partition.get_consumption_status(task_name)
     assert consumption_tensor is not None
-    assert consumption_tensor.shape[0] == partition.total_samples_num
+    assert consumption_tensor.shape[0] == partition.allocated_samples_num
 
     # Test 4: Production status update error conditions
     # Test with empty lists
@@ -371,80 +373,6 @@ def test_edge_cases_and_error_handling():
     print("Edge cases and error handling tests passed!\n")
 
 
-def test_backward_compatibility():
-    """Test backward compatibility with existing interfaces."""
-    print("Testing backward compatibility...")
-
-    from transfer_queue.controller import DataPartitionStatus
-
-    partition = DataPartitionStatus(partition_id="compat_test")
-
-    # Test 1: Basic workflow should work as before
-    sample_indices = [0, 1, 2, 3, 4]
-    field_names = ["input_ids", "attention_mask", "labels"]
-    dtypes = {
-        k: {"input_ids": "torch.int64", "attention_mask": "torch.bool", "labels": "torch.int64"} for k in sample_indices
-    }
-    shapes = {k: {"input_ids": (32,), "attention_mask": (32,), "labels": (32,)} for k in sample_indices}
-    success = partition.update_production_status(
-        sample_indices,
-        field_names,
-        dtypes=dtypes,
-        shapes=shapes,
-    )
-    assert success
-
-    # Traditional consumption tracking
-    task_name = "training_task"
-    ready_samples = partition.scan_data_status(field_names, task_name)
-    assert len(ready_samples) == 5
-
-    # Mark as consumed
-    partition.mark_consumed(task_name, ready_samples[:3])
-
-    # Should now return only unconsumed samples
-    remaining_ready = partition.scan_data_status(field_names, task_name)
-    assert len(remaining_ready) == 2
-
-    print("✓ Basic workflow maintains compatibility")
-
-    # Test 2: Field mapping should be consistent
-    for field in field_names:
-        assert field in partition.field_name_mapping
-        field_idx = partition.field_name_mapping[field]
-        assert field_idx >= 0
-        assert field_idx < partition.allocated_fields_num
-
-    print("✓ Field mapping consistency maintained")
-
-    # Test 3: Metadata access patterns
-    for sample_idx in sample_indices:
-        for field in field_names:
-            # These should return reasonable values or None
-            dtype = partition.get_field_dtype(sample_idx, field)
-            shape = partition.get_field_shape(sample_idx, field)
-            assert dtype is not None
-            assert shape is not None
-            # Should not crash even if metadata wasn't provided
-
-    print("✓ Metadata access patterns preserved")
-
-    # Test 4: Statistics format should be familiar
-    stats = partition.get_statistics()
-    familiar_keys = ["partition_id", "total_samples_num", "total_fields_num"]
-    for key in familiar_keys:
-        assert key in stats
-
-    assert isinstance(stats["total_samples_num"], int)
-    assert isinstance(stats["total_fields_num"], int)
-    assert stats["total_samples_num"] > 0
-    assert stats["total_fields_num"] == len(field_names)
-
-    print("✓ Statistics format maintains familiarity")
-
-    print("Backward compatibility tests passed!\n")
-
-
 def test_performance_characteristics():
     """Test performance characteristics of the refactored implementation."""
     print("Testing performance characteristics...")
@@ -512,65 +440,3 @@ def test_performance_characteristics():
     print("✓ Memory usage patterns reasonable")
 
     print("Performance characteristics tests passed!\n")
-
-
-def main():
-    """Run all tests."""
-    print("=== Comprehensive Testing of TransferQueue Controller ===\n")
-
-    test_functions = [
-        test_data_partition_status,
-        test_partition_interface,
-        test_dynamic_expansion_scenarios,
-        test_data_partition_status_advanced,
-        test_edge_cases_and_error_handling,
-        test_backward_compatibility,
-        test_performance_characteristics,
-    ]
-
-    passed_tests = 0
-    total_tests = len(test_functions)
-
-    try:
-        for test_func in test_functions:
-            try:
-                test_func()
-                passed_tests += 1
-            except Exception as e:
-                print(f"❌ {test_func.__name__} failed: {e}")
-                import traceback
-
-                traceback.print_exc()
-                print()
-
-        print("=" * 60)
-        print(f"TEST SUMMARY: {passed_tests}/{total_tests} test suites passed")
-
-        if passed_tests == total_tests:
-            print("🎉 ALL TESTS PASSED!")
-            print("\nThe refactored DataPartitionStatus demonstrates:")
-            print("1. ✅ Dynamic row and column expansion without pre-allocation")
-            print("2. ✅ Robust partition-controller interface design")
-            print("3. ✅ Self-contained state management in DataPartitionStatus")
-            print("4. ✅ Flexible consumption tracking per task")
-            print("5. ✅ Comprehensive scanning and query capabilities")
-            print("6. ✅ Advanced error handling and edge case management")
-            print("7. ✅ Backward compatibility with existing interfaces")
-            print("8. ✅ Good performance characteristics for large datasets")
-            print("\n🚀 DataPartitionStatus refactoring is ready for production!")
-        else:
-            print(f"⚠️  {total_tests - passed_tests} test suites failed.")
-            print("Please review the failures before deploying to production.")
-
-        print("=" * 60)
-
-    except Exception as e:
-        print(f"❌ Critical test failure: {e}")
-        import traceback
-
-        traceback.print_exc()
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()