run

#!/usr/bin/env python3

import argparse
import typing
import subprocess
import logging
import shutil
import shlex
import os

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("MaelstromInvoker")


DESCRIPTION = '''
Run the maelstrom tests for a given fly.io dist-sys challenge.

Currently, the following challenges are implemented:

- 01: Echo
- 02: Unique ID Generation
- 03a: Broadcast (Single Node)
- 03b: Broadcast (Multi-Node)
- 03c: Broadcast (Multi-Node Fault-tolerant)
- 03d: Broadcast (Multi-Node Efficient 1)
- 03e: Broadcast (Multi-Node Efficient 2)

'''


IMPLEMENTED_CHALLENGES = {
    "1": "Echo", 
    "2": "Unique ID Generation", 
    "3a": "Broadcast (Single Node)",
    "3b": "Broadcast (Multi-Node)",
    "3c": "Broadcast (Fault Tolerant)",
    "3d": "Broadcast (Efficient 1)",
    "3e": "Broadcast (Efficient 2)",
    "4": "Grow only counter",
    "5a": "Kafka style log (Single Node)"
}

CHALLENGE_COMMANDS = {
    "1": "./maelstrom test -w echo --bin solutions/maelstrom-echo --time-limit 5",
    "2": "./maelstrom test -w unique-ids --bin solutions/maelstrom-unique-id-generation --time-limit 30 --rate 1000 --node-count 3 --availability total --nemesis partition",
    "3a": "./maelstrom test -w broadcast --bin solutions/maelstrom-broadcast-single-node --node-count 1 --time-limit 20 --rate 10",
    "3b": "./maelstrom test -w broadcast --bin solutions/maelstrom-broadcast-multi-node --node-count 5 --time-limit 100 --rate 10",
    "3c": "./maelstrom test -w broadcast --bin solutions/maelstrom-broadcast-fault-tolerant --node-count 5 --time-limit 20 --rate 10 --nemesis partition",
    "3d": "./maelstrom test -w broadcast --bin solutions/maelstrom-broadcast-efficient-part-1 --node-count 25 --time-limit 20 --rate 100 --latency 100",
    "3e": "./maelstrom test -w broadcast --bin solutions/maelstrom-broadcast-efficient-part-2 --node-count 25 --time-limit 20 --rate 100 --latency 100",
    "4": "./maelstrom test -w g-counter --bin solutions/maelstrom-grow-only-counter --node-count 3 --rate 100 --time-limit 20 --nemesis partition",
    "5a": "./maelstrom test -w kafka --bin solutions/maelstrom-kafka-part-1 --node-count 1 --concurrency 2n --time-limit 20 --rate 1000"
}

def parse_args():
    parser = argparse.ArgumentParser(prog="MaelstromInvoker", description=DESCRIPTION)
    parser.add_argument(
        "challenge",
        metavar="CHALLENGE_ID",
        nargs="+",
        type=str,
        help="The ids of the challenges to run.",
    )
    parser.add_argument(
        "-l",
        "--language",
        type=str,
        choices=("go", "rust"),
        default="rust",
        help="The implementation to run the tests for. Defaults to 'rust'."
    )

    args = parser.parse_args()

    if any(
        challenge_id not in IMPLEMENTED_CHALLENGES for challenge_id in args.challenge
    ):
        raise ValueError(
            f"Invalid or unimplemented challenge id. Recognized challenge ids are: {str(IMPLEMENTED_CHALLENGES)}"
        )
    return args


def inv(command, check=True, *args, **kwargs):
    return subprocess.run(shlex.split(command), check=True, *args, **kwargs)


def build_and_copy_binaries(language: str = "rust"):
    """Build the binaries for the given language."""
    if language == "rust":
        inv("cargo build --release", cwd="solutions-rust")
        inv("mkdir -p maelstrom/solutions/rust")

        SRC_DIR = "solutions-rust/target/release"
        DEST_DIR = "maelstrom/solutions/rust"
        
        for file_path in os.listdir(SRC_DIR):
            src_file_path = os.path.join(SRC_DIR, file_path)
            # Copy over all the executables.

            if os.path.isfile(src_file_path) and os.access(src_file_path, os.X_OK):
                dest_file_path = os.path.join(DEST_DIR, file_path)
                shutil.copy(src_file_path, dest_file_path)
                logging.debug("Copied %s to %s", src_file_path, dest_file_path)
    elif language == "go":
        return inv("make all -j12", cwd="./solutions-go")


def run_test(challenge: str, lang: str):
    """Run maelstrom for a given challenge."""

    command = CHALLENGE_COMMANDS.get(challenge, None)

    if command is None:
        raise ValueError(f"Couldn't find a test for challenge: {str(challenge)}")

    command = command.replace("solutions/", f"solutions/{lang}/")

    logger.info(
        'Running maelstrom test for %s implementation of challenge "%s" (id: %s)',
        lang,
        IMPLEMENTED_CHALLENGES[challenge],
        challenge,
    )
    result = inv(command, cwd="maelstrom")
    return result


def try_setup_maelstrom():
    if os.path.exists("./maelstrom") and os.path.isdir("./maelstrom"):
        logger.debug("Maelstrom directory exists, so we'll assume maelstrom is already set up correctly.")
        return

    logger.info("Seems like this is the first time we're running tests. Setting up Maelstrom v0.2.2 ...")
    logger.info("Downloading Maelstrom v0.2.2 to ./maelstrom")
    inv('curl -L https://github.com/jepsen-io/maelstrom/releases/download/v0.2.2/maelstrom.tar.bz2 -o maelstrom.tar.bz2')
    inv("tar -xvf maelstrom.tar.bz2", capture_output=True)
    inv("rm maelstrom.tar.bz2", capture_output=True)
    inv("mkdir -p maelstrom/solutions", capture_output=True)
    logger.info(
        "Downloaded Maelstrom v0.2.2 to ./maelstrom and "
        "created a ./maelstrom/solutions directory to store our solutions."
    )


def main():
    args = parse_args()
    logger.debug("Arguments: %s", args)
    try_setup_maelstrom()
    build_and_copy_binaries(args.language)
    for challenge_id in args.challenge:
        run_test(challenge_id, args.language)


if __name__ == "__main__":
    main()