fix vllm compatibility issue (#417)

* Use vllm/sglang server wrapper for local launching and fix an NCCL issue with vllm

* fix

Author: garrett4wade

areal/engine/base_hf_engine.py

@@ -117,10 +117,6 @@ def parallelism_group(self) -> dist.ProcessGroup:
 
     def create_process_group(self, parallel_strategy: ParallelStrategy | None = None):
         backend = current_platform.communication_backend
-        if current_platform.communication_backend == "nccl":
-            # Required by NCCL weight update group for SGLang
-            os.environ["NCCL_CUMEM_ENABLE"] = "0"
-            os.environ["NCCL_NVLS_ENABLE"] = "0"
         if not dist.is_initialized():
             # TODO: Handle the condition when WORLD_SIZE and RANK is not set in launcher
             # NOTE: device_id **SHOULD NOT** be passed into init_process_group,

areal/engine/sglang_remote.py

@@ -26,13 +26,19 @@
 from areal.platforms import current_platform
 from areal.utils import logging, name_resolve, names
 from areal.utils.http import arequest_with_retry, get_default_connector
+from areal.utils.launcher import wait_llm_server_addrs
 
 RID_CACHE_SIZE = 128
 
 
 class RemoteSGLangEngine(InferenceEngine):
 
     def __init__(self, config: InferenceEngineConfig):
+        if current_platform.communication_backend == "nccl":
+            # Required by NCCL weight update group.
+            os.environ["NCCL_CUMEM_ENABLE"] = "0"
+            os.environ["NCCL_NVLS_ENABLE"] = "0"
+
         self.config = config
 
         self.rid_to_address = {}
@@ -83,9 +89,24 @@ def initialize(
 
         if addr:
             self.addresses = addr if isinstance(addr, list) else [addr]
+            self.logger.info(f"Get server addresses from the `addr` argument.")
         else:
+            if (
+                self.config.experiment_name is not None
+                and self.config.trial_name is not None
+            ):
+                try:
+                    self.addresses = wait_llm_server_addrs(
+                        experiment_name=self.config.experiment_name,
+                        trial_name=self.config.trial_name,
+                        timeout=1,
+                    )
+                    self.logger.info(f"Get server addresses from name_resolve.")
+                except TimeoutError:
+                    pass
+        if not self.addresses and os.getenv("AREAL_LLM_SERVER_ADDRS"):
             # When addr is not provided, fallback to reading addrs from env var
-            self.addresses = os.getenv("AREAL_LLM_SERVER_ADDRS").split(",")
+            self.addresses = os.environ["AREAL_LLM_SERVER_ADDRS"].split(",")
         if not self.addresses:
             raise RuntimeError(
                 "No configured SGLang servers. Please pass in SGLang server addresses by arguments "

areal/engine/vllm_remote.py

@@ -26,6 +26,7 @@
 from areal.platforms import current_platform
 from areal.utils import logging, name_resolve, names
 from areal.utils.http import arequest_with_retry, get_default_connector
+from areal.utils.launcher import wait_llm_server_addrs
 
 RID_CACHE_SIZE = 128
 
@@ -90,9 +91,24 @@ def initialize(
 
         if addr:
             self.addresses = addr if isinstance(addr, list) else [addr]
+            self.logger.info(f"Get server addresses from the `addr` argument.")
         else:
+            if (
+                self.config.experiment_name is not None
+                and self.config.trial_name is not None
+            ):
+                try:
+                    self.addresses = wait_llm_server_addrs(
+                        experiment_name=self.config.experiment_name,
+                        trial_name=self.config.trial_name,
+                        timeout=1,
+                    )
+                    self.logger.info(f"Get server addresses from name_resolve.")
+                except TimeoutError:
+                    pass
+        if not self.addresses and os.getenv("AREAL_LLM_SERVER_ADDRS"):
             # When addr is not provided, fallback to reading addrs from env var
-            self.addresses = os.getenv("AREAL_LLM_SERVER_ADDRS").split(",")
+            self.addresses = os.environ["AREAL_LLM_SERVER_ADDRS"].split(",")
         if not self.addresses:
             raise RuntimeError(
                 "No configured vLLM servers. Please pass in vLLM server addresses by arguments "

areal/launcher/local.py

@@ -22,8 +22,14 @@
 )
 from areal.platforms import current_platform
 from areal.utils import logging, name_resolve, names
-from areal.utils.launcher import JobException, JobInfo, JobState, get_env_vars
-from areal.utils.network import find_free_ports, gethostip
+from areal.utils.launcher import (
+    JobException,
+    JobInfo,
+    JobState,
+    get_env_vars,
+    wait_llm_server_addrs,
+)
+from areal.utils.network import find_free_ports
 from areal.utils.recover import check_if_recover
 
 logger = logging.getLogger("Local Scheduler")
@@ -136,7 +142,9 @@ def submit_array(
             )
             c = f"{c} 2>&1 | tee -a {self.log_path_of(job_name)}"
             logger.info("Starting local process with command: %s", c)
-            process = subprocess.Popen(c, shell=isinstance(c, str))
+            process = subprocess.Popen(
+                c, shell=isinstance(c, str), stdout=sys.stdout, stderr=sys.stdout
+            )
             self._jobs[f"{job_name}/{offset + i}"] = process
             self._job_counter[job_name] += 1
 
@@ -275,72 +283,64 @@ def local_main(config, run_id: int = 0):
         f"run_id={run_id}, is_recover_run={is_recover_run}"
     )
 
-    server_cmd = []
-    server_addrs = []
-    if alloc_mode.gen_backend == "sglang":
-        base_seed = config.sglang.random_seed
-        config.sglang = to_structured_cfg(config.sglang, SGLangConfig)
-        ports = find_free_ports(alloc_mode.gen.dp_size * 2, port_range=(10000, 50000))
-        host_ip = gethostip()
-        host = "localhost" if not config.sglang.enable_metrics else host_ip
-        for i in range(alloc_mode.gen.dp_size):
-            config.sglang.random_seed = base_seed + i
-            cmd = SGLangConfig.build_cmd(
-                config.sglang,
-                host=host,
-                tp_size=alloc_mode.gen.tp_size,
-                base_gpu_id=0,
-                port=ports[i * 2],
-                dist_init_addr=f"localhost:{ports[i*2+1]}",
-            )
-            server_cmd.append(cmd)
-            server_addrs.append(f"{host}:{ports[i * 2]}")
+    if alloc_mode.gen_backend in ("sglang", "vllm"):
+        # Launcher should launch llm servers according to allocation mode.
+        if alloc_mode.gen_backend == "sglang":
+            config.sglang = to_structured_cfg(config.sglang, SGLangConfig)
+            random_seed = config.sglang.random_seed
+        else:
+            config.vllm = to_structured_cfg(config.vllm, vLLMConfig)
+            random_seed = config.vllm.seed
+
+        backend_spec = {
+            "sglang": {
+                "module": "areal.launcher.sglang_server",
+                "seed_arg": "sglang.random_seed",
+                "set_device_env": False,
+            },
+            "vllm": {
+                "module": "areal.launcher.vllm_server",
+                "seed_arg": "vllm.seed",
+                "set_device_env": True,  # vLLM needs `device_control_env_var` to control GPU allocation
+            },
+        }
+
+        spec = backend_spec[alloc_mode.gen_backend]
+
+        base_seed = random_seed
+        seed_arg = spec["seed_arg"]
+        module = spec["module"]
+        server_cmd = (
+            f"python3 -m {module} {' '.join(sys.argv[2:])} {seed_arg}={base_seed}"
+        )
 
         # Launch inference servers.
         launcher.submit_array(
             job_name="llm_server",
             cmd=server_cmd,
-            count=alloc_mode.gen.dp_size,
-            gpu=alloc_mode.gen.pp_size * alloc_mode.gen.tp_size,
+            count=1,
+            gpu=alloc_mode.gen.pp_size
+            * alloc_mode.gen.tp_size
+            * alloc_mode.gen.dp_size,
             env_vars=get_env_vars(
                 config.cluster.cluster_name,
                 config.launcher.inference_server_env_vars,
             ),
         )
-        logger.info(
-            f"LLM inference server launched at: AREAL_LLM_SERVER_ADDRS={','.join(server_addrs)}"
-        )
-    elif alloc_mode.gen_backend == "vllm":
-        base_seed = config.vllm.seed
-        config.vllm = to_structured_cfg(config.vllm, vLLMConfig)
-        ports = find_free_ports(alloc_mode.gen.dp_size * 2, port_range=(10000, 50000))
-        host = "localhost"
-        for i in range(alloc_mode.gen.dp_size):
-            config.vllm.seed = base_seed + i
-            cmd = vLLMConfig.build_cmd(
-                config.vllm,
-                host=host,
-                tp_size=alloc_mode.gen.tp_size,
-                port=ports[i * 2],
-                dist_init_addr=f"localhost:{ports[i*2+1]}",
-            )
-            server_cmd.append(cmd)
-            server_addrs.append(f"{host}:{ports[i * 2]}")
 
-        # Launch inference servers.
-        launcher.submit_array(
-            job_name="llm_server",
-            cmd=server_cmd,
-            count=alloc_mode.gen.dp_size,
-            gpu=alloc_mode.gen.pp_size * alloc_mode.gen.tp_size,
-            env_vars=get_env_vars(
-                config.cluster.cluster_name,
-                config.launcher.inference_server_env_vars,
-            ),
+    # Get llm server addresses by name resolve
+    try:
+        server_addrs = wait_llm_server_addrs(
+            config.experiment_name,
+            config.trial_name,
+            n_rollout_servers=alloc_mode.gen.dp_size,
         )
         logger.info(
             f"LLM inference server launched at: AREAL_LLM_SERVER_ADDRS={','.join(server_addrs)}"
         )
+    except (TimeoutError, KeyboardInterrupt) as e:
+        launcher.stop_all(signal="SIGINT")
+        raise e
 
     # Launch trainer entrypoint
     if alloc_mode.type_ != AllocationType.LLM_SERVER_ONLY:

areal/launcher/vllm_server.py

@@ -37,6 +37,10 @@ def launch_server_cmd(command: str, custom_env: dict | None = None) -> subproces
     triton_cache_path = _env.get("TRITON_CACHE_PATH", TRITON_CACHE_PATH)
     unique_triton_cache_path = os.path.join(triton_cache_path, str(uuid.uuid4()))
     _env["TRITON_CACHE_PATH"] = unique_triton_cache_path
+    # To avoid vllm compile cache conflict
+    vllm_cache_path = _env.get("VLLM_CACHE_ROOT")
+    if vllm_cache_path:
+        _env["VLLM_CACHE_ROOT"] = os.path.join(vllm_cache_path, str(uuid.uuid4()))
     if custom_env is not None:
         _env.update(custom_env)
     return subprocess.Popen(

areal/utils/launcher.py

@@ -15,6 +15,7 @@
 PYTORCH_KERNEL_CACHE_PATH = (
     f"{LOCAL_CACHE_DIR}/.cache/{getpass.getuser()}/torch/kernels/"
 )
+VLLM_CACHE_ROOT = f"{LOCAL_CACHE_DIR}/.cache/{getpass.getuser()}/vllm/"
 TRITON_CACHE_PATH = f"{LOCAL_CACHE_DIR}/.cache/{getpass.getuser()}/triton/"
 PYTHONPATH = os.pathsep.join(
     filter(
@@ -26,11 +27,13 @@
     )
 )
 os.makedirs(PYTORCH_KERNEL_CACHE_PATH, exist_ok=True)
+os.makedirs(VLLM_CACHE_ROOT, exist_ok=True)
 os.makedirs(TRITON_CACHE_PATH, exist_ok=True)
 BASE_ENVIRONS = {
     "TOKENIZERS_PARALLELISM": "true",
     "PYTORCH_KERNEL_CACHE_PATH": PYTORCH_KERNEL_CACHE_PATH,
     "TRITON_CACHE_DIR": TRITON_CACHE_PATH,
+    "VLLM_CACHE_ROOT": VLLM_CACHE_ROOT,
     "CUDA_DEVICE_MAX_CONNECTIONS": "1",
     "PYTHONPATH": PYTHONPATH,
 }
@@ -48,7 +51,6 @@
     "NCCL_DEBUG": "WARN",
     "NCCL_DEBUG_SUBSYS": "INIT,TUNING,GRAPH",
 }
-LLM_SERVER_WAIT_TIMEOUT_SECONDS = 360
 
 
 def get_env_vars(
@@ -103,7 +105,8 @@ class JobInfo:
 def wait_llm_server_addrs(
     experiment_name: str,
     trial_name: str,
-    n_rollout_servers: int,
+    n_rollout_servers: int = 1,
+    timeout: int | None = 360,
 ):
     # Get rollout nodes, find the hosts
     name = names.gen_servers(experiment_name, trial_name)
@@ -117,7 +120,7 @@ def wait_llm_server_addrs(
             break
 
         time.sleep(1)
-        if time.perf_counter() - start > LLM_SERVER_WAIT_TIMEOUT_SECONDS:
+        if timeout is not None and time.perf_counter() - start > timeout:
             raise TimeoutError(
                 f"Timeout waiting for rollout servers to be ready. "
                 f"Expected {n_rollout_servers} servers, found {len(rollout_addrs)}."