From 8d08df4a208ce0086dcad610d0ba0a4349726fbf Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 01:45:23 -0700 Subject: [PATCH 1/9] SLURM Python launcher Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- examples/disaggregated/slurm/config.yaml | 36 + .../slurm/disagg_profiler/__init__.py | 13 + .../slurm/disagg_profiler/job_manager.py | 1137 +++++++++++++++++ examples/disaggregated/slurm/launcher.py | 89 ++ 4 files changed, 1275 insertions(+) create mode 100644 examples/disaggregated/slurm/config.yaml create mode 100644 examples/disaggregated/slurm/disagg_profiler/__init__.py create mode 100644 examples/disaggregated/slurm/disagg_profiler/job_manager.py create mode 100644 examples/disaggregated/slurm/launcher.py diff --git a/examples/disaggregated/slurm/config.yaml b/examples/disaggregated/slurm/config.yaml new file mode 100644 index 000000000000..3efb7622bd07 --- /dev/null +++ b/examples/disaggregated/slurm/config.yaml @@ -0,0 +1,36 @@ +exec: + config: + context: + tp: 4 + ep: 4 + pp: 1 + max_batch_size: 4 + max_num_tokens: 1024 + max_seq_len: 1024 + config: + kv_cache_config: + free_gpu_memory_fraction: 0.75 + enable_block_reuse: false + print_iter_log: true + dp: 1 + generation: + tp: 4 + ep: 4 + pp: 1 + max_batch_size: 1 + max_num_tokens: 4096 + max_seq_len: 2048 + config: + print_iter_log: true + kv_cache_config: + free_gpu_memory_fraction: 0.75 + enable_block_reuse: false + dp: 1 + model_path: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +profile: + isl: 1024 + osl: 1024 + use_benchmark_serving: true + concurrency: + - 128 + - 256 diff --git a/examples/disaggregated/slurm/disagg_profiler/__init__.py b/examples/disaggregated/slurm/disagg_profiler/__init__.py new file mode 100644 index 000000000000..e25f5b45807c --- /dev/null +++ b/examples/disaggregated/slurm/disagg_profiler/__init__.py @@ -0,0 +1,13 @@ +""" +Disaggregated serving profiler package. +This package contains the job management and parameter sweeping functionality +for the TRT-LLM disaggregated serving launcher. +""" + +from .job_manager import JobManager, calculate_nodes_needed, wait_for_server + +__all__ = [ + 'JobManager', 'calculate_nodes_needed', 'wait_for_server', + 'ParameterSweeper', 'AutoSweeper', 'get_slurm_allocation', + 'run_sweep_configuration' +] diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py new file mode 100644 index 000000000000..6d3ac837c5f7 --- /dev/null +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -0,0 +1,1137 @@ +import json +import os +import re +import subprocess +import time +import uuid + +import requests +import yaml + + +def get_slurm_allocation(account, partition, time_limit, job_name, num_nodes): + """ + Request a SLURM allocation and wait for it to be granted. + Args: + account (str): SLURM account + partition (str): SLURM partition + time_limit (str): Time limit for the allocation + job_name (str): Job name + num_nodes (int): Number of nodes to request + Returns: + str: Job ID of the SLURM allocation, or None if allocation failed + """ + print( + f"Requesting SLURM allocation for {num_nodes} nodes on partition {partition}..." + ) + + # capture SLURM_NODELIST inside allocation + cmd = [ + "salloc", "-A", account, "-p", partition, "-N", + str(num_nodes), "-t", time_limit, "-J", job_name, "--no-shell" + ] + + try: + # Run the command and stream the output in real-time + process = subprocess.Popen(cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + universal_newlines=True) + + output_lines = [] + job_id = None + + # Process the output in real-time + for line in process.stdout: + print(line.strip()) + output_lines.append(line) + + # Look for job ID in the output (e.g., "salloc: Granted job allocation 12345") + job_id_match = re.search(r"Granted job allocation (\d+)", line) + if job_id_match: + job_id = job_id_match.group(1) + print( + f"Successfully obtained SLURM allocation with job ID: {job_id}" + ) + break + + # Wait for the process to complete + process.wait() + + if job_id: + return job_id + else: + print( + "Failed to obtain SLURM allocation. Check the output above for details." + ) + return None + except Exception as e: + print(f"Error requesting SLURM allocation: {e}") + return None + + +def execute_cmd(cmd): + """Execute a command.""" + print(f"Executing command: {' '.join(cmd)}") + return subprocess.Popen(cmd) + + +def wait_for_server(host, port, max_retries=1000, delay=5): + """Wait for a server to become available.""" + for _ in range(max_retries): + try: + print(f"Checking server status at {host}:{port}") + state = requests.get(f"http://{host}:{port}/health") + if state.status_code == 200: + print("Server is running") + return True + except: + time.sleep(delay) + + print(f"Server did not start after {max_retries} attempts") + return False + + +def calculate_nodes_needed(config, num_gpus): + """ + Calculate the number of nodes needed based on the configuration. + + The calculation is based on: + 1. The number of context servers and their TP size + 2. The number of generation servers and their TP size + 3. The number of IFB servers and their TP size + 4. The number of GPUs per node (num_gpus) + + Args: + config (dict): The configuration dictionary + + Returns: + int: The number of nodes needed + """ + total_gpus_needed = 0 + + # Calculate GPUs needed for context servers + if 'context' in config['exec']['config']: + context_config = config['exec']['config']['context'] + context_tp = context_config['tp'] + context_dp = context_config['dp'] + total_gpus_needed += context_tp * context_dp + + # Calculate GPUs needed for generation servers + if 'generation' in config['exec']['config']: + generation_config = config['exec']['config']['generation'] + generation_tp = generation_config['tp'] + generation_dp = generation_config['dp'] + total_gpus_needed += generation_tp * generation_dp + + # Calculate GPUs needed for IFB servers + if 'ifb' in config['exec']['config']: + ifb_config = config['exec']['config']['ifb'] + ifb_tp = ifb_config['tp'] + ifb_dp = ifb_config['dp'] + total_gpus_needed += ifb_tp * ifb_dp + + # Calculate nodes needed (round up to ensure enough nodes) + nodes_needed = (total_gpus_needed + num_gpus - 1) // num_gpus + + # Ensure we request at least one node + return max(1, nodes_needed) + + +class JobManager: + """Manages SLURM jobs for the disaggregated serving setup.""" + + def __init__(self, + args, + load_config_only=False, + override_config=None, + override_job_id=None): + """ + Initialize the job manager. + + Args: + args (argparse.Namespace): Command line arguments. + load_config_only (bool): If True, only load the configuration without launching jobs. + override_config (dict): Override the base config with this config if provided. + override_job_id (str): Override the auto-generated job ID with this ID if provided. + """ + print( + f"DEBUG JobManager: __init__ called with load_config_only={load_config_only}, override_job_id={override_job_id}" + ) + + self.args = args + self.num_gpus = args.num_gpus + self.account = args.account + self.partition = args.partition + self.time = args.time + self.job_name = args.job_name + self.container_image = args.container_image + self.mounts = args.mounts + self.container_name = "trtllm-disagg-server" + self.workdir = os.path.abspath(os.path.join(os.getcwd())) + self.context_jobs = [] + self.generation_jobs = [] + self.ifb_jobs = [] + self.disagg_jobs = [] + self.experiment_path = args.experiment_path + + # Clean up any previous node info files + self.setup_shared_directories(override_job_id) + + # Generate configuration files + self.base_config = self.load_config() + + # If override_config is provided, use it instead of the base_config + if override_config: + self.config = override_config + else: + self.config = self.base_config + print(self.config) + + # If we only need to load the config, return here + if load_config_only: + print( + "DEBUG JobManager: load_config_only=True, returning without launching jobs" + ) + return + + self.output_folder = self.create_output_folder() + + # Check if results already exist and should be skipped + if hasattr(args, 'skip_existing') and args.skip_existing: + if self._check_existing_results(): + print( + f"Skipping execution - results already exist in {self.output_folder}" + ) + return + + # Run the job with the current configuration + print("DEBUG JobManager: About to call prepare_and_run_single_config") + self.prepare_and_run_single_config() + print("DEBUG JobManager: Completed prepare_and_run_single_config") + + def create_output_folder(self): + # The parent directory must be model name + model_name = self.config['exec']['model_path'].split('/')[-1] + if model_name == '': + model_name = self.config['exec']['model_path'].split('/')[-2] + model_name = os.path.join(self.experiment_path, model_name) + os.makedirs(model_name, exist_ok=True) + output_folder = os.path.join(model_name, ( + f"{self.config['profile']['isl']}_{self.config['profile']['osl']}")) + os.makedirs(output_folder, exist_ok=True) + ifb_tp = self.config['exec']['config']['ifb'][ + 'tp'] if 'ifb' in self.config['exec']['config'] else 0 + ifb_dp = self.config['exec']['config']['ifb'][ + 'dp'] if 'ifb' in self.config['exec']['config'] else 0 + ctx_tp = self.config['exec']['config']['context'][ + 'tp'] if 'context' in self.config['exec']['config'] else 0 + ctx_dp = self.config['exec']['config']['context'][ + 'dp'] if 'context' in self.config['exec']['config'] else 0 + gen_tp = self.config['exec']['config']['generation'][ + 'tp'] if 'generation' in self.config['exec']['config'] else 0 + gen_dp = self.config['exec']['config']['generation'][ + 'dp'] if 'generation' in self.config['exec']['config'] else 0 + + output_folder = os.path.join(output_folder, (f"{ctx_tp}_" + f"{ctx_dp}_" + f"{gen_tp}_" + f"{gen_dp}_" + f"{ifb_tp}_" + f"{ifb_dp}")) + + for server_type in ['generation', 'ifb']: + if f'{server_type}' in self.config['exec']['config']: + config = self.config['exec']['config'][f'{server_type}'][ + 'config'] + pytorch_config = config[ + 'pytorch_backend_config'] if 'pytorch_backend_config' in config else config + if 'disable_overlap_scheduler' not in pytorch_config or not pytorch_config[ + 'disable_overlap_scheduler']: + output_folder += '_overlap' + if 'use_cuda_graph' in pytorch_config: + if pytorch_config['use_cuda_graph']: + output_folder += '_cuda_graph' + if 'enable_attention_dp' in config: + if config['enable_attention_dp']: + output_folder += '_adp' + if 'speculative_config' in config: + if "num_nextn_predict_layers" in config[ + "speculative_config"]: + output_folder += f'_mtp{config["speculative_config"]["num_nextn_predict_layers"]}' + config = self.config['exec']['config'][f'{server_type}'] + if 'env' in config: + for key, value in config['env'].items(): + output_folder += f'_{key}_{value}' + + if not os.path.exists(output_folder): + os.makedirs(output_folder) + return output_folder + + def _check_existing_results(self): + """ + Check if results already exist for this configuration. + + Returns: + bool: True if all results exist and are valid, False otherwise + """ + if not os.path.exists(self.output_folder): + return False + + # Get concurrency levels from the configuration + concurrency_levels = self.config.get('profile', + {}).get('concurrency', [1]) + + for concurrency in concurrency_levels: + concurrency_path = os.path.join(self.output_folder, + str(concurrency)) + raw_data_path = os.path.join(concurrency_path, "raw_data.json") + + if not os.path.exists(raw_data_path): + return False + + # Check if the raw_data.json is valid and contains expected results + try: + with open(raw_data_path, 'r') as f: + data = json.load(f) + + if 'results' not in data or 'infbench_summary' not in data[ + 'results']: + return False + + if 'Requests/s' not in data['results']['infbench_summary']: + return False + except (json.JSONDecodeError, KeyError): + return False + + return True + + def setup_shared_directories(self, override_job_id=None): + """Create shared directories for coordination between nodes.""" + + # Use a unique directory for each job, or the override_job_id if provided + if override_job_id: + job_id = override_job_id + else: + job_id = str(uuid.uuid4()) + + os.makedirs(f"job_{job_id}", exist_ok=True) + self.job_id = job_id + + def load_config(self): + """Load the configuration file.""" + with open(self.args.config_file, "r") as f: + return yaml.safe_load(f) + + def create_context_config(self): + """Create the context server configuration file.""" + if 'context' in self.config['exec']['config']: + context_config = self.config['exec']['config']['context']['config'] + + # Create config in temporary job directory for execution + config_path = os.path.join(f"job_{self.job_id}", "context.yaml") + with open(config_path, 'w') as f: + yaml.dump(context_config, f) + + # Also save a copy to the output folder for reference + if hasattr(self, 'output_folder') and self.output_folder: + output_config_path = os.path.join(self.output_folder, + "context_config.yaml") + try: + with open(output_config_path, 'w') as f: + yaml.dump(context_config, + f, + default_flow_style=False, + indent=2) + print( + f"Context configuration saved to {output_config_path}") + except Exception as e: + print( + f"Warning: Could not save context config to output folder: {e}" + ) + + return os.path.abspath(config_path) + return None + + def create_generation_config(self): + """Create the generation server configuration file.""" + if 'generation' in self.config['exec']['config']: + generation_config = self.config['exec']['config']['generation'][ + 'config'] + + # Create config in temporary job directory for execution + config_path = os.path.join(f"job_{self.job_id}", "generation.yaml") + with open(config_path, 'w') as f: + yaml.dump(generation_config, f) + + # Also save a copy to the output folder for reference + if hasattr(self, 'output_folder') and self.output_folder: + output_config_path = os.path.join(self.output_folder, + "generation_config.yaml") + try: + with open(output_config_path, 'w') as f: + yaml.dump(generation_config, + f, + default_flow_style=False, + indent=2) + print( + f"Generation configuration saved to {output_config_path}" + ) + except Exception as e: + print( + f"Warning: Could not save generation config to output folder: {e}" + ) + + return os.path.abspath(config_path) + return None + + def create_ifb_config(self): + """Create the IFB server configuration file.""" + if 'ifb' in self.config['exec']['config']: + ifb_config = self.config['exec']['config']['ifb']['config'] + + # Create config in temporary job directory for execution + config_path = os.path.join(f"job_{self.job_id}", "ifb.yaml") + with open(config_path, 'w') as f: + yaml.dump(ifb_config, f) + + # Also save a copy to the output folder for reference + if hasattr(self, 'output_folder') and self.output_folder: + output_config_path = os.path.join(self.output_folder, + "ifb_config.yaml") + try: + with open(output_config_path, 'w') as f: + yaml.dump(ifb_config, + f, + default_flow_style=False, + indent=2) + print(f"IFB configuration saved to {output_config_path}") + except Exception as e: + print( + f"Warning: Could not save IFB config to output folder: {e}" + ) + + return os.path.abspath(config_path) + return None + + def get_node_hostnames(self): + """Get list of hostnames from SLURM allocation or fallback to localhost.""" + node_hostnames = [] + + # Try to get hostnames from SLURM + try: + # Check if we're running in a SLURM allocation + if "SLURM_JOB_ID" in os.environ: + cmd = [ + "srun", "--overlap", "--oversubscribe", "scontrol", "show", + "hostnames" + ] + print(f"Running command: {' '.join(cmd)}") + output = subprocess.check_output(cmd).decode().strip().split( + '\n') + for hostname in output: + if hostname.strip(): + node_hostnames.append(hostname.strip()) + + if not node_hostnames: + raise Exception("No hostnames returned by SLURM") + else: + raise Exception( + "Not running in a SLURM allocation (SLURM_JOB_ID not set)") + except Exception as e: + print( + f"Error getting node hostnames: {e}, falling back to localhost") + node_hostnames = ["localhost"] + + print(f"Using node hostnames: {node_hostnames}") + return node_hostnames + + def calculate_server_distribution(self): + """Calculate how to distribute context and generation servers across nodes.""" + distribution = [] + + total_ctx_servers = self.config['exec']['config']['context'][ + 'dp'] if 'context' in self.config['exec']['config'] else 0 + total_gen_servers = self.config['exec']['config']['generation'][ + 'dp'] if 'generation' in self.config['exec']['config'] else 0 + total_ifb_servers = self.config['exec']['config']['ifb'][ + 'dp'] if 'ifb' in self.config['exec']['config'] else 0 + port_base = 8001 + + gpu_index = 0 + server_id = 0 + j = 0 + for ctx_server in range(total_ctx_servers): + gpu_indices = [] + server_ids = [] + server_ids.append(self.node_hostnames[server_id]) + for i in range(self.config['exec']['config']['context']['tp']): + if i != 0 and i % self.num_gpus == 0: + server_id += 1 + print(f'Adding {self.node_hostnames[server_id]}') + server_ids.append(self.node_hostnames[server_id]) + gpu_indices.append(str(gpu_index % self.num_gpus)) + gpu_index += 1 + gpu_index = gpu_index % self.num_gpus + if len(gpu_indices) > self.num_gpus: + gpu_indices = gpu_indices[:self.num_gpus] + + distribution.append({ + "node_id": server_id, + "hostnames": server_ids, + "gpu_indices": gpu_indices, + "is_generation": False, + "is_ifb": False, + "port": port_base + j + }) + j += 1 + if gpu_index % self.num_gpus == 0: + gpu_index = 0 + server_id += 1 + + # If the last context server is not on a full node, add a generation server to the next node + if gpu_index % self.num_gpus != 0: + server_id += 1 + gpu_index = 0 + + for gen_server in range(total_gen_servers): + gpu_indices = [] + server_ids = [] + server_ids.append(self.node_hostnames[server_id]) + for i in range(self.config['exec']['config']['generation']['tp']): + if i != 0 and i % self.num_gpus == 0: + server_id += 1 + server_ids.append(self.node_hostnames[server_id]) + print(f'Adding {self.node_hostnames[server_id]}') + gpu_indices.append(str(gpu_index % self.num_gpus)) + gpu_index += 1 + gpu_index = gpu_index % self.num_gpus + + if len(gpu_indices) > self.num_gpus: + gpu_indices = gpu_indices[:self.num_gpus] + + distribution.append({ + "node_id": server_id, + "hostnames": server_ids, + "gpu_indices": gpu_indices, + "is_ifb": False, + "is_generation": True, + "port": port_base + j + }) + j += 1 + if gpu_index % self.num_gpus == 0: + gpu_index = 0 + server_id += 1 + for ifb_server in range(total_ifb_servers): + gpu_indices = [] + server_ids = [] + server_ids.append(self.node_hostnames[server_id]) + for i in range(self.config['exec']['config']['ifb']['tp']): + if i != 0 and i % self.num_gpus == 0: + server_id += 1 + server_ids.append(self.node_hostnames[server_id]) + print(f'Adding {self.node_hostnames[server_id]}') + gpu_indices.append(str(gpu_index % self.num_gpus)) + gpu_index += 1 + gpu_index = gpu_index % self.num_gpus + if len(gpu_indices) > self.num_gpus: + gpu_indices = gpu_indices[:self.num_gpus] + + distribution.append({ + "node_id": server_id, + "hostnames": server_ids, + "gpu_indices": gpu_indices, + "is_generation": False, + "is_ifb": True, + "port": port_base + j + }) + j += 1 + return distribution + + def run_in_trtllm_container(self, cmd, num_tasks=1): + """Run a command in the TRTLLM container.""" + cmd = [ + "srun", "--overlap", "--oversubscribe", "-A", self.account, "-p", + self.partition, "-N", "1", "-n", + str(num_tasks), "-t", self.time, "--container-image", + self.container_image, "--container-mounts", self.mounts, + "--container-workdir", self.workdir, "--mpi", "pmix", "bash", "-c", + " ".join(cmd) + ] + return subprocess.run(cmd) + + def spawn_in_trtllm_container(self, cmd, hostname): + """Spawn a command in the TRTLLM container.""" + cmd = [ + "srun", "--overlap", "--oversubscribe", "-A", self.account, "-p", + self.partition, "-N", "1", "-t", self.time, "--container-image", + self.container_image, "--container-mounts", self.mounts, + "--container-workdir", self.workdir, "--mpi", "pmix", "-w", + hostname, "bash", "-c", " ".join(cmd) + ] + return subprocess.Popen(cmd) + + def write_node_distribution(self, distribution): + """Write the server distribution to a file for the disaggregated server.""" + # Ensure the job_id directory exists + os.makedirs(f"job_{self.job_id}", exist_ok=True) + + with open(f"job_{self.job_id}/node_distribution.yaml", "w") as f: + yaml.dump(distribution, f) + + def install_trtllm(self): + """Install TensorRT-LLM into the container.""" + + cmd = [ + "srun", f"--container-name={self.container_name}", + f"--container-mounts={self.mounts}", "--mpi=pmix", "--overlap", + "-N", + os.environ.get('SLURM_NNODES', + '1'), "--ntasks-per-node=1", "bash", "-c", + f"echo 'Running install operation...' && pip install -e {self.trtllm_repo}" + ] + return execute_cmd(cmd) + + def launch_jobs(self): + """Launch SLURM jobs for all nodes.""" + print("DEBUG JobManager: launch_jobs started") + + # Install TensorRT-LLM into the container + self.install_trtllm() + + # Calculate server distribution + distribution = self.calculate_server_distribution() + self.write_node_distribution(distribution) + + # Track all launched jobs + # Launch context and generation servers on each node + for node in distribution: + node_id = node["node_id"] + hostnames = node["hostnames"] + gpu_indices = node["gpu_indices"] + print( + f"Launching context server on {hostnames[0]}:{node['port']}, GPUs: {gpu_indices}" + ) + + # Launch context servers + if node["is_ifb"]: + if 'use_trtllm_bench' not in self.config[ + 'profile'] or not self.config['profile'][ + 'use_trtllm_bench']: + log_file = f"ifb_server_{node_id}.log" + ifb_jobs = self.launch_ifb_server(node_id, hostnames, + gpu_indices, node["port"], + log_file) + self.ifb_jobs.extend(ifb_jobs) + elif not node["is_generation"]: + log_file = f"context_server_{node_id}.log" + ctx_jobs = self.launch_context_server(node_id, hostnames, + gpu_indices, node["port"], + log_file) + self.context_jobs.extend(ctx_jobs) + else: + log_file = f"generation_server_{node_id}.log" + gen_jobs = self.launch_generation_server( + node_id, hostnames, gpu_indices, node["port"], log_file) + self.generation_jobs.extend(gen_jobs) + # Wait for server to start + time.sleep(15) + + # Wait for all servers to start (give some time for them to initialize) + print("Waiting for all servers to start...") + + for node in distribution: + hostnames = node["hostnames"] + print(f"Waiting for {hostnames[0]}:{node['port']} to start...") + if 'use_trtllm_bench' not in self.config[ + 'profile'] or not self.config['profile']['use_trtllm_bench']: + wait_for_server(hostnames[0], node["port"]) + + # Collect server URLs and configure the disaggregated server + context_server_urls = [] + generation_server_urls = [] + ifb_server_urls = [] + + # Add all known server URLs based on the distribution + for node in distribution: + node_id = node["node_id"] + hostnames = node["hostnames"] + gpu_indices = node["gpu_indices"] + + # Add context server URLs + node_port = node["port"] + if node["is_ifb"]: + ifb_server_urls.append(f"{hostnames[0]}:{node_port}") + elif not node["is_generation"]: + context_server_urls.append(f"{hostnames[0]}:{node_port}") + else: + generation_server_urls.append(f"{hostnames[0]}:{node_port}") + + # Launch disaggregated server on the master node + distribution[0]["hostnames"][0] + port = 8000 + + if len(distribution) == 1 and distribution[0]["is_ifb"]: + distribution[0]["hostnames"][0] + port = distribution[0]["port"] + else: + for i in range(self.args.num_disagg_servers): + disagg_job = self.launch_disaggregated_server( + self.node_hostnames[i], context_server_urls, + generation_server_urls) + self.disagg_jobs.append(disagg_job) + + # Wait a bit for the disaggregated server to initialize + print("Waiting for disaggregated server to initialize...") + time.sleep(30) + + load_gen_jobs = [] + # Run the loadgen tests + print("DEBUG JobManager: About to call run_loadgen_tests") + for i in range(self.args.num_disagg_servers): + load_gen_jobs.append( + self.run_loadgen_tests(self.node_hostnames[i], port, + generation_server_urls)) + print("DEBUG JobManager: Completed run_loadgen_tests") + + for loadgen_job in load_gen_jobs: + loadgen_job.wait() + + self.terminate_jobs(self.context_jobs + self.generation_jobs + + self.ifb_jobs + self.disagg_jobs) + print("DEBUG JobManager: launch_jobs completed") + + def _build_server_launch_command(self, config, hostnames, gpu_indices, + node_port, config_path, log_file): + """Build the command to launch a server (context or generation). + + Returns: + cmd: The command to launch the server + """ + if not gpu_indices or not config_path: + return None + + # Get configuration based on server type + server_config = config + tp = server_config['tp'] + ep = server_config['ep'] + pp = server_config['pp'] + max_batch_size = server_config['max_batch_size'] + max_seq_len = server_config['max_seq_len'] + max_num_tokens = server_config['max_num_tokens'] + cuda_visible_devices = ",".join(gpu_indices) + + print( + f"Launching server on {hostnames[0]}:{node_port}, GPUs: {cuda_visible_devices}" + ) + + envs = "" + envs += "TRTLLM_USE_UCX_KVCACHE=1 " + envs += " CUDA_VISIBLE_DEVICES=" + cuda_visible_devices + envs += " UCX_CUDA_IPC_ENABLE_MNNVL=n " + if 'gen_only' in config and config['gen_only']: + envs += " TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1" + + if 'env' in config: + for key, value in config['env'].items(): + envs += f" {key}={value}" + unset_env = '' + if 'unset_env' in config: + for key in config['unset_env']: + unset_env += f" -u {key}" + nsys = False + nsys_prefix = '' + if 'nsys' in config: + nsys = True + if nsys: + envs += " TLLM_PROFILE_RECORD_GC=1" + envs += " TLLM_NVTX_DEBUG=1" + if 'TLLM_PROFILE_START_STOP' in config['nsys']: + envs += f" TLLM_PROFILE_START_STOP={config['nsys']['TLLM_PROFILE_START_STOP']}" + nsys_file = os.path.join(self.output_folder, + f"{log_file}.$SLURM_PROCID.nsys-rep") + nsys_prefix = f"nsys profile -e \"NSYS_MPI_STORE_TEAMS_PER_RANK=1\" -o {nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=none" + log_file = os.path.join(self.output_folder, log_file) + + # Build the command + cmd = [ + "srun", + "--oversubscribe", + "-A", + self.account, + "-p", + self.partition, + "-N", + f"{len(hostnames)}", + "-n", + str(tp), + "-t", + self.time, + "--mpi", + "pmix", + "-w", + ",".join(hostnames), # Specify the hostnames + f"--container-image={self.container_image}", + f"--container-name={self.container_name}", + f"--container-mounts={self.mounts}", + f"--container-workdir={self.workdir}", + "--export=\"'CUDA_VISIBLE_DEVICES=" + cuda_visible_devices + "'\"", + "bash", + "-c", + envs + " env " + unset_env + " " + nsys_prefix + + " trtllm-llmapi-launch trtllm-serve " + + self.config['exec']['model_path'] + " --host 0.0.0.0 --port " + + str(node_port) + " --backend pytorch --extra_llm_api_options " + + config_path + " --tp_size " + str(tp) + " --ep_size " + str(ep) + + " --max_batch_size " + str(max_batch_size) + " --max_seq_len " + + str(max_seq_len) + " --max_num_tokens " + str(max_num_tokens) + + " --pp_size " + str(pp) + " &> " + log_file + ] + + return cmd + + def launch_context_server(self, node_id, hostnames, gpu_indices, node_port, + log_file): + """Launch context server on a specific node.""" + launched_jobs = [] + + cmd = self._build_server_launch_command( + self.config['exec']['config']['context'], hostnames, gpu_indices, + node_port, self.context_config_path, log_file) + if not cmd: + return launched_jobs + + print(f"Running command: {' '.join(cmd)}") + proc = subprocess.Popen(cmd) + launched_jobs.append(proc) + + # Small delay to avoid race conditions in resource allocation + time.sleep(2) + + return launched_jobs + + def launch_generation_server(self, node_id, hostnames, gpu_indices, + node_port, log_file): + """Launch generation server on a specific node.""" + launched_jobs = [] + + cmd = self._build_server_launch_command( + self.config['exec']['config']['generation'], hostnames, gpu_indices, + node_port, self.generation_config_path, log_file) + if not cmd: + return launched_jobs + + print(f"Running command: {' '.join(cmd)}") + proc = subprocess.Popen(cmd) + launched_jobs.append(proc) + + return launched_jobs + + def launch_ifb_server(self, node_id, hostnames, gpu_indices, node_port, + log_file): + """Launch IFB server on a specific node.""" + launched_jobs = [] + + cmd = self._build_server_launch_command( + self.config['exec']['config']['ifb'], hostnames, gpu_indices, + node_port, self.ifb_config_path, log_file) + if not cmd: + return launched_jobs + + print(f"Running command: {' '.join(cmd)}") + proc = subprocess.Popen(cmd) + launched_jobs.append(proc) + + return launched_jobs + + def launch_disaggregated_server(self, hostname, context_servers, + generation_servers): + """Launch the disaggregated server on the master node.""" + if not generation_servers: + print( + "No context or generation servers available. Cannot launch disaggregated server." + ) + return None + + config = self.config['exec']['config'] + # Create the disaggregated server configuration + disagg_config = { + "hostname": hostname, + "port": 8000, + "backend": "pytorch", + "model": self.config['exec']['model_path'], + "context_servers": { + "num_instances": + len(context_servers), + "urls": + context_servers, + "tensor_parallel_size": + config['context']['tp'] if 'context' in config else 1, + }, + "generation_servers": { + "num_instances": + len(generation_servers), + "urls": + generation_servers, + "tensor_parallel_size": + self.config['exec']['config']['generation']['tp'], + } + } + + # Write the configuration to a file + disagg_config_path = os.path.join(f"job_{self.job_id}", + f"disagg_config_{hostname}.yaml") + with open(disagg_config_path, 'w') as f: + yaml.dump(disagg_config, f) + disagg_config_path = os.path.abspath(disagg_config_path) + + # Also save a copy to the output folder for reference + if hasattr(self, 'output_folder') and self.output_folder: + output_config_path = os.path.join(self.output_folder, + f"disagg_config_{hostname}.yaml") + try: + with open(output_config_path, 'w') as f: + yaml.dump(disagg_config, + f, + default_flow_style=False, + indent=2) + print( + f"Disaggregated server configuration saved to {output_config_path}" + ) + except Exception as e: + print( + f"Warning: Could not save disaggregated config to output folder: {e}" + ) + + print(f"Launching disaggregated server on {hostname}:8000") + envs = "" + gen_only = self.config['exec']['config']['generation'].get( + 'gen_only', False) + if gen_only: + envs += " TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1" + + # Store the log to the output folder + log_file = os.path.join(self.output_folder, + f"disagg_server_{hostname}.log") + # Build the command + cmd = [ + "srun", + "--overlap", + "--oversubscribe", + "-A", + self.account, + "-p", + self.partition, + "-N", + "1", + "-n", + "1", + "-t", + self.time, + "-w", + hostname, # Specify the hostname + "-J", + f"{self.job_name}_disagg", + "--mpi", + "pmix", + f"--container-image={self.container_image}", + f"--container-mounts={self.mounts}", + f"--container-workdir={self.workdir}", + "--export=OMPI_ALLOW_RUN_AS_ROOT=1", + "bash", + "-c", + f" {envs} trtllm-serve disaggregated -c " + disagg_config_path + + " --request_timeout " + " 1800000" + " &> " + log_file + ] + + print(f"Running command: {' '.join(cmd)}") + return subprocess.Popen(cmd) + + def run_loadgen_tests(self, hostname, port, generation_server_urls): + """Run loadgen tests on the master node.""" + print( + f"DEBUG JobManager: run_loadgen_tests called with hostname={hostname}, port={port}" + ) + output_folder = self.output_folder + + # Run loadgen tests for each concurrency level + if 'use_trtllm_bench' in self.config['profile'] and self.config[ + 'profile']['use_trtllm_bench']: + self.prepare_dataset(output_folder) + + concurrency_levels = self.config['profile']['concurrency'] + print( + f"DEBUG JobManager: Will run loadgen for concurrency levels: {concurrency_levels}" + ) + + for concurrency in concurrency_levels: + print( + f"DEBUG JobManager: Starting loadgen for concurrency {concurrency}" + ) + concurrency_output_folder = f"{output_folder}/{concurrency}" + if not os.path.exists(concurrency_output_folder): + os.makedirs(concurrency_output_folder) + + # Build and run the loadgen command + if 'use_trtllm_bench' in self.config['profile'] and self.config[ + 'profile']['use_trtllm_bench']: + self.run_trtllm_bench_for_concurrency( + hostname, port, concurrency, concurrency_output_folder) + elif 'use_benchmark_serving' in self.config[ + 'profile'] and self.config['profile'][ + 'use_benchmark_serving']: + pid = self.run_benchmark_serving_for_concurrency( + hostname, port, concurrency, concurrency_output_folder) + pid.wait() + print( + f"DEBUG JobManager: Completed loadgen for concurrency {concurrency}" + ) + + print("DEBUG JobManager: run_loadgen_tests completed") + return pid + + def record_kv_cache_stats(self, generation_server_urls, output_folder): + """Record the kv cache stats for a specific concurrency level.""" + print( + f"DEBUG JobManager: Recording kv cache stats for generation servers" + ) + # Get the kv cache stats + # Send a request to "/metrics" and store the output in a file + for url in generation_server_urls: + response = requests.get(f"http://{url}/metrics") + with open( + os.path.join(output_folder, + f"{url.split(':')[0]}_kv_cache_stats.txt"), + "w") as f: + f.write(response.text) + print( + f"DEBUG JobManager: Recorded kv cache stats for generation servers") + + def prepare_dataset(self, output_folder): + """Prepare the dataset for the loadgen tests.""" + input_tokens = str(self.config['profile']['isl']) + output_tokens = str(self.config['profile']['osl']) + num_prompts = str( + 8 * self.config['exec']['config']['ifb']['max_batch_size']) + + # If the dataset file already exists, skip the preparation + dataset_file = os.path.join(output_folder, "dataset.json") + if os.path.exists(dataset_file): + return + + cmd = [ + "python3", self.args.prepare_dataset_script, "--stdout", + "--tokenizer", self.config['exec']['model_path'], "token-norm-dist", + "--input-mean", input_tokens, "--output-mean", output_tokens, + "--input-stdev", "0", "--output-stdev", "0", "--num-requests", + num_prompts, ">", dataset_file + ] + print(f"Running command: {' '.join(cmd)}") + self.run_in_trtllm_container(cmd) + + def run_trtllm_bench_for_concurrency(self, hostname, port, concurrency, + output_folder): + """Run the trtllm bench for a specific concurrency level.""" + config = self.config['exec']['config']['ifb'] + tp = config['tp'] if 'tp' in config else 0 + ep = config['ep'] if 'ep' in config else 0 + max_batch_size = config[ + 'max_batch_size'] if 'max_batch_size' in config else 0 + max_num_tokens = config[ + 'max_num_tokens'] if 'max_num_tokens' in config else 0 + gpu_fraction = config['config']['kv_cache_config'][ + 'free_gpu_memory_fraction'] if 'config' in config and 'kv_cache_config' in config[ + 'config'] else 0 + num_requests = concurrency * 8 + # Parent of output_folder is the dataset file + dataset_file = os.path.join(os.path.dirname(output_folder), + "dataset.json") + # Redirect the output to a file + cmd = [ + "trtllm-llmapi-launch", "trtllm-bench", "-m", + "deepseek-ai/DeepSeek-R1", "--model_path", + f"{self.config['exec']['model_path']}", "throughput", "--tp", + f"{tp}", "--ep", f"{ep}", "--warmup", "0", "--dataset", + f"{dataset_file}", "--backend", "pytorch", "--max_batch_size", + f"{max_batch_size}", "--max_num_tokens", f"{max_num_tokens}", + "--kv_cache_free_gpu_mem_fraction", f"{gpu_fraction}", + "--extra_llm_api_options", f"{self.ifb_config_path}", + "--num_requests", f"{num_requests}", "--concurrency", + f"{concurrency}", "--report_json", + f"{output_folder}/trtllm_bench_throughput.json", "--streaming", + "&>" + os.path.join(output_folder, "output.log") + ] + print(f"Running command: {' '.join(cmd)}") + self.run_in_trtllm_container(cmd, tp) + + def run_benchmark_serving_for_concurrency(self, hostname, port, concurrency, + output_folder): + """Run the benchmark serving test for a specific concurrency level.""" + model_name = self.config['exec']['model_path'] + dataset_path = self.config['profile']['dataset_path'] + num_prompts = self.config['profile']['num_prompts'] + cmd = [ + "/usr/bin/python3", "-m", + "tensorrt_llm.serve.scripts.benchmark_serving", "--model", + model_name, "--tokenizer", model_name, "--dataset-name", + "trtllm_custom", "--dataset-path", dataset_path, "--num-prompts", + str(num_prompts), "--host", hostname, "--port", + str(port), "--max-concurrency", + str(concurrency) + ] + if 'ignore_eos' in self.config['profile'] and self.config['profile'][ + 'ignore_eos']: + cmd += ["--ignore-eos"] + cmd += ["&> " + os.path.join(output_folder, "benchmark_serving.log")] + print(f"Running command: {' '.join(cmd)}") + return self.spawn_in_trtllm_container(cmd, hostname) + + def terminate_jobs(self, jobs): + """Terminate all running jobs.""" + + # Then terminate the jobs + for job in jobs: + if job and job.poll() is None: + try: + job.terminate() + job.wait(timeout=5) + except: + job.kill() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.terminate_jobs(self.context_jobs + self.generation_jobs + + self.ifb_jobs + self.disagg_jobs) + + def save_config_to_output_folder(self): + """Save the complete configuration used for this run to the output folder.""" + if not hasattr(self, 'output_folder') or not self.output_folder: + return + + config_file_path = os.path.join(self.output_folder, "config.yaml") + try: + with open(config_file_path, 'w') as f: + yaml.dump(self.config, f, default_flow_style=False, indent=2) + print(f"Configuration saved to {config_file_path}") + except Exception as e: + print( + f"Warning: Could not save configuration to output folder: {e}") + + def prepare_and_run_single_config(self): + """Prepare and run a single configuration.""" + print("DEBUG JobManager: prepare_and_run_single_config started") + + # Create the configuration files + self.context_config_path = self.create_context_config() + self.generation_config_path = self.create_generation_config() + self.ifb_config_path = self.create_ifb_config() + + # Save the complete configuration to the output folder + self.save_config_to_output_folder() + + # Get hostnames of available nodes + self.node_hostnames = self.get_node_hostnames() diff --git a/examples/disaggregated/slurm/launcher.py b/examples/disaggregated/slurm/launcher.py new file mode 100644 index 000000000000..f44f51d2d66f --- /dev/null +++ b/examples/disaggregated/slurm/launcher.py @@ -0,0 +1,89 @@ +import argparse +import os + +from disagg_profiler.job_manager import JobManager, get_slurm_allocation + + +def setup_argparse(): + """Set up command line argument parsing.""" + parser = argparse.ArgumentParser( + description="TensorRT-LLM Disaggregated Serving Launcher") + + # SLURM job arguments + parser.add_argument("--account", + type=str, + required=True, + help="SLURM account") + parser.add_argument("--partition", + type=str, + required=True, + help="SLURM partition") + parser.add_argument("--time", + type=str, + default="04:00:00", + help="SLURM time limit") + parser.add_argument("--job-name", + type=str, + required=True, + help="SLURM job name") + parser.add_argument("--container-image", + type=str, + required=True, + help="Container image") + parser.add_argument("--mounts", + type=str, + required=True, + help="Container mount points") + parser.add_argument("--num-gpus", + type=int, + required=True, + help="Number of GPUs on each node") + parser.add_argument("--num-nodes", + type=int, + required=True, + help="Number of nodes") + parser.add_argument("--ntasks", + type=int, + required=True, + help="Number of tasks") + parser.add_argument("--ntasks-per-node", + type=int, + required=True, + help="Number of tasks per node") + + parser.add_argument("--trtllm-repo", + type=str, + default=None, + help="TensorRT-LLM repository") + parser.add_argument("--config-file", + type=str, + default="./config.yaml", + help="Path to the configuration file") + parser.add_argument("--experiment-path", + type=str, + default=".", + help="Path to the configuration file") + parser.add_argument("--num-disagg-servers", + type=int, + default=1, + help="Number of disaggregated servers to use.") + + return parser + + +def main(): + parser = setup_argparse() + args = parser.parse_args() + + slurm_job_id = get_slurm_allocation(args.account, + args.partition, + args.time, + f"{args.job_name}", + num_nodes=args.num_nodes) + os.environ["SLURM_JOB_ID"] = slurm_job_id + with JobManager(args) as job_manager: + job_manager.launch_jobs() + + +if __name__ == "__main__": + main() From 6029a4876c392142ca3b330915ee92f736cd3ed7 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 02:36:37 -0700 Subject: [PATCH 2/9] Update Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- .../disaggregated/slurm/disagg_profiler/job_manager.py | 4 +++- examples/disaggregated/slurm/launcher.py | 7 ++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py index 6d3ac837c5f7..f275ac01c97f 100644 --- a/examples/disaggregated/slurm/disagg_profiler/job_manager.py +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -9,7 +9,7 @@ import yaml -def get_slurm_allocation(account, partition, time_limit, job_name, num_nodes): +def get_slurm_allocation(account, partition, time_limit, job_name, num_nodes, gres=None): """ Request a SLURM allocation and wait for it to be granted. Args: @@ -30,6 +30,8 @@ def get_slurm_allocation(account, partition, time_limit, job_name, num_nodes): "salloc", "-A", account, "-p", partition, "-N", str(num_nodes), "-t", time_limit, "-J", job_name, "--no-shell" ] + if gres: + cmd.append(f"--gres={gres}") try: # Run the command and stream the output in real-time diff --git a/examples/disaggregated/slurm/launcher.py b/examples/disaggregated/slurm/launcher.py index f44f51d2d66f..d85a941357ee 100644 --- a/examples/disaggregated/slurm/launcher.py +++ b/examples/disaggregated/slurm/launcher.py @@ -50,6 +50,10 @@ def setup_argparse(): type=int, required=True, help="Number of tasks per node") + parser.add_argument("--gres", + type=str, + default=None, + help="GPUs per node") parser.add_argument("--trtllm-repo", type=str, @@ -79,7 +83,8 @@ def main(): args.partition, args.time, f"{args.job_name}", - num_nodes=args.num_nodes) + num_nodes=args.num_nodes, + gres=args.gres) os.environ["SLURM_JOB_ID"] = slurm_job_id with JobManager(args) as job_manager: job_manager.launch_jobs() From ae86a523f171d359c33753ad5f42c40f676f1cfa Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 02:48:20 -0700 Subject: [PATCH 3/9] Fix Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- .../slurm/disagg_profiler/job_manager.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py index f275ac01c97f..ca9623c82fb8 100644 --- a/examples/disaggregated/slurm/disagg_profiler/job_manager.py +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -172,6 +172,7 @@ def __init__(self, self.container_image = args.container_image self.mounts = args.mounts self.container_name = "trtllm-disagg-server" + self.trtllm_repo = args.trtllm_repo self.workdir = os.path.abspath(os.path.join(os.getcwd())) self.context_jobs = [] self.generation_jobs = [] @@ -586,15 +587,16 @@ def write_node_distribution(self, distribution): def install_trtllm(self): """Install TensorRT-LLM into the container.""" - cmd = [ - "srun", f"--container-name={self.container_name}", - f"--container-mounts={self.mounts}", "--mpi=pmix", "--overlap", - "-N", - os.environ.get('SLURM_NNODES', - '1'), "--ntasks-per-node=1", "bash", "-c", - f"echo 'Running install operation...' && pip install -e {self.trtllm_repo}" - ] - return execute_cmd(cmd) + if self.trtllm_repo: + cmd = [ + "srun", f"--container-name={self.container_name}", + f"--container-mounts={self.mounts}", "--mpi=pmix", "--overlap", + "-N", + os.environ.get('SLURM_NNODES', + '1'), "--ntasks-per-node=1", "bash", "-c", + f"echo 'Running install operation...' && pip install -e {self.trtllm_repo}" + ] + return execute_cmd(cmd) def launch_jobs(self): """Launch SLURM jobs for all nodes.""" From 519d4bdb03be9c450e26868b48ab8f6fc08db7d4 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 02:54:57 -0700 Subject: [PATCH 4/9] Update Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- examples/disaggregated/slurm/disagg_profiler/job_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py index ca9623c82fb8..199c60c707da 100644 --- a/examples/disaggregated/slurm/disagg_profiler/job_manager.py +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -594,7 +594,7 @@ def install_trtllm(self): "-N", os.environ.get('SLURM_NNODES', '1'), "--ntasks-per-node=1", "bash", "-c", - f"echo 'Running install operation...' && pip install -e {self.trtllm_repo}" + f"echo 'Running install operation...' && pip install -e {self.trtllm_repo} 2>&1 | tee install.log" ] return execute_cmd(cmd) @@ -603,7 +603,7 @@ def launch_jobs(self): print("DEBUG JobManager: launch_jobs started") # Install TensorRT-LLM into the container - self.install_trtllm() + self.install_trtllm().wait() # Calculate server distribution distribution = self.calculate_server_distribution() From 37e5a0b0fcc0ca19a4fc25e6ab54b36a1310461b Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 02:58:42 -0700 Subject: [PATCH 5/9] Update Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- .../slurm/disagg_profiler/job_manager.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py index 199c60c707da..e15a1e6278fd 100644 --- a/examples/disaggregated/slurm/disagg_profiler/job_manager.py +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -584,6 +584,16 @@ def write_node_distribution(self, distribution): with open(f"job_{self.job_id}/node_distribution.yaml", "w") as f: yaml.dump(distribution, f) + def start_container(self): + """Start the container.""" + cmd = [ + "srun", "-l", f"--container-image={self.container_image}", + f"--container-name={self.container_name}", + f"--container-mounts={self.mounts}", "--mpi=pmix", + "echo 'Container up.'" + ] + return execute_cmd(cmd) + def install_trtllm(self): """Install TensorRT-LLM into the container.""" @@ -602,6 +612,9 @@ def launch_jobs(self): """Launch SLURM jobs for all nodes.""" print("DEBUG JobManager: launch_jobs started") + # Start the container + self.start_container().wait() + # Install TensorRT-LLM into the container self.install_trtllm().wait() From 9ca31bc499339671e9c26cf0431c58a966380ca5 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 03:06:18 -0700 Subject: [PATCH 6/9] Update Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- .../disaggregated/slurm/disagg_profiler/job_manager.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py index e15a1e6278fd..abb0c667ead2 100644 --- a/examples/disaggregated/slurm/disagg_profiler/job_manager.py +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -165,6 +165,7 @@ def __init__(self, self.args = args self.num_gpus = args.num_gpus + self.num_nodes = args.num_nodes self.account = args.account self.partition = args.partition self.time = args.time @@ -589,7 +590,7 @@ def start_container(self): cmd = [ "srun", "-l", f"--container-image={self.container_image}", f"--container-name={self.container_name}", - f"--container-mounts={self.mounts}", "--mpi=pmix", + f"--container-mounts={self.mounts}", "--mpi=pmix", "bash", "-c", "echo 'Container up.'" ] return execute_cmd(cmd) @@ -601,9 +602,7 @@ def install_trtllm(self): cmd = [ "srun", f"--container-name={self.container_name}", f"--container-mounts={self.mounts}", "--mpi=pmix", "--overlap", - "-N", - os.environ.get('SLURM_NNODES', - '1'), "--ntasks-per-node=1", "bash", "-c", + "-N", str(self.num_nodes), "--ntasks-per-node=1", "bash", "-c", f"echo 'Running install operation...' && pip install -e {self.trtllm_repo} 2>&1 | tee install.log" ] return execute_cmd(cmd) From a9bbb5b0c5450ca2a5f586acee6935fd5b8d3d5e Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 04:56:48 -0700 Subject: [PATCH 7/9] Update Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- examples/disaggregated/slurm/config.yaml | 2 ++ .../slurm/disagg_profiler/job_manager.py | 12 +++++++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/examples/disaggregated/slurm/config.yaml b/examples/disaggregated/slurm/config.yaml index 3efb7622bd07..310de724c485 100644 --- a/examples/disaggregated/slurm/config.yaml +++ b/examples/disaggregated/slurm/config.yaml @@ -28,9 +28,11 @@ exec: dp: 1 model_path: TinyLlama/TinyLlama-1.1B-Chat-v1.0 profile: + dataset_path: /tmp/ShareGPT_V3_unfiltered_cleaned_split.json isl: 1024 osl: 1024 use_benchmark_serving: true concurrency: - 128 - 256 + ignore_eos: true diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py index abb0c667ead2..1f574d309efe 100644 --- a/examples/disaggregated/slurm/disagg_profiler/job_manager.py +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -1085,17 +1085,23 @@ def run_trtllm_bench_for_concurrency(self, hostname, port, concurrency, def run_benchmark_serving_for_concurrency(self, hostname, port, concurrency, output_folder): """Run the benchmark serving test for a specific concurrency level.""" - model_name = self.config['exec']['model_path'] dataset_path = self.config['profile']['dataset_path'] + cmd=["wget", "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json", "-O", dataset_path] + execute_cmd(cmd).wait() + + model_name = self.config['exec']['model_path'] + isl = self.config['profile']['isl'] + osl = self.config['profile']['osl'] num_prompts = self.config['profile']['num_prompts'] cmd = [ "/usr/bin/python3", "-m", "tensorrt_llm.serve.scripts.benchmark_serving", "--model", model_name, "--tokenizer", model_name, "--dataset-name", - "trtllm_custom", "--dataset-path", dataset_path, "--num-prompts", + "random", "--dataset-path", dataset_path, "--random-input-len", isl, + "--random-output-len", osl, "--random-prefix-len", 0, "--num-prompts", str(num_prompts), "--host", hostname, "--port", str(port), "--max-concurrency", - str(concurrency) + str(concurrency), "--no-test-input" ] if 'ignore_eos' in self.config['profile'] and self.config['profile'][ 'ignore_eos']: From e2887f24c866a94a645555fc7f793b363a92a6e3 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 05:10:52 -0700 Subject: [PATCH 8/9] Update Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- examples/disaggregated/slurm/config.yaml | 1 + .../slurm/disagg_profiler/job_manager.py | 86 ++++--------------- examples/disaggregated/slurm/launcher.py | 5 +- 3 files changed, 19 insertions(+), 73 deletions(-) diff --git a/examples/disaggregated/slurm/config.yaml b/examples/disaggregated/slurm/config.yaml index 310de724c485..bc6f2805d9c3 100644 --- a/examples/disaggregated/slurm/config.yaml +++ b/examples/disaggregated/slurm/config.yaml @@ -31,6 +31,7 @@ profile: dataset_path: /tmp/ShareGPT_V3_unfiltered_cleaned_split.json isl: 1024 osl: 1024 + num_prompts: 8192 use_benchmark_serving: true concurrency: - 128 diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py index 1f574d309efe..746618cf8e7c 100644 --- a/examples/disaggregated/slurm/disagg_profiler/job_manager.py +++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py @@ -9,7 +9,12 @@ import yaml -def get_slurm_allocation(account, partition, time_limit, job_name, num_nodes, gres=None): +def get_slurm_allocation(account, + partition, + time_limit, + job_name, + num_nodes, + gres=None): """ Request a SLURM allocation and wait for it to be granted. Args: @@ -96,52 +101,6 @@ def wait_for_server(host, port, max_retries=1000, delay=5): return False -def calculate_nodes_needed(config, num_gpus): - """ - Calculate the number of nodes needed based on the configuration. - - The calculation is based on: - 1. The number of context servers and their TP size - 2. The number of generation servers and their TP size - 3. The number of IFB servers and their TP size - 4. The number of GPUs per node (num_gpus) - - Args: - config (dict): The configuration dictionary - - Returns: - int: The number of nodes needed - """ - total_gpus_needed = 0 - - # Calculate GPUs needed for context servers - if 'context' in config['exec']['config']: - context_config = config['exec']['config']['context'] - context_tp = context_config['tp'] - context_dp = context_config['dp'] - total_gpus_needed += context_tp * context_dp - - # Calculate GPUs needed for generation servers - if 'generation' in config['exec']['config']: - generation_config = config['exec']['config']['generation'] - generation_tp = generation_config['tp'] - generation_dp = generation_config['dp'] - total_gpus_needed += generation_tp * generation_dp - - # Calculate GPUs needed for IFB servers - if 'ifb' in config['exec']['config']: - ifb_config = config['exec']['config']['ifb'] - ifb_tp = ifb_config['tp'] - ifb_dp = ifb_config['dp'] - total_gpus_needed += ifb_tp * ifb_dp - - # Calculate nodes needed (round up to ensure enough nodes) - nodes_needed = (total_gpus_needed + num_gpus - 1) // num_gpus - - # Ensure we request at least one node - return max(1, nodes_needed) - - class JobManager: """Manages SLURM jobs for the disaggregated serving setup.""" @@ -602,7 +561,8 @@ def install_trtllm(self): cmd = [ "srun", f"--container-name={self.container_name}", f"--container-mounts={self.mounts}", "--mpi=pmix", "--overlap", - "-N", str(self.num_nodes), "--ntasks-per-node=1", "bash", "-c", + "-N", + str(self.num_nodes), "--ntasks-per-node=1", "bash", "-c", f"echo 'Running install operation...' && pip install -e {self.trtllm_repo} 2>&1 | tee install.log" ] return execute_cmd(cmd) @@ -1008,23 +968,6 @@ def run_loadgen_tests(self, hostname, port, generation_server_urls): print("DEBUG JobManager: run_loadgen_tests completed") return pid - def record_kv_cache_stats(self, generation_server_urls, output_folder): - """Record the kv cache stats for a specific concurrency level.""" - print( - f"DEBUG JobManager: Recording kv cache stats for generation servers" - ) - # Get the kv cache stats - # Send a request to "/metrics" and store the output in a file - for url in generation_server_urls: - response = requests.get(f"http://{url}/metrics") - with open( - os.path.join(output_folder, - f"{url.split(':')[0]}_kv_cache_stats.txt"), - "w") as f: - f.write(response.text) - print( - f"DEBUG JobManager: Recorded kv cache stats for generation servers") - def prepare_dataset(self, output_folder): """Prepare the dataset for the loadgen tests.""" input_tokens = str(self.config['profile']['isl']) @@ -1086,7 +1029,11 @@ def run_benchmark_serving_for_concurrency(self, hostname, port, concurrency, output_folder): """Run the benchmark serving test for a specific concurrency level.""" dataset_path = self.config['profile']['dataset_path'] - cmd=["wget", "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json", "-O", dataset_path] + cmd = [ + "wget", + "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json", + "-O", dataset_path + ] execute_cmd(cmd).wait() model_name = self.config['exec']['model_path'] @@ -1096,9 +1043,10 @@ def run_benchmark_serving_for_concurrency(self, hostname, port, concurrency, cmd = [ "/usr/bin/python3", "-m", "tensorrt_llm.serve.scripts.benchmark_serving", "--model", - model_name, "--tokenizer", model_name, "--dataset-name", - "random", "--dataset-path", dataset_path, "--random-input-len", isl, - "--random-output-len", osl, "--random-prefix-len", 0, "--num-prompts", + model_name, "--tokenizer", model_name, "--dataset-name", "random", + "--dataset-path", dataset_path, "--random-input-len", isl, + "--random-output-len", osl, "--random-prefix-len", 0, + "--num-prompts", str(num_prompts), "--host", hostname, "--port", str(port), "--max-concurrency", str(concurrency), "--no-test-input" diff --git a/examples/disaggregated/slurm/launcher.py b/examples/disaggregated/slurm/launcher.py index d85a941357ee..396af4a51d18 100644 --- a/examples/disaggregated/slurm/launcher.py +++ b/examples/disaggregated/slurm/launcher.py @@ -50,10 +50,7 @@ def setup_argparse(): type=int, required=True, help="Number of tasks per node") - parser.add_argument("--gres", - type=str, - default=None, - help="GPUs per node") + parser.add_argument("--gres", type=str, default=None, help="GPUs per node") parser.add_argument("--trtllm-repo", type=str, From 35031e46631ef403cc889a7f6f1cc48ec94154a5 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 26 Aug 2025 05:12:44 -0700 Subject: [PATCH 9/9] Fix Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- examples/disaggregated/slurm/disagg_profiler/__init__.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/examples/disaggregated/slurm/disagg_profiler/__init__.py b/examples/disaggregated/slurm/disagg_profiler/__init__.py index e25f5b45807c..8aba3b2e5724 100644 --- a/examples/disaggregated/slurm/disagg_profiler/__init__.py +++ b/examples/disaggregated/slurm/disagg_profiler/__init__.py @@ -4,10 +4,8 @@ for the TRT-LLM disaggregated serving launcher. """ -from .job_manager import JobManager, calculate_nodes_needed, wait_for_server +from .job_manager import JobManager, wait_for_server __all__ = [ - 'JobManager', 'calculate_nodes_needed', 'wait_for_server', - 'ParameterSweeper', 'AutoSweeper', 'get_slurm_allocation', - 'run_sweep_configuration' + 'JobManager', 'wait_for_server', ]