fix interchange server boot

Author: binary-husky

ajet/default_config/ajet_default.yaml

@@ -192,7 +192,7 @@ ajet:
     fix_retokenization_drift: True
 
     # log tool format check results
-    log_tool_format_check: True
+    log_tool_format_check: False
 
     # log tool format check results
     log_tool_format_error_detail: False
@@ -281,7 +281,7 @@ ajet:
 
 
   # the experimental reverse proxy feature that allows `tuner.as_oai_baseurl_apikey` feature
-  enable_experimental_interchange_server: False
+  enable_experimental_interchange_server: True
   interchange_server:
     interchange_method: 'ipc' # options: 'tcp' (multi-nodes) or  'ipc' (1 node)
     interchange_server_port: 'auto'

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_server.py

@@ -102,6 +102,11 @@ def _begin_handle_chat_completion(episode_address, int_req: InterchangeCompletio
         return result_object
 
 
+    @app.get("/health")
+    async def health():
+        return {"status": "ok"}
+
+
     @app.post("/v1/chat/completions")
     async def chat_completions(request: Request, authorization: str = Header(None)):
         """
@@ -167,6 +172,7 @@ def __init__(self, experiment_dir: str, port: int, num_fastapi_process: int = 2,
         self.max_fastapi_threads = max_fastapi_threads
 
     def run(self):
+        logger.info(f"Starting Interchange Server on port {self.port} with {self.num_fastapi_process} processes and {self.max_fastapi_threads} threads per process.")
         app = get_app(self.max_fastapi_threads)
         async def serve_with_monitor():
             # Start the server
@@ -215,15 +221,16 @@ def start_interchange_server(config) -> int:
             logger.error(f"Interchange server subprocess failed to start. Return code: {interchange_server.exitcode}")
             raise RuntimeError("Interchange server subprocess failed to start.")
         if time.time() - start_time > 30:
-            logger.error("Interchange server subprocess failed to start within 30 seconds.")
-            raise RuntimeError("Interchange server subprocess failed to start within 30 seconds.")
+            msg = f"Interchange server subprocess failed to start within {time.time() - start_time} seconds."
+            logger.error(msg)
+            raise RuntimeError(msg)
         try:
             if httpx.get(health_url, timeout=0.5).status_code == 200:
                 break
         except Exception:
             # keep waiting
             pass
-        time.sleep(0.5)
+        time.sleep(1)
 
     # register a termination handler
     if DEBUG: logger.info(f"Interchange server subprocess started on port {port} (pid: {interchange_server.pid})")

docs/en/installation.md

@@ -83,7 +83,10 @@ AgentJet supports multiple backbones. Currently we have `verl` and `trinity` (re
     ```
 
     !!! warning "flash-attn Installation"
-        `flash-attn` must be installed after other dependencies. To build faster, export `MAX_JOBS=${N_CPU}`, or ensure a healthy connection to GitHub to install pre-compiled wheels.
+        - `flash-attn` must be installed **after** other dependencies.
+        - Ensure a healthy connection to GitHub to install pre-compiled wheels.
+        - If you find your machine spend a long time installing flash-attn, ensure a healthy connection to GitHub.
+        - To build faster, export `MAX_JOBS=${N_CPU}`.
 
 
 === "Trinity"

tutorial/example_math_agent/math_agentscope_urlkey.py

@@ -0,0 +1,74 @@
+import re
+from loguru import logger
+from agentscope.message import Msg
+from agentscope.agent import ReActAgent
+from agentscope.formatter import OpenAIChatFormatter
+from agentscope.model import OpenAIChatModel
+from agentscope.memory import InMemoryMemory
+from agentscope.tool import Toolkit, execute_python_code
+from ajet import AjetTuner, Workflow, WorkflowOutput, WorkflowTask
+
+
+def extract_final_answer(result) -> str:
+    """Extract the final answer from the agent's response."""
+    try:
+        if (
+            hasattr(result, "metadata")
+            and isinstance(result.metadata, dict)
+            and "result" in result.metadata
+        ):
+            return result.metadata["result"]
+        if hasattr(result, "content"):
+            if isinstance(result.content, dict) and "result" in result.content:
+                return result.content["result"]
+            return str(result.content)
+        return str(result)
+    except Exception as e:
+        logger.warning(f"Extract final answer error: {e}. Raw: {result}")
+        return str(result)
+
+
+system_prompt = """
+You are an agent specialized in solving math problems with tools.
+Please solve the math problem given to you.
+You can write and execute Python code to perform calculation or verify your answer.
+You should return your final answer within \\boxed{{}}.
+"""
+
+
+class MathToolWorkflow(Workflow): # ✨✨ inherit `Workflow` class
+    name: str = "math_agent_workflow"
+
+    async def execute(self, workflow_task: WorkflowTask, tuner: AjetTuner) -> WorkflowOutput:
+        # run agentscope
+        query = workflow_task.task.main_query
+        self.toolkit = Toolkit()
+        self.toolkit.register_tool_function(execute_python_code)
+
+        url_and_apikey = tuner.as_oai_baseurl_apikey()
+        base_url = url_and_apikey.base_url
+        api_key = url_and_apikey.api_key    # the api key contain information, do not discard it
+        model = OpenAIChatModel(
+            model_name="whatever",
+            client_args={"base_url": base_url},
+            api_key=api_key,
+            stream=False,
+        )
+        self.agent = ReActAgent(
+            name="math_react_agent", sys_prompt=system_prompt,
+            model=model,  # ✨✨ compared with a normal agentscope agent, here is the difference!
+            formatter=OpenAIChatFormatter(),
+            toolkit=self.toolkit,
+            memory=InMemoryMemory(), max_iters=2,
+        )
+        self.agent.set_console_output_enabled(False)
+        msg = Msg("user", query, role="user")
+        result = await self.agent.reply(msg)
+        final_answer = extract_final_answer(result)
+
+        # compute reward
+        reference_answer = workflow_task.task.metadata["answer"].split("####")[-1].strip()
+        match = re.search(r"\\boxed\{([^}]*)\}", final_answer)
+        if match: is_success = (match.group(1) == reference_answer)
+        else:     is_success = False
+        return WorkflowOutput(reward=(1.0 if is_success else 0.0), metadata={"final_answer": final_answer})