Merge pull request NVIDIA#942 from NVIDIA/ipod/llm-sigterm

[LLM] Gracefully shutdown processes with scancel

Author: podkidyshev

conf/experimental/vllm/test_scenario/vllm.toml

@@ -16,6 +16,27 @@
 
 name = "vllm"
 
+[[Tests]]
+id = "vllm.disagg.4nodes"
+test_name = "vllm"
+num_nodes = 4
+time_limit = "00:30:00"
+
+  [Tests.cmd_args.prefill]
+  num_nodes = 2
+  enforce_eager = ""
+  tensor_parallel_size = 8
+  max_num_batched_tokens = 1024
+
+  [Tests.cmd_args.decode]
+  num_nodes = 2
+  enforce_eager = ""
+  tensor_parallel_size = 8
+  max_num_batched_tokens = 1024
+
+  [Tests.extra_env_vars]
+  CUDA_VISIBLE_DEVICES = "0,1,2,3"
+
 [[Tests]]
 id = "vllm.agg.1node"
 test_name = "vllm"

src/cloudai/workloads/common/llm_serving.py

@@ -565,41 +565,103 @@ def generate_wait_for_health_function(self) -> str:
     return 1
 }}"""
 
-    def generate_cleanup_function(self, pid_vars: list[str], timeout: int = 15) -> str:
-        if len(pid_vars) == 1:
-            pid_var = pid_vars[0]
-            return f"""\
-cleanup() {{
-    echo "Cleaning up PIDs: {pid_var}=${pid_var}"
-    kill -TERM "${pid_var}" 2>/dev/null
-    i=0
-    while kill -0 "${pid_var}" 2>/dev/null; do
-        [ "$i" -ge {timeout} ] && echo "PID did not exit in time" && return 1
-        sleep 1
-        i=$((i+1))
-    done
-}}
-trap cleanup EXIT"""
+    def slurm_step_name(self, role: str) -> str:
+        return f"cloudai-{self.workload_slug}-{role.replace('_', '-')}"
 
-        pid_values = " ".join(f"{pid_var}=${pid_var}" for pid_var in pid_vars)
-        pid_array = " ".join(f'"${p}"' for p in pid_vars)
+    @staticmethod
+    def slurm_step_id_var(role: str) -> str:
+        return f"{role.replace('-', '_').upper()}_STEP_IDS"
+
+    def cleanup_step_id_vars(self, pid_vars: list[str]) -> list[str]:
+        return [f"{pid_var.removesuffix('_PID')}_STEP_IDS" for pid_var in pid_vars]
+
+    def cleanup_prelude(self) -> str:
+        return ""
+
+    def cleanup_guard_var(self) -> str:
+        run_id = re.sub(r"\W+", "_", f"{self.test_run.name}_{self.test_run.current_iteration}").strip("_")
+        return f"CLOUDAI_{run_id.upper()}_CLEANUP_DONE"
+
+    def _with_slurm_step_name(self, srun_prefix: str, role: str) -> str:
+        return f"{srun_prefix} --job-name={self.slurm_step_name(role)}"
+
+    def render_step_discovery(self, role: str, step_id_var: str | None = None, expected_count: int = 1) -> str:
+        step_id_var = step_id_var or self.slurm_step_id_var(role)
+        step_name = self.slurm_step_name(role)
+        squeue_cmd = (
+            'squeue --noheader --steps --job "$SLURM_JOB_ID" --format="%i %j" 2>/dev/null '
+            f"| awk '$2 == \"{step_name}\" {{ print $1 }}'"
+        )
         return f"""\
-cleanup() {{
-    echo "Cleaning up PIDs: {pid_values}"
+{step_id_var}=
+for _ in {{1..10}}; do
+    {step_id_var}=$({squeue_cmd})
+    CLOUDAI_STEP_COUNT=$(printf '%s\\n' "${{{step_id_var}}}" | wc -w | tr -d ' ')
+    if [ "$CLOUDAI_STEP_COUNT" -ge {expected_count} ]; then break; fi
+    sleep 1
+done
+echo "Slurm step IDs for {role}: ${{{step_id_var}:-unknown}}"
+"""
 
-    for pid in {pid_array}; do
-        [ -n "$pid" ] && kill -TERM "$pid" 2>/dev/null
-    done
+    def _render_cleanup_signal_block(self, pid_var: str, step_id_var: str, signal_name: str) -> str:
+        pid_signal = "KILL" if signal_name == "KILL" else "TERM"
+        return f"""\
+    if [ -n "${{{step_id_var}:-}}" ]; then
+        for step_id in ${{{step_id_var}}}; do
+            scancel --signal={signal_name} "$step_id" 2>/dev/null || true
+        done
+    elif [ -n "${{{pid_var}:-}}" ]; then
+        kill -{pid_signal} "${{{pid_var}}}" 2>/dev/null || true
+    fi"""
 
-    for pid in {pid_array}; do
-        [ -z "$pid" ] && continue
+    def _render_cleanup_wait_block(self, pid_var: str, step_id_var: str, timeout: int) -> str:
+        force_block = self._render_cleanup_signal_block(pid_var, step_id_var, "KILL")
+        force_block = "\n".join(f"            {line}" for line in force_block.splitlines())
+        return f"""\
+    if [ -n "${{{pid_var}:-}}" ]; then
         i=0
-        while kill -0 "$pid" 2>/dev/null; do
-            [ "$i" -ge {timeout} ] && echo "PID $pid did not exit in time" && return 1
+        while kill -0 "${{{pid_var}}}" 2>/dev/null; do
+            if [ "$i" -ge {timeout} ]; then
+                echo "PID ${{{pid_var}}} did not exit in time"
+                cleanup_status=1
+{force_block}
+                break
+            fi
             sleep 1
             i=$((i+1))
         done
-    done
+    fi"""
+
+    def generate_cleanup_function(self, pid_vars: list[str], timeout: int = 60) -> str:
+        step_id_vars = self.cleanup_step_id_vars(pid_vars)
+        guard_var = self.cleanup_guard_var()
+        pid_values = " ".join(f"{pid_var}=${pid_var}" for pid_var in pid_vars)
+        step_values = " ".join(f"{step_id_var}=${step_id_var}" for step_id_var in step_id_vars)
+        prelude = self.cleanup_prelude()
+        if prelude:
+            prelude = f"{prelude}\n"
+        term_blocks = "\n".join(
+            self._render_cleanup_signal_block(pid_var, step_id_var, "TERM")
+            for pid_var, step_id_var in zip(pid_vars, step_id_vars, strict=True)
+        )
+        wait_blocks = "\n".join(
+            self._render_cleanup_wait_block(pid_var, step_id_var, timeout)
+            for pid_var, step_id_var in zip(pid_vars, step_id_vars, strict=True)
+        )
+        return f"""\
+cleanup() {{
+    if [ "${{{guard_var}:-0}}" = "1" ]; then
+        return 0
+    fi
+    {guard_var}=1
+    cleanup_status=0
+{prelude}    echo "Cleaning up PIDs: {pid_values}"
+    echo "Cleaning up Slurm step IDs: {step_values}"
+
+{term_blocks}
+
+{wait_blocks}
+    return "$cleanup_status"
 }}
 trap cleanup EXIT"""
 
@@ -728,12 +790,15 @@ def render_serve_launch(
         head_node_var: str,
         nodelist_var: str,
     ) -> str:
-        del role, node_count, nodelist_var
+        del node_count, nodelist_var
+        srun_prefix = self._with_slurm_step_name(self._single_role_srun_prefix(head_node_var), role)
+        step_id_var = self.slurm_step_id_var(role)
         return f"""\
-{self._single_role_srun_prefix(head_node_var)} \\
+{srun_prefix} \\
     --output={self.test_run.output_path.absolute()}/{log_file} \\
     {self._with_custom_bash(command_tail)} &
-{pid_var}=$!"""
+{pid_var}=$!
+{self.render_step_discovery(role, step_id_var)}"""
 
     def _expand_semantic_eval_args(self, args: str, *, host: str) -> str:
         replacements = {
@@ -792,11 +857,15 @@ def _gen_aggregated_script(self, serve_cmd: list[str], bench_cmd: str) -> str:
         node_setup = self.generate_aggregated_node_setup(serve_node_count)
         preamble = self.aggregated_script_preamble()
         if legacy_single_node:
+            serve_srun_prefix = self._with_slurm_step_name(
+                f"{srun_prefix} --overlap --ntasks-per-node=1 --ntasks=1", "serve"
+            )
             serve_launch = f"""\
-{srun_prefix} --overlap --ntasks-per-node=1 --ntasks=1 \\
+{serve_srun_prefix} \\
     --output={(self.test_run.output_path / self.serve_log_file).absolute()} \\
     {self._with_custom_bash(serve_cmd_with_env)} &
-{self.serve_pid_var}=$!"""
+{self.serve_pid_var}=$!
+{self.render_step_discovery("serve", self.slurm_step_id_var("serve"))}"""
         else:
             serve_launch = self.render_serve_launch(
                 "serve",
@@ -889,10 +958,11 @@ def _gen_disaggregated_script(self, serve_commands: list[list[str]], bench_cmd:
 {wait_block}
 
 echo "Starting {self.proxy_router_name}..."
-{prefill_srun_prefix} \\
+{self._with_slurm_step_name(prefill_srun_prefix, "helper")} \\
     --output={self.test_run.output_path.absolute()}/{self.proxy_router_log_file} \\
     {self._with_custom_bash(" ".join(helper_cmd))} &
 {self.proxy_router_pid_var}=$!
+{self.render_step_discovery("helper", self.slurm_step_id_var("helper"))}
 
 {wait_block_helper}
 

src/cloudai/workloads/vllm/slurm_command_gen_strategy.py

@@ -174,6 +174,22 @@ def disaggregated_cleanup_pid_vars(self) -> list[str]:
             pid_vars.insert(insert_at, "DECODE_RAY_PID")
         return pid_vars
 
+    def cleanup_step_id_vars(self, pid_vars: list[str]) -> list[str]:
+        step_id_vars: list[str] = []
+        for pid_var in pid_vars:
+            base = pid_var.removesuffix("_PID")
+            if base.endswith("_RAY"):
+                step_id_vars.append(f"{base}_WORKER_STEP_IDS")
+                continue
+
+            role = base.lower()
+            if role in {"serve", "prefill", "decode"} and self._needs_ray(role):
+                step_id_vars.append(f"{base}_RAY_HEAD_STEP_IDS")
+                continue
+
+            step_id_vars.append(f"{base}_STEP_IDS")
+        return step_id_vars
+
     @property
     def proxy_router_healthcheck(self) -> str:
         fields_set = self.tdef.cmd_args.model_fields_set
@@ -216,12 +232,8 @@ def _ray_stop_cleanup_block(self) -> str:
             )
         return "\n".join(lines)
 
-    def generate_cleanup_function(self, pid_vars: list[str], timeout: int = 15) -> str:
-        cleanup = super().generate_cleanup_function(pid_vars, timeout)
-        ray_stop_block = self._ray_stop_cleanup_block()
-        if not ray_stop_block:
-            return cleanup
-        return cleanup.replace("cleanup() {\n", f"cleanup() {{\n{ray_stop_block}\n", 1)
+    def cleanup_prelude(self) -> str:
+        return self._ray_stop_cleanup_block()
 
     def render_serve_launch(
         self,
@@ -246,8 +258,12 @@ def render_serve_launch(
         ray_worker_log = f"{self.workload_slug}-{role}-ray-worker-%N.log"
         serve_log = f"{self.test_run.output_path.absolute()}/{log_file}"
         head_node_expr = f"${{{head_node_var}}}"
-        worker_prefix = self._role_srun_prefix("$node")
-        head_prefix = self._single_role_srun_prefix(head_node_var)
+        worker_role = f"{role}-ray-worker"
+        head_role = f"{role}-ray-head"
+        worker_step_id_var = f"{role_prefix}_RAY_WORKER_STEP_IDS"
+        head_step_id_var = f"{role_prefix}_RAY_HEAD_STEP_IDS"
+        worker_prefix = self._with_slurm_step_name(self._role_srun_prefix("$node"), worker_role)
+        head_prefix = self._with_slurm_step_name(self._single_role_srun_prefix(head_node_var), head_role)
         serve_cmd = self._with_custom_bash(f'env RAY_ADDRESS="{head_node_expr}:${{{ray_port_var}}}" {command_tail}')
         ray_head_args = self._ray_start_args(role, "head", {"head": True, "port": f'"${{{ray_port_var}}}"'})
         ray_worker_args = self._ray_start_args(
@@ -301,11 +317,13 @@ def render_serve_launch(
     wait
 ) &
 {ray_pid_var}=$!
+{self.render_step_discovery(worker_role, worker_step_id_var, node_count - 1)}
 {head_prefix} \\
     --output={serve_log} \\
     --error={self.test_run.output_path.absolute()}/{ray_head_log} \\
     bash -lc {ray_head_command} &
-{pid_var}=$!"""
+{pid_var}=$!
+{self.render_step_discovery(head_role, head_step_id_var)}"""
 
     def disaggregated_role_env(self, role: str, gpu_ids: list[int]) -> dict[str, str]:
         env = super().disaggregated_role_env(role, gpu_ids)