sharing httpx client

Author: binary-husky

.gitignore

@@ -171,3 +171,4 @@ prompts
 swarmexp
 swarmlog
 werewolves_swarm
+.claude

ajet/task_rollout/native_parallel_worker.py

@@ -233,7 +233,7 @@ def enough_finished_task_stop_condition(completed_task_id_map_ct) -> bool:
                     f"Too many cached episodes [{total_completed_episodes}] has exceeded the max cached episodes [{self.config.ajet.swarm_mode_sample_collection_max_cached_episodes}] "
                     f"Deleting cached episodes to release memory..."
                 )
-                completed_task_id_map_ct = {}
+                completed_task_id_map_ct.clear()
             return (total_completed_tasks >= n_batch_task)
 
         def enough_non_dummy_task_stop_condition(completed_task_id_map_ct) -> bool:
@@ -258,7 +258,7 @@ def enough_non_dummy_task_stop_condition(completed_task_id_map_ct) -> bool:
                     f"Too many cached episodes [{total_completed_episodes}] has exceeded the max cached episodes [{self.config.ajet.swarm_mode_sample_collection_max_cached_episodes}] "
                     f"Deleting cached episodes to release memory..."
                 )
-                completed_task_id_map_ct = {}
+                completed_task_id_map_ct.clear()
             return (total_completed_non_dummy_tasks >= n_batch_task)
 
         # select stop condition function based on config
@@ -387,6 +387,7 @@ def update_rollout_result_array_preview(observation_window, completed_task_id_ma
         logger.info('Collecting results...')
         for ct_list in completed_task_id_map_ct.values():
             tracker_array.extend(ct_list)
+        completed_task_id_map_ct.clear()
 
         # TODO: support multi-step reward
         task_success_rate = np.mean(
@@ -402,7 +403,6 @@ def update_rollout_result_array_preview(observation_window, completed_task_id_ma
 
         update_rollout_result_array_preview(observation_window, completed_task_id_map_ct)
         self._write_swarm_rollout_dynamic_log(observation_window)
-
         return tracker_array
 
 

ajet/tuner_lib/experimental/as_oai_model_server.py

@@ -428,6 +428,7 @@ def start_interchange_server(config, blocking=False, env={}) -> int:
 
     # polling for server ready
     start_time = time.time()
+    _httpx_client = httpx.Client(timeout=0.5)
     while True:
         if interchange_server and interchange_server.exitcode is not None:
             logger.error(f"Interchange server subprocess failed to start. Return code: {interchange_server.exitcode}")
@@ -437,7 +438,7 @@ def start_interchange_server(config, blocking=False, env={}) -> int:
             logger.error(msg)
             raise RuntimeError(msg)
         try:
-            if httpx.get(health_url, timeout=0.5).status_code == 200:
+            if _httpx_client.get(health_url).status_code == 200:
                 break
         except Exception:
             # keep waiting
@@ -462,7 +463,7 @@ def start_interchange_server(config, blocking=False, env={}) -> int:
                 interchange_server.join()
         except KeyboardInterrupt:
             logger.info("Shutting down interchange server...")
-            try: httpx.post(f"http://127.0.0.1:{port}/stop_engine", timeout=8).status_code
+            try: _httpx_client.post(f"http://127.0.0.1:{port}/stop_engine", timeout=8).status_code
             except Exception: pass
 
             if interchange_server:

ajet/tuner_lib/experimental/as_swarm_client.py

@@ -71,6 +71,8 @@ def __init__(self, server_url: str):
         self._agent_jet_job = None
         # throttle
         self._recent_seen_tasks = []
+        # reuse httpx client to avoid creating SSL context repeatedly
+        self._http_client = httpx.Client(timeout=GENERAL_TIMEOUT)
 
     def logger_info(self, message):
         # logger with de-duplication within 1 second to prevent log flooding
@@ -252,10 +254,9 @@ def _begin_episode_auto_retry(self, discard_episode_timeout=240, episode_type="t
                     discard_episode_timeout=discard_episode_timeout,
                     throttle_policy=throttle_policy
                 )
-                resp = httpx.post(
+                resp = self._http_client.post(
                     f"{self.server_url}/claim_episode",
-                    json=req_obj.model_dump(),
-                    timeout=GENERAL_TIMEOUT
+                    json=req_obj.model_dump()
                 )
                 raise_for_status_with_detail(resp)
                 data = ClaimEpisodeResponse.model_validate(resp.json())
@@ -337,10 +338,9 @@ def end_episode(self, task:Task, episode_uuid: str, workflow_output: WorkflowOut
             task_id=task_id
         )
 
-        resp = httpx.post(
+        resp = self._http_client.post(
             f"{self.server_url}/end_episode",
-            json=req_obj.model_dump(),
-            timeout=GENERAL_TIMEOUT
+            json=req_obj.model_dump()
         )
         raise_for_status_with_detail(resp)
         data = EndEpisodeResponse.model_validate(resp.json())
@@ -366,10 +366,9 @@ def abort_episode(self, episode_uuid: str):
                 task_id=""
             )
 
-            resp = httpx.post(
+            resp = self._http_client.post(
                 f"{self.server_url}/abort_episode",
-                json=req_obj.model_dump(),
-                timeout=GENERAL_TIMEOUT
+                json=req_obj.model_dump()
             )
             raise_for_status_with_detail(resp)
             data = EndEpisodeResponse.model_validate(resp.json())
@@ -399,10 +398,9 @@ def sync_train_config(self, agent_jet_job: AgentJetJob):
 
             req_obj = SyncTrainConfigRequest(yaml_as_string=yaml_str)
 
-            resp = httpx.post(
+            resp = self._http_client.post(
                 f"{self.server_url}/sync_train_config",
-                json=req_obj.model_dump(),
-                timeout=GENERAL_TIMEOUT
+                json=req_obj.model_dump()
             )
             raise_for_status_with_detail(resp)
             self.logger_info("Synced train config to Swarm server")
@@ -422,7 +420,7 @@ def start_engine(self):
             raise RuntimeError(f"Cannot start engine when engine is NOT ENGINE.OFFLINE. (current status: {current_status})")
 
         # Send start engine request
-        resp = httpx.post(
+        resp = self._http_client.post(
             f"{self.server_url}/start_engine",
             json={},
             timeout=600
@@ -487,7 +485,7 @@ def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=
     @cache_with_ttl(ttl=0.5)
     def get_engine_status(self) -> Tuple[str, dict]:
         try:
-            resp = httpx.get(
+            resp = self._http_client.get(
                 f"{self.server_url}/get_engine_status",
                 timeout=10
             )
@@ -512,7 +510,7 @@ def can_continue_episode(self, episode_uuid: str) -> bool:
                 client_uuid=self.client_uuid,
                 episode_uuid=episode_uuid
             )
-            resp = httpx.post(
+            resp = self._http_client.post(
                 f"{self.server_url}/can_continue_episode",
                 json=req_obj.model_dump(),
                 timeout=10
@@ -526,7 +524,7 @@ def can_continue_episode(self, episode_uuid: str) -> bool:
 
     def get_episode_buffer(self) -> List[EpisodeStatus]:
         try:
-            resp = httpx.post(
+            resp = self._http_client.post(
                 f"{self.server_url}/get_episode_buffer",
                 json={},
                 timeout=10
@@ -585,7 +583,7 @@ def stop_engine(self):
             self.logger_info("Engine is already OFFLINE. No action needed.")
             return
 
-        resp = httpx.post(
+        resp = self._http_client.post(
             f"{self.server_url}/stop_engine",
             json={},
             timeout=600
@@ -605,7 +603,7 @@ def get_rollout_stat(self) -> CurrentBatchRolloutPoolInformation:
         Returns statistics about completed episodes, tasks, and progress.
         """
         try:
-            resp = httpx.get(
+            resp = self._http_client.get(
                 f"{self.server_url}/get_current_batch_rollout_pool_information",
                 timeout=10
             )

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -338,7 +338,8 @@ async def start_engine():
             experiment_dir = config_dict.get("ajet", {}).get("experiment_dir", DEFAULT_DIR)
             if experiment_dir == "auto":
                 exp_base_dir = DEFAULT_DIR
-            exp_base_dir = os.path.dirname(os.path.abspath(experiment_dir))
+            else:
+                exp_base_dir = os.path.dirname(os.path.abspath(experiment_dir))
 
             # Save YAML to temporary file
             with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".yaml") as temp_file:

ajet/tuner_lib/experimental/interchange_utils.py

@@ -109,6 +109,8 @@ class UpdateEngineStatusRequest(BaseModel):
 
 VERBOSE = True
 
+shared_http_client = httpx.Client(timeout=10.0)
+
 def get_interchange_server_url(config):
     port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
     if isinstance(config, dict):
@@ -127,7 +129,7 @@ def http_change_engine_status(config, new_status: str, new_status_detail: str|No
     if new_status not in VALID_STATUSES:
         raise ValueError(f"Invalid engine status: {new_status}")
 
-    resp = httpx.post(
+    resp = shared_http_client.post(
         f"{get_interchange_server_url(config)}/update_engine_status",
         json={"engine_status": new_status, "engine_status_detail": new_status_detail, "global_step": global_step},
         timeout=10
@@ -137,7 +139,7 @@ def http_change_engine_status(config, new_status: str, new_status_detail: str|No
 
 
 def is_episode_claimed(config, episode_uuid: str, unregister_if_not_claimed: bool) -> bool:
-    resp = httpx.post(
+    resp = shared_http_client.post(
         f"{get_interchange_server_url(config)}/is_episode_claimed",
         json={"episode_uuid": episode_uuid, "unregister_if_not_claimed": unregister_if_not_claimed},
         timeout=5
@@ -168,7 +170,7 @@ def http_register_episode(config,
         zmq_listen_result_addr=zmq_listen_result_addr,
     )
     # send http request to swarm server to register episode
-    response = httpx.post(
+    response = shared_http_client.post(
         f"{interchange_http_addr}/register_episode",
         json=rer.model_dump(),  # 或者 rer.model_dump() 如果使用 Pydantic v2
         timeout=2

ajet/utils/swarm_overwatch.py

@@ -37,11 +37,12 @@ def __init__(self, server_url: str, refresh_interval: float = 2.0):
         self.last_update_time = None
         self.error_count = 0
         self.total_requests = 0
+        self._httpx_client = httpx.Client(timeout=5.0)
 
     def fetch_pool_info(self) -> Optional[CurrentBatchRolloutPoolInformation]:
         """Fetch current batch rollout pool information from server"""
         try:
-            response = httpx.get(
+            response = self._httpx_client.get(
                 f"{self.server_url}/get_current_batch_rollout_pool_information",
                 timeout=5.0,
             )

ajet/utils/tokenizer.py

@@ -19,6 +19,9 @@ def cleanup_messages(messages: List[Dict]) -> List[Dict]:
                     pass
     return messages_copied
 
+# Cache storage
+_cache = {}
+
 
 def ajet_apply_chat_template(
     tokenizer,
@@ -28,16 +31,39 @@ def ajet_apply_chat_template(
     tokenize: bool = True,
 ):
     conversation = cleanup_messages(conversation)
+
+    # Create cache key by hashing all inputs
+    cache_key = (
+        id(tokenizer),
+        hash(json.dumps(conversation, sort_keys=True)),
+        hash(json.dumps(tools, sort_keys=True)) if tools else 0,
+        add_generation_prompt,
+        tokenize,
+    )
+
+    # Check cache
+    if cache_key in _cache:
+        return _cache[cache_key]
+
+    # Compute result
     if tools:
-        return tokenizer.apply_chat_template(
+        result = tokenizer.apply_chat_template(
             conversation,
             tools,
             add_generation_prompt=add_generation_prompt,
             tokenize=tokenize,
         )
     else:
-        return tokenizer.apply_chat_template(
+        result = tokenizer.apply_chat_template(
             conversation,
             tokenize=tokenize,
             add_generation_prompt=add_generation_prompt,
         )
+
+    # Store in cache (implement LRU eviction if cache gets too large)
+    if len(_cache) >= 1024:
+        # Remove oldest item (first inserted)
+        _cache.pop(next(iter(_cache)))
+
+    _cache[cache_key] = result
+    return result