keyurl multinode training support

Author: binary-husky

ajet/backbone/main_trinity.py

@@ -23,7 +23,8 @@ def get_ajet_config_from_trinity_side():
 
 def patch_runtime_env_to_get_actor():
     """Patch the classmethod of Explorer and Trainer to pass in the runtime env."""
-    runtime_env = get_runtime_env(is_trinity=True)
+    ajet_config = get_ajet_config_from_trinity_side()
+    runtime_env = get_runtime_env(ajet_config, is_trinity=True)
 
     def patched_explorer_get_actor(cls, config: Config):
         return (
@@ -50,7 +51,6 @@ def patched_trainer_get_actor(cls, config: Config):
     Explorer.get_actor = classmethod(patched_explorer_get_actor)
     Trainer.get_actor = classmethod(patched_trainer_get_actor)
 
-    ajet_config = get_ajet_config_from_trinity_side()
     if ajet_config.ajet.enable_experimental_interchange_server:
         from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
         start_interchange_server(ajet_config)

ajet/backbone/main_verl.py

@@ -55,7 +55,7 @@ def run_ppo(config) -> None:
     # Check if Ray is not initialized
     if not ray.is_initialized():
         # this is for local ray cluster
-        runtime_env = get_runtime_env()
+        runtime_env = get_runtime_env(config)
         print_dict(runtime_env["env_vars"], "runtime_env")
         ray.init(
             runtime_env=runtime_env,

ajet/backbone/main_vllm.py

@@ -180,7 +180,7 @@ def run(config):
 def main(config):
     from omegaconf import OmegaConf
     OmegaConf.resolve(config)
-    runtime_env = get_runtime_env()
+    runtime_env = get_runtime_env(config)
     os.environ.update(runtime_env["env_vars"])
     # atexit.register(lambda: print("Process exiting, performing cleanup..."))
 

ajet/backbone/warm_up.py

@@ -43,6 +43,42 @@ def init_parallel_rollout_logger(experiment_name):
 
 
 
+def warm_up_task_judge_when_needed(config):
+    if config.ajet.task_judge.judge_type == "rubrics_auto_grader":
+        from ajet.task_judge.rm_auto_grader_judge import AutoGraderJudge
+
+        judge = AutoGraderJudge(config)
+        asyncio.run(judge.generate_rubrics_from_samples())
+        asyncio.run(judge.load_rubrics_from_cache())
+
+
+def clean_up_tmp_ajet_dir(config):
+    """Clean up old IPC socket files in /tmp/ajet directory."""
+    import time
+    if config.ajet.enable_experimental_interchange_server is False:
+        return
+
+    tmp_dir = "/tmp/ajet"
+    if not os.path.exists(tmp_dir):
+        return
+    current_time = time.time()
+    ttl = 4 * 3600
+    try:
+        for filename in os.listdir(tmp_dir):
+            if not filename.endswith(".sock"):
+                continue
+
+            file_path = os.path.join(tmp_dir, filename)
+            try:
+                print(current_time - os.path.getmtime(file_path))
+                if current_time - os.path.getmtime(file_path) > ttl:
+                    os.remove(file_path)
+            except OSError:
+                pass
+    except OSError:
+        pass
+
+
 def warm_up_process(config):
     """
     Process level warm up
@@ -65,12 +101,4 @@ def warm_up_process(config):
     experiment_name = config.ajet.experiment_name
     init_parallel_rollout_logger(experiment_name)
     warm_up_task_judge_when_needed(config)
-
-
-def warm_up_task_judge_when_needed(config):
-    if config.ajet.task_judge.judge_type == "rubrics_auto_grader":
-        from ajet.task_judge.rm_auto_grader_judge import AutoGraderJudge
-
-        judge = AutoGraderJudge(config)
-        asyncio.run(judge.generate_rubrics_from_samples())
-        asyncio.run(judge.load_rubrics_from_cache())
+    clean_up_tmp_ajet_dir(config)

ajet/default_config/ajet_default.yaml

@@ -280,7 +280,7 @@ ajet:
     save_trajectory_as_json_file: False
 
 
-  # the experimental reverse proxy feature that allows `tuner.as_oai_baseurl_apikey` feature
+  # the experimental ZeroMQ interchange server feature that allows `tuner.as_oai_baseurl_apikey` feature
   enable_experimental_interchange_server: True
   interchange_server:
     interchange_method: 'ipc' # options: 'tcp' (multi-nodes) or  'ipc' (1 node)

ajet/tuner.py

@@ -26,7 +26,8 @@ def __init__(
         self.context_tracker = context_tracker
         self.llm_inference_fn = llm_inference_fn
         self.target2proxy_registry: dict[str, dict[str,TunerTypeUnion]] = {}
-        if config.ajet.enable_experimental_interchange_server:
+        self.enable_interchange_server = config.ajet.enable_experimental_interchange_server
+        if self.enable_interchange_server:
             self.proxy_client_started = False
 
 
@@ -104,7 +105,7 @@ def as_oai_baseurl_apikey(
             ```
         """
 
-        assert self.config.ajet.enable_experimental_interchange_server, "Please enable `ajet.enable_experimental_interchange_server` in yaml config to use `as_oai_baseurl_apikey` feature."
+        assert self.enable_interchange_server, "Please enable `ajet.enable_experimental_interchange_server` in yaml config to use `as_oai_baseurl_apikey` feature."
         if self.proxy_client_started is False:
             self.proxy_client_started = True
             self._enable_experimental_interchange_server(self.llm_inference_fn)
@@ -171,7 +172,7 @@ def get_context_tracker(self) -> MultiAgentContextTracker:
 
     def _enable_experimental_interchange_server(self, llm_inference_fn):
         # experimental reverse proxy start
-        if self.config.ajet.enable_experimental_interchange_server:
+        if self.enable_interchange_server:
             from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_client import InterchangeClient
             self.interchange_client = InterchangeClient(
                 episode_uuid=self.context_tracker.episode_uuid,
@@ -184,6 +185,6 @@ def _enable_experimental_interchange_server(self, llm_inference_fn):
 
     def terminate_episode(self):
         # experimental reverse proxy cleanup
-        if self.config.ajet.enable_experimental_interchange_server:
+        if self.enable_interchange_server:
             if (self.proxy_client_started is True) and hasattr(self, "interchange_client"):
                 self.interchange_client._should_terminate = True

ajet/tuner_deprecated.py

@@ -12,7 +12,7 @@
 from openai.resources.chat.chat import Chat, AsyncChat
 from openai.resources.completions import AsyncCompletions
 from openai import OpenAI, AsyncOpenAI
-from ajet.utils.free_port import find_free_port
+from ajet.utils.networking import find_free_port
 from .experimental.as_oai_model_client import generate_auth_token
 
 if TYPE_CHECKING:
@@ -47,9 +47,12 @@ def __init__(
         episode_contect_address: str,
         **kwargs,
     ):
+
         port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
         assert port is not None, "AJET_DAT_INTERCHANGE_PORT env var must be set"
-        base_url = f"http://localhost:{port}/v1"
+        master_node_ip = os.getenv("MASTER_NODE_IP", "localhost")
+
+        base_url = f"http://{master_node_ip}:{port}/v1"
         api_key = generate_auth_token(
             agent_name=agent_name,
             target_tag=target_tag,