Support multi episode with vllm_v2. (#187)

* Support multi episode.

chatlearn/schedule/model_manager.py

@@ -190,8 +190,7 @@ def sync_parameters(self, episode_offset=0, requires_grad=None, validate=False):
     def set_func_decorator(self, model):
         if is_decorated(model.name):
             return
-        # decorate async method here will raise cannot serialize coroutine object error
-        call_funcs = model.call_funcs if not isinstance(model, VLLMModuleV2) else []
+        call_funcs = model.call_funcs
 
         model_cls = model.__class__
         for func_name in call_funcs:

examples/megatron/models/vllm_policy_inference.py

@@ -15,20 +15,21 @@
 """vllm policy inference"""
 
 import copy
-import os
 import random
 
 import torch
 import torch.nn.functional as F
 
+from chatlearn.models.vllm import is_vllm_v2
 from examples.megatron.data.prompt_dataset import VLLMPromptPipeline
 from .utils import get_loss_mask
 
-if os.environ.get("ENABLE_VLLM_V2"):
+# pylint: disable=ungrouped-imports
+if is_vllm_v2():
     from chatlearn import VLLMModuleV2 as VLLMModule
 else:
     from chatlearn import VLLMModule
-
+# pylint: enable=ungrouped-imports
 
 
 class VLLMPolicyInference(VLLMModule):