remove weight_tuner directory layer

binary-husky · binary-husky · commit 9637c9cc5597 · 2026-02-11T17:31:18.000+08:00
diff --git a/ajet/backbone/main_trinity.py b/ajet/backbone/main_trinity.py
@@ -53,7 +53,7 @@ def patched_trainer_get_actor(cls, config: Config):
     Trainer.get_actor = classmethod(patched_trainer_get_actor)
 
     if ajet_config.ajet.enable_experimental_interchange_server:
-        from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
+        from ajet.tuner_lib.experimental.as_oai_model_server import start_interchange_server
         start_interchange_server(ajet_config)
 
 
diff --git a/ajet/backbone/main_verl.py b/ajet/backbone/main_verl.py
@@ -21,7 +21,6 @@
 
 import hydra
 import ray
-from beast_logger import print_dict
 from omegaconf import DictConfig, OmegaConf
 from verl.trainer.ppo.reward import load_reward_manager
 from verl.utils.device import is_cuda_available
@@ -243,7 +242,7 @@ def run(self, config):
         from ajet.backbone.trainer_verl import AjetRayPPOTrainer
 
         if config.ajet.enable_experimental_interchange_server:
-            from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
+            from ajet.tuner_lib.experimental.as_oai_model_server import start_interchange_server
             start_interchange_server(config)
 
         # Initialize the PPO trainer.
diff --git a/ajet/backbone/main_vllm.py b/ajet/backbone/main_vllm.py
@@ -187,10 +187,10 @@ def main(config):
     # atexit.register(lambda: print("Process exiting, performing cleanup..."))
 
     if config.ajet.enable_experimental_interchange_server:
-        from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
+        from ajet.tuner_lib.experimental.as_oai_model_server import start_interchange_server
         start_interchange_server(config)
         if config.ajet.enable_swarm_mode:
-            from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
+            from ajet.tuner_lib.experimental.interchange_utils import http_change_engine_status
             http_change_engine_status(config, "ENGINE.ROLLING")
 
     def companion_launch():
diff --git a/ajet/backbone/trainer_verl.py b/ajet/backbone/trainer_verl.py
@@ -459,7 +459,7 @@ def init_workers(self):
     def _update_interchange_server_status_flag(self, status: str):
         if self.config.ajet.enable_experimental_interchange_server:
             if self.config.ajet.enable_swarm_mode:
-                from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
+                from ajet.tuner_lib.experimental.interchange_utils import http_change_engine_status
                 http_change_engine_status(self.config, status, global_step=self.global_steps)
 
     # #######################################
@@ -859,7 +859,7 @@ def fit(self):  # noqa: C901
 
                 # # when enabled oai request interchange, we need to clear the cache from time to time
                 # if self.config.ajet.enable_experimental_interchange_server:
-                #     from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import ensure_dat_interchange_server_cache_clear
+                #     from ajet.tuner_lib.experimental.as_oai_model_server import ensure_dat_interchange_server_cache_clear
                 #     ensure_dat_interchange_server_cache_clear()
 
                 if is_last_step:
diff --git a/ajet/launcher.py b/ajet/launcher.py
@@ -1,7 +1,5 @@
 import argparse
 import os
-import subprocess
-from types import SimpleNamespace
 
 from dotenv import load_dotenv
 from loguru import logger
@@ -154,7 +152,7 @@ def start_swarm_server(env, config):
     assert config.ajet.enable_experimental_interchange_server, (
         "Please enable_experimental_interchange_server in config to start swarm server."
     )
-    from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import (
+    from ajet.tuner_lib.experimental.as_oai_model_server import (
         start_interchange_server,
     )
 
diff --git a/ajet/task_rollout/native_parallel_worker.py b/ajet/task_rollout/native_parallel_worker.py
@@ -20,7 +20,7 @@
 from ajet.schema.trajectory import Sample
 from ajet.task_rollout.single_worker import BaseRolloutManager
 from ajet.context_tracker.basic_tracker import BaseContextTracker
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import (
+from ajet.tuner_lib.experimental.interchange_utils import (
     http_change_engine_status,
     http_update_rollout_pool_information,
     CurrentBatchRolloutPoolInformation,
diff --git a/ajet/task_runner/base_runner.py b/ajet/task_runner/base_runner.py
@@ -11,7 +11,7 @@
 from ajet.utils.async_utils import run_async_coroutine_with_timeout
 from ajet.utils.dynamic_import import dynamic_import
 from ajet.workflow import Workflow
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import is_episode_claimed
+from ajet.tuner_lib.experimental.interchange_utils import is_episode_claimed
 
 gc_lock = Lock()
 
diff --git a/ajet/task_runner/swarm_runner.py b/ajet/task_runner/swarm_runner.py
@@ -11,8 +11,8 @@
 from ajet.schema.trajectory import Reward
 from ajet.task_runner.base_runner import BaseAgentRunner
 from ajet.utils.retry import SwarmReceiveAbortException
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_register_episode, get_zmq_socket, is_episode_claimed
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG
+from ajet.tuner_lib.experimental.interchange_utils import http_register_episode, get_zmq_socket
+from ajet.tuner_lib.experimental.interchange_utils import DEBUG
 from loguru import logger
 from ajet import Workflow
 from typing import Callable
@@ -66,11 +66,11 @@ def register_episode_and_wait_output(
 
             while True:
                 # <wait for 1/2>:
-                #   <from_sourcefile>: ajet/tuner_lib/weight_tuner/experimental/as_swarm_server.py
+                #   <from_sourcefile>: ajet/tuner_lib/experimental/as_swarm_server.py
                 #   <from_code>: socket.send_string(workflow_output.model_dump_json())
                 #   <expect>: workflow_output: WorkflowOutput
                 # <wait for 2/2>:
-                #   <from_sourcefile>: ajet/tuner_lib/weight_tuner/experimental/as_swarm_server.py
+                #   <from_sourcefile>: ajet/tuner_lib/experimental/as_swarm_server.py
                 #   <from_code>: socket.send_string("RUNNER.SPECIAL.RESET_CONTEXT_TRACKER")
                 #   <expect>: "RUNNER.SPECIAL.RESET_CONTEXT_TRACKER"
                 try:
diff --git a/ajet/tuner.py b/ajet/tuner.py
@@ -1,9 +1,9 @@
 from typing import TYPE_CHECKING, Callable, Union, Type
 
 from ajet.context_tracker.multiagent_tracking import MultiAgentContextTracker
-from ajet.tuner_lib.weight_tuner import AgentScopeModelTuner
-from ajet.tuner_lib.weight_tuner import OpenaiClientModelTuner
-from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiClientBaseUrlTuner
+from ajet.tuner_lib import AgentScopeModelTuner
+from ajet.tuner_lib import OpenaiClientModelTuner
+from ajet.tuner_lib.as_oai_baseurl_apikey import OpenaiClientBaseUrlTuner
 if TYPE_CHECKING:
     from ajet import Workflow
 
@@ -171,7 +171,7 @@ def get_context_tracker(self) -> MultiAgentContextTracker:
     def _enable_experimental_interchange_server(self, llm_inference_fn):
         # experimental reverse proxy start
         if self.enable_interchange_server:
-            from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_client import InterchangeClient
+            from ajet.tuner_lib.experimental.as_oai_model_client import InterchangeClient
             self.interchange_client = InterchangeClient(
                 episode_uuid=self.context_tracker.episode_uuid,
                 context_tracker=self.context_tracker,
diff --git a/ajet/tuner_lib/__init__.py b/ajet/tuner_lib/__init__.py
@@ -0,0 +1,3 @@
+
+from ajet.tuner_lib.as_agentscope_model import AgentScopeModelTuner
+from ajet.tuner_lib.as_oai_sdk_model import OpenaiClientModelTuner
diff --git a/ajet/tuner_lib/as_agentscope_model.py b/ajet/tuner_lib/as_agentscope_model.py
diff --git a/ajet/tuner_lib/as_oai_baseurl_apikey.py b/ajet/tuner_lib/as_oai_baseurl_apikey.py
diff --git a/ajet/tuner_lib/as_oai_sdk_model.py b/ajet/tuner_lib/as_oai_sdk_model.py
diff --git a/ajet/tuner_lib/experimental/as_oai_model_client.py b/ajet/tuner_lib/experimental/as_oai_model_client.py
@@ -12,10 +12,10 @@
 from typing import TYPE_CHECKING
 from vllm.entrypoints.openai.protocol import ChatCompletionRequest
 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.tuner_lib.experimental.as_oai_model_server import InterchangeCompletionRequest, API_KEY_PREFIX
 from ajet.utils.thread_executors import SharedInferenceTrackerThreadExecutor, SharedInterchangeThreadExecutor
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import get_zmq_socket, is_episode_claimed
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG
+from ajet.tuner_lib.experimental.interchange_utils import get_zmq_socket
+from ajet.tuner_lib.experimental.interchange_utils import DEBUG
 
 context = zmq.Context()
 atexit.register(context.term)
@@ -215,4 +215,3 @@ def _begin_service_threading(self):
                 if os.path.exists(self.ipc_path):
                     os.remove(self.ipc_path)
                     if DEBUG: logger.info(f"[client] {self.episode_uuid} | IPC socket file {self.ipc_path} removed.")
-
diff --git a/ajet/tuner_lib/experimental/as_oai_model_server.py b/ajet/tuner_lib/experimental/as_oai_model_server.py
@@ -33,9 +33,9 @@
 from vllm.entrypoints.openai.protocol import ChatCompletionRequest
 from openai.types.chat.chat_completion import ChatCompletion
 
-from ajet.utils.networking import find_free_port, get_host_ip
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import EpisodeStatus
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG, VERBOSE
+from ajet.utils.networking import get_host_ip
+from ajet.tuner_lib.experimental.interchange_utils import EpisodeStatus
+from ajet.tuner_lib.experimental.interchange_utils import DEBUG, VERBOSE
 
 API_KEY_PREFIX = "sk-ajet-"
 
@@ -158,7 +158,7 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
 
         # enable_swarm_mode
         if enable_swarm_mode:
-            from ajet.tuner_lib.weight_tuner.experimental.as_swarm_server import ep_key
+            from ajet.tuner_lib.experimental.as_swarm_server import ep_key
             assert shared_mem_dict is not None
             assert shared_mem_dict_lock is not None
 
@@ -189,7 +189,7 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
 
 
     if enable_swarm_mode:
-        from ajet.tuner_lib.weight_tuner.experimental.as_swarm_server import register_enable_swarm_mode_routes
+        from ajet.tuner_lib.experimental.as_swarm_server import register_enable_swarm_mode_routes
         assert shared_mem_dict is not None, "shared_mem_dict must not be None when enable_swarm_mode is True."
         assert shared_mem_dict_lock is not None, "shared_mem_dict_lock must not be None when enable_swarm_mode is True."
         app, additional_coro = register_enable_swarm_mode_routes(app, zmq_context=context, shared_mem_dict=shared_mem_dict, shared_mem_dict_lock=shared_mem_dict_lock)
@@ -354,6 +354,6 @@ def start_interchange_server(config, blocking=False, env={}) -> int:
             if interchange_server:
                 interchange_server.terminate()
             if enable_swarm_mode:
-                from ajet.tuner_lib.weight_tuner.experimental.as_swarm_server import kill_process_tree
+                from ajet.tuner_lib.experimental.as_swarm_server import kill_process_tree
                 kill_process_tree(None, None)
-        return -1
+        return -1
diff --git a/ajet/tuner_lib/experimental/as_swarm_client.py b/ajet/tuner_lib/experimental/as_swarm_client.py
@@ -9,8 +9,8 @@
 from ajet.schema.task import WorkflowOutput, Task
 from ajet.copilot.job import AgentJetJob
 from ajet.utils.thread_executors import BoundedThreadPoolExecutor
-from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import (
+from ajet.tuner_lib.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
+from ajet.tuner_lib.experimental.interchange_utils import (
     SyncTrainConfigRequest,
     ClaimEpisodeRequest,
     ClaimEpisodeResponse,
diff --git a/ajet/tuner_lib/experimental/as_swarm_server.py b/ajet/tuner_lib/experimental/as_swarm_server.py
@@ -11,8 +11,8 @@
 from multiprocessing.managers import DictProxy
 from typing import Coroutine, Optional, Tuple, List
 from ajet.utils.process_killer import kill_process_tree
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG, VERBOSE
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import (
+from ajet.tuner_lib.experimental.interchange_utils import DEBUG, VERBOSE
+from ajet.tuner_lib.experimental.interchange_utils import (
     SyncTrainConfigRequest,
     ClaimEpisodeRequest,
     ClaimEpisodeResponse,
@@ -689,7 +689,6 @@ async def get_current_batch_rollout_pool_information():
                 CurrentBatchRolloutPoolInformation(),
             )
             # Fetch additional fields from shared_mem_dict
-            pool_info.sample_collection_method = shared_mem_dict.get("sample_collection_method", "")
             pool_info.engine_status = shared_mem_dict.get("engine_status", None)
             pool_info.global_step = shared_mem_dict.get("global_step", None)
 
diff --git a/ajet/tuner_lib/experimental/interchange_utils.py b/ajet/tuner_lib/experimental/interchange_utils.py
diff --git a/ajet/tuner_lib/weight_tuner/__init__.py b/ajet/tuner_lib/weight_tuner/__init__.py
diff --git a/ajet/utils/swarm_overwatch.py b/ajet/utils/swarm_overwatch.py
@@ -16,7 +16,7 @@
 from rich.text import Text
 from loguru import logger
 
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import (
+from ajet.tuner_lib.experimental.interchange_utils import (
     CurrentBatchRolloutPoolInformation,
 )
 
diff --git a/ajet_swarm_threading.py b/ajet_swarm_threading.py
@@ -1,6 +1,5 @@
 import re
 import time
-import threading
 import requests
 from loguru import logger
 from textwrap import dedent
@@ -9,8 +8,8 @@
 from ajet.task_reader import RouterTaskReader
 from ajet.utils.retry import retry_with_backoff
 from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
-from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
-from ajet.tuner_lib.weight_tuner.experimental.as_swarm_client import SwarmClient
+from ajet.tuner_lib.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
+from ajet.tuner_lib.experimental.as_swarm_client import SwarmClient
 from concurrent.futures import ThreadPoolExecutor
 
 # --------- configurations that take effect locally -------------
diff --git a/tutorial/example_academic_trans/trans.py b/tutorial/example_academic_trans/trans.py
@@ -1,25 +1,18 @@
 
 import re
 import os
-import time
 import asyncio
-import threading
-from loguru import logger
-from textwrap import dedent
 from openai import OpenAI
 
 from ajet import WorkflowOutput
 from ajet.schema.task import Task
-from ajet.copilot.job import AgentJetJob
-from ajet.task_reader import RouterTaskReader
 from ajet.utils.retry import retry_with_backoff
-from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
-from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
+from ajet.tuner_lib.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
 from beast_logger import print_listofdict
 
 # Import reward computation from trans_reward.py
 from openjudge.models import OpenAIChatModel
-from .trans_reward import TranslationQualityGrader, build_translation_quality_messages, examples
+from .trans_reward import TranslationQualityGrader, examples
 
 
 
diff --git a/tutorial/example_academic_trans/trans_roll.py b/tutorial/example_academic_trans/trans_roll.py
@@ -1,14 +1,7 @@
-from loguru import logger
-from textwrap import dedent
 from ajet.copilot.job import AgentJetJob
-from ajet.tuner_lib.weight_tuner.experimental.as_swarm_client import SwarmClient, run_episodes_until_all_complete, auto_train_with_dataset
+from ajet.tuner_lib.experimental.as_swarm_client import SwarmClient, run_episodes_until_all_complete
 from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
-from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
-from ajet.utils.thread_executors import BoundedThreadPoolExecutor
-from ajet.schema.task import Task
 from ajet.task_reader import RouterTaskReader
-from ajet.utils.retry import retry_with_backoff
-from concurrent.futures import ThreadPoolExecutor
 from tutorial.example_academic_trans.trans import execute_agent
 
 # python -m tutorial.example_academic_trans.trans_roll
@@ -80,5 +73,3 @@ def rollout(task) -> float | None:
 
 if __name__ == "__main__":
     main()
-
-
diff --git a/tutorial/example_academic_trans/trans_roll_backup.py b/tutorial/example_academic_trans/trans_roll_backup.py
@@ -1,14 +1,8 @@
 from loguru import logger
-from textwrap import dedent
 from ajet.copilot.job import AgentJetJob
-from ajet.tuner_lib.weight_tuner.experimental.as_swarm_client import SwarmClient, run_episodes_until_all_complete, auto_train_with_dataset
+from ajet.tuner_lib.experimental.as_swarm_client import SwarmClient, run_episodes_until_all_complete
 from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
-from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
-from ajet.utils.thread_executors import BoundedThreadPoolExecutor
-from ajet.schema.task import Task
 from ajet.task_reader import RouterTaskReader
-from ajet.utils.retry import retry_with_backoff
-from concurrent.futures import ThreadPoolExecutor
 from tutorial.example_academic_trans.trans import execute_agent
 
 # python -m tutorial.example_academic_trans.trans_roll
@@ -82,5 +76,3 @@ def rollout(task) -> float | None:
 
 if __name__ == "__main__":
     main()
-
-

Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,5 @@`
`1`	`1`	`import argparse`
`2`	`2`	`import os`
`3`		`-import subprocess`
`4`		`-from types import SimpleNamespace`
`5`	`3`
`6`	`4`	`from dotenv import load_dotenv`
`7`	`5`	`from loguru import logger`
`@@ -154,7 +152,7 @@ def start_swarm_server(env, config):`
`154`	`152`	`assert config.ajet.enable_experimental_interchange_server, (`
`155`	`153`	`"Please enable_experimental_interchange_server in config to start swarm server."`
`156`	`154`	`)`
`157`		`- from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import (`
	`155`	`+ from ajet.tuner_lib.experimental.as_oai_model_server import (`
`158`	`156`	`start_interchange_server,`
`159`	`157`	`)`
`160`	`158`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+`
	`2`	`+from ajet.tuner_lib.as_agentscope_model import AgentScopeModelTuner`
	`3`	`+from ajet.tuner_lib.as_oai_sdk_model import OpenaiClientModelTuner`