[Feat] support vllm with slurm launcher (#404)

* [Feat] support vllm with slurm launcher

Author: fishcrap

areal/api/cli_args.py

@@ -455,7 +455,7 @@ class vLLMConfig:
     skip_tokenizer_init: bool = False
     enforce_eager: bool = True
     dtype: str = "bfloat16"
-    distributed_executor_backend = "mp"
+    distributed_executor_backend: str = "mp"
     # original
     max_num_seqs: int = 256
     # kv_cache_type: str = "auto"
@@ -479,6 +479,7 @@ class vLLMConfig:
         "areal.thirdparty.vllm.vllm_worker_extension.VLLMWorkerExtension"
     )
     enable_sleep_mode: bool = False
+    uvicorn_log_level: str = "warning"
 
     @staticmethod
     def build_args(

areal/engine/vllm_remote.py

@@ -432,7 +432,6 @@ def update_weights_from_distributed(
     ]
 
     async def _fn():
-        tik = time.perf_counter()
         if init_group:
             await asyncio.gather(
                 *[
@@ -472,8 +471,6 @@ async def _fn():
             ]
         )
 
-        logger.info(f"Distributed update weights done in {time.perf_counter() - tik}s")
-
     return uvloop.run(_fn())
 
 

areal/launcher/slurm.py

@@ -16,6 +16,7 @@
     SGLangConfig,
     parse_cli_args,
     to_structured_cfg,
+    vLLMConfig,
 )
 from areal.platforms import current_platform
 from areal.utils import logging, name_resolve, names
@@ -431,58 +432,111 @@ def slurm_main(config, run_id: int = 0):
     n_gpus_per_node = config.cluster.n_gpus_per_node
     allocation_mode = config.allocation_mode
     allocation_mode = AllocationMode.from_str(allocation_mode)
-    sglang_cmds = []
-    sglang_addrs = []
-    n_sglang_nodes = 0
-    if allocation_mode.gen_backend == "sglang":
-        # Launcher should launch SGLang servers according to allocation mode.
-        config.sglang = to_structured_cfg(config.sglang, SGLangConfig)
-        n_sglang_servers = allocation_mode.gen.dp_size
-        n_sglang_nodes = allocation_mode.gen.world_size // n_gpus_per_node
-        node_group_size = max(1, allocation_mode.gen_instance_size // n_gpus_per_node)
-        n_servers_per_node = max(n_sglang_servers // n_sglang_nodes, 1)
-
-        cross_nodes = allocation_mode.gen_instance_size > n_gpus_per_node
-        env_vars = get_env_vars(
-            config.cluster.cluster_name,
-            config.launcher.inference_server_env_vars,
-        )
-        env_vars = [copy.deepcopy(env_vars) for _ in range(n_sglang_nodes)]
-        base_seed = config.sglang.random_seed
-        sglang_server_cmd_template = f"python3 -m areal.launcher.sglang_server {' '.join(sys.argv[2:])} sglang.random_seed={{seed}}"
-        for i in range(n_sglang_nodes):
-            sglang_cmd = sglang_server_cmd_template.format(
-                seed=base_seed + i * n_servers_per_node
+    n_backend_nodes = 0
+
+    if allocation_mode.gen_backend in ("sglang", "vllm"):
+        # Launcher should launch llm servers according to allocation mode.
+        if allocation_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",
+                "prefix": "AREAL_SGLANG",
+                "set_device_env": False,
+            },
+            "vllm": {
+                "module": "areal.launcher.vllm_server",
+                "seed_arg": "vllm.seed",
+                "prefix": "AREAL_VLLM",
+                "set_device_env": True,  # vLLM needs `device_control_env_var` to control GPU allocation
+            },
+        }
+
+        def _build_llm_server_plan(backend: str, spec: Dict):
+            # Returns: cmds, env_vars_list, n_nodes, n_servers
+
+            if backend not in backend_spec:
+                raise NotImplementedError(f"Unknown backend: {backend}")
+
+            spec = backend_spec[backend]
+
+            n_backend_servers = allocation_mode.gen.dp_size
+            n_backend_nodes = allocation_mode.gen.world_size // n_gpus_per_node
+            node_group_size = max(
+                1, allocation_mode.gen_instance_size // n_gpus_per_node
+            )
+            n_servers_per_node = max(n_backend_servers // n_backend_nodes, 1)
+
+            cross_nodes = allocation_mode.gen_instance_size > n_gpus_per_node
+            base_env_bars = get_env_vars(
+                config.cluster.cluster_name,
+                config.launcher.inference_server_env_vars,
+            )
+            if spec["set_device_env"]:
+                base_env_bars[current_platform.device_control_env_var] = ",".join(
+                    list(map(str, range(n_gpus_per_node)))
+                )
+            env_list = [copy.deepcopy(base_env_bars) for _ in range(n_backend_nodes)]
+
+            base_seed = random_seed
+            seed_arg = spec["seed_arg"]
+            module = spec["module"]
+            backend_server_cmd_template = (
+                f"python3 -m {module} {' '.join(sys.argv[2:])} {seed_arg}={{seed}}"
             )
-            sglang_cmds.append(sglang_cmd)
-            if cross_nodes:
-                # master_addrs and master_ports are the IP addresses and free ports of the all nodes in the job array, obtained in the SBATCH script.
-                env_vars[i] |= dict(
-                    AREAL_SGLANG_MULTI_NODE_RANK=i % node_group_size,
-                    AREAL_SGLANG_MULTI_NODE_MASTER_ADDR=f"${{master_addrs[{i // node_group_size * node_group_size}]}}",
-                    AREAL_SGLANG_MULTI_NODE_MASTER_PORT=f"${{master_ports[{i // node_group_size * node_group_size}]}}",
+
+            backend_cmds = []
+            for i in range(n_backend_nodes):
+                backend_cmd = backend_server_cmd_template.format(
+                    seed=base_seed + i * n_servers_per_node
                 )
+                backend_cmds.append(backend_cmd)
+                if cross_nodes:
+                    # master_addrs and master_ports are the IP addresses and free ports of the all nodes in the job array, obtained in the SBATCH script.
+                    prefix = spec["prefix"]
+                    env_list[i] |= dict(
+                        **{
+                            f"{prefix}_MULTI_NODE_RANK": i % node_group_size,
+                            f"{prefix}_MULTI_NODE_MASTER_ADDR": f"${{master_addrs[{i // node_group_size * node_group_size}]}}",
+                            f"{prefix}_MULTI_NODE_MASTER_PORT": f"${{master_ports[{i // node_group_size * node_group_size}]}}",
+                        }
+                    )
+
+            return backend_cmds, env_list, n_backend_nodes, n_backend_servers
+
+        backend_cmds, env_list, n_backend_nodes, n_backend_servers = (
+            _build_llm_server_plan(
+                allocation_mode.gen_backend,
+                random_seed,
+            )
+        )
 
         launcher.submit_array(
             job_name="llm_server",
-            cmd=sglang_cmds,
-            count=n_sglang_nodes,
-            nodes=n_sglang_nodes,
-            n_gpus_per_node=config.cluster.n_gpus_per_node,
+            cmd=backend_cmds,
+            count=n_backend_nodes,
+            nodes=n_backend_nodes,
+            n_gpus_per_node=n_gpus_per_node,
             cpus_per_task=config.launcher.inference_server_cpus_per_gpu
             * n_gpus_per_node,
             mem_per_task=config.launcher.inference_server_mem_per_gpu * n_gpus_per_node,
             srun_additional_args=config.launcher.slurm.srun_additional_args,
             container_image=config.launcher.slurm.inference_server_image,
             container_mounts=config.launcher.slurm.mount,
-            env_vars=env_vars,
+            env_vars=env_list,
         )
-        # Get SGLang server addresses by name resolve
+        # Get llm server addresses by name resolve
         try:
-            sglang_addrs = wait_llm_server_addrs(
+            llm_addrs = wait_llm_server_addrs(
                 config.experiment_name,
                 config.trial_name,
-                n_sglang_servers,
+                n_backend_servers,
             )
         except (TimeoutError, KeyboardInterrupt) as e:
             launcher.stop_all(force=True)
@@ -492,7 +546,7 @@ def slurm_main(config, run_id: int = 0):
         trainer_n_nodes = 1
         gpus_per_node = 0
     else:
-        trainer_n_nodes = n_nodes - n_sglang_nodes
+        trainer_n_nodes = n_nodes - n_backend_nodes
         gpus_per_node = config.cluster.n_gpus_per_node
 
     # Here $head_node_ip is the IP address of the first node in the job array.
@@ -534,7 +588,7 @@ def slurm_main(config, run_id: int = 0):
                     config.cluster.cluster_name,
                     config.launcher.trainer_env_vars,
                 ),
-                AREAL_LLM_SERVER_ADDRS=",".join(sglang_addrs),
+                AREAL_LLM_SERVER_ADDRS=",".join(llm_addrs),
                 AREAL_RECOVER_RUN=str(int(is_recover_run)),
             ),
         )

areal/launcher/vllm_server.py

@@ -24,18 +24,21 @@
 logger = logging.getLogger("vLLMServer Wrapper")
 
 
-def launch_server_cmd(command: str) -> subprocess.Popen:
+def launch_server_cmd(command: str, custom_env: dict | None = None) -> subprocess.Popen:
     """
     Execute a shell command and return its process handle.
     """
     # Replace newline continuations and split the command string.
     command = command.replace("\\\n", " ").replace("\\", " ")
+    logger.info(f"Launch command: {command}")
     parts = command.split()
     _env = os.environ.copy()
     # To avoid DirectoryNotEmpty error caused by triton
     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
+    if custom_env is not None:
+        _env.update(custom_env)
     return subprocess.Popen(
         parts,
         text=True,
@@ -94,13 +97,10 @@ def run(self):
         device_control_env_var = current_platform.device_control_env_var
         if device_control_env_var in os.environ:
             visible = os.getenv(device_control_env_var).split(",")
-            ordered = ",".join(sorted(visible, key=int))
-            os.environ[device_control_env_var] = ordered
             n_visible_devices = len(visible)
             n_servers_per_proc = max(1, n_visible_devices // gpus_per_server)
-            server_idx_offset = int(visible[0]) // gpus_per_server
+            server_idx_offset = min(list(map(int, visible))) // gpus_per_server
         else:
-            n_visible_devices = self.n_gpus_per_node
             n_servers_per_proc = n_servers_per_node
             server_idx_offset = 0
 
@@ -109,6 +109,7 @@ def run(self):
         ports_per_server = 40000 // n_servers_per_node
         launch_server_args = []
         server_addresses = []
+        base_random_seed = self.config.seed
         for server_local_idx in range(
             server_idx_offset, server_idx_offset + n_servers_per_proc
         ):
@@ -121,15 +122,21 @@ def run(self):
             dist_init_addr = f"localhost:{dist_init_port}"
             host_ip = gethostip()
 
-            (server_local_idx - server_idx_offset) * gpus_per_server
+            base_gpu_id = (server_local_idx - server_idx_offset) * gpus_per_server
+            custom_env = {
+                device_control_env_var: ",".join(
+                    map(str, range(base_gpu_id, base_gpu_id + gpus_per_server))
+                )
+            }
+            self.config.seed = base_random_seed + server_local_idx
             cmd = vLLMConfig.build_cmd(
                 self.config,
                 tp_size=self.allocation_mode.gen.tp_size,
                 host=host_ip,
                 port=server_port,
                 dist_init_addr=dist_init_addr,
             )
-            launch_server_args.append((cmd, host_ip, server_port))
+            launch_server_args.append((cmd, host_ip, server_port, custom_env))
             server_addresses.append(f"http://{host_ip}:{server_port}")
 
         with ThreadPoolExecutor(max_workers=n_servers_per_proc) as executor:
@@ -159,8 +166,10 @@ def run(self):
 
             time.sleep(1)
 
-    def launch_one_server(self, cmd, host_ip, server_port):
-        server_process = launch_server_cmd(cmd)
+    def launch_one_server(
+        self, cmd: str, host_ip: str, server_port: int, custom_env: dict | None = None
+    ):
+        server_process = launch_server_cmd(cmd, custom_env)
         wait_for_server(f"http://{host_ip}:{server_port}")
         name = names.gen_servers(self.experiment_name, self.trial_name)
         name_resolve.add_subentry(name, f"{host_ip}:{server_port}")