Skip to content
This repository has been archived by the owner on Mar 3, 2023. It is now read-only.

Use click instead of argparse in tracker #3599

Merged
merged 1 commit into from
Oct 30, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions heron/tools/tracker/src/python/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@ pex_library(
exclude = ["main.py"],
),
reqs = [
"protobuf==3.8.0",
"tornado==4.0.2",
"click==7.1.2",
"javaobj-py3==0.4.1",
"networkx==2.4",
"protobuf==3.8.0",
"tornado==4.0.2",
],
deps = [
"//heron/common/src/python:common-py",
Expand Down
257 changes: 98 additions & 159 deletions heron/tools/tracker/src/python/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
# under the License.

''' main.py '''
import argparse
import logging
import os
import signal
import sys
Expand All @@ -29,14 +29,16 @@
from tornado.options import define
from tornado.httpclient import AsyncHTTPClient

import heron.tools.common.src.python.utils.config as common_config
import heron.common.src.python.utils.log as log
from heron.tools.common.src.python.utils import config as common_config
from heron.common.src.python.utils import log
from heron.tools.tracker.src.python import constants
from heron.tools.tracker.src.python import handlers
from heron.tools.tracker.src.python import utils
from heron.tools.tracker.src.python.config import Config, STATEMGRS_KEY
from heron.tools.tracker.src.python.tracker import Tracker

import click

Log = log.Log

class Application(tornado.web.Application):
Expand Down Expand Up @@ -93,173 +95,110 @@ def __init__(self, config):
def stop(self):
self.tracker.stop_sync()

# pylint: disable=protected-access
class _HelpAction(argparse._HelpAction):
""" HelpAction """
def __call__(self, parser, namespace, values, option_string=None):
parser.print_help()

# retrieve subparsers from parser
subparsers_actions = [
action for action in parser._actions
if isinstance(action, argparse._SubParsersAction)]

# there will probably only be one subparser_action,
# but better save than sorry
for subparsers_action in subparsers_actions:
# get all subparsers and print help
for choice, subparser in list(subparsers_action.choices.items()):
print("Subparser '{}'".format(choice))
print(subparser.format_help())

parser.exit()

# pylint: disable=bad-super-call
class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter):
""" Subcommand help formatter """
def _format_action(self, action):
parts = super(argparse.RawDescriptionHelpFormatter, self)._format_action(action)
if action.nargs == argparse.PARSER:
parts = "\n".join(parts.split("\n")[1:])
return parts


def add_titles(parser):
""" add titles """
parser._positionals.title = "Required arguments"
parser._optionals.title = "Optional arguments"
return parser


def add_arguments(parser):
""" add arguments """
default_config_file = os.path.join(
utils.get_heron_tracker_conf_dir(), constants.DEFAULT_CONFIG_FILE)

parser.add_argument(
'--config-file',
metavar='(a string; path to config file; default: "' + default_config_file + '")',
default=default_config_file)

parser.add_argument(
'--type',
metavar='(an string; type of state manager (zookeeper or file, etc.); example: ' \
+ str(constants.DEFAULT_STATE_MANAGER_TYPE) + ')',
choices=["file", "zookeeper"])

parser.add_argument(
'--name',
metavar='(an string; name to be used for the state manager; example: ' \
+ str(constants.DEFAULT_STATE_MANAGER_NAME) + ')')

parser.add_argument(
'--rootpath',
metavar='(an string; where all the states are stored; example: ' \
+ str(constants.DEFAULT_STATE_MANAGER_ROOTPATH) + ')')

parser.add_argument(
'--tunnelhost',
metavar='(an string; if ssh tunneling needs to be established to connect to it; example: ' \
+ str(constants.DEFAULT_STATE_MANAGER_TUNNELHOST) + ')')

parser.add_argument(
'--hostport',
metavar='(an string; only used to connect to zk, must be of the form \'host:port\';'\
' example: ' + str(constants.DEFAULT_STATE_MANAGER_HOSTPORT) + ')')

parser.add_argument(
'--port',
metavar='(an integer; port to listen; default: ' + str(constants.DEFAULT_PORT) + ')',
type=int,
default=constants.DEFAULT_PORT)

parser.add_argument(
'--verbose',
action='store_true')

return parser

def create_parsers():
""" create argument parser """
parser = argparse.ArgumentParser(
epilog='For detailed documentation, go to http://github.com/apache/incubator-heron',
usage="%(prog)s [options] [help]",
add_help=False)

parser = add_titles(parser)
parser = add_arguments(parser)

ya_parser = argparse.ArgumentParser(
parents=[parser],
formatter_class=SubcommandHelpFormatter,
add_help=False)

subparsers = ya_parser.add_subparsers(
title="Available commands")

help_parser = subparsers.add_parser(
'help',
help='Prints help',
add_help=False)

help_parser.set_defaults(help=True)

subparsers.add_parser(
'version',
help='Prints version',
add_help=True)

return parser, ya_parser

def define_options(port, config_file):
def define_options(port: int, config_file: str) -> None:
""" define Tornado global variables """
define("port", default=port)
define("config_file", default=config_file)

def create_tracker_config(namespace):

def create_tracker_config(config_file: str, stmgr_override: dict) -> dict:
# try to parse the config file if we find one
config_file = namespace["config_file"]
config = utils.parse_config_file(config_file)
if config is None:
Log.debug("Config file does not exists: %s" % config_file)
Log.debug(f"Config file does not exists: {config_file}")
config = {STATEMGRS_KEY:[{}]}

# update the config if we have any flags
config_flags = ["type", "name", "rootpath", "tunnelhost", "hostport"]
config_to_update = config[STATEMGRS_KEY][0]
for flag in config_flags:
value = namespace.get(flag, None)
if value is not None:
config_to_update[flag] = value

# update non-null options
config[STATEMGRS_KEY][0].update(
(k, v)
for k, v in stmgr_override.items()
if v is not None
)
return config

def main():
""" main """
# create the parser and parse the arguments
(parser, _) = create_parsers()
(args, remaining) = parser.parse_known_args()

if remaining == ['help']:
parser.print_help()
parser.exit()

elif remaining == ['version']:
def show_version(_, __, value):
if value:
common_config.print_build_info()
parser.exit()

elif remaining != []:
Log.error('Invalid subcommand')
sys.exit(1)

namespace = vars(args)

log.set_logging_level(namespace)
sys.exit(0)


@click.command()
@click.option(
"--version",
is_flag=True,
is_eager=True,
expose_value=False,
callback=show_version,
)
@click.option('--verbose', is_flag=True)
@click.option(
'--config-file',
help="path to a tracker config file",
default=os.path.join(utils.get_heron_tracker_conf_dir(), constants.DEFAULT_CONFIG_FILE),
show_default=True,
)
@click.option(
'--port',
type=int,
default=constants.DEFAULT_PORT,
show_default=True,
help="local port to serve on",
)
@click.option(
'--type',
"stmgr_type",
help=f"statemanager type e.g. {constants.DEFAULT_STATE_MANAGER_TYPE}",
type=click.Choice(choices=["file", "zookeeper"]),
)
@click.option(
'--name',
help=f"statemanager name e.g. {constants.DEFAULT_STATE_MANAGER_NAME}",
)
@click.option(
'--rootpath',
help=f"statemanager rootpath e.g. {constants.DEFAULT_STATE_MANAGER_ROOTPATH}",
)
@click.option(
'--tunnelhost',
help=f"statemanager tunnelhost e.g. {constants.DEFAULT_STATE_MANAGER_TUNNELHOST}",
)
@click.option(
'--hostport',
help=f"statemanager hostport e.g. {constants.DEFAULT_STATE_MANAGER_HOSTPORT}",
)
def cli(
config_file: str,
stmgr_type: str,
name: str,
rootpath: str,
tunnelhost: str,
hostport: str,
port: int,
verbose: bool,
) -> None:
"""
A HTTP service for serving data about clusters.

The statemanager's config from the given config file can be overrided using
options on this executable.

"""

log.configure(logging.DEBUG if verbose else logging.INFO)

# set Tornado global option
define_options(namespace['port'], namespace['config_file'])
define_options(port, config_file)

config = Config(create_tracker_config(namespace))
stmgr_override = {
"type": stmgr_type,
"name": name,
"rootpath": rootpath,
"tunnelhost": tunnelhost,
"hostport": hostport,
}
config = Config(create_tracker_config(config_file, stmgr_override))

# create Tornado application
application = Application(config)
Expand All @@ -278,15 +217,15 @@ def signal_handler(signum, frame):
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

Log.info("Running on port: %d", namespace['port'])
if namespace["config_file"]:
Log.info("Using config file: %s", namespace['config_file'])
Log.info("Using state manager:\n" + str(config))
Log.info("Running on port: %d", port)
if config_file:
Log.info("Using config file: %s", config_file)
Log.info(f"Using state manager:\n{config}")

http_server = tornado.httpserver.HTTPServer(application)
http_server.listen(namespace['port'])
http_server.listen(port)

tornado.ioloop.IOLoop.instance().start()

if __name__ == "__main__":
main()
cli() # pylint: disable=no-value-for-parameter
1 change: 0 additions & 1 deletion heron/tools/tracker/src/python/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,6 @@ def parse_config_file(config_file):
if not os.path.lexists(expanded_config_file_path):
return None

configs = {}
# Read the configuration file
with open(expanded_config_file_path, 'r') as f:
configs = yaml.load(f)
Expand Down