| 108 | raise e |
| 109 | |
| 110 | def create_workers(self, worker_cls, worker_kwargs): |
| 111 | llm_args = worker_kwargs.get("llm_args") |
| 112 | placement_config = getattr(llm_args, 'ray_placement_config', |
| 113 | None) if llm_args else None |
| 114 | |
| 115 | # When set to be a fraction, it allows Ray to schedule |
| 116 | # multiple actors on a single GPU for colocate use cases. |
| 117 | num_gpus = float(os.getenv("TRTLLM_RAY_PER_WORKER_GPUS", "1.0")) |
| 118 | if placement_config and placement_config.per_worker_gpu_share is not None: |
| 119 | num_gpus = placement_config.per_worker_gpu_share |
| 120 | |
| 121 | logger.debug(f"{num_gpus=} for each worker.") |
| 122 | |
| 123 | runtime_env = ray.runtime_env.RuntimeEnv() |
| 124 | runtime_env["env_vars"] = os.environ.copy() |
| 125 | runtime_env["env_vars"].update({ |
| 126 | "TLLM_DISABLE_MPI": "1", |
| 127 | "MASTER_ADDR": self.master_address, # head-IP for NCCL/Gloo |
| 128 | }) |
| 129 | |
| 130 | placement_groups, self.bundle_indices = self._get_placement_group( |
| 131 | tp_size=self.tp_size, worker_kwargs=worker_kwargs) |
| 132 | |
| 133 | if isinstance(placement_groups, list): |
| 134 | self.placement_group = None |
| 135 | else: |
| 136 | self.placement_group = placement_groups |
| 137 | |
| 138 | self.workers = [] |
| 139 | for rank in range(self.world_size): |
| 140 | pg = placement_groups[rank] if isinstance( |
| 141 | placement_groups, list) else placement_groups |
| 142 | worker = RayWorkerWrapper.options( |
| 143 | num_gpus=num_gpus, |
| 144 | runtime_env=runtime_env, |
| 145 | scheduling_strategy=PlacementGroupSchedulingStrategy( |
| 146 | placement_group=pg, |
| 147 | placement_group_bundle_index=self.bundle_indices[rank], |
| 148 | )).remote(worker_cls, worker_kwargs, self.world_size, rank) |
| 149 | self.workers.append(worker) |
| 150 | |
| 151 | def init_workers_sync(self): |
| 152 | self.create_workers(RayGPUWorker, self.worker_kwargs) |