(
self, config: InferenceModelConfig, role: str, engine_id: int
)
| 77 | self.logger.info(" > Bundle %s: Actor %s", bundle_id, actor_name) |
| 78 | |
| 79 | async def create_engine( |
| 80 | self, config: InferenceModelConfig, role: str, engine_id: int |
| 81 | ) -> ModelWrapper: |
| 82 | config = deepcopy(config) |
| 83 | config.engine_id = engine_id |
| 84 | |
| 85 | actor_bundle_lists = [] |
| 86 | for node_id in range(config.nnodes): |
| 87 | actor_name = self.get_actor_name(role, engine_id, node_id) |
| 88 | actor_bundle_lists.append((actor_name, self.bundle_result.actor_bundle_map[actor_name])) |
| 89 | |
| 90 | model_cls = None |
| 91 | if config.engine_type.startswith("vllm"): |
| 92 | from trinity.common.models.vllm_model import vLLMRolloutModel |
| 93 | |
| 94 | model_cls = vLLMRolloutModel |
| 95 | elif config.engine_type == "sglang": |
| 96 | from trinity.common.models.sglang_model import SGLangRolloutModel |
| 97 | |
| 98 | model_cls = SGLangRolloutModel |
| 99 | elif config.engine_type == "tinker": |
| 100 | from trinity.common.models.tinker_model import TinkerModel |
| 101 | |
| 102 | model_cls = TinkerModel |
| 103 | elif config.engine_type == "external": |
| 104 | return await get_external_model_wrapper(config=config) |
| 105 | else: |
| 106 | raise ValueError(f"Unsupported engine type: {config.engine_type}") |
| 107 | |
| 108 | self.logger.info( |
| 109 | f"Creating inference_model {self.get_actor_name(role, engine_id, 0)} in {config.ray_namespace}." |
| 110 | ) |
| 111 | return await get_model_wrapper(model_cls, config, self.pg, actor_bundle_lists) |
| 112 | |
| 113 | async def create_all_models(self) -> Tuple[List[ModelWrapper], List[List[ModelWrapper]]]: |
| 114 | """Create all model actors for the rollout model and auxiliary models based on the configuration.""" |
no test coverage detected