modify launch argparse for micro_perf.

byte_micro_perf/launch.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import argparse
+import pathlib
 import json
 import logging
 import os
@@ -21,12 +22,13 @@
 import subprocess
 import sys
 
-BYTE_MLPERF_ROOT = os.path.dirname(os.path.abspath(__file__))
-os.chdir(BYTE_MLPERF_ROOT)
+BYTE_MLPERF_ROOT = pathlib.Path(__file__).parent.absolute()
+CUR_DIR = pathlib.Path.cwd().absolute()
+os.chdir(str(BYTE_MLPERF_ROOT))
 sys.path.insert(0, BYTE_MLPERF_ROOT)
 
 logging.basicConfig(level=logging.INFO)
-log = logging.getLogger("lanuch")
+logger = logging.getLogger("lanuch")
 
 
 def parse_task(task_dir):
@@ -39,31 +41,40 @@ def parse_task(task_dir):
     return tasks
 
 
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--task", default="", help="The task going to be evaluted, refs to workloads/"
-    )
-    parser.add_argument(
-        "--task_dir", default="", help="The direcotry of tasks going to be evaluted, e.g., set to workloads"
-    )
+
+    # hardware config
     parser.add_argument(
         "--hardware_type",
-        default="GPU",
+        default="GPU", 
         help="The backend going to be evaluted, refs to backends/",
     )
     parser.add_argument(
         "--vendor_path",
-        required=False,
-        help="The hardware configs need to be loaded, refs to vendor_zoo/NVIDIA/A100-PCIe.json",
+        required=False, 
+        help="The hardware configs need to be loaded, refs to vendor_zoo/",
     )
+
+    # task config
     parser.add_argument(
-        "--compile_only",
-        action="store_true",
-        help="Task will stoped after compilation finished",
+        "--task_dir", 
+        default=str(BYTE_MLPERF_ROOT.joinpath("workloads").absolute()), 
+        help="The direcotry of tasks going to be evaluted, e.g., set to workloads"
     )
     parser.add_argument(
-        "--show_task_list", action="store_true", help="Print all task names"
+        "--task", 
+        default="all", 
+        help="The task going to be evaluted, refs to workloads/, default use all tasks in workloads/"
+    )
+
+    # list all supported task and hardware
+    parser.add_argument(
+        "--show_task_list", 
+        action="store_true", 
+        help="Print all available task names"
     )
     parser.add_argument(
         "--show_hardware_list",
@@ -72,37 +83,59 @@ def parse_task(task_dir):
     )
     args = parser.parse_args()
 
+
+    # show tasks
+    task_list = [file.stem for file in pathlib.Path(args.task_dir).iterdir()]
+    task_list.sort()
     if args.show_task_list:
-        logging.info("******************* Supported Task *******************")
-        for file in os.listdir("workloads"):
-            print(file[:-5])
+        logger.info("******************* Supported Task *******************")
+        print(task_list)        
+        exit(0)
+
+    # show hardwares
+    hardware_list = []
+    for file in BYTE_MLPERF_ROOT.joinpath("backends").iterdir():
+        if file.is_dir() and file.stem.startswith("_") is False:
+            hardware_list.append(file.stem)
     if args.show_hardware_list:
-        log.info("***************** Supported Hardware Backend *****************")
-        for file in os.listdir("backends"):
-            if not file.endswith(".py") and not file.startswith("_"):
-                print(file)
-    if args.task or args.task_dir:
-        log.info("******************* Pip Package Installing *******************")
-        subprocess.call(
-            ["python3", "-m", "pip", "install", "pip", "--upgrade", "--quiet"]
-        )
+        logger.info("***************** Supported Hardware Backend *****************")
+        print(hardware_list)
+        exit(0)
+
+    # check task
+    tasks = task_list if args.task == "all" else args.task.split(",")
+    for task in tasks:
+        if task not in task_list:
+            logger.error(f"Task {task} not found in {args.task_dir}")
+            exit(1)
+
+    logger.info(f"******************* Tasks: *****************")
+    logger.info(f"{tasks}\n")
+
+    # check hardware
+    hardware = args.hardware_type
+    if hardware not in hardware_list:
+        logger.error(f"Hardware {hardware} not found in {BYTE_MLPERF_ROOT.joinpath('backends')}")
+        exit(1)
+
+    logger.info(f"******************* hardware: *****************")
+    logger.info(f"{hardware}\n")
+
+
+    logger.info("******************* Pip Package Installing *******************")
+    subprocess.run(
+        ["python3", "-m", "pip", "install", "-r", "requirements.txt", "--quiet"],
+        capture_output=True
+    )
+    subprocess.run(
+        ["python3", "-m", "pip", "install", "-r", "requirements.txt", "--quiet"], 
+        capture_output=True
+    )
 
-        subprocess.call(
-            ["python3", "-m", "pip", "install", "-r", "requirements.txt", "--quiet"]
+    for task in tasks:
+        cmd = "python3 core/perf_engine.py --hardware_type {} --task {} --vendor_path {}".format(
+            args.hardware_type, task, args.vendor_path
         )
+        exit_code = subprocess.call(cmd, shell=True)
 
-        if args.task:
-            if args.task_dir:
-                log.warning("task and task_dir are both set, task_dir will be ignored")
-            tasks = args.task.split(',')
-        elif args.task_dir:
-            tasks = parse_task(args.task_dir)
-        logging.info(f"******************* Tasks: {tasks}")
-        exit_code = 0
-        for task in tasks:
-            cmd = "python3 core/perf_engine.py --hardware_type {} --task {} --vendor_path {}".format(
-                args.hardware_type, task, args.vendor_path
-            )
-            exit_code = subprocess.call(cmd, shell=True)
-
-        sys.exit(exit_code)
+    sys.exit(exit_code)