improve tinkerscript

Author: binary-husky

ajet/backbone/main_vllm.py

@@ -189,6 +189,9 @@ def main(config):
     if config.ajet.enable_experimental_interchange_server:
         from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
         start_interchange_server(config)
+        if config.ajet.enable_tinkerscript_mode:
+            from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
+            http_change_engine_status(config, "ROLLING")
 
     def companion_launch():
         import torch

ajet/backbone/trainer_verl.py

@@ -443,6 +443,13 @@ def init_workers(self):
             tokenizer=self.tokenizer,
         )
 
+    def _update_interchange_server_status_flag(self, status: str):
+        # if interchange server is enabled, change engine status to ROLLING
+        if self.config.ajet.enable_experimental_interchange_server:
+            if self.config.ajet.enable_tinkerscript_mode:
+                from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
+                http_change_engine_status(self.config, status)
+
     # #######################################
     # training loop
     # #######################################
@@ -552,6 +559,7 @@ def fit(self):  # noqa: C901
                         assert self.async_rollout_mode
                         logger.info("=== wake up begin ===")
                         self.async_rollout_manager.wake_up()
+                        self._update_interchange_server_status_flag("ROLLING")
                         logger.info("=== wake up end ===")
                         tasks: List[Task] = [
                             dict_to_ajet_task(dict(
@@ -577,6 +585,7 @@ def fit(self):  # noqa: C901
                             tasks, mode="sample", epoch=f"train.{epoch}"
                         )
                         logger.info("=" * 10 + "end fit rollout" + "=" * 10)
+                        self._update_interchange_server_status_flag("UPDATE_WEIGHT")
                         logger.info("begin to convert context_tracker_arr to dataproto")
                         gen_batch_output = self.parallel_env.to_dataproto(context_tracker_arr)
                         logger.info("end convertion")

ajet/task_runner/tinkerscript_runner.py

@@ -1,10 +1,8 @@
 
 import atexit
 import json
-import requests
 import zmq
 import os
-import time
 from ajet import AjetTuner
 from ajet import WorkflowOutput
 from ajet.context_tracker.multiagent_tracking import (
@@ -14,77 +12,38 @@
 from ajet.schema.task import WorkflowTask
 from ajet.schema.trajectory import Reward
 from ajet.task_runner.base_runner import BaseAgentRunner
-from ajet.utils.networking import find_free_port
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_register_episode, get_zmq_socket
 from loguru import logger
 from ajet import Workflow
 
 context = zmq.Context()
 atexit.register(context.term)
+DEBUG = True
 
 class TinkerScriptRunner(BaseAgentRunner):
 
-    def get_zmq_socket(self, episode_uuid: str):
-        interchange_method = self.config.ajet.interchange_server.interchange_method
-        if interchange_method == 'tcp':
-            master_node_ip = os.getenv("MASTER_NODE_IP", "localhost")
-            episode_contect_address = f"tcp://{master_node_ip}:{find_free_port()}"
-        elif interchange_method == 'ipc':
-            ipc_path = f"/tmp/ajet/{episode_uuid}-workflow.sock"
-            episode_contect_address = f"ipc://{ipc_path}"
-        else:
-            raise RuntimeError(f"Unknown interchange_method: {interchange_method}")
-        return episode_contect_address
-
-
-    def get_interchange_server_url(self):
-        port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
-        if self.config.ajet.interchange_server.interchange_server_port != 'auto':
-            port = str(int(self.config.ajet.interchange_server.interchange_server_port))
-        assert port is not None, "AJET_DAT_INTERCHANGE_PORT env var must be set"
-        master_node_ip = os.getenv("MASTER_NODE_IP", "localhost")
-        base_url = f"http://{master_node_ip}:{port}"
-        return base_url
-
-
     def register_episode_and_wait_output(self, episode_uuid: str, openai_base_url: str, openai_api_key: str) -> WorkflowOutput:
         """Register the episode as ready in the TinkerScript data interchange center."""
-        from ajet.tuner_lib.weight_tuner.experimental.as_tinkerscript_server import RegisterEpisodeRequest
-
         # parse episode_uuid, openai_base_url, openai_api_key
-        zmq_listen_result_addr = self.get_zmq_socket(episode_uuid)
-        interchange_http_addr = self.get_interchange_server_url()
-        rer = RegisterEpisodeRequest(
+        zmq_listen_result_addr, ipc_path = get_zmq_socket(self.config, episode_uuid, tag="workflow")
+        http_register_episode(
+            self.config,
             episode_uuid=episode_uuid,
             openai_base_url=openai_base_url,
             openai_api_key=openai_api_key,
             zmq_listen_result_addr=zmq_listen_result_addr,
         )
-        logger.info(f"zmq_listen_result_addr: {zmq_listen_result_addr}, interchange_http_addr: {interchange_http_addr}")
-
-        # send http request to tinkerscript server to register episode
-        while True:
-            try:
-                response = requests.post(
-                    f"{interchange_http_addr}/register_episode",
-                    json=rer.model_dump(),  # 或者 rer.model_dump() 如果使用 Pydantic v2
-                    timeout=30
-                )
-                response.raise_for_status()
-                result = response.json()
-                if not result.get('success'):
-                    raise RuntimeError(f"Failed to register episode {episode_uuid}")
-                logger.info(f"Successfully registered episode {episode_uuid}")
-                break
-            except requests.RequestException as e:
-                logger.error(f"Error registering episode {episode_uuid}: {e}. Retrying...")
-                time.sleep(5)
+        logger.info(f"zmq_listen_result_addr: {zmq_listen_result_addr}")
 
         # begin wait for result
         zmq_socket = zmq.Context().socket(zmq.REP)
         zmq_socket.bind(zmq_listen_result_addr)
         message = zmq_socket.recv_string()
         logger.success(f"Received workflow output for episode {episode_uuid}")
         zmq_socket.send_string("ack")
+        zmq_socket.close()
+        if ipc_path and os.path.exists(ipc_path): os.remove(ipc_path)
+
         return WorkflowOutput(**json.loads(message))
 
 

ajet/tuner_lib/weight_tuner/as_oai_baseurl_apikey.py

@@ -27,6 +27,7 @@ class OpenaiBaseUrlAndApiKey(BaseModel):
     base_url: str = Field(default="http://localhost:27788/v1", description="The base URL for the Ajet's fake OpenAI API")
     api_key: str = Field(default="invalid_apikey", description="The Ajet's fake key, which is not a real key, it is a encoded string contain episode_uuid and other stuff.")
     model: str = Field(default="reserved_field", description="reserved field.")
+    episode_uuid: str = Field(default="episode_id", description="reserved field.")
 
 
 class OpenaiClientBaseUrlTuner(BaseModel):

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_client.py

@@ -14,7 +14,7 @@
 from openai.types.chat.chat_completion import ChatCompletion
 from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import InterchangeCompletionRequest, API_KEY_PREFIX
 from ajet.utils.thread_executors import SharedInferenceTrackerThreadExecutor, SharedInterchangeThreadExecutor
-from ajet.utils.networking import find_free_port
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import get_zmq_socket
 
 context = zmq.Context()
 atexit.register(context.term)
@@ -67,17 +67,11 @@ def __init__(self, episode_uuid: str, context_tracker: "MultiAgentContextTracker
         self.llm_inference_fn = llm_inference_fn
         self.config = config
         self._should_terminate = False
-
+        self.episode_contect_address, ipc_path = get_zmq_socket(config, episode_uuid, tag="llm")
+        self.ipc_path = ipc_path
         self.interchange_method = config.ajet.interchange_server.interchange_method
-        if self.interchange_method == 'tcp':
-            master_node_ip = os.getenv("MASTER_NODE_IP", "localhost")
-            self.episode_contect_address = f"tcp://{master_node_ip}:{find_free_port()}"
-        elif self.interchange_method == 'ipc':
-            self.ipc_path = f"/tmp/ajet/{self.episode_uuid}.sock"
-            self.episode_contect_address = f"ipc://{self.ipc_path}"
         self.max_inference_tracker_threads = config.ajet.interchange_server.max_inference_tracker_threads
 
-
     async def llm_infer(
             self,
             req: ChatCompletionRequest,

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_server.py

@@ -33,6 +33,8 @@
 from vllm.entrypoints.openai.protocol import ChatCompletionRequest
 from openai.types.chat.chat_completion import ChatCompletion
 
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import EpisodeStatus
+
 API_KEY_PREFIX = "sk-ajet-"
 
 class InterchangeCompletionRequest(BaseModel):
@@ -151,6 +153,21 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
         # Create timeline UUID
         timeline_uuid = uuid.uuid4().hex
 
+        # enable_tinkerscript_mode
+        if enable_tinkerscript_mode:
+            assert shared_mem_dict is not None
+            assert shared_mem_dict_lock is not None
+            if shared_mem_dict['engine_status'] != "ROLLING":
+                logger.error(f"The server is not in ROLLING status (current status: [{shared_mem_dict['engine_status']}]), cannot accept new requests.")
+                raise HTTPException(status_code=503, detail="The server is not in ROLLING status, cannot accept new requests.")
+            if (f"episodes-{episode_uuid}") not in shared_mem_dict:
+                raise HTTPException(status_code=404, detail=f"Episode {episode_uuid} not found.")
+            # update activate timestamp
+            with shared_mem_dict_lock:
+                es:EpisodeStatus = shared_mem_dict[f"episodes-{episode_uuid}"]
+                es.latest_activity_timestamp = time.time()
+                shared_mem_dict[f"episodes-{episode_uuid}"] = es
+
         # Add to received queue
         int_req = InterchangeCompletionRequest(
             completion_request = new_req,

ajet/tuner_lib/weight_tuner/experimental/as_tinkerscript_client.py

@@ -2,47 +2,32 @@
 import time
 import httpx
 import yaml
+from typing import List, Tuple
 from loguru import logger
-from pydantic import BaseModel
 from ajet.schema.task import WorkflowOutput
 from ajet.copilot.job import AgentJetJob
 from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import (
+    SyncTrainConfigRequest,
+    ClaimEpisodeRequest,
+    ClaimEpisodeResponse,
+    CanContinueEpisodeRequest,
+    CanContinueEpisodeResponse,
+    EndEpisodeRequest,
+    EndEpisodeResponse,
+    EpisodeStatus,
+    EpisodeBufferResponse,
+)
 
-# --- Schema Definitions ---
-
-class SyncTrainConfigRequest(BaseModel):
-    yaml_as_string: str
-
-class ClaimEpisodeRequest(BaseModel):
-    client_uuid: str
-    episode_type: str
-
-class ClaimEpisodeResponse(BaseModel):
-    success: bool
-    client_uuid: str
-    episode_uuid: str
-    openai_base_url: str = ""
-    openai_api_key: str = ""
-    fail_cause: str = ""
-
-class EndEpisodeRequest(BaseModel):
-    client_uuid: str
-    episode_uuid: str
-    workflow_output: WorkflowOutput
-
-class EndEpisodeResponse(BaseModel):
-    success: bool
 
 class TinkerScriptClient(object):
 
     def __init__(self, server_url: str):
         self.server_url = server_url
         self.client_uuid = str(uuid.uuid4())
-        self.episode_uuid = None
-        self.openai_base_url = None
-        self.openai_api_key = None
 
-    def begin_episode(self) -> OpenaiBaseUrlAndApiKey:
+
+    def begin_episode(self, allow_discard_timeout=60) -> Tuple[str, OpenaiBaseUrlAndApiKey]:
         """
         Block until an episode is claimed.
         Return (episode_uuid, openai_base_url, openai_api_key)
@@ -51,7 +36,8 @@ def begin_episode(self) -> OpenaiBaseUrlAndApiKey:
             try:
                 req_obj = ClaimEpisodeRequest(
                     client_uuid=self.client_uuid,
-                    episode_type="default"
+                    episode_type="default",
+                    allow_discard_timeout=allow_discard_timeout,
                 )
                 resp = httpx.post(
                     f"{self.server_url}/claim_episode",
@@ -60,15 +46,17 @@ def begin_episode(self) -> OpenaiBaseUrlAndApiKey:
                 )
                 resp.raise_for_status()
                 data = ClaimEpisodeResponse.model_validate(resp.json())
+                episode_uuid = data.episode_uuid
 
                 if data.success:
-                    self.episode_uuid = data.episode_uuid
-                    self.openai_base_url = data.openai_base_url
-                    self.openai_api_key = data.openai_api_key
-                    logger.info(f"Claimed episode {self.episode_uuid}")
-                    return OpenaiBaseUrlAndApiKey(
-                        base_url=self.openai_base_url,
-                        api_key=self.openai_api_key,
+                    episode_uuid = data.episode_uuid
+                    openai_base_url = data.openai_base_url
+                    openai_api_key = data.openai_api_key
+                    logger.info(f"Claimed episode {episode_uuid}")
+                    return episode_uuid, OpenaiBaseUrlAndApiKey(
+                        base_url=openai_base_url,
+                        api_key=openai_api_key,
+                        episode_uuid=episode_uuid
                     )
                 else:
                     logger.info(f"Failed to claim episode: {data.fail_cause}. Retrying in 5s...")
@@ -77,15 +65,15 @@ def begin_episode(self) -> OpenaiBaseUrlAndApiKey:
                 logger.error(f"Error claiming episode: {e}. Retrying in 5s...")
                 time.sleep(5)
 
-    def end_episode(self, workflow_output: WorkflowOutput):
-        if not self.episode_uuid:
+    def end_episode(self, episode_uuid: str, workflow_output: WorkflowOutput):
+        if not episode_uuid:
             logger.error("No episode to end.")
             return
 
         try:
             req_obj = EndEpisodeRequest(
                 client_uuid=self.client_uuid,
-                episode_uuid=self.episode_uuid,
+                episode_uuid=episode_uuid,
                 workflow_output=workflow_output
             )
 
@@ -98,10 +86,9 @@ def end_episode(self, workflow_output: WorkflowOutput):
             data = EndEpisodeResponse.model_validate(resp.json())
 
             if data.success:
-                logger.info(f"Ended episode {self.episode_uuid}")
-                self.episode_uuid = None
+                logger.info(f"Ended episode {episode_uuid}")
             else:
-                 logger.error(f"Failed to end episode {self.episode_uuid}")
+                logger.error(f"Failed to end episode {episode_uuid}")
 
         except Exception as e:
             logger.error(f"Error ending episode: {e}")
@@ -122,3 +109,50 @@ def sync_train_config(self, agent_jet_job: AgentJetJob):
             logger.info("Synced train config")
         except Exception as e:
             logger.error(f"Error syncing train config: {e}")
+
+    def get_engine_status(self) -> str:
+        try:
+            resp = httpx.get(
+                f"{self.server_url}/get_engine_status",
+                timeout=10
+            )
+            resp.raise_for_status()
+            return resp.json().get("engine_status", "unknown")
+        except Exception as e:
+            logger.error(f"Error getting engine status: {e}")
+            return "unknown"
+
+    def can_continue_episode(self, episode_uuid: str) -> bool:
+        if not episode_uuid:
+            return False
+
+        try:
+            req_obj = CanContinueEpisodeRequest(
+                client_uuid=self.client_uuid,
+                episode_uuid=episode_uuid
+            )
+            resp = httpx.post(
+                f"{self.server_url}/can_continue_episode",
+                json=req_obj.model_dump(),
+                timeout=10
+            )
+            resp.raise_for_status()
+            data = CanContinueEpisodeResponse.model_validate(resp.json())
+            return data.can_continue
+        except Exception as e:
+            logger.error(f"Error checking can_continue_episode: {e}")
+            return False
+
+    def get_episode_buffer(self) -> List[EpisodeStatus]:
+        try:
+            resp = httpx.post(
+                f"{self.server_url}/get_episode_buffer",
+                json={},
+                timeout=10
+            )
+            resp.raise_for_status()
+            data = EpisodeBufferResponse.model_validate(resp.json())
+            return data.buffer
+        except Exception as e:
+            logger.error(f"Error getting episode buffer: {e}")
+            return []