"""
this is atk base module
"""
import importlib
import logging
import os
import time
import sys
from types import SimpleNamespace
import copy
from collections import OrderedDict
from typing import List
import click
from atk.common.log import log_init
from atk.common.output_manager import OutputManager
from atk.common.file_check import FileCheck
from atk.common.utils import chmod_output_path
class AtkCommand(click.Command):
"""Customized command for atk."""
def format_options(self, ctx, formatter):
"""Write all the options into the formatter if they exist."""
opts = OrderedDict()
for param in self.get_params(ctx):
ret_val = param.get_help_record(ctx)
if ret_val is not None:
if hasattr(param, "help_group") and param.help_group:
opts.setdefault(str(param.help_group), []).append(ret_val)
else:
opts.setdefault("Options", []).append(ret_val)
for name, opts_group in opts.items():
with formatter.section(name):
formatter.write_dl(opts_group)
class AtkOption(click.Option):
"""Customized option for atk."""
def __init__(self, *args, **kwargs):
"""Initialize a Celery option."""
self.help_group = kwargs.pop("help_group", None)
self.default_value_from_context = kwargs.pop("default_value_from_context", None)
super().__init__(*args, **kwargs)
def get_default(self, ctx, *args, **kwargs):
if self.default_value_from_context:
self.default = ctx.obj[self.default_value_from_context]
return super().get_default(ctx, *args, **kwargs)
class DynamicCLI(click.Group):
TERMINAL_COMMANDS = {"pytorch", "aclnn"}
def list_commands(self, ctx):
commands = []
import atk
commands_dir = os.path.join(os.path.dirname(atk.__file__), "bin")
for file in os.listdir(commands_dir):
if "base.py" in file:
continue
if file.endswith(".py") and not file.startswith("_"):
commands.append(file[:-3])
commands.sort()
return commands
def get_command(self, ctx, name):
try:
module = importlib.import_module(f"atk.bin.{name}")
if name in self.TERMINAL_COMMANDS:
ctx.meta["atk_terminal_command"] = name
return module.cli
except ImportError as exc:
logging.exception(exc)
raise click.UsageError(f"atk 子命令 '{name}' 不存在")
def parse_args(self, ctx, args):
terminal_commands = self.TERMINAL_COMMANDS.intersection(args)
if terminal_commands:
original_chain = self.chain
self.chain = False
try:
return super().parse_args(ctx, args)
finally:
self.chain = original_chain
return super().parse_args(ctx, args)
def resolve_command(self, ctx, args):
cmd_name, cmd, rest = super().resolve_command(ctx, args)
if cmd_name in self.TERMINAL_COMMANDS:
return cmd_name, cmd, rest
return cmd_name, cmd, rest
def invoke(self, ctx):
protected_args = getattr(ctx, "_protected_args", None)
if protected_args is None:
protected_args = getattr(ctx, "protected_args", [])
args = [*protected_args, *ctx.args]
if args and args[0] in self.TERMINAL_COMMANDS:
ctx.args = []
if hasattr(ctx, "_protected_args"):
ctx._protected_args = []
elif hasattr(ctx, "protected_args"):
ctx.protected_args = []
with ctx:
cmd_name, cmd, sub_args = self.resolve_command(ctx, args)
if cmd is None:
raise ValueError("Command not found")
ctx.invoked_subcommand = cmd_name
super(click.Group, self).invoke(ctx)
sub_ctx = cmd.make_context(cmd_name, sub_args, parent=ctx)
with sub_ctx:
return sub_ctx.command.invoke(sub_ctx)
return super().invoke(ctx)
class OptionStrToList:
@staticmethod
def list_to_type(value: List, dist_type):
try:
value = [dist_type(_value) for _value in value]
except Exception as e:
raise click.BadParameter(f"Value must be List[{str(dist_type)}]") from e
return value
@classmethod
def str_to_str_list(cls, ctx, param, value):
if not value:
return value
value = value.split(",")
return value
@classmethod
def str_to_int_list(cls, ctx, param, value):
if not value:
return value
value = value.split(",")
value = cls.list_to_type(value, int)
return value
class Task:
def __init__(self, ctx, **kwargs):
params = copy.deepcopy(kwargs)
if ctx.obj and ctx.obj.get("nodes_list"):
params["nodes_list"] = ctx.obj["nodes_list"]
else:
params["nodes_list"] = {}
args = SimpleNamespace(**params)
self.args = args
self.file_check()
self.output_manager = None
self.init()
logging.info(f"start atk task: {args}")
def get_output_head_name(self):
if self.args.case:
return os.path.splitext(os.path.basename(self.args.case))[0]
elif self.args.fast_case:
return os.path.splitext(os.path.basename(self.args.fast_case))[0]
else:
return None
def init(self):
"""
初始化
"""
output_path = OutputManager.make_output_path(
nodes_list=self.args.nodes_list, nodes_yaml=self.args.nodes
)
case_file_name = self.get_output_head_name()
task_file_name = OutputManager.get_task_file_name(head_name=case_file_name)
self.output_manager = OutputManager(
output_path=output_path, task_file_name=task_file_name
)
log_init(
file_path=self.output_manager.default_log_file,
level=self.args.log,
)
chmod_output_path(self.output_manager.default_log_file)
logging.info("task_output_path: %s", self.output_manager.task_output_path)
def file_check(self):
FileCheck.check(self.args.case)
FileCheck.check(self.args.nodes)
FileCheck.check(self.args.plugin_path)
FileCheck.check(self.args.input_data)
for node in self.args.nodes_list:
FileCheck.check(node.get("bm_file"))
FileCheck.check(node.get("output_path"))
def run(self):
"""
执行入口
"""
self.run_task()
def run_task(self):
"""
执行任务
"""
start_time = time.time()
try:
from atk.tasks.main import MainProcess
process = MainProcess(self.args)
try:
process.init_all()
process.run()
except KeyboardInterrupt:
logging.error("received Ctrl+C. start stop!")
logging.error("atk task failed!")
except BaseException as e:
logging.exception(f"Error: {str(e)}")
logging.error("atk task failed!")
else:
logging.info("atk task success!")
finally:
process.clear()
finally:
logging.info(f"atk task finish. cost time: %d s", time.time() - start_time)