class FacilitySet(Application): """Set facility profile data.""" interface = Interface(PROGRAM, USAGE, HELP) source: str = '-' interface.add_argument('source') profile: Optional[str] = None interface.add_argument('--profile', default=profile) debug: bool = False verbose: bool = False logging_interface = interface.add_mutually_exclusive_group() logging_interface.add_argument('-d', '--debug', action='store_true') logging_interface.add_argument('-v', '--verbose', action='store_true') syslog: bool = False interface.add_argument('--syslog', action='store_true') exceptions = { RuntimeError: functools.partial(log_and_exit, logger=log.critical, status=exit_status.runtime_error), FacilityNotFound: functools.partial(log_and_exit, logger=log.critical, status=exit_status.runtime_error), DatabaseError: functools.partial(log_and_exit, logger=log.critical, status=exit_status.runtime_error), } def run(self) -> None: """Set facility profiles.""" if self.source == '-': data = sys.stdin.read() else: with open(self.source, mode='r') as source: data = source.read() profile = Facility.from_dict(json.loads(data)) profile.to_database() def __enter__(self) -> FacilitySet: """Initialize resources.""" cli_setup(self) database.connect(profile=self.profile) return self def __exit__(self, *exc) -> None: """Release resources.""" database.disconnect()
class RefittControllerApp(Application): """Application class for the refitt daemon controller.""" interface = Interface(PROGRAM, USAGE, HELP) interface.add_argument('-v', '--version', version=__version__, action='version') action: str = None interface.add_argument('action', choices=ACTIONS) exceptions = { RuntimeError: functools.partial(log_exception, logger=log.critical, status=exit_status.runtime_error), ConnectionRefusedError: daemon_unavailable, } def run(self) -> None: """Delegate action.""" if self.action in ['start', 'status']: action = getattr(self, f'run_{self.action}') action() else: self.run_action() @staticmethod def run_start() -> None: """Start the daemon.""" subprocess.run(['refittd', '--all', '--daemon']) @staticmethod def run_status() -> None: """Show the status of daemon services.""" # retrieve status from daemon (returns dict) with DaemonClient() as daemon: info = daemon.status for service, status in info.items(): alive = status.pop('alive') color = ANSI_GREEN if alive else ANSI_RED state = 'alive' if alive else 'dead' pid = status.pop('pid') print(f'{color}● {service}: {state} ({pid}){ANSI_RESET}') for key, value in status.items(): print(f'{key:>10}: {value}') def run_action(self) -> None: """Run the action.""" with DaemonClient() as daemon: daemon.request(self.action)
class RunApp(Application): """Run benchmark.""" interface = Interface('pybench run', run_usage, run_help) name: str = None interface.add_argument('name', choices=list(benchmark.listing)) args: List[Any] = None interface.add_argument('args', nargs='*', type=coerce_type, default=[]) repeat: int = 1 interface.add_argument('-n', '--repeat', type=int, default=repeat) spacing: float = 1.0 interface.add_argument('-s', '--spacing', type=float, default=spacing) monitor_cpu: bool = False interface.add_argument('-c', '--monitor-cpu', action='store_true') monitor_memory: bool = False interface.add_argument('-m', '--monitor-memory', action='store_true') resolution: float = 1.0 interface.add_argument('-r', '--resolution', type=float, default=resolution) cpu_thread: Optional[CPUResource] = None mem_thread: Optional[MemoryResource] = None exceptions = { BenchmarkError: functools.partial(log_exception, status=exit_status.runtime_error) } def run(self) -> None: """Run requested benchmark.""" self.setup_telemetry() self.run_benchmark() if self.cpu_thread or self.mem_thread: time.sleep(self.spacing) # NOTE: pause to get one more metric def setup_telemetry(self) -> None: """Start telemetry threads if requested.""" self.cpu_thread = None if not self.monitor_cpu else CPUResource.new( self.resolution) self.mem_thread = None if not self.monitor_memory else MemoryResource.new( self.resolution) def run_benchmark(self) -> None: """Setup and initiate benchmark.""" benchmark_type = benchmark.listing.get(self.name) benchmark_type(self.repeat, self.spacing, *self.args).run()
class ConfigApp(ApplicationGroup): """Application class for config command group.""" interface = Interface(PROGRAM, USAGE, HELP) interface.add_argument('command') command = None commands = {'get': get.GetConfigApp, 'set': set.SetConfigApp, 'edit': edit.EditConfigApp, 'which': which.WhichConfigApp, }
class DatabaseApp(ApplicationGroup): """Application class for database command group.""" interface = Interface(PROGRAM, USAGE, HELP) interface.add_argument('command') command = None commands = {'init': init.InitDatabaseApp, 'check': check.CheckDatabaseApp, 'query': query.QueryDatabaseApp, }
class ServiceApp(ApplicationGroup): """Application class for service command group.""" interface = Interface(PROGRAM, USAGE, HELP) interface.add_argument('command') command = None commands = { 'api': api.WebApp, 'stream': stream.StreamApp, 'test': test.TestApp, }
class CPUMemory(Application): """Monitor CPU memory utilization.""" ALLOW_NOARGS = True # no usage behavior interface = Interface(PROGRAM, USAGE, HELP) sample_rate: float = 1 interface.add_argument('-s', '--sample-rate', type=float, default=sample_rate) human_readable: bool = False interface.add_argument('-H', '--human-readable', action='store_true') memory_actual: bool = False memory_percent: bool = True memory_interface = interface.add_mutually_exclusive_group() memory_interface.add_argument('--actual', action='store_true', dest='memory_actual') memory_interface.add_argument('--percent', action='store_true', dest='memory_percent') format_csv: bool = False format_plain: bool = True format_interface = interface.add_mutually_exclusive_group() format_interface.add_argument('--csv', action='store_true', dest='format_csv') format_interface.add_argument('--plain', action='store_true', dest='format_plain') no_header: bool = False interface.add_argument('--no-header', action='store_true') exceptions = { RuntimeError: functools.partial(log_and_exit, logger=log.critical, status=exit_status.runtime_error), } def run(self) -> None: """Run monitor.""" if not self.format_csv and self.no_header: raise ArgumentError('--no-header only applies to --csv mode.') mem_attr = 'used' if self.memory_actual else 'percent' if not self.memory_actual and self.human_readable: raise ArgumentError('"--human-readable" only applies to "--actual" values.') log.handlers[0] = PLAIN_HANDLER if self.format_csv: log.handlers[0] = CSV_HANDLER if not self.no_header: print(f'timestamp,hostname,resource,memory_{mem_attr}') formatter = functools.partial(format_size, scale_units=self.human_readable) while True: time.sleep(self.sample_rate) value = getattr(psutil.virtual_memory(), mem_attr) log.debug(formatter(value))
class Group(ApplicationGroup): interface = Interface(APP, APP_USAGE, APP_HELP) commands = {CMD1: Command1, CMD2: Command2} command: str = None interface.add_argument('command') version: str = '1.2.3' interface.add_argument('-v', '--version', action='version', version=version)
class GPUPercent(Application): """Monitor GPU percent utilization.""" ALLOW_NOARGS = True interface = Interface(PROGRAM, USAGE, HELP) sample_rate: float = 1 interface.add_argument('-s', '--sample-rate', type=float, default=sample_rate) format_plain: bool = True format_csv: bool = False format_interface = interface.add_mutually_exclusive_group() format_interface.add_argument('--plain', action='store_true', dest='format_plain') format_interface.add_argument('--csv', action='store_true', dest='format_csv') no_header: bool = False interface.add_argument('--no-header', action='store_true') exceptions = { RuntimeError: functools.partial(log_and_exit, logger=log.critical, status=exit_status.runtime_error), } def run(self) -> None: """Run monitor.""" if not self.format_csv and self.no_header: raise ArgumentError('--no-header only applies to --csv mode.') log.handlers[0] = PLAIN_HANDLER if self.format_csv: log.handlers[0] = CSV_HANDLER if not self.no_header: print('timestamp,hostname,resource,gpu_id,gpu_percent') smi = SMIData() while True: time.sleep(self.sample_rate) for gpu_id, gpu_percent in smi.percent.items(): log.debug(f'[{gpu_id}] {gpu_percent}')
class NamesGeneratorApp(Application): """Top-level application class for `generate_name` console application.""" interface = Interface(PROGRAM, USAGE, HELP) interface.add_argument('-v', '--version', action='version', version=__version__) style: str = 'underscore' interface.add_argument('-s', '--style', default=style, choices=list(_formatting_methods)) # run even without arguments (do not print usage) ALLOW_NOARGS = True def run(self) -> None: """Generate a random name and print it.""" print(generate_name(style=self.style))
class Command2(Application): interface = Interface(CMD2, CMD2_USAGE, CMD2_HELP) arg: str = None interface.add_argument('arg') option: int = 42 # arbitrary interface.add_argument('-o', '--option', type=int, default=option) debug: bool = False interface.add_argument('-d', '--debug', action='store_true') def run(self) -> None: print(f'arg: {self.arg}, option: {self.option}, debug: {self.debug}')
class PyBenchApp(ApplicationGroup): """Run benchmark.""" interface = Interface('pybench', app_usage, app_help) interface.add_argument('-v', '--version', action='version', version=__version__) interface.add_argument('command') command = None commands = { 'run': RunApp, 'list': ListApp, 'graph': GraphApp, }
class RefittApp(ApplicationGroup): """Top-level application class for Refitt.""" interface = Interface(PROGRAM, USAGE, HELP) interface.add_argument('command') interface.add_argument('-v', '--version', action='version', version=__version__) interface.add_argument('--ascii-art', action='version', version=__ascii_art__) command = None commands = {'auth': auth.AuthApp, 'login': login.LoginApp, 'config': config.ConfigApp, 'database': database.DatabaseApp, 'service': service.ServiceApp, 'api': api.WebApp, }
class SharedGroup(ApplicationGroup): interface = Interface(APP, APP_USAGE, APP_HELP) commands = {CMD1: Command1, CMD2: Command3} ALLOW_PARSE = True command: str = None interface.add_argument('command') version: str = '1.2.3' interface.add_argument('-v', '--version', action='version', version=version) opt: bool = False interface.add_argument('-g', '--global', action='store_true', dest='opt')
class StreamKit(ApplicationGroup): """Application class for streamkit entry-point.""" interface = Interface('streamkit', USAGE, HELP) interface.add_argument('-v', '--version', version=__version__, action='version') interface.add_argument('command') command = None commands = { 'config': config.ConfigApp, 'publish': publish.PublisherApp, 'subscribe': subscribe.SubscriberApp, 'database': database.DatabaseApp }
class SubscriberApp(Application): """Application class for subscriber.""" interface = Interface(PROGRAM, USAGE, HELP) name: str = None interface.add_argument('name') topics: List[str] = [] interface.add_argument('topics', nargs='+') batch_size: int = DEFAULT_BATCHSIZE interface.add_argument('-b', '--batch-size', type=int, default=batch_size) poll_interval: float = DEFAULT_POLL interface.add_argument('-p', '--poll-interval', type=float, default=poll_interval) timeout: float = DEFAULT_TIMEOUT interface.add_argument('-t', '--timeout', type=float, default=timeout) exceptions = { RuntimeError: partial(log_exception, log=log.critical, status=exit_status.runtime_error), ConfigurationError: partial(log_exception, log=log.critical, status=exit_status.bad_config), } def run(self) -> None: """Print messages as they arrive.""" with Subscriber(self.name, self.topics, batchsize=self.batch_size, poll=self.poll_interval, timeout=self.timeout) as stream: for message in stream: print( f'{message.time} {message.host} {message.topic} {message.level} {message.text}' )
class Slack(Application): """Post slack messages and files.""" interface = Interface(PROGRAM, USAGE, HELP) channel: str = None interface.add_argument('channel') message: str = '@-' interface.add_argument('message', default=message) botname: str = 'refitt' interface.add_argument('-f', '--from', default=botname, dest='botname') attachment: List[str] = None interface.add_argument('-a', '--attach', dest='attachment') debug: bool = False verbose: bool = False logging_interface = interface.add_mutually_exclusive_group() logging_interface.add_argument('-d', '--debug', action='store_true') logging_interface.add_argument('-v', '--verbose', action='store_true') syslog: bool = False interface.add_argument('--syslog', action='store_true') exceptions = { RuntimeError: functools.partial(log_and_exit, logger=log.critical, status=exit_status.runtime_error), } def run(self) -> None: """Create and send email.""" raise RuntimeError('slack integration is not implemented') def __enter__(self) -> Slack: """Initialize resources.""" cli_setup(self) return self def __exit__(self, *exc) -> None: """Release resources."""
class GPUDevice(Application): """Monitor GPU resources.""" interface = Interface(PROGRAM, USAGE, HELP) resource: str = None interface.add_argument('resource') exceptions = { CompletedCommand: (lambda exc: int(exc.args[0])), } def run(self) -> None: """Show usage/help/version or defer to group.""" if self.resource in RESOURCES: status = RESOURCES[self.resource].main(sys.argv[3:]) raise CompletedCommand(status) else: raise ArgumentError(f'"{self.resource}" is not a GPU resource.')
class InitDatabaseApp(Application): """Application class for database init entry-point.""" interface = Interface(PROGRAM, USAGE, HELP) ALLOW_NOARGS = True drop_tables: bool = False interface.add_argument('--drop', action='store_true', dest='drop_tables') load_core: bool = False load_test: bool = False load_interface = interface.add_mutually_exclusive_group() load_interface.add_argument('--core', action='store_true', dest='load_core') load_interface.add_argument('--test', action='store_true', dest='load_test') exceptions = { RuntimeError: partial(log_exception, logger=log.critical, status=exit_status.runtime_error), DatabaseError: partial(log_exception, logger=log.critical, status=exit_status.runtime_error), ConfigurationError: partial(log_exception, logger=log.critical, status=exit_status.bad_config), } def run(self) -> None: """Business logic of command.""" if self.drop_tables: drop_database() init_database() if self.load_core: load_records('core') if self.load_test: load_records('test')
class EditConfigApp(Application): """Application class for config edit command.""" interface = Interface(PROGRAM, USAGE, HELP) local: bool = False user: bool = False system: bool = False site_interface = interface.add_mutually_exclusive_group() site_interface.add_argument('--local', action='store_true') site_interface.add_argument('--user', action='store_true') site_interface.add_argument('--system', action='store_true') exceptions = { RuntimeError: functools.partial(log_exception, log=log.critical, status=exit_status.runtime_error), ConfigurationError: functools.partial(log_exception, log=log.critical, status=exit_status.bad_config), } def run(self) -> None: """Open editor for configuration.""" site = SITE config_path = None for key in ('local', 'user', 'system'): if getattr(self, key) is True: config_path = CONF_PATH[key] site = key if not os.path.exists(config_path): log.info(f'{config_path} does not exist - initializing') init_config(site) if 'EDITOR' not in os.environ: raise RuntimeError('EDITOR must be set') editor = os.environ['EDITOR'] subprocess.run([editor, config_path])
class EditConfigApp(Application): """Application class for config edit command.""" interface = Interface(PROGRAM, USAGE, HELP) local: bool = False user: bool = False system: bool = False site_interface = interface.add_mutually_exclusive_group() site_interface.add_argument('--local', action='store_true') site_interface.add_argument('--user', action='store_true') site_interface.add_argument('--system', action='store_true') exceptions = { RuntimeError: partial(log_exception, logger=log.critical, status=exit_status.runtime_error), ConfigurationError: partial(log_exception, logger=log.critical, status=exit_status.bad_config), } def run(self) -> None: """Open editor for configuration.""" site = SITE path = PATH[site].config for key in ('local', 'user', 'system'): if getattr(self, key) is True: site = key path = PATH[site].config dirname = os.path.dirname(path) if not os.path.exists(dirname): os.makedirs(dirname, exist_ok=True) if 'EDITOR' not in os.environ: raise RuntimeError('EDITOR must be set') editor = os.environ['EDITOR'] run([editor, path])
class ResourceMonitor(Application): """Application class for resource-monitor.""" interface = Interface(PROGRAM, USAGE, HELP) interface.add_argument('-v', '--version', version=__version__, action='version') device: str = None interface.add_argument('device') exceptions = { CompletedCommand: (lambda exc: int(exc.args[0])), } def run(self) -> None: """Show usage/help/version or defer to group.""" if self.device in DEVICES: status = DEVICES[self.device].main(sys.argv[2:3]) raise CompletedCommand(status) else: raise ArgumentError(f'"{self.device}" is not a device.')
class FacilityGroup(Application): """Manage facility profiles.""" interface = Interface(PROGRAM, USAGE, HELP) command: str = None interface.add_argument('command') exceptions = { CompletedCommand: (lambda exc: int(exc.args[0])), } def run(self) -> None: """Show usage/help/version or defer to command.""" if self.command in COMMANDS: status = COMMANDS[self.command].main(sys.argv[4:]) raise CompletedCommand(status) else: raise ArgumentError(f'"{self.command}" is not a command.')
class ListApp(Application): """List available benchmarks.""" interface = Interface('pybench list', list_usage, list_help) ALLOW_NOARGS = True pattern: re.Pattern = re.compile(f'.*') interface.add_argument('pattern', nargs='?', type=re.compile, default=pattern) long_mode: bool = False interface.add_argument('-l', '--long', action='store_true', dest='long_mode') def run(self) -> None: """List benchmarks.""" for name, benchmark_type in benchmark.listing.items(): if self.pattern.match(name): self.output(name, benchmark_type) @property def output(self) -> Callable[[str, Type[Benchmark]], None]: return self.detailed_output if self.long_mode else self.basic_output @staticmethod def basic_output(name: str, benchmark_type: Type[Benchmark]) -> None: name = f'{name:<17}' if not sys.stdout.isatty( ) else f'{ANSI_BLUE}{name:<18}{ANSI_RESET}' print(f'{name} {benchmark_type.__doc__}') @staticmethod def detailed_output(name: str, benchmark_type: Type[Benchmark]) -> None: name = f'{name:<18}' if not sys.stdout.isatty( ) else f'{ANSI_BLUE}{name:<18}{ANSI_RESET}' print( f'{name} {benchmark_type.annotation:>27} {benchmark_type.__doc__}' )
class PublisherApp(Application): """Application class for publisher.""" interface = Interface(PROGRAM, USAGE, HELP) topic: str = None interface.add_argument('topic') level: str = None interface.add_argument('level') source: TextIOWrapper = stdin interface.add_argument('source', nargs='?', type=FileType(mode='r'), default=source) batch_size: int = DEFAULT_BATCHSIZE interface.add_argument('-b', '--batch-size', type=int, default=batch_size) timeout: float = DEFAULT_TIMEOUT interface.add_argument('-t', '--timeout', type=float, default=timeout) exceptions = { RuntimeError: partial(log_exception, log=log.critical, status=exit_status.runtime_error), ConfigurationError: partial(log_exception, log=log.critical, status=exit_status.bad_config), } def run(self) -> None: """Send source lines of text to publisher.""" with Publisher(topic=self.topic, level=self.level, batchsize=self.batch_size, timeout=self.timeout) as publisher: for message in self.source: publisher.write(message.strip())
class InitConfigApp(Application): """Application class for config initialization/.""" interface = Interface(PROGRAM, USAGE, HELP) user: bool = False system: bool = False site_interface = interface.add_mutually_exclusive_group() site_interface.add_argument('--user', action='store_true') site_interface.add_argument('--system', action='store_true') exceptions = { RuntimeError: functools.partial(log_exception, log=log.critical, status=exit_status.runtime_error), PermissionError: functools.partial(log_exception, log=log.critical, status=exit_status.runtime_error), ConfigurationError: functools.partial(log_exception, log=log.critical, status=exit_status.bad_config), } def run(self) -> None: """Business logic of command.""" self.check_exists() init_config(self.cfg_site) @property def cfg_site(self) -> str: """Either 'system' or 'user'.""" return 'system' if self.system else 'user' @property def cfg_path(self) -> str: """The path to the relevant configuration file.""" return CONF_PATH[self.cfg_site] def check_exists(self) -> None: """Check to see if the configuration file already exists.""" if os.path.exists(self.cfg_path): raise RuntimeError(f'{self.cfg_path} exists')
class LoginApp(Application): """Application class for web login method.""" interface = Interface(PROGRAM, USAGE, HELP) ALLOW_NOARGS = True force: bool = False interface.add_argument('--force', action='store_true') exceptions = { RuntimeError: functools.partial(log_exception, logger=log.critical, status=exit_status.runtime_error), } def run(self) -> None: """Run login method.""" if 'key' in config.api and 'secret' in config.api and not self.force: log.info('Already logged in, use --force to get new credentials') else: request.login(force=self.force)
class Forecast(Application): """Run forecast on candidate.""" interface = Interface(PROGRAM, USAGE, HELP) # input file containing list of candidates/alerts source: str = '-' interface.add_argument('source') debug: bool = False verbose: bool = False logging_interface = interface.add_mutually_exclusive_group() logging_interface.add_argument('-d', '--debug', action='store_true') logging_interface.add_argument('-v', '--verbose', action='store_true') syslog: bool = False interface.add_argument('--syslog', action='store_true') exceptions = { RuntimeError: functools.partial(log_and_exit, logger=log.critical, status=exit_status.runtime_error), } def run(self) -> None: """Run Refitt pipeline.""" raise RuntimeError('not implemented') def __enter__(self) -> Forecast: """Initialize resources.""" cli_setup(self) return self def __exit__(self, *exc) -> None: """Release resources."""
class DemoApp(Application): """Fixture for Application class to run unit tests against.""" interface = Interface(DEMO_NAME, DEMO_USAGE, DEMO_HELP) arg_1: str = None interface.add_argument('arg_1') option_1: int = 4 # arbitrary interface.add_argument('-o', '--option', dest='option_1', type=int, default=option_1) debug_mode: bool = False interface.add_argument('-d', '--debug', dest='debug_mode', action='store_true') def run(self) -> None: """Business logic of application.""" pass
class InitDatabaseApp(Application): """Application class for database init entry-point.""" interface = Interface(PROGRAM, USAGE, HELP) ALLOW_NOARGS = True include_test_data: bool = False interface.add_argument('--test', action='store_true', dest='include_test_data') include_extensions: bool = False interface.add_argument('--ext', action='store_true', dest='include_extensions') echo: bool = False interface.add_argument('--echo', action='store_true') exceptions = { RuntimeError: functools.partial(log_exception, log=log.critical, status=exit_status.runtime_error), DatabaseError: functools.partial(log_exception, log=log.critical, status=exit_status.runtime_error), ConfigurationError: functools.partial(log_exception, log=log.critical, status=exit_status.bad_config), } def run(self) -> None: """Business logic of command.""" self.check_backend() self.config_engine() init() if self.include_extensions: init_extensions() if self.include_test_data: init_test_data() @staticmethod def check_backend() -> None: """Check if we are touching an external database.""" backend = db_config['backend'] if backend != 'sqlite': response = input( f'Connecting to {backend} database, proceed [Y/n]: ') if response in ('Yes', 'yes', 'Y', 'y'): pass elif response in ('No', 'no', 'N', 'n'): raise RuntimeError('Stopping now') else: raise RuntimeError('Response not understood') def config_engine(self) -> None: """Set `echo` parameter for engine.""" if self.echo: from ...database.core.engine import engine engine.echo = True for name, value in db_config.items(): log.debug(f'{name}: {value}')