run.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import subprocess
import threading
import argparse
import sys
import os
import signal
from datetime import datetime
import time
import json
import glob
import shutil
# import matplotlib.pyplot as plt
# import numpy as np
try:
    from ruamel.yaml import YAML
    from ruamel.yaml.scalarstring import DoubleQuotedScalarString
    from ruamel.yaml.comments import CommentedSeq
except ImportError:
    print("[set-yaml] The ruamel.yaml library is required, please run first: pip install ruamel.yaml", file=sys.stderr)
    sys.exit(3)

# Import analysis functions
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from analysis import plot_exp_range_results
from analysis_cnp_samples import analyze_cnp_samples_files
from analysis_roce_stats import analyze_roce_stats_files
from analysis_detailed_stats import analyze_detailed_stats_file
from analysis_wqe_completion import analyze_wqe_completion_files

# --- Configuration area ---
# Directly list the host + bond combinations to run, each record independent
HOST_CONFIGS = [
    {
        "id": "33.255.69.193:bond0",
        "control_ip": "33.255.69.193",
        "data_ip": "33.255.69.193",
        "pcie_addr": "0001:06:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.193:bond1",
        "control_ip": "33.255.69.193",
        "data_ip": "26.248.37.98",
        "pcie_addr": "0000:2a:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.194:bond0",
        "control_ip": "33.255.69.194",
        "data_ip": "33.255.69.194",
        "pcie_addr": "0001:06:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.194:bond1",
        "control_ip": "33.255.69.194",
        "data_ip": "26.248.37.99",
        "pcie_addr": "0000:2a:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.195:bond0",
        "control_ip": "33.255.69.195",
        "data_ip": "33.255.69.195",
        "pcie_addr": "0001:06:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.195:bond1",
        "control_ip": "33.255.69.195",
        "data_ip": "26.248.41.98",
        "pcie_addr": "0000:2a:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.196:bond0",
        "control_ip": "33.255.69.196",
        "data_ip": "33.255.69.196",
        "pcie_addr": "0001:06:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.196:bond1",
        "control_ip": "33.255.69.196",
        "data_ip": "26.248.38.98",
        "pcie_addr": "0000:2a:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.197:bond0",
        "control_ip": "33.255.69.197",
        "data_ip": "33.255.69.197",
        "pcie_addr": "0001:06:00.0",
        "lcores": "2-6",
    },
    {
        "id": "33.255.69.197:bond1",
        "control_ip": "33.255.69.197",
        "data_ip": "26.248.38.99",
        "pcie_addr": "0000:2a:00.0",
        "lcores": "2-6",
    },
    # More host:bond combinations can be added here
]

def _deduplicate_preserve_order(seq):
    seen = set()
    deduped = []
    for item in seq:
        if item in seen:
            continue
        deduped.append(item)
        seen.add(item)
    return deduped

HOSTS = _deduplicate_preserve_order([cfg["control_ip"] for cfg in HOST_CONFIGS])
BASE_DIR = "/home/zhangzhaochen.zzc/anytest/"

# Repeat count configuration for the exp-range feature
EXP_RANGE_REPEAT_COUNT = 1

FORCE_QUIT = False  # Whether to force-quit the script, defaults to False

def build_host_run_entries(program_basename="PROGRAM", config_basename="CONFIG"):
    """
    Build the commands and associated metadata needed to run, based on the current HOST_CONFIGS.
    program_basename/config_basename are the file names after being copied to BASE_DIR.
    """
    entries = []
    program_path = os.path.join(BASE_DIR, program_basename)
    config_path = os.path.join(BASE_DIR, config_basename)
    libs_path = os.path.join(BASE_DIR, "libs")
    
    for host_cfg in HOST_CONFIGS:
        host_ip = host_cfg["data_ip"]
        lcores = host_cfg["lcores"]
        pcie_addr = host_cfg["pcie_addr"]
        file_prefix = host_cfg["id"]
        
        cmd = (
            f"sudo LD_LIBRARY_PATH={libs_path}:$LD_LIBRARY_PATH "
            f"{program_path} -n 8 -l {lcores} -a {pcie_addr},dv_flow_en=2 "
            f"--file-prefix={file_prefix} -- -c {config_path} -l {host_ip}"
        )
        entry = dict(host_cfg)
        entry["command"] = cmd
        entries.append(entry)
    
    return entries


def select_host_entries_from_config(config_data, host_entries):
    """
    Select the host entries on which to execute commands based on the servers.peer_ip field in the config file.
    If peer_ip is not configured, returns all host_entries by default.
    """
    servers_cfg = config_data.get('servers') if isinstance(config_data, dict) else {}
    peer_ips = servers_cfg.get('peer_ip') if isinstance(servers_cfg, dict) else None
    
    if not peer_ips:
        return host_entries
    
    if isinstance(peer_ips, str):
        peer_ip_list = [peer_ips]
    else:
        peer_ip_list = list(peer_ips)
    
    filtered = []
    missing_ips = []
    seen_ids = set()
    
    for ip in peer_ip_list:
        matches = [entry for entry in host_entries if entry["data_ip"] == ip]
        if not matches:
            missing_ips.append(ip)
            continue
        for entry in matches:
            if entry["id"] in seen_ids:
                continue
            filtered.append(entry)
            seen_ids.add(entry["id"])
    
    if missing_ips:
        print(f"[run] Warning: the following peer_ip did not match any HOST_CONFIGS: {', '.join(missing_ips)}", file=sys.stderr)

    if not filtered:
        print("[run] No host matched, falling back to all HOST_CONFIGS", file=sys.stderr)
        return host_entries

    return filtered


def augment_config_with_peer_ips(config_data):
    """
    If servers.peer_id is provided in the config, automatically populate the peer_ip list based on HOST_CONFIGS.
    """
    if not isinstance(config_data, dict):
        return False
    servers_cfg = config_data.get('servers')
    if not isinstance(servers_cfg, dict):
        return False
    peer_ids = servers_cfg.get('peer_id')
    if not peer_ids:
        return False
    if isinstance(peer_ids, str):
        peer_id_list = [peer_ids]
    else:
        peer_id_list = list(peer_ids)
    
    resolved_ips = []
    missing_ids = []
    for peer_id in peer_id_list:
        match = next((cfg for cfg in HOST_CONFIGS if cfg["id"] == peer_id), None)
        if not match:
            missing_ids.append(peer_id)
            continue
        resolved_ips.append(DoubleQuotedScalarString(str(match["data_ip"])))
    
    if missing_ids:
        print(f"[run] Warning: the following peer_id did not match any HOST_CONFIGS: {', '.join(missing_ids)}", file=sys.stderr)
    if not resolved_ips:
        return False
    
    dedup_ips = _deduplicate_preserve_order(resolved_ips)
    if not dedup_ips:
        return False

    peer_ip_seq = CommentedSeq(dedup_ips)
    try:
        peer_ip_seq.fa.set_flow_style()
    except AttributeError:
        # Older versions of ruamel may not have .fa, fall back to a plain list but still keep the bracket format
        peer_ip_seq = dedup_ips

    if 'peer_ip' in servers_cfg:
        del servers_cfg['peer_ip']

    insert_pos = None
    if hasattr(servers_cfg, 'insert'):
        print(f"[run] Using the insert method to insert peer_ip")
        try:
            peer_id_index = list(servers_cfg).index('peer_id')
            insert_pos = peer_id_index + 1
        except ValueError:
            insert_pos = len(servers_cfg)
        servers_cfg.insert(insert_pos, 'peer_ip', peer_ip_seq)
    else:
        print(f"[run] Using the dict approach to insert peer_ip")
        servers_cfg['peer_ip'] = peer_ip_seq

    print(f"[run] Generated peer_ip list from peer_id: {dedup_ips}")
    return True

# --- Global variables, used to track and manage child processes ---
# List used to store all active Popen objects
ACTIVE_PROCESSES = []
# Used to protect concurrent access to the ACTIVE_PROCESSES list
PROCESS_LOCK = threading.Lock()


# --- Core script functionality ---
def kill_all_child_processes():
    """
    Kill all child processes, but do not exit the main script.
    """
    print("\n\n[kill-all-child-processes] Killing all child processes...", file=sys.stderr)
    with PROCESS_LOCK:
        for process_info in list(ACTIVE_PROCESSES):
            proc = process_info['proc']
            host = process_info['host']
            command = process_info['command']
            try:
                print(f"[kill-all-child-processes]  -> Remotely terminating the process on {host}: {command}", file=sys.stderr)
                kill_cmd = f"sudo pkill -f \"{command}\""
                subprocess.Popen(["ssh", host, kill_cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            except ProcessLookupError:
                pass # Process no longer exists
            except Exception as e:
                print(f"[kill-all-child-processes]  -> Error terminating process {proc.pid} (host: {host}): {e}", file=sys.stderr)
    print("[kill-all-child-processes] Termination signal sent to all child processes.", file=sys.stderr)


def cleanup_processes(signum, frame):
    """
    Signal handler function. Called when SIGINT (Ctrl+C) is received.
    Terminates processes in different ways depending on whether the host is local or remote.
    """
    print("\n\n[MAIN] Detected Ctrl+C! Terminating all child processes...", file=sys.stderr)

    FORCE_QUIT = True  # Set the global variable, indicating the script needs to force-quit
    kill_all_child_processes()
    # print("[MAIN] Termination signal sent to all child processes. Exiting.", file=sys.stderr)
    # sys.exit(130)

def clean_line(line):
    """Clean control characters from an output line"""
    # Remove carriage returns
    line = line.replace('\r\n', '\n').replace('\r', '\n')
    # Remove other common control characters
    import re
    line = re.sub(r'\x1b\[[0-9;]*[mGKH]', '', line)  # Remove ANSI color codes, etc.
    return line

def execute_command(host, command, log_dir):
    """
    Execute a command on the specified host and capture output in real time.
    host: 'localhost' or a remote hostname 'user@hostname'
    command: the command string to execute
    log_dir: the directory to store logs
    """
    # Sanitize the hostname so it can be used as a file name
    sanitized_host = host.replace('@', '_').replace('.', '_')
    log_file_path = os.path.join(log_dir, f"{sanitized_host}.log")

    print(f"--- Starting execution on {host} ---")

    command_to_exec = "exec " + command

    # Construct the command
    if host == "localhost":
        # Execute directly on the local machine
        cmd_list = ["/bin/bash", "-c", command_to_exec]
    else:
        # Execute via ssh on the remote machine
        # Wrap the command in single quotes to prevent the local shell from interpreting special characters in the command
        # cmd_list = ["ssh", "-t", host, command_to_exec]
        cmd_list = ["ssh", host, command_to_exec]

    try:
        # Use Popen to start the child process so we can read output in real time
        # stderr=subprocess.STDOUT: merge stderr into stdout, so reading only stdout below
        #   captures the child process's stderr as well (the program emits a lot of errors via
        #   fprintf(stderr,...); previously stderr=PIPE was used but never read, so real errors
        #   were discarded and not visible in the logs).
        # bufsize=1: line buffering, ensures we can read line by line
        # Popen call section (removed text=True and encoding)
        process = subprocess.Popen(
            cmd_list,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            bufsize=1
        )
        
        # Pack the process info into a dict and store it in the global list
        process_info = {
            'proc': process,
            'host': host,
            'command': command  # Store the original command, used for later pkill
        }
        with PROCESS_LOCK:
            ACTIVE_PROCESSES.append(process_info)

        # Read loop section (manual decoding)
        with open(log_file_path, 'w') as log_file:
            # Read the child process's byte stream line by line and decode manually
            for line_bytes in iter(process.stdout.readline, b''):
                line = line_bytes.decode('utf-8', errors='ignore')

                line = clean_line(line)
                # Skip the line if it is empty
                if not line.strip():
                    continue
                # Print to terminal in real time, with the hostname prefix
                sys.stdout.write(f"[{host}] {line}")

                # Write to the log file
                log_file.write(line)

        # Wait for the process to finish and get the return code
        process.wait()
        if process.returncode == 0:
            print(f"--- {host} executed successfully ---")
        else:
            print(f"--- !!! {host} execution failed, return code: {process.returncode} !!! ---")

    except FileNotFoundError:
        # If the ssh command does not exist
        error_msg = f"Error: command 'ssh' not found. Please make sure it is installed and in your PATH."
        print(error_msg, file=sys.stderr)
        with open(log_file_path, 'w') as log_file:
            log_file.write(error_msg)
    except Exception as e:
        error_msg = f"An unknown error occurred while executing on {host}: {e}"
        print(error_msg, file=sys.stderr)
        with open(log_file_path, 'w') as log_file:
            log_file.write(error_msg)
    finally:
        # *** When the process ends (whether success, failure, or exception), remove it from the global list ***
        if process_info:
            with PROCESS_LOCK:
                if process_info in ACTIVE_PROCESSES:
                    ACTIVE_PROCESSES.remove(process_info)


def distribute_files(source_path, dest_path):
    """
    Distribute a file or directory to all remote hosts using scp.
    """
    print(f"Starting to distribute '{source_path}' to all remote hosts...")
    for host in HOSTS:
        if host == "localhost":
            continue  # Skip the local machine, since it doesn't need to copy via scp
        print(f"--> Copying to {host}:{dest_path}")
        # Use the -r option to support recursive directory copying
        # command = ["rsync", "-avz", "--progress", source_path, f"{host}:{dest_path}"]
        # command = ["scp", source_path, f"{host}:{dest_path}"]
        command = ["rsync", "-avz", "-e", "ssh", source_path, f"{host}:{dest_path}"]
        # tar czf - /local/directory | ssh user@remote_host "tar xzf - -C /remote/path"
        # command = ["tar", "czf", "-", source_path, "|", "ssh", host, '"tar', "xzf", "-", "-C", f'{dest_path[:-8]}"']
        # command = f"tar czf - {source_path} | ssh {host} 'tar xzf - -C {dest_path[:-8]}'"

        # Print the command for debugging
        # print(f"Executing command: {command}")

        try:
            # subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            print(f"    Successfully copied to {host}")
        except subprocess.CalledProcessError as e:
            print(f"    !!! Copy to {host} failed !!!", file=sys.stderr)
            print(f"    Error message: {e.stderr.decode('utf-8', errors='ignore')}", file=sys.stderr)

            print("File distribution complete.")

def collect_experiment_results(log_dir, file_name_temp="cnp_samples_*.csv"):
    """
    Collect experiment result files from all hosts into the specified log directory.
    Copy files matching the given template from each host's /tmp/ directory into the local log directory,
    rename them to "machineIP_originalfilename", then delete the remote original files.
    After collection, analyze all CSV files and aggregate the results into a result.json file.

    log_dir: local log directory path
    file_name_temp: file name template to collect, defaults to "cnp_samples_*.csv"
    """
    print(f"[collect-results] Starting to collect experiment result files into {log_dir}")
    print(f"[collect-results] File template: {file_name_temp}")

    collected_files = []  # Record the list of successfully collected files

    # Remote host processing
    for host in HOSTS:
        try:
            # First look for files on the remote host
            list_cmd = ["ssh", host, f"ls /tmp/{file_name_temp} 2>/dev/null || true"]
            result = subprocess.run(list_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)

            if result.stdout.strip():
                remote_files = result.stdout.strip().split('\n')

                for remote_file in remote_files:
                    if remote_file:  # Make sure the file name is not empty
                        filename = os.path.basename(remote_file)
                        new_filename = f"{host}_{filename}"
                        dest_path = os.path.join(log_dir, new_filename)

                        # Use scp to copy the file to the local machine
                        # ssh_cmd = ["ssh", "h194", "echo 'Establishing connection...'"]
                        # subprocess.run(ssh_cmd, check=True)
                        scp_cmd = ["scp", "-o", "ControlPath=none", f"{host}:{remote_file}", dest_path]
                        # scp_cmd = ["rsync", "-az", "-e", "ssh",  f"{host}:{remote_file}", dest_path]
                        print(f"[collect-results] {host}: {filename} -> {new_filename}")
                        subprocess.run(scp_cmd, check=True)
                        collected_files.append(dest_path)  # Record the successfully collected file
                        # success = transfer_file_with_dd_chunks(remote_file, dest_path, source_host=host, dest_host=None)
                        # transfer_file_with_http_chunks(remote_file, dest_path, source_host=host, dest_host=None)

                # Delete the original files on the remote host
                cleanup_cmd = ["ssh", host, f"sudo rm -f /tmp/{file_name_temp}"]
                subprocess.run(cleanup_cmd, check=True)
                print(f"[collect-results] Cleaned up the original files on {host}")
            else:
                print(f"[collect-results] No {file_name_temp} files found on {host}")

        except subprocess.CalledProcessError as e:
            print(f"[collect-results] Error while processing {host}: {e}", file=sys.stderr)
        except Exception as e:
            print(f"[collect-results] An unknown error occurred while processing {host}: {e}", file=sys.stderr)

    print("[collect-results] Experiment result file collection complete")

    if not collected_files:
        print("[collect-results] No CSV files collected, skipping the analysis step")
        return

    # Classify by type, supporting mixed-template scenarios
    cnp_files = [f for f in collected_files if "cnp_samples_" in os.path.basename(f)]
    roce_files = [f for f in collected_files if "roce_stats_" in os.path.basename(f)]
    tx_files = [f for f in collected_files if "tx_detailed_stats_" in os.path.basename(f)]
    rx_files = [f for f in collected_files if "rx_detailed_stats_" in os.path.basename(f)]
    wqe_files = [f for f in collected_files if "wqe_completion_latency_" in os.path.basename(f)]

    analyzed_any = False
    if cnp_files:
        analyze_cnp_samples_files(log_dir, cnp_files)
        analyzed_any = True
    if roce_files:
        analyze_roce_stats_files(log_dir, roce_files, True)
        analyzed_any = True
    if rx_files and tx_files:
        analyze_detailed_stats_file(log_dir, rx_files, tx_files)
        analyzed_any = True
    if wqe_files:
        # Directly use the default parameters defined in analysis_wqe_completion.py
        analyze_wqe_completion_files(log_dir, wqe_files)
        analyzed_any = True

    if not analyzed_any:
        print(f"[collect-results] No supported CSV type found, skipping the analysis step")


def execute_launch_commands(launch_config, config_type):
    """
    Execute the commands in the pre_launch or post_launch configuration.
    launch_config: the pre_launch or post_launch configuration dict
    config_type: "pre_launch" or "post_launch", used for log identification
    """
    if not launch_config or not launch_config.get('enable', False):
        print(f"[{config_type}] Not enabled, skipping")
        return

    host = launch_config.get('host')
    commands = launch_config.get(f'{config_type}_cmd', [])

    if not host:
        print(f"[{config_type}] No host specified, skipping", file=sys.stderr)
        return

    if not commands:
        print(f"[{config_type}] No commands specified, skipping")
        return

    print(f"[{config_type}] Executing {len(commands)} commands on host {host}")

    for i, cmd in enumerate(commands):
        print(f"[{config_type}] Executing command {i+1}/{len(commands)}: {cmd}")
        try:
            if host == "localhost" or host == "127.0.0.1":
                # Execute locally
                result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
            else:
                # Execute on the remote host
                ssh_cmd = ["ssh", host, cmd]
                result = subprocess.run(ssh_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)

            if result.returncode == 0:
                print(f"[{config_type}] Command executed successfully")
                if result.stdout.strip():
                    print(f"[{config_type}] Output: {result.stdout.strip()}")
            else:
                print(f"[{config_type}] Command execution failed, return code: {result.returncode}", file=sys.stderr)
                if result.stderr.strip():
                    print(f"[{config_type}] Error: {result.stderr.strip()}", file=sys.stderr)

        except Exception as e:
            print(f"[{config_type}] An exception occurred while executing the command: {e}", file=sys.stderr)

    print(f"[{config_type}] All commands executed")


def run_main_command(program_path=None, config_path=None, wait=True, log_dir=None):
    """
    Encapsulate the main logic of the run subcommand, taking program_path and config_path arguments directly.
    Execute commands in parallel on the local machine and all remote hosts.
    wait: whether to wait for all commands to finish executing
    log_dir: log directory path, automatically generated if None
    """
    program_path = program_path if program_path else "./anytest"
    config_path = config_path if config_path else "./cnp4.yaml"
    program_basename = os.path.basename(program_path)
    config_basename = os.path.basename(config_path)

    yaml = YAML()
    yaml.preserve_quotes = True
    yaml.indent(mapping=2, sequence=4, offset=2)
    config_data = {}
    file_name_template = "cnp_samples_*.csv"

    try:
        with open(config_path, 'r') as f:
            config_data = yaml.load(f)
        if augment_config_with_peer_ips(config_data):
            with open(config_path, 'w') as f:
                yaml.dump(config_data, f)
        exp_mode = config_data.get('experiment', {}).get('exp_mode', 'RECEIVER')
        if exp_mode == 'SENDER':
            # In sender mode, match roce_stats_*.csv, rx_detailed_stats_*, tx_detailed_stats_*, wqe_completion_latency_*.csv all at once
            # Use brace expansion (interpreted by the remote shell)
            file_name_template = "{roce_stats_*.csv,rx_detailed_stats_*.csv,tx_detailed_stats_*.csv,wqe_completion_latency_*.csv}"
        elif exp_mode == 'RECEIVER':
            file_name_template = "cnp_samples_*.csv"
        elif exp_mode == 'ROCE':
            file_name_template = "wqe_completion_latency_*.csv"
        print(f"[run] Detected experiment mode: {exp_mode}")
        print(f"[run] Will collect file template: {file_name_template}")
    except Exception as e:
        print(f"[run] Failed to read config file, using default template: {e}", file=sys.stderr)
        config_data = {}

    dest_dir = BASE_DIR
    distribute_files(config_path, dest_dir)
    host_run_entries = build_host_run_entries(program_basename, config_basename)
    time.sleep(1)
    
    # If no log_dir is specified, create a timestamped directory to store the logs for this run
    if log_dir is None:
        log_dir = f"logs/{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    os.makedirs(log_dir, exist_ok=True)
    print(f"All output logs will be saved in directory: '{log_dir}'")

    # Execute pre_launch commands
    pre_launch_config = config_data.get('pre_launch', {})
    execute_launch_commands(pre_launch_config, 'pre_launch')

    threads = []
    target_entries = select_host_entries_from_config(config_data, host_run_entries)
    print(f"[run] Will execute commands on {len(target_entries)} hosts")

    # Create a thread for each host entry
    for entry in target_entries:
        host = entry["control_ip"]
        command = entry["command"]
        print(f"Running command on {host} ({entry['id']}): {command}")
        thread = threading.Thread(target=execute_command, args=(host, command, log_dir))
        threads.append(thread)
        thread.start()

    # Copy the config file to the log directory
    config_filename = os.path.basename(config_path)
    config_dest_path = os.path.join(log_dir, config_filename)
    try:
        shutil.copy2(config_path, config_dest_path)
        print(f"[run] Copied config file to: {config_dest_path}")
    except Exception as e:
        print(f"[run] Failed to copy config file: {e}", file=sys.stderr)

    if wait:
        # Wait for all threads to finish executing
        for thread in threads:
            thread.join()
        print("\nAll tasks completed.")

        # Execute post_launch commands
        post_launch_config = config_data.get('post_launch', {})
        execute_launch_commands(post_launch_config, 'post_launch')

        # Collect experiment result files
        # wait == True means this is a single run
        collect_experiment_results(log_dir, file_name_template)

    return threads, log_dir, file_name_template


def generate_ibtest_commands():
    """
    Generate all the commands needed for ibtest along with their host and port info.
    Returns: List[dict], each dict containing host, command, port, role
    """
    base_port = 8259
    # num_servers = 32
    # num_clients_per_host = 16
    server_host = "h200"
    client_hosts = ["h194", "h195"]
    num_servers = len(client_hosts) 
    ib_path = "/usr/bin/ib_write_bw"
    server_cmds = []
    client_cmds = []

    # Server-side commands
    for i in range(num_servers):
        port = base_port + i
        cmd = f"{ib_path} -d mlx5_bond_1 -p {port}"
        server_cmds.append({
            'host': server_host,
            'command': cmd,
            'port': port,
            'role': 'server',
        })

    # Client-side commands
    for idx, chost in enumerate(client_hosts):
        # for i in range(num_clients_per_host):
        # port = base_port + idx * num_clients_per_host + i
        port = base_port + idx
        # The target is always h194
        cmd = f"sleep 1; {ib_path} -d mlx5_bond_1 -D 30 -p {port} 33.255.69.200 -q 8"
        client_cmds.append({
            'host': chost,
            'command': cmd,
            'port': port,
            'role': 'client',
        })

    return server_cmds + client_cmds


def run_ibtest_commands():
    """
    Execute all ibtest commands concurrently; log names include the host and port.
    """
    from datetime import datetime
    import threading
    import os

    commands = generate_ibtest_commands()
    log_dir = f"logs/ibtest_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    os.makedirs(log_dir, exist_ok=True)
    print(f"All ibtest output logs will be saved in directory: '{log_dir}'")

    threads = []
    for info in commands:
        host = info['host']
        command = info['command']
        port = info['port']
        role = info['role']
        # Log name includes host, port, and role
        sanitized_host = host.replace('@', '_').replace('.', '_')
        log_file_path = os.path.join(log_dir, f"{sanitized_host}_{role}_{port}.log")
        # Wrap execute_command to pass a custom log path
        def run_cmd(host=host, command=command, log_file_path=log_file_path):
            # command_to_exec = "exec " + command
            command_to_exec = command
            if host == "localhost":
                cmd_list = ["/bin/bash", "-c", command_to_exec]
            else:
                cmd_list = ["ssh", host, command_to_exec]
            try:
                print(f"Executing command on {host}:{port}: {cmd_list}")
                process = subprocess.Popen(
                    cmd_list,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    bufsize=1
                )
                process_info = {
                    'proc': process,
                    'host': host,
                    'command': command,
                    'port': port,
                    'role': role,
                }
                with PROCESS_LOCK:
                    ACTIVE_PROCESSES.append(process_info)
                import threading
                def stream_output(stream, prefix, log_file, is_stderr=False):
                    for line_bytes in iter(stream.readline, b''):
                        line = line_bytes.decode('utf-8', errors='ignore')
                        line = clean_line(line)
                        if not line.strip():
                            continue
                        tag = f"[{host}:{port}]"
                        if is_stderr:
                            tag += "[STDERR] "
                        else:
                            tag += " "
                        sys.stdout.write(f"{tag}{line}")
                        log_file.write(line)
                with open(log_file_path, 'w') as log_file:
                    t_out = threading.Thread(target=stream_output, args=(process.stdout, 'STDOUT', log_file, False))
                    t_err = threading.Thread(target=stream_output, args=(process.stderr, 'STDERR', log_file, True))
                    t_out.start()
                    t_err.start()
                    t_out.join()
                    t_err.join()
                process.wait()
                if process.returncode == 0:
                    print(f"--- {host}:{port} executed successfully ---")
                else:
                    print(f"--- !!! {host}:{port} execution failed, return code: {process.returncode} !!! ---")
            except Exception as e:
                error_msg = f"An unknown error occurred while executing on {host}:{port}: {e}"
                print(error_msg, file=sys.stderr)
                with open(log_file_path, 'w') as log_file:
                    log_file.write(error_msg)
            finally:
                if process_info:
                    with PROCESS_LOCK:
                        if process_info in ACTIVE_PROCESSES:
                            ACTIVE_PROCESSES.remove(process_info)
        thread = threading.Thread(target=run_cmd)
        threads.append(thread)
        thread.start()
    for thread in threads:
        thread.join()
    print("\nAll ibtest tasks completed.")


def pingmesh():
    """
    Batch-ping target IPs on all hosts to refresh the arp cache.
    """
    ping_ips = HOSTS[1:] # Exclude the local machine
    for ip in ping_ips:
        cmd = ["timeout", "2", "clush", "-g", "allt", "-b", f"ping {ip}"]
        print(f"\n[pingmesh] Running: {' '.join(cmd)}")
        try:
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if result.stdout:
                print(result.stdout.decode('utf-8', errors='ignore'))
            if result.stderr:
                print(result.stderr.decode('utf-8', errors='ignore'), file=sys.stderr)
        except Exception as e:
            print(f"[pingmesh] Failed to execute command: {e}", file=sys.stderr)

def killall():
    """
    Kill all anytest processes on the local machine and all hosts, but exclude the Python script itself.
    """
    # Remote kill
    remote_cmd = ["clush", "-g", "allt", "-x", "h193", "-b", "sudo pkill -f anytest"]
    print(f"[killall] Running: {' '.join(remote_cmd)}")
    try:
        result = subprocess.run(remote_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.stdout:
            print(result.stdout.decode('utf-8', errors='ignore'))
        if result.stderr:
            print(result.stderr.decode('utf-8', errors='ignore'), file=sys.stderr)
    except Exception as e:
        print(f"[killall] Remote kill failed: {e}", file=sys.stderr)

    # Local kill - use more precise matching, excluding python processes
    print(f"[killall] Running local kill command...")
    try:
        # First get the PID of the current python process
        current_pid = os.getpid()
        # Use the ps command for precise matching, excluding the python process and the current script
        local_cmd = ["sudo", "bash", "-c", f"ps aux | grep '[c]np_test' | grep -v python | grep -v run.py | grep -v {current_pid} | awk '{{print $2}}' | xargs -r kill"]
        result = subprocess.run(local_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.stdout:
            print(result.stdout.decode('utf-8', errors='ignore'))
        if result.stderr:
            print(result.stderr.decode('utf-8', errors='ignore'), file=sys.stderr)
    except Exception as e:
        print(f"[killall] Local kill failed: {e}", file=sys.stderr)

def set_cnp_interval(ip, cnp_interval):
    """
    Set cnp_interval on the specified machine.
    """
    set_cmd = f"sudo bash -c 'for f in /sys/class/net/eth*/ecn/roce_np/min_time_between_cnps; do echo {cnp_interval} > \$f; done'"
    if ip == "localhost" or ip == "127.0.0.1":
        print(f"[set-cnp-interval] Executing locally: {set_cmd}")
        try:
            result = subprocess.run(set_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if result.stdout:
                print(result.stdout.decode('utf-8', errors='ignore'))
            if result.stderr:
                print(result.stderr.decode('utf-8', errors='ignore'), file=sys.stderr)
        except Exception as e:
            print(f"[set-cnp-interval] Local execution failed: {e}", file=sys.stderr)
    else:
        ssh_cmd = f'ssh {ip} "{set_cmd}"'
        print(f"[set-cnp-interval] Executing on remote {ip}: {ssh_cmd}")
        try:
            result = subprocess.run(ssh_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if result.stdout:
                print(result.stdout.decode('utf-8', errors='ignore'))
            if result.stderr:
                print(result.stderr.decode('utf-8', errors='ignore'), file=sys.stderr)
        except Exception as e:
            print(f"[set-cnp-interval] Remote execution failed: {e}", file=sys.stderr)


def init_hosts():
    """
    Execute the following operations in sequence on HOSTS, printing the return values:
    1) Set RoCE NP traffic_class=136
    2) Set the number of 2MB hugepages to 8192
    3) Start the nusad service
    """
    commands = [
        "sudo bash -c 'for f in /sys/class/infiniband/*/tc/1/traffic_class; do echo 136 > $f; done'",
        "sudo bash -c 'echo 8192 > /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages'",
        "sudo systemctl start nusad",
        "sudo ip link set dev bond0 mtu 9000",
        "sudo ip link set dev bond1 mtu 9000",
        # Stop nusad after waiting 10s
        "sleep 10; sudo systemctl stop nusad",
        # Fully detach hpcc_doca from the current SSH/stdout, otherwise subprocess.run will keep waiting for it to exit
        "sudo nohup /usr/local/bin/hpcc_doca -d mlx5_bond_0 >/dev/null 2>&1 </dev/null &",
        "sudo nohup /usr/local/bin/hpcc_doca -d mlx5_bond_1 >/dev/null 2>&1 </dev/null &"
    ]

    print("[init] Starting initialization operations on all HOSTS (in parallel)")
    threads = []

    def run_on_host(host):
        print(f"[init] Host {host}: starting to execute {len(commands)} commands")
        for i, cmd in enumerate(commands, start=1):
            print(f"[init] Host {host} command {i}/{len(commands)}: {cmd}")
            try:
                if host == "localhost" or host == "127.0.0.1":
                    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
                else:
                    ssh_cmd = ["ssh", host, cmd]
                    result = subprocess.run(ssh_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)

                print(f"[init] Host {host} return code: {result.returncode}")
                if result.stdout and result.stdout.strip():
                    print(f"[init] Host {host} output: {result.stdout.strip()}")
                if result.stderr and result.stderr.strip():
                    print(f"[init] Host {host} error: {result.stderr.strip()}", file=sys.stderr)
            except Exception as e:
                print(f"[init] Host {host} encountered an exception while executing the command: {e}", file=sys.stderr)
        print(f"[init] Host {host}: initialization commands executed")

    for host in HOSTS:
        t = threading.Thread(target=run_on_host, args=(host,))
        threads.append(t)
        t.start()

    for t in threads:
        t.join()
    print("[init] All hosts initialized")

def set_yaml_value(file_path, key, value):
    """
    Modify the specified key-value pair in a YAML file, supporting key paths of arbitrary depth and list indices.
    Key path format: level1.level2.list_index.key
    For example: cnp_state_machine.states.0.duration_us
    List indices are represented as numbers; other keys use strings.

    Additional support: when the final target is a string and the last-level key is a number, it means
    replacing by index within the list of numbers extracted from that string.
    For example: pre_launch.pre_launch_cmd.3.1 means in the 3rd command string, replace the 1st extracted number (0-based) with the new value.
    """
    yaml = YAML()
    yaml.preserve_quotes = True
    yaml.indent(mapping=2, sequence=4, offset=2)
    yaml.default_flow_style = False

    try:
        with open(file_path, 'r') as f:
            data = yaml.load(f)
    except Exception as e:
        print(f"[set-yaml] Failed to read file: {e}", file=sys.stderr)
        sys.exit(1)

    # Helper function: replace the index-th number in the string with new_num_str (numbers in order of appearance, 0-based)
    def replace_nth_number_in_string(s, index, new_num_str):
        import re
        # Match standalone integers (with an optional sign)
        pattern = re.compile(r'(?<!\d)[+-]?\d+(?!\d)')
        matches = list(pattern.finditer(s))
        if index < 0 or index >= len(matches):
            return None  # Indicates index out of range
        parts = []
        last_end = 0
        for i, m in enumerate(matches):
            parts.append(s[last_end:m.start()])
            if i == index:
                parts.append(str(new_num_str))
            else:
                parts.append(m.group(0))
            last_end = m.end()
        parts.append(s[last_end:])
        return ''.join(parts)

    # Parse the key path
    keys = key.split('.')
    current = data
    parent = None
    parent_key = None

    # Traverse to the parent node of the last level
    for i, k in enumerate(keys[:-1]):
        # Check whether it is a numeric index (list access)
        if k.isdigit():
            k = int(k)
            if not isinstance(current, list):
                print(f"[set-yaml] Path error: expected a list but got {type(current).__name__}, path: {'.'.join(keys[:i])}", file=sys.stderr)
                sys.exit(2)
            if k >= len(current):
                print(f"[set-yaml] List index out of range: {k}, list length: {len(current)}, path: {'.'.join(keys[:i+1])}", file=sys.stderr)
                sys.exit(2)
            parent = current
            parent_key = k
            current = current[k]
        else:
            if not isinstance(current, dict):
                print(f"[set-yaml] Path error: expected a dict but got {type(current).__name__}, path: {'.'.join(keys[:i])}", file=sys.stderr)
                sys.exit(2)
            if k not in current:
                print(f"[set-yaml] Key does not exist: {'.'.join(keys[:i+1])}", file=sys.stderr)
                sys.exit(2)
            parent = current
            parent_key = k
            current = current[k]

    # Handle the last-level key
    final_key = keys[-1]

    # Check whether the last level is a numeric index
    if final_key.isdigit():
        final_key = int(final_key)
        # Case A: parent node is a list => regular list-index assignment
        if isinstance(current, list):
            if final_key >= len(current):
                print(f"[set-yaml] List index out of range: {final_key}, list length: {len(current)}, path: {key}", file=sys.stderr)
                sys.exit(2)
            old_value = current[final_key]
            old_type = type(old_value)
        # Case B: parent node is a string => replace numbers extracted from the string by index
        elif isinstance(current, str):
            # On the string current, replace the final_key-th number with value (used as a string)
            new_str = replace_nth_number_in_string(current, final_key, value)
            if new_str is None:
                print(f"[set-yaml] String number index out of range: {final_key}, the string contains only a limited number of numbers, path: {'.'.join(keys[:-1])}", file=sys.stderr)
                sys.exit(2)
            # Write the new string back to its parent container
            if parent is None:
                print(f"[set-yaml] Internal error: parent container not found, cannot write back the string", file=sys.stderr)
                sys.exit(1)
            try:
                if isinstance(parent, list) and isinstance(parent_key, int):
                    parent[parent_key] = new_str
                elif isinstance(parent, dict) and isinstance(parent_key, str):
                    parent[parent_key] = new_str
                else:
                    print(f"[set-yaml] Internal error: parent container type mismatch, cannot write back the string", file=sys.stderr)
                    sys.exit(1)
            except Exception as e:
                print(f"[set-yaml] Failed to write back the string: {e}", file=sys.stderr)
                sys.exit(1)
            # Write back to the file and finish
            try:
                with open(file_path, 'w') as f:
                    yaml.dump(data, f)
                print(f"[set-yaml] Modified {file_path}: {'.'.join(keys[:-1])} (string number #{final_key}) = {value}")
            except Exception as e:
                print(f"[set-yaml] Failed to write file: {e}", file=sys.stderr)
                sys.exit(1)
            return
        else:
            print(f"[set-yaml] Path error: expected a list or string but got {type(current).__name__}, path: {'.'.join(keys[:-1])}", file=sys.stderr)
            sys.exit(2)
    else:
        if not isinstance(current, dict):
            print(f"[set-yaml] Path error: expected a dict but got {type(current).__name__}, path: {'.'.join(keys[:-1])}", file=sys.stderr)
            sys.exit(2)
        if final_key not in current:
            print(f"[set-yaml] Key does not exist: {key}", file=sys.stderr)
            sys.exit(2)

        # Get the type of the original value and convert the new value
        old_value = current[final_key]
        old_type = type(old_value)

    # Type conversion
    try:
        if old_type is int:
            value = int(value)
        elif old_type is float:
            value = float(value)
        elif old_type is bool:
            value = value.lower() in ('true', '1', 'yes', 'on')
        elif old_type is list:
            import ast
            value = ast.literal_eval(value)
        elif old_type is dict:
            import ast
            value = ast.literal_eval(value)
        # String type stays as-is, no conversion needed
    except Exception as e:
        print(f"[set-yaml] Type conversion failed: {e}, will use string type", file=sys.stderr)
        # If conversion fails, keep it as a string

    # Set the new value
    current[final_key] = value

    # Write back to the file
    try:
        with open(file_path, 'w') as f:
            yaml.dump(data, f)
        print(f"[set-yaml] Modified {file_path}: {key} = {value}")
    except Exception as e:
        print(f"[set-yaml] Failed to write file: {e}", file=sys.stderr)
        sys.exit(1)


def _iter_exp_range_values(start_value: int, stop_value: int, step: int, sweep_mode: str):
    """
    Generate the parameter sequence for exp-range.
    - sweep_mode="linear": arithmetic progression: start, start+step, ... <= stop
    - sweep_mode="multiply": exponential/multiplicative progression: start, start*step, start*step^2, ... <= stop
    """
    if sweep_mode not in ("linear", "multiply"):
        raise ValueError(f"invalid sweep_mode: {sweep_mode}")

    if sweep_mode == "linear":
        if step == 0:
            raise ValueError("step cannot be 0 (linear mode)")
        if start_value <= stop_value:
            for v in range(start_value, stop_value + 1, step):
                yield v
        else:
            # Allow a reverse sweep (step must be negative)
            for v in range(start_value, stop_value - 1, step):
                yield v
        return

    # multiply mode
    if step <= 1:
        raise ValueError("in multiply mode step must be > 1, otherwise it loops forever or is meaningless")
    if start_value < 0 or stop_value < 0:
        raise ValueError("multiply mode only supports non-negative integer ranges")
    if start_value == 0:
        # 0 multiplied by anything is still 0, avoid an infinite loop
        yield 0
        return
    if start_value > stop_value:
        # Multiplicative progression cannot "converge" from a larger value down to a smaller stop, return an empty sequence directly
        return

    v = start_value
    while v <= stop_value:
        yield v
        v *= step


def exp_range_main(program, file, key, start_value, stop_value, step, log_dir=None, data_path=None, sweep_mode="linear"):
    """
    Batch-modify the yaml over a range and run the experiment multiple times, killing the process after 30 seconds each time.
    The type of files collected is determined by experiment.exp_mode in the config file:
    - RECEIVER: collect cnp_samples_*.csv
    - SENDER: collect roce_stats_*.csv
    log_dir: log directory path, automatically generated if None
    """
    start_value = int(start_value)
    stop_value = int(stop_value)
    step = int(step)

    # If no log_dir is specified, create a timestamped directory
    if log_dir is None:
        # Get the config file name (without extension), drop the first dir segment (usually config) and keep the remaining dir segments
        config_basename = os.sep.join(os.path.splitext(file)[0].split(os.sep)[1:])
        # config_basename = os.path.splitext(os.path.basename(file))[0]
        # Get the part of the key after the dot
        key_suffix = key.split('.')[-1] if '.' in key else key
        # Generate the range string
        range_str = f"{start_value}-{stop_value}"
        # Generate the timestamp
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        # Combine into the new log directory name
        log_dir = f"logs/{config_basename}-{key_suffix}-{range_str}-{timestamp}"

    print(f"[exp-range] Using log directory: {log_dir}")

    try:
        values = list(_iter_exp_range_values(start_value, stop_value, step, sweep_mode))
    except Exception as e:
        print(f"[exp-range] Failed to generate parameter sequence: {e}", file=sys.stderr)
        sys.exit(2)

    if not values:
        print(f"[exp-range] No parameter values generated (start={start_value}, stop={stop_value}, step={step}, mode={sweep_mode}), exiting", file=sys.stderr)
        sys.exit(2)

    print(f"[exp-range] sweep_mode={sweep_mode}, will run parameter sequence: {values}")

    for value in values:
        print(f"[exp-range] Setting {key} = {value} and starting the experiment")
        set_yaml_value(file, key, value)
        for i in range(EXP_RANGE_REPEAT_COUNT):
            print(f"[exp-range] Experiment {i+1}/{EXP_RANGE_REPEAT_COUNT}, value: {value}")
            # Create a subdirectory for each experiment
            experiment_log_dir = os.path.join(log_dir, f"value_{value}_run_{i+1}")

            threads, experiment_log_dir, file_name_template = run_main_command(program, file, wait=False, log_dir=experiment_log_dir)
            # Wait at most 30 seconds; if the threads finish early, skip the remaining wait
            max_wait_seconds = 30
            check_interval = 1
            elapsed = 0
            all_done = False
            print(f"[exp-range] Waiting at most {max_wait_seconds} seconds; if finished early, skip the remaining wait")
            while elapsed < max_wait_seconds:
                if all(not t.is_alive() for t in threads):
                    all_done = True
                    print(f"[exp-range] All threads completed within {elapsed} seconds, entering the collection phase early")
                    break
                time.sleep(check_interval)
                elapsed += check_interval
                if elapsed % 5 == 0:
                    print(f"[exp-range] Waited {elapsed} seconds, still running...")

            if not all_done:
                print(f"[exp-range] Timed out after {max_wait_seconds} seconds, terminating all child processes")
                kill_all_child_processes()
                kill_ts = time.time()
            else:
                kill_ts = None

            # Wait for all threads to finish, then collect the results
            print(f"[exp-range] Waiting for all threads to finish...")
            for thread in threads:
                thread.join()
            print(f"[exp-range] Collecting experiment results...")
            collect_experiment_results(experiment_log_dir, file_name_template)

            if FORCE_QUIT:
                print("[exp-range] Received force-quit signal, terminating all experiments - exiting the inner loop")
                break

            # Only when a kill was triggered, ensure we wait another 10 seconds after the kill
            if kill_ts is not None:
                while time.time() - kill_ts < 10:
                    time.sleep(1)
        if FORCE_QUIT:
            print("[exp-range] Received force-quit signal, ending experiments early - exiting the outer loop")
            return

        # Plot the experiment result charts
    print("[exp-range] Starting to plot the experiment result charts...")
    plot_exp_range_results(log_dir, key, data_path=data_path)

    # After the last experiment finishes, clean up all processes
    print("[exp-range] All experiments complete, cleaning up all processes")
    killall()


def plot_exp_range_main(log_dir, key, start_value, stop_value, data_path=None):
    """
    Plot the exp-range experiment result charts within the specified range
    """
    print(f"[plot-exp-range] Starting to plot the experiment result charts...")
    print(f"[plot-exp-range] Log directory: {log_dir}")
    print(f"[plot-exp-range] Parameter key: {key}")
    print(f"[plot-exp-range] Range: {start_value} to {stop_value}")

    # Check whether the log directory exists
    if not os.path.exists(log_dir):
        print(f"[plot-exp-range] Error: log directory does not exist: {log_dir}")
        sys.exit(1)

    # Call the analysis function to plot the charts
    plot_exp_range_results(log_dir, key, start_value, stop_value, data_path=data_path)
    print(f"[plot-exp-range] Chart plotting complete")


def main():
    """
    Main function, parses command-line arguments.
    """
    signal.signal(signal.SIGINT, cleanup_processes)

    parser = argparse.ArgumentParser(
        description="A tool for distributing files and executing commands across multiple machines.",
        epilog="Examples: \n"
               "  python deploy_tool.py distribute ./myapp.zip /tmp/\n"
               "  python deploy_tool.py run",
        formatter_class=argparse.RawTextHelpFormatter
    )

    subparsers = parser.add_subparsers(dest="command", help="Available subcommands")

    # Create the 'distribute' subcommand parser
    parser_distribute = subparsers.add_parser("distribute", help="Distribute a file or directory to all remote hosts")
    parser_distribute.add_argument("source", help="Local file or directory path to copy")
    parser_distribute.add_argument("destination", help="Target path on the remote hosts")

    # Create the 'run' subcommand parser
    parser_run = subparsers.add_parser("run", help="Execute a command on the local machine and all remote hosts")
    parser_run.add_argument("--program", type=str, default=None, help="Path to the program to execute")
    parser_run.add_argument("--config", type=str, default=None, help="Path to the config file to distribute to all hosts")

    # Create the 'ibtest' subcommand parser
    parser_ibtest = subparsers.add_parser("ibtest", help="Batch-start ib_write_bw servers on h194, and batch-start clients on localhost and h196")

    # Create the 'pingmesh' subcommand parser
    parser_pingmesh = subparsers.add_parser("pingmesh", help="Batch-ping target IPs on all hosts to refresh the arp cache")

    # Create the 'killall' subcommand parser
    parser_killall = subparsers.add_parser("killall", help="Kill all anytest processes on the local machine and all hosts")

    # Create the 'init' subcommand parser
    parser_init = subparsers.add_parser("init", help="Set traffic_class and hugepages on HOSTS and start nusad")

    # Create the 'set-cnp-interval' subcommand parser
    parser_set_cnp_interval = subparsers.add_parser("set-cnp-interval", help="Set cnp_interval on the specified machine")
    parser_set_cnp_interval.add_argument("ip", type=str, help="IP address of the target machine")
    parser_set_cnp_interval.add_argument("cnp_interval", type=str, help="The cnp_interval value to set")

    # Add the 'set-yaml' subcommand
    parser_set_yaml = subparsers.add_parser("set-yaml", help="Modify a specific key-value pair in a yaml file")
    parser_set_yaml.add_argument("config", type=str, help="Path to the yaml file")
    parser_set_yaml.add_argument("key", type=str, help="The key to modify (supports a or a.b)")
    parser_set_yaml.add_argument("value", type=str, help="The new value to set")

    # Add the 'exp-range' subcommand
    parser_exp_range = subparsers.add_parser("exp-range", help="Batch-modify yaml and run the experiment multiple times")
    parser_exp_range.add_argument("program", type=str, help="Program path or name")
    parser_exp_range.add_argument("config", type=str, help="Path to the yaml file")
    parser_exp_range.add_argument("key", type=str, help="The key to modify (supports a or a.b)")
    parser_exp_range.add_argument("start_value", type=int, help="Start value")
    parser_exp_range.add_argument("stop_value", type=int, help="Stop value")
    parser_exp_range.add_argument("step", type=int, help="Step")
    parser_exp_range.add_argument(
        "--sweep-mode",
        type=str,
        default="linear",
        choices=["linear", "multiply"],
        help="Parameter sweep mode: linear=arithmetic progression; multiply=exponential/multiplicative progression (value*=step)",
    )
    parser_exp_range.add_argument("--data-path", type=str, default=None, help="Dotted path: resolve a numeric field from the top level of result.json for plotting")
    parser_exp_range.add_argument("--log-dir", type=str, default=None, help="Log directory path, automatically generated if not specified")

    # Add the 'plot-exp-range' subcommand
    parser_plot_exp_range = subparsers.add_parser("plot-exp-range", help="Plot the exp-range experiment result charts within the specified range")
    parser_plot_exp_range.add_argument("log_dir", type=str, help="Log directory path")
    parser_plot_exp_range.add_argument("key", type=str, help="The parameter key to analyze")
    parser_plot_exp_range.add_argument("start_value", type=int, help="Start value")
    parser_plot_exp_range.add_argument("stop_value", type=int, help="End value")
    parser_plot_exp_range.add_argument("--data-path", type=str, default=None, help="Dotted path: resolve a numeric field from the top level of result.json for plotting")

    args = parser.parse_args()
    
    if args.command is None:
        parser.print_help()
        sys.exit(1)

    if args.command == "distribute":
        distribute_files(args.source, args.destination)
        # distribute_files_with_chunks(args.source, args.destination)
    elif args.command == "run":
        run_main_command(args.program, args.config)
    elif args.command == "ibtest":
        run_ibtest_commands()
    elif args.command == "pingmesh":
        pingmesh()
    elif args.command == "killall":
        killall()
    elif args.command == "init":
        init_hosts()
    elif args.command == "set-cnp-interval":
        set_cnp_interval(args.ip, args.cnp_interval)
    elif args.command == "set-yaml":
        set_yaml_value(args.config, args.key, args.value)
    elif args.command == "exp-range":
        exp_range_main(
            args.program,
            args.config,
            args.key,
            args.start_value,
            args.stop_value,
            args.step,
            args.log_dir,
            args.data_path,
            args.sweep_mode,
        )
    elif args.command == "plot-exp-range":
        plot_exp_range_main(args.log_dir, args.key, args.start_value, args.stop_value, args.data_path)
    

if __name__ == "__main__":
    main()