optimize further

Author: binary-husky

ajet/context_tracker/multiagent_tracking.py

@@ -107,6 +107,7 @@ def extract_text_content_from_content_dict(self, msg):
         should_skip_message = False
         return str_content, should_skip_message
 
+
     def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_toolcalls: bool = False) -> List[ExtendedMessage]:
         """Spawn a timeline from messages.
 

ajet/task_rollout/async_llm_bridge.py

@@ -535,14 +535,11 @@ async def run_infer(
 
         # run llm inference ✨
         if self.config.ajet.llm_infer_submit_method == "sync":
-            llm_output = await asyncio.wait_for(
-                asyncio.to_thread(
-                    self.llm_inference_fn, converted_message, custom_sampling_params, tools
-                ),
-                timeout=1800,
+            llm_output = await asyncio.to_thread(
+                self.llm_inference_fn, converted_message, custom_sampling_params, tools
             )
         else:
-            llm_output = await asyncio.wait_for(self.llm_inference_fn(converted_message, custom_sampling_params, tools), timeout=1800)
+            llm_output = await self.llm_inference_fn(converted_message, custom_sampling_params, tools)
 
 
         # begin context tracking

ajet/task_runner/base_runner.py

@@ -68,9 +68,16 @@ def generated_token_callback_fn(token_array):
     async def wrapper_type_asyncio(self, workflow_cls: Type[Workflow], workflow_task: WorkflowTask, tuner: AjetTuner) -> WorkflowOutput:
         user_workflow: Workflow = workflow_cls(name="ajet-workflow")
         result = await user_workflow.execute(workflow_task, tuner)
+
+        # malloc garbage collection
         del user_workflow
-        with gc_lock:
-            gc.collect()    # force garbage collection
+
+        # run gc in a thread-safe way
+        if gc_lock.acquire(blocking=False):
+            try:
+                gc.collect()
+            finally:
+                gc_lock.release()
         return result
 
 

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_client.py

@@ -13,7 +13,7 @@
 from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import InterchangeCompletionRequest
 from redis.exceptions import TimeoutError
 from ajet.utils.free_port import find_free_port
-from ajet.utils.sington import ThreadExecutorLlmInferSingleton, ThreadExecutorSingleton
+from ajet.utils.sington import ThreadExecutorContextTrackerSingleton, ThreadExecutorSingleton
 from functools import cache
 
 import pickle
@@ -141,14 +141,19 @@ def begin_service(self):
         if DEBUG: logger.info(f"[client] {self.episode_uuid} | Starting InterchangeClient service loop...")
         self.socket = context.socket(zmq.REP)
         self.socket.bind(f"{self.episode_contect_address}")
-        self.socket.setsockopt(zmq.RCVTIMEO, 2*1000)  # 60 秒超时
+        self.socket.setsockopt(zmq.RCVTIMEO, 3*1000)  # 60 秒超时
 
         self.executor = ThreadExecutorSingleton().get_executor()
         if DEBUG: logger.info(f"[client] {self.episode_uuid} | Submitting _begin_service_threading to executor...")
         future = self.executor.submit(self._begin_service_threading)
-        time.sleep(1)
+
+        # wait till service begin running
+        time.sleep(0.5)
+        w_time = 1
         while future._state == 'PENDING':
-            time.sleep(1)
+            time.sleep(min(w_time * 2, 10))
+            w_time += 1
+
         if DEBUG: logger.info(f"[client] {self.episode_uuid} | Future ready...")
 
         # t = threading.Thread(target=self._begin_service_threading, daemon=True)
@@ -189,9 +194,9 @@ def _begin_service_threading(self):
                     loop = asyncio.get_running_loop()
                 except:
                     loop = asyncio.new_event_loop()
-                executor = ThreadExecutorLlmInferSingleton().get_executor()
+                context_tracker_executor = ThreadExecutorContextTrackerSingleton().get_executor()
                 future = loop.run_in_executor(
-                    executor,  # executor
+                    context_tracker_executor,
                     asyncio.run,
                     self.llm_infer(
                         req=parsed_msg.completion_request,

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_server.py

@@ -186,7 +186,7 @@ async def serve_with_monitor():
                 host="0.0.0.0",
                 port=self.port,
                 log_level="error",
-                # workers=4
+                workers=2
             )
             server = uvicorn.Server(config)
             await server.serve()

ajet/utils/sington.py

@@ -4,16 +4,16 @@
 @singleton
 class ThreadExecutorSingleton:
     def __init__(self):
-        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=16)
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=64)
 
     def get_executor(self) -> concurrent.futures.ThreadPoolExecutor:
         return self.executor
 
 
 @singleton
-class ThreadExecutorLlmInferSingleton:
+class ThreadExecutorContextTrackerSingleton:
     def __init__(self):
-        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=16)
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=64)
 
     def get_executor(self) -> concurrent.futures.ThreadPoolExecutor:
         return self.executor