tools/vcluster/main.py

#!/usr/bin/env python
from argparse import ArgumentParser
from argparse import RawDescriptionHelpFormatter
import os

from config_generator import load_config
from vcluster import VCluster

DATE = "2017-09-13"
AUTHOR = "Chunyang Shen"

LOGO = """
  Created by %s on %s.
  Copyright Uber Compute Platform. All rights reserved.

USAGE
"""


def parse_arguments():
    program_license = LOGO % (AUTHOR, DATE)

    parser = ArgumentParser(
        description=program_license,
        formatter_class=RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "-c",
        "--config-file",
        dest="config_file",
        help="Path to vcluster config file. Default is config/default.yaml",
    )

    parser.add_argument(
        "-z",
        "--zookeeper-server",
        dest="zookeeper_server",
        help="the DNS of zookeeper server of the physical Peloton cluster",
    )

    parser.add_argument(
        "-n", "--name", dest="label_name", help="the name of the vcluster"
    )

    parser.add_argument(
        "-p",
        "--resource-pool",
        dest="respool_path",
        help="the path of the resource pool",
    )

    parser.add_argument(
        "--auth-type",
        dest="auth_type",
        help="auth type of the physical Peloton cluster",
    )

    parser.add_argument(
        "--auth-config-file",
        dest="auth_config_file",
        help="auth config file used to talk to the physical Peloton cluster",
    )

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

    # Subparser for the 'mesos' command
    parser_mesos = subparsers.add_parser(
        "mesos",
        help="set up a virtual cluster with Mesos master and Mesos slave",
    )

    parser_mesos.add_argument(
        "-s",
        "--agent-number",
        type=int,
        dest="agent_number",
        help="the number of Mesos agent in the vcluster",
    )

    # Subparser for the 'mesos-slave' command
    parser_mesos_slave = subparsers.add_parser(
        "mesos-slave", help="adding Mesos slaves giving a zookeeper"
    )

    parser_mesos_slave.add_argument(
        "-s",
        "--agent-number",
        type=int,
        dest="agent_number",
        help="the number of Mesos slaves add into the vcluster",
    )

    parser_mesos_slave.add_argument(
        "-k", "--zk", dest="zk", help="the host of virtual zk"
    )

    # Subparser for the 'mesos-master' command
    parser_mesos_slave = subparsers.add_parser(
        "mesos-master",
        help="set up a virtual cluster with Mesos master and Mesos slave",
    )

    parser_mesos_slave.add_argument(
        "-k", "--zk", dest="zk", help="the host of virtual zk"
    )

    # Subparser for the 'setup' command
    parser_setup = subparsers.add_parser(
        "setup", help="set up a virtual cluster"
    )

    parser_setup.add_argument(
        "-v",
        "--peloton-version",
        nargs="?",
        dest="peloton_version",
        help="The image version for Peloton",
    )

    parser_setup.add_argument(
        "-c",
        "--peloton-app-config",
        dest="peloton_apps_config_path",
        help="the path of the peloton apps config",
    )

    parser_setup.add_argument(
        "-i",
        "--peloton-image",
        dest="peloton_image",
        default=None,
        help="Docker image to use for Peloton. "
        + "If specified, overrides option -v",
    )

    parser_setup.add_argument(
        "-s",
        "--agent-number",
        type=int,
        dest="agent_number",
        help="the number of Mesos agent in the vcluster",
    )

    parser_setup.add_argument(
        "--no-respool",
        action="store_true",
        dest="skip_respool",
        help="If set, default resource-pool will not be created",
    )
    parser_setup.add_argument(
        "--clean",
        action="store_true",
        dest="clean_setup",
        help="Clean up old instance(s) of vcluster before creating a new one",
    )

    # Subparser for the 'teardown' command
    parser_teardown = subparsers.add_parser(
        "teardown", help="shut down a virtual cluster"
    )
    parser_teardown.add_argument(
        "-o", "--option", dest="option", help="option of action"
    )
    parser_teardown.add_argument(
        "--remove",
        action="store_true",
        dest="remove",
        help="Delete Peloton jobs as part of tearing down vcluster",
    )
    # Subparser for the 'peloton' command
    parser_peloton = subparsers.add_parser(
        "peloton", help="start peloton on a Mesos cluster"
    )

    parser_peloton.add_argument(
        "-k", "--zk", dest="zk", help="the host of virtual zk"
    )

    parser_peloton.add_argument(
        "-s",
        "--agent-number",
        type=int,
        dest="agent_number",
        help="the number of Mesos agent in the vcluster",
    )

    parser_peloton.add_argument(
        "-v",
        "--peloton-version",
        nargs="?",
        dest="peloton_version",
        help="The image version for Peloton",
    )

    parser_peloton.add_argument(
        "-i",
        "--peloton-image",
        dest="peloton_image",
        default=None,
        help="Docker image to use for Peloton. "
        + "If specified, overrides option -v",
    )

    subparsers.add_parser(
        "parameters", help="get the parameters of the vCluster"
    )

    return parser.parse_args()


def main():
    args = parse_arguments()
    if not args.config_file:
        args.config_file = os.path.join(
            os.path.dirname(os.path.abspath(__file__)),
            "config",
            "default.yaml",
        )
    config = load_config(args.config_file)

    vcluster = VCluster(
        config,
        args.label_name,
        args.zookeeper_server,
        args.respool_path,
        auth_type=args.auth_type,
        auth_file=args.auth_config_file,
    )

    command = args.command

    if command == "mesos":
        agent_number = args.agent_number
        vcluster.start_mesos(agent_number)

    elif command == "mesos-slave":
        agent_number = args.agent_number
        zk = args.zk.split(":")
        if len(zk) != 2:
            raise Exception("Invalid zk")
        vcluster.start_mesos_slave(args.zk, agent_number)

    elif command == "mesos-master":
        zk = args.zk.split(":")
        if len(zk) != 2:
            raise Exception("Invalid zk")
        vcluster.start_mesos_master(args.zk)

    elif command == "peloton":
        zk = args.zk.split(":")
        if len(zk) != 2:
            raise Exception("Invalid zk")
        vcluster.start_peloton(
            args.zk,
            args.agent_number,
            args.peloton_version,
            skip_respool=args.skip_respool,
            peloton_image=args.peloton_image,
            peloton_apps_config=args.peloton_apps_config_path,
        )

    elif command == "setup":
        agent_number = args.agent_number
        peloton_version = args.peloton_version
        if args.clean_setup:
            vcluster.teardown()
        vcluster.start_all(
            agent_number,
            peloton_version,
            skip_respool=args.skip_respool,
            peloton_image=args.peloton_image,
            peloton_apps_config=args.peloton_apps_config_path,
        )

    elif command == "teardown":
        option = args.option
        if option == "slave":
            vcluster.teardown_slave(remove=args.remove)
        elif option == "peloton":
            vcluster.teardown_peloton(remove=args.remove)
        else:
            vcluster.teardown(remove=args.remove)


if __name__ == "__main__":
    main()