imp raw http training

Author: binary-husky

ajet/backbone/main_vllm.py

@@ -9,19 +9,21 @@
 from ajet.task_rollout.native_parallel_worker import VerlRolloutManager
 from ajet.utils.launch_utils import set_loguru_default_color
 from ajet.schema.logprob import TokenAndProb
+from ajet.utils.core_env_vars import get_runtime_env
 
 set_loguru_default_color()
 
 
 class TokenAndProbVllmDebug(TokenAndProb):
     def __init__(self, t):
         # ChatCompletionTokenLogprob(token='token_id:73594', bytes=[96, 96, 96], logprob=-1.9073468138230965e-06, top_logprobs=[])
-        self.token_id = int(t.token.split("token_id:")[-1])
-        self.logprob = t.logprob
+        token_id = int(t.token.split("token_id:")[-1])
+        logprob = t.logprob
         try:
-            self.decoded_string = bytes(t.bytes).decode("utf-8")
+            decoded_string = bytes(t.bytes).decode("utf-8")
         except Exception:
-            self.decoded_string = "<cannot decode>" + str(t.bytes)
+            decoded_string = "<cannot decode>" + str(t.bytes)
+        super().__init__(token_id=token_id, logprob=logprob, decoded_string=decoded_string)
 
 
 class ChatCompletionScheduler:
@@ -87,6 +89,8 @@ def submit_chat_completions(self, messages, sampling_params, request_id, tools=[
 
 
 def run(config):
+    from ajet.task_reader import RouterTaskReader
+
     # --------- fast adjustment for debugging ---------
     warm_up_process(config)
     max_parallel = config.ajet.debug.debug_max_parallel
@@ -106,7 +110,6 @@ def run(config):
         tokenizer=async_rollout_manager.tokenizer,
     )
 
-    from ajet.task_reader import RouterTaskReader
 
     task_reader = RouterTaskReader(
         config.ajet.task_reader.type,
@@ -132,6 +135,12 @@ def main(config):
     OmegaConf.resolve(config)
     print("*" * 20)
 
+    runtime_env = get_runtime_env()
+    os.environ.update(runtime_env["env_vars"])
+    if config.ajet.enable_experimental_reverse_proxy:
+        from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
+        start_interchange_server()
+
     def companion_launch():
         import torch
 

ajet/backbone/trainer_verl.py

@@ -832,6 +832,11 @@ def fit(self):  # noqa: C901
                 progress_bar.update(1)
                 self.global_steps += 1
 
+                # when enabled oai request interchange, we need to clear the cache from time to time
+                if self.config.ajet.enable_experimental_reverse_proxy:
+                    from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import ensure_dat_interchange_server_cache_clear
+                    ensure_dat_interchange_server_cache_clear()
+
                 if is_last_step:
                     pprint(f"Final validation metrics: {last_val_metrics}")
                     progress_bar.close()

ajet/context_tracker/multiagent_tracking.py

@@ -64,17 +64,33 @@ def __init__(
         self.episode_uuid = episode_uuid
 
 
-    def step_prepare(self, messages: List[dict], tools: List = [], timeline_uuid: str = ""):
+    def preprocess_tools_field(self, tools: List[dict] = [], disable_toolcalls: bool = False):
+        if disable_toolcalls:
+            tools = []
+        else:
+            if tools is not None:
+                # rerank tool parameters to improve compatibility
+                for i in range(len(tools)):
+                    tools[i]["function"]["parameters"] = tools[i]["function"].pop("parameters")
+        return tools
+
+
+    def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_toolcalls: bool = False) -> List[ExtendedMessage]:
+        """Spawn a timeline from messages.
+
+        Args:
+            messages: List of message dictionaries
+            tools: List of tool dictionaries
+            disable_toolcalls: Whether to disable tool calls
+
+        Returns:
+            List of ExtendedMessage objects representing the timeline
+        """
         timeline = []
+
         consider_roles = ["user", "assistant", "system", "tool"]
-        disable_toolcalls = self.config.ajet.rollout.force_disable_toolcalls
         if disable_toolcalls:
             consider_roles.remove("tool")
-            tools = []
-        else:
-            # rerank tool parameters to improve compatibility
-            for i in range(len(tools)):
-                tools[i]["function"]["parameters"] = tools[i]["function"].pop("parameters")
 
         for i, msg in enumerate(messages):
             if (disable_toolcalls) and (not isinstance(msg["content"], str)):
@@ -132,6 +148,14 @@ def step_prepare(self, messages: List[dict], tools: List = [], timeline_uuid: st
                 )
             ]
 
+        return timeline
+
+
+    def step_prepare(self, messages: List[dict], tools: List = [], timeline_uuid: str = ""):
+        disable_toolcalls = self.config.ajet.rollout.force_disable_toolcalls
+        tools = self.preprocess_tools_field(tools, disable_toolcalls=disable_toolcalls)
+        timeline = self.step_spawn_timeline(messages, tools, disable_toolcalls)
+
         # check token overflow
         converted_message = self.to_role_content(timeline)
         timeline = ExtendedMessage.check_and_merge_chained_tool_response(

ajet/schema/logprob.py

@@ -5,8 +5,10 @@
 # from verl import DataProto
 
 
-class TokenAndProb:
-    def __init__(self, token_id, logprob, decoded_string):
-        self.token_id = token_id
-        self.logprob = logprob
-        self.decoded_string = decoded_string
+from pydantic import BaseModel
+
+
+class TokenAndProb(BaseModel):
+    token_id: int
+    logprob: float
+    decoded_string: str

ajet/task_rollout/async_llm_bridge.py

@@ -5,8 +5,8 @@
 import uuid
 from typing import Any, Callable, Dict, List, Literal, Type, Union
 
-from agentscope._utils._common import _json_loads_with_repair
-from agentscope.message import TextBlock, ToolUseBlock
+
+
 from loguru import logger
 from omegaconf import DictConfig
 from pydantic import BaseModel
@@ -33,15 +33,15 @@
 class AjetStandardLlmBridgeRequest(BaseModel):
     messages: List[Dict[str, str]]
     custom_sampling_params: dict = {}
-    tools=[]
+    tools: List = []
     request_id: str = ""
 
 class AjetStandardLlmBridgeResponse(BaseModel):
     role: str = "assistant"
     request_id: str = ""
     content: str = ""
     tool_calls: List[Dict] = []
-    tokens: List[TokenAndProb]
+    tokens: List[TokenAndProb] = []
 
 
 # -------------------------------------------------------------------------------------

ajet/tuner_lib/weight_tuner/as_oai_baseurl_apikey.py

@@ -1,3 +1,4 @@
+import os
 import asyncio
 from typing import TYPE_CHECKING, Any, List, Callable, Literal, Type, Union
 from loguru import logger
@@ -45,10 +46,19 @@ def __init__(
         episode_uuid: str,
         **kwargs,
     ):
-        self.base_url = "http://localhost:27788/v1"
-        self.api_key = generate_auth_token(
+        port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
+        assert port is not None, "AJET_DAT_INTERCHANGE_PORT env var must be set"
+        base_url = f"http://localhost:{port}/v1"
+        api_key = generate_auth_token(
             agent_name=agent_name,
             target_tag=target_tag,
             episode_uuid=episode_uuid,
         )
-        self.model = "reserved_field"
+        model = "reserved_field"
+
+        # Properly initialize the Pydantic BaseModel
+        super().__init__(
+            base_url=base_url,
+            api_key=api_key,
+            model=model,
+        )

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_client.py

@@ -108,7 +108,7 @@ async def _service_loop(self):
         This design is for efficiency
         """
 
-        from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import TypeCompletionRequest
+        from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import InterchangeCompletionRequest
 
         port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
         assert port is not None, "AJET_DAT_INTERCHANGE_PORT env var must be set"
@@ -118,13 +118,13 @@ async def _service_loop(self):
             try:
                 # Send initialization parameters
                 # Sending as a list [agent_name, target_tag, episode_uuid] to match "input (a,b,c)" structure
-                await websocket.send(f"episode_uuid:{self.episode_uuid}")
+                await websocket.send(pickle.dumps(f"episode_uuid:{self.episode_uuid}"))
 
                 while not self.should_terminate:
 
                     try:
                         # wait message from ajet/tuner_lib/weight_tuner/experimental/as_oai_model_server.py
-                        parsed_msg: TypeCompletionRequest = pickle.loads(
+                        parsed_msg: InterchangeCompletionRequest = pickle.loads(
                             await asyncio.wait_for(websocket.recv(decode=False), timeout=0.25)
                         )