piccolo_driver.py

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the Apache 2.0 License.
import argparse
import os
import infra.e2e_args
import infra.remote_client
import infra.jwt_issuer
from infra.perf import PERF_COLUMNS
from random import seed
import getpass
from loguru import logger as LOG
import time
import hashlib
import json
from piccolo import generator
from piccolo import analyzer
import infra.bencher
import infra.proc


def get_command_args(args, network, get_command):
    client_ident = network.users[0]
    command_args = [
        "--cert",
        client_ident.cert_path,
        "--key",
        client_ident.key_path,
        "--cacert",
        network.cert_path,
    ]
    return get_command(*command_args)


def minimum_number_of_local_nodes(args):
    if args.send_tx_to == "backups":
        return 2

    return 1


def filter_nodes(primary, backups, filter_type):
    if filter_type == "primary":
        return [primary]
    elif filter_type == "backups":
        assert backups, "--send-tx-to backups but no backup was found"
        return backups
    else:
        return [primary] + backups


def my_configure_remote_client(args, client_id, client_host, node, command_args):
    client_host = infra.net.expand_localhost()

    try:
        remote_client = infra.remote_client.CCFRemoteClient(
            f"client_{client_id}",
            client_host,
            args.client,
            node.get_public_rpc_host(),
            node.get_public_rpc_port(),
            args.workspace,
            args.label,
            args.config,
            command_args,
            piccolo_run=True,
        )
        remote_client.setup()
        return remote_client
    except Exception:
        LOG.exception("Failed to start client {}".format(client_host))
        raise


def run(get_command, args):
    if args.fixed_seed:
        seed(getpass.getuser())

    hosts = args.nodes
    if not hosts:
        hosts = infra.e2e_args.nodes(args, minimum_number_of_local_nodes(args))

    args.initial_user_count = 3
    args.sig_ms_interval = 100
    args.ledger_chunk_bytes = "5MB"  # Set to node default value

    LOG.info("Starting nodes on {}".format(hosts))

    with infra.network.network(
        hosts, args.binary_dir, args.debug_nodes, pdb=args.pdb
    ) as network:
        network.start_and_open(args)

        primary, backups = network.find_nodes()

        command_args = get_command_args(args, network, get_command)

        additional_headers = {}
        if args.use_jwt:
            jwt_issuer = infra.jwt_issuer.JwtIssuer("https://example.issuer")
            jwt_issuer.register(network)
            jwt = jwt_issuer.issue_jwt()
            additional_headers["Authorization"] = f"Bearer {jwt}"

        LOG.info(f"Generating {args.repetitions} parquet requests")
        msgs = generator.Messages()
        for i in range(args.repetitions):
            body = {
                "id": i % 100,
                "msg": f"Unique message: {hashlib.sha256(str(i).encode()).hexdigest()}",
            }
            msgs.append(
                "/app/log/private",
                "POST",
                additional_headers=additional_headers,
                body=json.dumps(body),
            )

        filename_prefix = "piccolo_driver"
        path_to_requests_file = os.path.join(
            network.common_dir, f"{filename_prefix}_requests.parquet"
        )
        LOG.info(f"Writing generated requests to {path_to_requests_file}")
        msgs.to_parquet_file(path_to_requests_file)

        path_to_send_file = os.path.join(
            network.common_dir, f"{filename_prefix}_send.parquet"
        )

        path_to_response_file = os.path.join(
            network.common_dir, f"{filename_prefix}_response.parquet"
        )

        # Add filepaths in commands
        command_args += ["--send-filepath", path_to_send_file]
        command_args += ["--response-filepath", path_to_response_file]
        command_args += ["--generator-filepath", path_to_requests_file]

        nodes_to_send_to = filter_nodes(primary, backups, args.send_tx_to)
        clients = []
        client_hosts = []
        if args.one_client_per_backup:
            assert backups, "--one-client-per-backup was set but no backup was found"
            client_hosts = ["localhost"] * len(backups)
        else:
            if args.client_nodes:
                client_hosts.extend(args.client_nodes)

        if args.num_localhost_clients:
            client_hosts.extend(["localhost"] * int(args.num_localhost_clients))

        if not client_hosts:
            client_hosts = ["localhost"]

        for client_id, client_host in enumerate(client_hosts):
            node = nodes_to_send_to[client_id % len(nodes_to_send_to)]

            remote_client = my_configure_remote_client(
                args, client_id, client_host, node, command_args
            )
            clients.append(remote_client)

        if args.network_only:
            for remote_client in clients:
                LOG.info(f"Client can be run with: {remote_client.remote.get_cmd()}")
            while True:
                time.sleep(60)
        else:
            for remote_client in clients:
                remote_client.start()

            hard_stop_timeout = 90
            format_width = len(str(hard_stop_timeout)) + 3

            try:
                start_time = time.time()
                while True:
                    stop_waiting = True
                    for i, remote_client in enumerate(clients):
                        done = remote_client.check_done()
                        # all the clients need to be done
                        LOG.info(
                            f"Client {i} has {'completed' if done else 'not completed'} running ({time.time() - start_time:>{format_width}.2f}s / {hard_stop_timeout}s)"
                        )
                        stop_waiting = stop_waiting and done
                    if stop_waiting:
                        break
                    if time.time() > start_time + hard_stop_timeout:
                        raise TimeoutError(
                            f"Client still running after {hard_stop_timeout}s"
                        )

                    time.sleep(5)

                perf_label = args.perf_label

                for remote_client in clients:
                    analysis = analyzer.Analyze()

                    LOG.info(
                        f"Analyzing results from {path_to_send_file} and {path_to_response_file}"
                    )
                    df_sends = analyzer.get_df_from_parquet_file(path_to_send_file)
                    df_responses = analyzer.get_df_from_parquet_file(
                        path_to_response_file
                    )
                    time_spent = analysis.total_time_in_sec(df_sends, df_responses)

                    perf_result = round(len(df_sends.index) / time_spent, 1)
                    LOG.success(f"{args.label}/{remote_client.name}: {perf_result}")

                    # Throughput from only one client, preserved for legacy reason
                    # see basicperf.py for a better, cross-client approach.
                    bf = infra.bencher.Bencher()
                    bf.set(
                        perf_label,
                        infra.bencher.Throughput(perf_result),
                    )

                primary, _ = network.find_primary()
                mem = infra.proc.get_proc_memory_stats(primary.remote.remote.proc.pid)
                if mem is not None:
                    bf = infra.bencher.Bencher()
                    bf.set_memory(perf_label, mem)

                for remote_client in clients:
                    remote_client.stop()

            except Exception as e:
                LOG.error(f"Stopping clients due to exception: {e}")
                for remote_client in clients:
                    remote_client.stop()
                raise


def cli_args(add=lambda x: None, accept_unknown=False):
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--client", help="Client binary", required=True)
    parser.add_argument(
        "-n",
        "--nodes",
        help="List of hostnames[,pub_hostnames:ports]. If empty, spawn minimum working number of local nodes (minimum depends on consensus and other args)",
        action="append",
    )
    client_args_group = parser.add_mutually_exclusive_group()
    client_args_group.add_argument(
        "-cn",
        "--client-nodes",
        help="List of hostnames for spawning client(s). If empty, one client is spawned locally",
        action="append",
    )
    client_args_group.add_argument(
        "--one-client-per-backup",
        help="If set, allocates one (local) client per backup",
        action="store_true",
    )
    parser.add_argument(
        "-nlc",
        "--num-localhost-clients",
        help="The number of localhost clients. \
        This argument is cumulative with the client-nodes and one-client-per-backup and arguments",
    )
    parser.add_argument(
        "--send-tx-to",
        choices=["primary", "backups", "all"],
        default="all",
        help="Send client requests only to primary, only to backups, or to all nodes",
    )
    parser.add_argument(
        "--metrics-file",
        default="metrics.json",
        help="Path to json file where the transaction rate metrics will be saved to",
    )
    parser.add_argument(
        "-f",
        "--fixed-seed",
        help="Set a fixed seed for port and IP generation.",
        action="store_true",
    )
    parser.add_argument(
        "--use-jwt",
        help="Use JWT with a temporary issuer as authentication method.",
        action="store_true",
    )
    parser.add_argument(
        "--repetitions",
        help="Number of requests to send",
        type=int,
        default=100,
    )
    parser.add_argument("--config", help="Path to config for client binary", default="")

    return infra.e2e_args.cli_args(
        add=add, parser=parser, accept_unknown=accept_unknown
    )


def generic_run(*args, **kwargs):
    infra.path.mk_new("perf_summary.csv", PERF_COLUMNS)

    run(*args, **kwargs)


if __name__ == "__main__":
    args, unknown_args = cli_args(accept_unknown=True)

    unknown_args = [term for arg in unknown_args for term in arg.split(" ")]

    def get_command(*args):
        return [*args] + unknown_args

    run(get_command, args)