provision-script.py

#!/bin/env python3
import os
import subprocess
from multiprocessing import Process
import time
import sys
from util import *

from params import *

# To get bash
# chsh -s /bin/bash


# adds a new user with passwordless sudo, assumes root ssh-key access
def adduser(server, username=REMOTE_USER):
    if username is None:
        username = REMOTE_USER
    run_remote(server, "", f"useradd -m {username}", user="root")
    run_remote(server, "", f"usermod -aG sudo {username}", user="root")
    run_remote(
        server,
        "",
        f'echo "{username} ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/{username}',
        user="root",
    )
    run_remote(server, "", f"mkdir /home/{username}/.ssh", user="root")
    copy_to_remote(
        server,
        f"~/.ssh/{PUBKEY}",
        f"/home/{username}/.ssh/authorized_keys",
        user="root",
    )
    run_remote(server, "", f"chown -R {username} /home/{username}/.ssh", user="root")
    run_remote(server, "", "chsh -s /bin/bash")
    run_remote(server, "", f"sudo passwd {username}")
    run_remote(server, "", f'sudo echo "done making user {username}"')


def setup_mlx(server):
    copy_to_remote(server, OFED_DIR + OFED_VERSION + ".tgz", "~")
    run_remote(server, "~", f"tar -xf {OFED_VERSION}.tgz")
    run_remote(server, "~/" + OFED_VERSION, "sudo ./mlnxofedinstall --all --force")
    run_remote(server, "~/" + OFED_VERSION, "sudo /etc/init.d/openibd restart")


# Takes down bond and sets up a new interface
def update_netplan(server, port1=False):
    if port1:
        copy_to_remote(
            server, "netplan-port1.yaml", "/etc/netplan/01-netcfg.yaml", user="root"
        )
    else:
        copy_to_remote(
            server, "netplan-port0.yaml", "/etc/netplan/01-netcfg.yaml", user="root"
        )
    print("go drink a coffee, this script waits for 10 minutes now")
    run_remote(server, "~", "sudo netplan apply")


def setup_caladan(server):
    os.system(f"cd ~/dev/sw && rm -rf {REPO_NAME} && git clone {CALADAN_REPO_LINK}")
    run_remote(server, "~", f"rm -rf ~/dev/sw/{REPO_NAME}")
    copy_to_remote(server, f"~/dev/sw/{REPO_NAME}", "~", rec=True)
    run_remote(
        server,
        "~",
        "sudo apt install make gcc cmake pkg-config libnl-3-dev libnl-route-3-dev libnuma-dev uuid-dev libssl-dev libaio-dev libcunit1-dev libclang-dev libncurses-dev meson python3-pyelftools -y",
    )
    make_initial(server)


def create_config_file(ip_address, num_cores):
    # Split the IP address to get the first three octets
    octets = ip_address.split(".")
    # assert this really is an ip
    for octet in octets:
        int(octet)

    # The gateway is usually the first IP in the subnet, which typically ends with .1
    gateway_ip = f"{octets[0]}.{octets[1]}.{octets[2]}.1"

    # Create the config file content
    config_content = f"""# an example runtime config file
host_addr {ip_address}
host_netmask 255.255.255.0
host_gateway {gateway_ip}
host_mtu 9000
#nosteal y
runtime_kthreads {num_cores}
runtime_guaranteed_kthreads {num_cores}
runtime_spinning_kthreads {num_cores}
runtime_priority lc
enable_directpath fs
"""
    filename = f"tmp.config_{ip_address}"
    print(f"Created CONFIG (server {ip_address}):\n{config_content}")

    # Write the content to tmp.config file
    with open(filename, "w") as file:
        file.write(config_content)
        file.flush()
    return filename


def set_cpus_to_performance(server):
    copy_to_remote(server, "set-cpus-to-performance.py", "~/set-cpus-to-performance.py")
    run_remote(
        server,
        "~",
        "chmod +x set-cpus-to-performance.py && sudo ./set-cpus-to-performance.py",
    )


def make_config_file(server, num_cores=NUM_CONFIG_CORES):
    ip = get_interface_ip(server, get_caladan_iface(server))
    filename = create_config_file(ip, num_cores)
    copy_to_remote(server, filename, f"~/{REPO_NAME}/client.config")


def per_server_after_reboot(server):
    set_cpus_to_performance(server)
    run_remote(
        server,
        f"~/{REPO_NAME}",
        "make && cd ksched && make && cd .. && sudo ./scripts/setup_machine.sh",
    )


def update_caladan(server):
    setup_caladan(server)
    make_config_file(server)


def per_user_once(server):
    adduser(server)
    update_caladan(server)


# must call adduser before
def per_server_once(server):
    update_netplan(server)  # Probably unnecessary for cloudlab
    setup_mlx(server)


def init_new_server(server, rerun=0):
    if rerun != 0 and rerun != 1:
        print("Pass whether this is the first run (0) or a rerun (1)")
        exit(1)
    if not rerun:
        adduser(server)
    per_server_once(server)
    update_caladan(server)

    per_server_after_reboot(server)


def update_install_essentials(server):
    run_remote(server, "~", "sudo apt update")
    run_remote(
        server,
        "~",
        "sudo apt install ethtool htop make gcc cmake pkg-config libnl-3-dev libnl-route-3-dev libnuma-dev uuid-dev libssl-dev libaio-dev libcunit1-dev libclang-dev libncurses-dev meson python3-pyelftools -y",
    )


def copy_github_keys(server):
    copy_to_remote(server, f"/home/{LOCAL_USER}/.ssh/id_rsa", f"~/.ssh")
    copy_to_remote(server, f"/home/{LOCAL_USER}/.ssh/id_rsa.pub", f"~/.ssh")
    run_remote(server, "~", "ssh-keyscan github.com >> ~/.ssh/known_hosts")
    run_remote(server, "~", "chmod 600 ~/.ssh/id_rsa")


def clone_and_make_sird_caladan(server):
    run_remote(server, "~", f"git clone {CALADAN_REPO_LINK}")
    run_remote(server, f"~/{REPO_NAME}", "make submodules")
    per_server_after_reboot(server)

def config_nic(server):
    print(f"Configuring NIC with caladan address {get_caladan_ip(server)} and iface: {get_caladan_iface(server)}")
    cmd = f"sudo ip link set dev {get_caladan_iface(server)} mtu 9000"
    print(f"Setting {server}'s NIC mtu: {cmd}")
    run_remote(server, "~", cmd)

    run_remote(server, "~", f"sudo ip addr flush dev {get_caladan_iface(server)}")
    cmd = f"sudo ip addr add {get_caladan_ip(server)}/24 dev {get_caladan_iface(server)}"
    run_remote(server, "~", cmd)
    run_remote(server, "~", f"sudo ip link set dev {get_caladan_iface(server)} up")

def make_apps(server):
    run_remote(server, f"~/{REPO_NAME}", "cd apps/sird_tester && make clean && make")



def setup_cloudlab(server):
    # Remember to add a mapping in .ssh/config
    print(f"Setting up {server}")
    copy_github_keys(server)
    update_install_essentials(server)
    setup_mlx(server)
    config_nic(server)
    clone_and_make_sird_caladan(server)
    make_config_file(server)
    make_apps(server)

def main():
    if len(sys.argv) < 3:
        print(f"correct usage: {sys.argv[0]} [mode] [ip/hostname]")
        return
    mode = sys.argv[1]
    server = sys.argv[2]
    if mode == "server":
        init_new_server(server, int(sys.argv[3]))
        # Note: if after provisioning, the shell of the mgmnt terminal is not bash, do
        # chsh -s /bin/bash
    elif mode == "user":
        per_user_once(server)
    elif mode == "update":
        update_caladan(server)
    elif mode == "postreboot":
        per_server_after_reboot(server)
    elif mode == "netplan":
        update_netplan(server)
    elif mode == "setup-mlx":
        setup_mlx(server)
    elif mode == "setup-cloudlab":
        servers = server.split(",")
        for s in servers:
            p = Process(target=setup_cloudlab, args=(s,))
            p.start()
    elif mode == "runtime-config":
        servers = server.split(",")
        for s in servers:
            p = Process(target=make_config_file, args=(s,))
            p.start()
    else:
        print("valid modes: server, user, update, postreboot")


main()