MCPcopy Create free account
hub / github.com/agentscope-ai/Trinity-RFT / asyncSetUp

Method asyncSetUp

tests/common/vllm_test.py:135–170  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

133)
134class ModelWrapperTest(VLLMTestBase):
135 async def asyncSetUp(self):
136 # configure the model
137 self.config = get_template_config()
138 self.config.mode = "explore"
139 self.config.model.model_path = (
140 get_model_path() if not self.enable_return_routed_experts else get_moe_model_path()
141 )
142 if self.enable_return_routed_experts:
143 self.text_config = _get_text_config(self.config.model.model_path)
144 self.expected_routed_experts_layers = _count_moe_layers(self.text_config)
145 self.expected_routed_experts_topk = int(self.text_config.num_experts_per_tok)
146 else:
147 self.expected_routed_experts_layers = 0
148 self.expected_routed_experts_topk = 0
149 self.config.model.custom_chat_template = CHAT_TEMPLATE
150 self.config.explorer.rollout_model.engine_num = self.engine_num
151 self.config.explorer.rollout_model.nnodes = self.nnodes
152 self.config.explorer.rollout_model.tensor_parallel_size = self.tensor_parallel_size
153 self.config.explorer.rollout_model.data_parallel_size = self.data_parallel_size
154 self.config.explorer.rollout_model.pipeline_parallel_size = self.pipeline_parallel_size
155 self.config.explorer.rollout_model.enable_expert_parallel = (
156 self.enable_return_routed_experts
157 )
158 self.config.algorithm.repeat_times = self.repeat_times
159 self.config.explorer.rollout_model.enable_history = self.enable_history
160 self.config.explorer.rollout_model.enable_openai_api = self.enable_return_routed_experts
161 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE
162 self.config.explorer.rollout_model.extra_engine_args = {"max_num_seqs": 24}
163 if self.enable_return_routed_experts:
164 self.config.explorer.rollout_model.extra_engine_args["moe_backend"] = "triton"
165 self.config.explorer.rollout_model.extra_engine_args["gdn_prefill_backend"] = "triton"
166 self.config.algorithm.enable_router_replay = self.enable_return_routed_experts
167 self.config.check_and_update()
168
169 self.engines, self.auxiliary_engines = await create_test_models(self.config)
170 self.model_wrapper = self.engines[0]
171
172 def _assert_openai_response_routed_experts(self, response, expected_choices: int):
173 self.assertEqual(len(response.choices), expected_choices)

Callers

nothing calls this directly

Calls 7

get_template_configFunction · 0.90
get_model_pathFunction · 0.90
get_moe_model_pathFunction · 0.90
_count_moe_layersFunction · 0.85
create_test_modelsFunction · 0.85
check_and_updateMethod · 0.80
_get_text_configFunction · 0.70

Tested by

no test coverage detected