class WhichConfigApp(Application):
"""Application class for config which command."""
interface = Interface(PROGRAM, USAGE, HELP)
varpath: str = None
interface.add_argument('varpath', metavar='VAR')
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:
"""Business logic for `config which`."""
try:
site = config.which(*self.varpath.split('.'))
except KeyError:
log.critical(f'"{self.varpath}" not found')
return
if site in ('default', 'env'):
print(site)
else:
path = CONF_PATH[site]
print(f'{site}: {path}')
class SetConfigApp(Application):
"""Application class for config set command."""
interface = Interface(PROGRAM, USAGE, HELP)
varpath: str = None
interface.add_argument('varpath', metavar='VAR')
value: str = None
interface.add_argument('value', type=smart_type)
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:
"""Business logic for `config set`."""
site = SITE
config_path = get_site()
for key in ('local', 'user', 'system'):
if getattr(self, key) is True:
site = key
init_config(key)
config_path = get_site(key)
if not os.path.exists(config_path):
raise RuntimeError(f'{config_path} does not exist')
# parse variable specification
if '.' not in self.varpath:
raise ArgumentError('missing section in variable path')
section, *subsections, variable = self.varpath.split('.')
config = {section: {}}
config_section = config[section]
for subsection in subsections:
if subsection not in config_section:
config_section[subsection] = dict()
config_section = config_section[subsection]
config_section[variable] = self.value
update_config(site, config)
class ConfigApp(ApplicationGroup):
"""Application class for config command group."""
interface = Interface(PROGRAM, USAGE, HELP)
interface.add_argument('command')
command = None
commands = COMMANDS
class GraphApp(Application):
"""Graph benchmark data."""
interface = Interface('pybench graph', graph_usage, graph_help)
source: str = '-'
interface.add_argument('source')
outpath: Optional[str] = None
interface.add_argument('-o', '--output', dest='outpath', default=outpath)
label_benchmark: Optional[str] = None
interface.add_argument('--label-benchmark', default=label_benchmark)
label_build: Optional[str] = None
interface.add_argument('--label-build', default=label_build)
label_version: Optional[str] = None
interface.add_argument('--label-version', default=label_version)
data: LogData
graph: PerfChart
def run(self) -> None:
"""List benchmarks."""
self.load_data()
self.graph = PerfChart(self.data, **self.get_labels())
self.graph.draw()
self.graph.save(self.outpath)
def load_data(self) -> None:
"""Load data from source."""
if self.source == '-':
self.data = LogData.from_io(sys.stdin)
else:
self.data = LogData.from_local(self.source)
def get_labels(self) -> Dict[str, str]:
"""Derive labels from data or from explicit assignment."""
return {
'label_benchmark': self.label_benchmark
or self.data.benchmark_name,
'label_build': self.label_build or self.get_python_info(),
'label_version': self.label_version or self.get_numpy_info()
}
@staticmethod
def get_python_info() -> str:
"""Load Python build information."""
implementation = python_implementation()
version = python_version()
return f'{implementation} {version}'
@staticmethod
def get_numpy_info() -> str:
"""Load NumPy build information."""
return f'NumPy {np.__version__}'
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 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 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 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 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 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 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 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 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)