Method create

trinity/utils/distributed.py:114–133

Factory method to create the appropriate weight transfer engine based on the rollout engine type.

(
        engine_type: str, master_address: str, master_port: int, world_size: int, group_name: str
    )

    @staticmethod
    def create(
        engine_type: str, master_address: str, master_port: int, world_size: int, group_name: str
    ):
        """Factory method to create the appropriate weight transfer engine based on the rollout engine type."""
        if engine_type == "vllm":
            return VLLMWeightTransferEngine(
                master_address=master_address,
                master_port=master_port,
                world_size=world_size,
                group_name=group_name,
            )
        elif engine_type == "sglang":
            return SGLangWeightTransferEngine(
                master_address=master_address,
                master_port=master_port,
                world_size=world_size,
                group_name=group_name,
            )
        else:
            raise ValueError(f"Unsupported engine type: {engine_type}")


class VLLMWeightTransferEngine(WeightTransferEngine):