set_frcc_params.py

#!/usr/bin/env python3

# NOTE: generated by Chat GPT

import os
import argparse
import subprocess

# Define all available parameters and their types
PARAMS = {
    "p_ub_rtprop_us": int,
    "p_ub_rtterr_us": int,
    "p_ub_flow_count": int,

    "p_lb_cwnd_pkts": int,
    "p_contract_min_qdel_us": int,
    "p_probe_duration_us": int,
    "p_probe_multiplier_unit": int,
    "p_cwnd_averaging_factor_unit": int,
    "p_cwnd_clamp_hi_unit": int,
    "p_cwnd_clamp_lo_unit": int,
    "p_slot_load_factor_unit": int,
    "p_ub_slots_per_round": int,
    "p_rprobe_interval_us": int,
    "p_probe_wait_rtts": int,

    "f_use_rtprop_probe": bool,
    "f_wait_rtt_after_probe": bool,
    "f_use_stable_cwnd_update": bool,
    "f_probe_wait_in_max_rtts": bool,
    "f_probe_duration_max_rtt": bool,
    "f_drain_over_rtt": bool,
    "f_probe_over_rtt": bool,
    "f_slot_greater_than_rtprop": bool,
}

P_SCALE = 8
P_UNIT = 1 << P_SCALE

DEFAULT_VALUES = {
    "p_ub_rtprop_us": 100000,  # 100 ms
    "p_ub_rtterr_us": 10000,  # 10 ms
    "p_ub_flow_count": 3,

    "p_lb_cwnd_pkts": 4,
    "p_contract_min_qdel_us": 10000,  # p_ub_rtprop_us / 2
    "p_probe_duration_us": 10000,  # 10 ms
    "p_probe_multiplier_unit": 4 * P_UNIT,  # P_UNIT * 4
    "p_cwnd_averaging_factor_unit": P_UNIT * 1,  # P_UNIT * 1
    "p_cwnd_clamp_hi_unit": (P_UNIT * 13) // 10,  # P_UNIT * 13 / 10
    "p_cwnd_clamp_lo_unit": (P_UNIT * 10) // 13,  # P_UNIT * 10 / 13
    "p_slot_load_factor_unit": P_UNIT * 2,  # P_UNIT * 1
    "p_ub_slots_per_round": 20,
    "p_rprobe_interval_us": 30000000,  # 30 seconds
    "p_probe_wait_rtts": 2,

    "f_use_rtprop_probe": "Y",  # True
    "f_wait_rtt_after_probe": "Y",  # True
    "f_use_stable_cwnd_update": "Y",  # True
    "f_probe_wait_in_max_rtts": "Y",
    "f_probe_duration_max_rtt": "Y",
    "f_drain_over_rtt": "Y",
    "f_probe_over_rtt": "Y",
    "f_slot_greater_than_rtprop": "Y",
}

MODULE_PATH = "/sys/module/tcp_frcc/parameters/"


def get_param_name(param):
    return f"static_{param}"


def set_kernel_param(param, value):
    """Writes the given value to the corresponding kernel module parameter."""
    name = get_param_name(param)
    param_path = os.path.join(MODULE_PATH, name)

    try:
        # Use sudo tee to write values
        subprocess.run(
            ["sudo", "tee", param_path],
            input=str(value).encode(),
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        print(f"Set {param} = {value}")
    except subprocess.CalledProcessError:
        print(f"Error setting {param}. Check if the module is loaded and path exists.")


def get_kernel_param(param):
    name = get_param_name(param)
    param_path = os.path.join(MODULE_PATH, name)
    try:
        with open(param_path, "r") as f:
            value = f.read().strip()
            return value
    except FileNotFoundError:
        print(f"Error: {param} not found. Check if the module is loaded.")


def print_kernel_params():
    """Reads and prints all kernel module parameters."""
    for param in PARAMS.keys():
        value = get_kernel_param(param)
        if value != str(DEFAULT_VALUES[param]):
            print(f"{param} = {value} (default: {DEFAULT_VALUES[param]})")
        else:
            print(f"{param} = {value}")


def get_write_value(args, param):
    param_type = PARAMS[param]
    env_var = f"TCP_FRCC_{param.upper()}"
    value = getattr(args, param, None)  # CLI argument

    if value is None and env_var in os.environ:
        # Convert environment variable to correct format
        env_value = os.environ[env_var]
        if param_type is bool:
            value = env_value.lower() in ("1", "true", "yes")
        else:
            value = param_type(env_value)

    # For boolean params, convert to 'Y' or 'N'
    if value is not None:
        if param_type is bool:
            value = "Y" if value else "N"

    return value


def check_params(args):
    """Checks that the params have been set as specified in the env vars"""
    for param in PARAMS.keys():
        read_value = get_kernel_param(param)
        write_value = get_write_value(args, param)
        if write_value is not None:
            write_value = str(write_value)
            assert read_value == write_value, f"Mismatch for {param}: read {read_value}, expected {write_value}"
        else:
            default_value = str(DEFAULT_VALUES[param])
            assert read_value == default_value, f"Mismatch for {param}: read {read_value}, expected {default_value}"
    print("All parameters are set correctly.")


def reset_all_params():
    """Resets all kernel parameters to their default values."""
    for param, default_value in DEFAULT_VALUES.items():
        set_kernel_param(param, default_value)
        print(f"Reset {param} to default value {default_value}")


def main():
    parser = argparse.ArgumentParser(description="Set kernel module parameters for tcp_frcc.")
    parser.add_argument("--print", action="store_true", help="Print all current kernel parameters")
    parser.add_argument("--reset", action="store_true", help="Reset all parameters to default values")
    parser.add_argument("--check", action="store_true", help="Check that all values are set according to cli or env vars")

    for param, param_type in PARAMS.items():
        parser.add_argument(f"--{param}", type=param_type, help=f"Set {param}")

    args = parser.parse_args()

    if args.print:
        print_kernel_params()
        return

    if args.reset:
        reset_all_params()
        return

    if args.check:
        check_params(args)
        return

    # import pdb; pdb.set_trace()
    for param, param_type in PARAMS.items():
        value = get_write_value(args, param)
        if value is not None:
            set_kernel_param(param, value)


if __name__ == "__main__":
    main()