get worker of different device
(fd_config: FDConfig, local_rank: int, rank: int)
| 72 | |
| 73 | |
| 74 | def get_worker(fd_config: FDConfig, local_rank: int, rank: int) -> WorkerBase: |
| 75 | """ |
| 76 | get worker of different device |
| 77 | """ |
| 78 | if fd_config.model_config.enable_logprob and not current_platform.is_cuda() and not current_platform.is_xpu(): |
| 79 | raise NotImplementedError("Only CUDA and XPU platforms support logprob.") |
| 80 | if current_platform.is_dcu(): |
| 81 | from fastdeploy.worker.dcu_worker import DcuWorker |
| 82 | |
| 83 | return DcuWorker(fd_config=fd_config, local_rank=local_rank, rank=rank) |
| 84 | if current_platform.is_cuda(): |
| 85 | from fastdeploy.worker.gpu_worker import GpuWorker |
| 86 | |
| 87 | return GpuWorker(fd_config=fd_config, local_rank=local_rank, rank=rank) |
| 88 | if current_platform.is_xpu(): |
| 89 | from fastdeploy.worker.xpu_worker import XpuWorker |
| 90 | |
| 91 | return XpuWorker(fd_config=fd_config, local_rank=local_rank, rank=rank) |
| 92 | if current_platform.is_iluvatar(): |
| 93 | from fastdeploy.worker.iluvatar_worker import IluvatarWorker |
| 94 | |
| 95 | return IluvatarWorker(fd_config=fd_config, local_rank=local_rank, rank=rank) |
| 96 | if current_platform.is_gcu(): |
| 97 | from fastdeploy.worker.gcu_worker import GcuWorker |
| 98 | |
| 99 | return GcuWorker(fd_config=fd_config, local_rank=local_rank, rank=rank) |
| 100 | if current_platform.is_maca(): |
| 101 | from fastdeploy.worker.metax_worker import MetaxWorker |
| 102 | |
| 103 | return MetaxWorker(fd_config=fd_config, local_rank=local_rank, rank=rank) |
| 104 | if current_platform.is_intel_hpu(): |
| 105 | from fastdeploy.worker.hpu_worker import HpuWorker |
| 106 | |
| 107 | return HpuWorker(fd_config=fd_config, local_rank=local_rank, rank=rank) |
| 108 | |
| 109 | |
| 110 | def init_distributed_environment(seed: int = 20) -> Tuple[int, int]: |
no test coverage detected