launcher.py

import os
import shutil
import subprocess
import sys
from argparse import REMAINDER, ArgumentParser


class Launcher:
    def parse_args():
        """
        Helper function parsing the command line options
        @retval ArgumentParser
        """
        parser = ArgumentParser(description="command-launch a Python script with quantization auto-enabled")

        parser.add_argument("-o", "--opt", type=str, default="", help="optimization feature to enable")

        parser.add_argument("-a", "--approach", type=str, default="auto", help="quantization approach (strategy)")

        parser.add_argument("--config", type=str, default="", help="quantization configuration file path")

        parser.add_argument(
            "-b", "--bench", default=False, action="store_true", help="conduct auto_quant benchmark instead of enable"
        )

        parser.add_argument(
            "-e", "--enable", default=False, action="store_true", help="only do enable, not overwrite or run program"
        )

        # positional
        parser.add_argument(
            "script",
            type=str,
            help="The full path to the script to be launched. " "followed by all the arguments for the script",
        )

        # script args
        parser.add_argument("script_args", nargs=REMAINDER)
        return parser.parse_args()

    def execute(
        args,
        use_modular=False,
        modular_pattern={},
        use_inc=True,
    ):
        # copy user entry script (main.py -> main_optimized.py)
        script_copied = args.script[:-3] + "_optimized.py"
        shutil.copy(args.script, script_copied)

        if not args.bench:  # "enable and run" or "only enable"
            # optimize on copied script with Neural Coder
            from neural_coder import enable

            if args.opt == "":
                if args.approach == "static":
                    args.opt = "pytorch_inc_static_quant_fx"
                if args.approach == "static_ipex":
                    args.opt = "pytorch_inc_static_quant_ipex"
                if args.approach == "static_ipex_xpu":
                    args.opt = "pytorch_inc_static_quant_ipex_xpu"
                if args.approach == "dynamic":
                    args.opt = "pytorch_inc_dynamic_quant"
                if args.approach == "auto":
                    args.opt = "inc_auto"
                features = [args.opt]
            else:
                features = args.opt.split(",")

            # modular design
            modular_item = ""
            if use_modular:
                modular_item = modular_pattern[args.opt]

            # execute optimization enabling
            enable(
                code=script_copied,
                features=features,
                overwrite=True,
                use_modular=use_modular,
                modular_item=modular_item,
                use_inc=use_inc,
            )

            if not args.enable:  # enable and run
                # execute on copied script, which has already been optimized
                cmd = []

                cmd.append(sys.executable)  # "/xxx/xxx/python"
                cmd.append("-u")
                cmd.append(script_copied)
                cmd.extend(args.script_args)

                cmd = " ".join(cmd)  # list convert to string

                process = subprocess.Popen(cmd, env=os.environ, shell=True)  # nosec
                process.wait()
        else:  # auto_quant
            from neural_coder import auto_quant

            auto_quant(
                code=script_copied,
                args=" ".join(args.script_args),  # convert list of strings to a single string
                use_inc=use_inc,
            )