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

Method asyncSetUp

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

Source from the content-addressed store, hash-verified

1229 """
1230
1231 async def asyncSetUp(self):
1232 self.config = get_template_config()
1233 self.config.mode = "explore"
1234 self.config.model.model_path = get_model_path()
1235 self.config.explorer.rollout_model.engine_type = "vllm"
1236 self.config.explorer.rollout_model.engine_num = 1
1237 self.config.explorer.rollout_model.tensor_parallel_size = 4
1238 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE
1239 self.config.explorer.rollout_model.enable_openai_api = True
1240 self.config.explorer.rollout_model.enable_history = True
1241 # Use the checkpoint-based weight sync. The checkpoint holds the *same*
1242 # weights as the running model, so the swap is a semantic no-op and the
1243 # test isolates the concurrency behavior rather than the weight values.
1244 self.config.explorer.rollout_model.sync_method = SyncMethod.CHECKPOINT
1245 self.config.explorer.rollout_model.extra_engine_args = {
1246 "attention_backend": "FLASHINFER",
1247 }
1248 # A throwaway checkpoint root for this test run.
1249 self.config.checkpoint_root_dir = get_checkpoint_path()
1250 self.config.check_and_update()
1251
1252 self.engines, self.auxiliary_engines = await create_test_models(self.config)
1253 self.model_wrapper = self.engines[0]
1254 self.openai_client = self.model_wrapper.get_openai_async_client()
1255 self.model_id = self.openai_client.model_path
1256 master_addr, master_port = await self.model_wrapper.get_available_address_async()
1257 # Stand up the weight-transfer process group for the single inference rank
1258 # (this mirrors the single-engine deployment path in Explorer).
1259 await self.model_wrapper.init_process_group(
1260 master_address=master_addr,
1261 master_port=master_port,
1262 rank_offset=0,
1263 world_size=4,
1264 group_name=ROLLOUT_WEIGHT_SYNC_GROUP_NAME,
1265 )
1266
1267 # Materialise an identical checkpoint at the path the engine expects:
1268 # <checkpoint_job_dir>/global_step_<version>/actor/huggingface/
1269 self._target_version = 1
1270 huggingface_dir = os.path.join(
1271 self.config.get_checkpoint_job_dir(),
1272 f"global_step_{self._target_version}",
1273 "actor",
1274 "huggingface",
1275 )
1276 os.makedirs(huggingface_dir, exist_ok=True)
1277 for entry in os.listdir(self.config.model.model_path):
1278 link = os.path.join(huggingface_dir, entry)
1279 if not os.path.exists(link):
1280 os.symlink(os.path.join(self.config.model.model_path, entry), link)
1281
1282 async def asyncTearDown(self):
1283 try:

Callers

nothing calls this directly

Calls 9

get_template_configFunction · 0.90
get_model_pathFunction · 0.90
get_checkpoint_pathFunction · 0.90
create_test_modelsFunction · 0.85
check_and_updateMethod · 0.80
init_process_groupMethod · 0.45

Tested by

no test coverage detected