Fix WPIWeightTransferEngine instantiation by implementing missing abstract lifecycle methods

Signed-off-by: Bill Du <yangspirit@google.com>

Author: yangspirit

tests/distributed/test_weight_transfer.py

@@ -267,6 +267,20 @@ def test_create_engine_sparse_nccl(self):
         )
         assert isinstance(engine, SparseNCCLWeightTransferEngine)
 
+    def test_create_engine_wpi(self):
+        from vllm.distributed.weight_transfer.wpi_engine import (
+            WPIWeightTransferEngine,
+        )
+
+        config = WeightTransferConfig(backend="wpi")
+        engine = WeightTransferEngineFactory.create_engine(
+            config,
+            create_mock_vllm_config(),
+            torch.device("cuda"),
+            MagicMock(spec=torch.nn.Module),
+        )
+        assert isinstance(engine, WPIWeightTransferEngine)
+
     def test_create_engine_invalid_backend(self):
         config = WeightTransferConfig(backend="invalid")
         with pytest.raises(ValueError, match="Invalid weight transfer backend"):

vllm/distributed/weight_transfer/wpi_engine.py

@@ -316,6 +316,22 @@ def init_transfer_engine(self, init_info: WPIWeightTransferInitInfo) -> None:
             total_shards,
         )
 
+    def start_weight_update(self) -> None:
+        """Initialize layerwise reloading for the incoming checkpoint weights."""
+        from vllm.model_executor.model_loader.reload import (
+            initialize_layerwise_reload,
+        )
+
+        initialize_layerwise_reload(self.model)
+
+    def finish_weight_update(self) -> None:
+        """Finalize layerwise reloading after all weights have been received."""
+        from vllm.model_executor.model_loader.reload import (
+            finalize_layerwise_reload,
+        )
+
+        finalize_layerwise_reload(self.model, self.model_config)
+
     def receive_weights(
         self,
         update_info: WPIWeightTransferUpdateInfo,