TensorRT-LLMs/examples/disaggregated/slurm/benchmark/submit.py

#!/usr/bin/env python3

import argparse
import glob
import os
import subprocess
import sys

import yaml


def parse_args():
    parser = argparse.ArgumentParser(
        description='Submit disaggregated benchmark job')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-c',
                       '--config',
                       type=str,
                       help='Path to the configuration YAML file')
    group.add_argument('-d',
                       '--dir',
                       type=str,
                       help='Directory containing YAML configuration files')
    return parser.parse_args()


def load_config(config_path):
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)


def save_worker_config(config, output_path, worker_type):
    """Save worker config to a separate YAML file."""
    # Get just the worker configuration without the wrapper
    worker_config = config['worker_config'][worker_type]

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, 'w') as f:
        yaml.dump(worker_config, f, default_flow_style=False)


def calculate_nodes(tp_size, num_servers, gpus_per_node):
    """Calculate required nodes based on tensor parallel size and server count."""
    return (tp_size + gpus_per_node - 1) // gpus_per_node * num_servers


def submit_job(config):
    # Extract configurations
    slurm_config = config['slurm']
    hw_config = config['hardware']
    env_config = config['environment']

    # Calculate nodes based on tensor parallel sizes
    ctx_tp_size = config['worker_config']['ctx']['tensor_parallel_size']
    gen_tp_size = config['worker_config']['gen']['tensor_parallel_size']

    # Get number of servers from config
    ctx_num = hw_config['num_ctx_servers']
    gen_num = hw_config['num_gen_servers']

    # Get mtp_size from gen config's speculative_config
    gen_config = config['worker_config']['gen']
    mtp_size = gen_config.get('speculative_config',
                              {}).get('num_nextn_predict_layers', 0)

    ctx_nodes = calculate_nodes(ctx_tp_size, ctx_num,
                                hw_config['gpus_per_node'])
    gen_nodes = calculate_nodes(gen_tp_size, gen_num,
                                hw_config['gpus_per_node'])
    total_nodes = ctx_nodes + gen_nodes
    total_tasks = total_nodes * hw_config['gpus_per_node']

    # Generate log directory path based on configuration
    isl = config['sequence']['input_length']
    osl = config['sequence']['output_length']
    gen_batch_size = config['worker_config']['gen']['max_batch_size']
    gen_enable_attention_dp = config['worker_config']['gen'][
        'enable_attention_dp']

    # Create base log directory path
    log_base = os.path.join(env_config['work_dir'], f"{isl}-{osl}")

    # Determine directory suffix based on attention_dp
    if gen_enable_attention_dp:
        dir_suffix = f"ctx{ctx_num}_gen{gen_num}_dep{gen_tp_size}_batch{gen_batch_size}_eplb{config['worker_config']['eplb_num_slots']}_mtp{mtp_size}"
    else:
        dir_suffix = f"ctx{ctx_num}_gen{gen_num}_tep{gen_tp_size}_batch{gen_batch_size}_eplb{config['worker_config']['eplb_num_slots']}_mtp{mtp_size}"

    # Create full log directory path
    log_dir = os.path.join(log_base, dir_suffix)
    os.makedirs(log_dir, exist_ok=True)

    # Setup config file paths and save worker configs
    ctx_config_path = os.path.join(log_dir, 'ctx_config.yaml')
    gen_config_path = os.path.join(log_dir, 'gen_config.yaml')
    save_worker_config(config, ctx_config_path, 'ctx')
    save_worker_config(config, gen_config_path, 'gen')

    # Prepare sbatch command
    cmd = [
        'sbatch',
        f'--partition={slurm_config["partition"]}',
        f'--gres=gpu:{hw_config["gpus_per_node"]}',
        f'--account={slurm_config["account"]}',
        f'--time={slurm_config["job_time"]}',
        f'--job-name={slurm_config["job_name"]}',
        f'--nodes={total_nodes}',
        f'--ntasks={total_tasks}',
        f'--ntasks-per-node={hw_config["gpus_per_node"]}',
        f'--segment={total_nodes}',
        slurm_config['script_file'],
        # Hardware configuration
        str(hw_config['gpus_per_node']),
        str(slurm_config['numa_bind']).lower(),
        str(ctx_nodes),  # Number of nodes needed for ctx workers
        str(gen_nodes),  # Number of nodes needed for gen workers
        str(ctx_tp_size),  # Tensor parallel size for ctx workers
        str(gen_tp_size),  # Tensor parallel size for gen workers

        # Worker configuration
        str(ctx_num),
        ctx_config_path,
        str(gen_num),
        gen_config_path,
        config['benchmark']['concurrency_list'],

        # Sequence and benchmark parameters
        str(config['sequence']['input_length']),
        str(config['sequence']['output_length']),
        str(config['benchmark']['multi_round']),
        str(config['benchmark']['benchmark_ratio']),
        str(config['benchmark']['streaming']).lower(),
        str(config['benchmark']['use_nv_sa_benchmark']).lower(),
        config['benchmark']['mode'],
        str(config['worker_config']['gen']['cache_transceiver_config']
            ['max_tokens_in_buffer']),

        # Environment and paths
        env_config['dataset_file'],
        env_config['model_path'],
        env_config['trtllm_repo'],
        env_config['work_dir'],
        log_dir,  # Pass the generated log directory
        env_config['container_mount'],
        env_config['container_image'],
        str(env_config['build_wheel']).lower(),

        # Profiling
        str(config['profiling']['nsys_on']).lower()
    ]

    # Submit the job
    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error submitting job: {e}", file=sys.stderr)
        sys.exit(1)


def main():
    args = parse_args()

    # Determine which mode to use
    if args.config:
        # Single config file mode
        config_files = [args.config]
    else:
        # Directory mode - find all YAML files
        yaml_pattern = os.path.join(args.dir, '*.yaml')
        config_files = sorted(glob.glob(yaml_pattern))

        if not config_files:
            print(f"No YAML files found in directory: {args.dir}",
                  file=sys.stderr)
            sys.exit(1)

        print(f"Found {len(config_files)} YAML file(s) in {args.dir}")

    # Process each config file
    for config_file in config_files:
        print(f"\nProcessing: {config_file}")
        try:
            config = load_config(config_file)
            submit_job(config)
            print(f"Successfully submitted job for: {config_file}")
        except Exception as e:
            print(f"Error processing {config_file}: {e}", file=sys.stderr)
            # Continue processing other files even if one fails
            continue


if __name__ == '__main__':
    main()