use-case-resources: Enable user provided metadata

* An optional argument has been added to the
  `set_up_default_resources.py` Python script to allow
  passing of a user defined use case resources metadata JSON file.
  This shortens the build time by only downloading the resources the
  end user is interested in. It also shortens the optimization part
  which takes additional minutes as it is done for all models and for
  all the specified NPU configurations.
* Adding changes to comply with Pylint
* Adding --use-case argument in `set_up_default_resources.py` to
  restrict setting up resources to the specified use cases.

Signed-off-by: Hugues Kamba-Mpiana <hugues.kambampiana@arm.com>
Change-Id: I8d38249d8a0b52e66c26e5e74c03657e29f979b0
Signed-off-by: Alex Tawse <alex.tawse@arm.com>

2 files changed, 170 insertions, 105 deletions

diff --git a/build_default.py b/build_default.py
index 907bf4d..5adff22 100755
--- a/build_default.py
+++ b/build_default.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-#  SPDX-FileCopyrightText:  Copyright 2021-2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
+#  SPDX-FileCopyrightText:  Copyright 2021-2024 Arm Limited and/or its affiliates <open-source-office@arm.com>
 #  SPDX-License-Identifier: Apache-2.0
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -25,25 +25,36 @@ import sys
 import threading
 from argparse import ArgumentDefaultsHelpFormatter
 from argparse import ArgumentParser
-from collections import namedtuple
+from dataclasses import dataclass
 from pathlib import Path
 
+from set_up_default_resources import SetupArgs
 from set_up_default_resources import default_npu_config_names
 from set_up_default_resources import get_default_npu_config_from_name
 from set_up_default_resources import set_up_resources
 from set_up_default_resources import valid_npu_config_names
 
-BuildArgs = namedtuple(
-    "BuildArgs",
-    [
-        "toolchain",
-        "download_resources",
-        "run_vela_on_models",
-        "npu_config_name",
-        "make_jobs",
-        "make_verbose",
-    ],
-)
+
+@dataclass(frozen=True)
+class BuildArgs:
+    """
+    Args used to build the project.
+
+    Attributes:
+        toolchain (str)            : Specifies if 'gnu' or 'arm' toolchain needs to be used.
+        download_resources (bool)  : Specifies if 'Download resources' step is performed.
+        run_vela_on_models (bool)  : Only if `download_resources` is True, specifies whether to
+                                     run Vela on downloaded models.
+        npu_config_name (str)      : Ethos-U NPU configuration name. See "valid_npu_config_names"
+        make_jobs (int)            : The number of make jobs to use (`-j` flag).
+        make_verbose (bool)        : Runs make with VERBOSE=1.
+    """
+    toolchain: str
+    download_resources: bool
+    run_vela_on_models: bool
+    npu_config_name: str
+    make_jobs: int
+    make_verbose: bool
 
 
 class PipeLogging(threading.Thread):
@@ -182,16 +193,7 @@ def run(args: BuildArgs):
 
     Parameters:
     ----------
-    args (BuildArgs)    :   Parsed set of build args expecting:
-                            - toolchain
-                            - download_resources
-                            - run_vela_on_models
-                            - np_config_name
-    toolchain (str)             :   Specifies if 'gnu' or 'arm' toolchain needs to be used.
-    download_resources (bool)   :   Specifies if 'Download resources' step is performed.
-    run_vela_on_models (bool)   :   Only if `download_resources` is True, specifies if
-                                    run vela on downloaded models.
-    npu_config_name(str)        :   Ethos-U NPU configuration name. See "valid_npu_config_names"
+    args (BuildArgs)    : Arguments used to build the project
     """
 
     current_file_dir = Path(__file__).parent.resolve()
@@ -202,11 +204,13 @@ def run(args: BuildArgs):
     # 2. Download models if specified
     if args.download_resources is True:
         logging.info("Downloading resources.")
-        env_path = set_up_resources(
+        setup_args = SetupArgs(
             run_vela_on_models=args.run_vela_on_models,
-            additional_npu_config_names=(args.npu_config_name,),
-            additional_requirements_file=current_file_dir / "scripts" / "py" / "requirements.txt"
+            additional_npu_config_names=[args.npu_config_name],
+            additional_requirements_file=current_file_dir / "scripts" / "py" / "requirements.txt",
+            use_case_resources_file=current_file_dir / "scripts" / "py" / "use_case_resources.json",
         )
+        env_path = set_up_resources(setup_args)
 
     # 3. Build default configuration
     logging.info("Building default configuration.")
diff --git a/set_up_default_resources.py b/set_up_default_resources.py
index f5cd0ac..7ed9e97 100755
--- a/set_up_default_resources.py
+++ b/set_up_default_resources.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-#  SPDX-FileCopyrightText:  Copyright 2021-2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
+#  SPDX-FileCopyrightText:  Copyright 2021-2024 Arm Limited and/or its affiliates <open-source-office@arm.com>
 #  SPDX-License-Identifier: Apache-2.0
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -31,7 +31,6 @@ import urllib.request
 import venv
 from argparse import ArgumentParser
 from argparse import ArgumentTypeError
-from collections import namedtuple
 from dataclasses import dataclass
 from pathlib import Path
 from urllib.error import URLError
@@ -56,18 +55,29 @@ valid_npu_config_names = [
 # Default NPU configurations (these are always run when the models are optimised)
 default_npu_config_names = [valid_npu_config_names[2], valid_npu_config_names[4]]
 
-# NPU config named tuple
-NPUConfig = namedtuple(
-    "NPUConfig",
-    [
-        "config_name",
-        "memory_mode",
-        "system_config",
-        "ethos_u_npu_id",
-        "ethos_u_config_id",
-        "arena_cache_size",
-    ],
-)
+# The internal SRAM size for Corstone-300 implementation on MPS3 specified by AN552
+# The internal SRAM size for Corstone-310 implementation on MPS3 specified by AN555
+# is 4MB, but we are content with the 2MB specified below.
+MPS3_MAX_SRAM_SZ = 2 * 1024 * 1024  # 2 MiB (2 banks of 1 MiB each)
+
+default_use_case_resources_path = (Path(__file__).parent.resolve()
+                                   / 'scripts' / 'py' / 'use_case_resources.json')
+
+default_requirements_path = (Path(__file__).parent.resolve()
+                             / 'scripts' / 'py' / 'requirements.txt')
+
+
+@dataclass(frozen=True)
+class NpuConfig:
+    """
+    Represent an NPU configuration for Vela
+    """
+    config_name: str
+    memory_mode: str
+    system_config: str
+    ethos_u_npu_id: str
+    ethos_u_config_id: str
+    arena_cache_size: str
 
 
 @dataclass(frozen=True)
@@ -90,36 +100,70 @@ class UseCase:
     resources: typing.List[UseCaseResource]
 
 
-# The internal SRAM size for Corstone-300 implementation on MPS3 specified by AN552
-# The internal SRAM size for Corstone-310 implementation on MPS3 specified by AN555
-# is 4MB, but we are content with the 2MB specified below.
-MPS3_MAX_SRAM_SZ = 2 * 1024 * 1024  # 2 MiB (2 banks of 1 MiB each)
-
-
-def load_use_case_resources(current_file_dir: Path) -> typing.List[UseCase]:
+@dataclass(frozen=True)
+class SetupArgs:
+    """
+    Args used to set up the project.
+
+    Attributes:
+        run_vela_on_models (bool)           :   Whether to run Vela on the downloaded models
+        additional_npu_config_names (list)  :   List of strings of Ethos-U NPU configs.
+        use_case_names (list)               :   List of names of use cases to set up resources for
+                                                (default is all).
+        arena_cache_size (int)              :   Specifies arena cache size in bytes. If a value
+                                                greater than 0 is provided, this will be taken
+                                                as the cache size. If 0, the default values, as per
+                                                the NPU config requirements, are used.
+        check_clean_folder (bool)           :   Indicates whether the resources folder needs to
+                                                be checked for updates and cleaned.
+        additional_requirements_file (str)  :   Path to a requirements.txt file if
+                                                additional packages need to be
+                                                installed.
+        use_case_resources_file (str)       :   Path to a JSON file containing the use case
+                                                metadata resources.
+    """
+    run_vela_on_models: bool = False
+    additional_npu_config_names: typing.List[str] = ()
+    use_case_names: typing.List[str] = ()
+    arena_cache_size: int = 0
+    check_clean_folder: bool = False
+    additional_requirements_file: Path = ""
+    use_case_resources_file: Path = ""
+
+
+def load_use_case_resources(
+        use_case_resources_file: Path,
+        use_case_names: typing.List[str] = ()
+) -> typing.List[UseCase]:
     """
     Load use case metadata resources
 
     Parameters
     ----------
-    current_file_dir:   Directory of the current script
-
+    use_case_resources_file :   Path to a JSON file containing the use case
+                                metadata resources.
+    use_case_names          :   List of named use cases to restrict
+                                resource loading to.
     Returns
     -------
     The use cases resources object parsed to a dict
     """
 
-    resources_path = current_file_dir / "scripts" / "py" / "use_case_resources.json"
-    with open(resources_path, encoding="utf8") as f:
-        use_cases = json.load(f)
-        return [
+    with open(use_case_resources_file, encoding="utf8") as f:
+        parsed_use_cases = json.load(f)
+        use_cases = (
             UseCase(
                 name=u["name"],
                 url_prefix=u["url_prefix"],
                 resources=[UseCaseResource(**r) for r in u["resources"]],
             )
-            for u in use_cases
-        ]
+            for u in parsed_use_cases
+        )
+
+        if len(use_case_names) == 0:
+            return list(use_cases)
+
+        return [uc for uc in use_cases if uc.name in use_case_names]
 
 
 def call_command(command: str, verbose: bool = True) -> str:
@@ -147,7 +191,7 @@ def call_command(command: str, verbose: bool = True) -> str:
 
 def get_default_npu_config_from_name(
         config_name: str, arena_cache_size: int = 0
-) -> typing.Optional[NPUConfig]:
+) -> typing.Optional[NpuConfig]:
     """
     Gets the file suffix for the TFLite file from the
     `accelerator_config` string.
@@ -190,7 +234,7 @@ def get_default_npu_config_from_name(
     for i, string_id in enumerate(strings_ids):
         if config_name.startswith(string_id):
             npu_config_id = config_name.replace(string_id, prefix_ids[i])
-            return NPUConfig(
+            return NpuConfig(
                 config_name=config_name,
                 memory_mode=memory_modes[i],
                 system_config=system_configs[i],
@@ -332,7 +376,7 @@ def download_resources(
 
 
 def run_vela(
-        config: NPUConfig,
+        config: NpuConfig,
         env_activate_cmd: str,
         model: Path,
         config_file: Path,
@@ -555,56 +599,50 @@ def set_up_python_venv(
 def update_metadata(
         metadata_dict: typing.Dict,
         setup_script_hash: str,
-        json_uc_res: typing.List[UseCase],
+        use_case_resources: typing.List[UseCase],
         metadata_file_path: Path
 ):
     """
     Update the metadata file
 
-    @param metadata_dict:       The metadata dictionary to update
-    @param setup_script_hash:   The setup script hash
-    @param json_uc_res:         The use case resources metadata
-    @param metadata_file_path   The metadata file path
+    @param metadata_dict        :   The metadata dictionary to update
+    @param setup_script_hash    :   The setup script hash
+    @param use_case_resources   :   The use case resources metadata
+    @param metadata_file_path   :   The metadata file path
     """
     metadata_dict["ethosu_vela_version"] = VELA_VERSION
     metadata_dict["set_up_script_md5sum"] = setup_script_hash.strip("\n")
-    metadata_dict["resources_info"] = [dataclasses.asdict(uc) for uc in json_uc_res]
+    metadata_dict["resources_info"] = [dataclasses.asdict(uc) for uc in use_case_resources]
 
     with open(metadata_file_path, "w", encoding="utf8") as metadata_file:
         json.dump(metadata_dict, metadata_file, indent=4)
 
 
-def set_up_resources(
-        run_vela_on_models: bool = False,
-        additional_npu_config_names: tuple = (),
-        arena_cache_size: int = 0,
-        check_clean_folder: bool = False,
-        additional_requirements_file: Path = ""
-) -> Path:
+def get_default_use_cases_names() -> typing.List[str]:
+    """
+    Get the names of the default use cases
+
+    :return :   List of use case names as strings
+    """
+    use_case_resources = load_use_case_resources(default_use_case_resources_path)
+    return [uc.name for uc in use_case_resources]
+
+
+def set_up_resources(args: SetupArgs) -> Path:
     """
     Helpers function that retrieve the output from a command.
 
     Parameters:
     ----------
-    run_vela_on_models (bool):  Specifies if run vela on downloaded models.
-    additional_npu_config_names(list):  list of strings of Ethos-U NPU configs.
-    arena_cache_size (int): Specifies arena cache size in bytes. If a value
-                            greater than 0 is provided, this will be taken
-                            as the cache size. If 0, the default values, as per
-                            the NPU config requirements, are used.
-    check_clean_folder (bool): Indicates whether the resources folder needs to
-                               be checked for updates and cleaned.
-    additional_requirements_file (str): Path to a requirements.txt file if
-                                        additional packages need to be
-                                        installed.
+    args (SetupArgs)        :   Arguments used to set up the project.
 
     Returns
     -------
 
-    Tuple of pair of Paths: (download_directory_path,  virtual_env_path)
+    Tuple of pairs of Paths: (download_directory_path,  virtual_env_path)
 
-    download_directory_path: Root of the directory where the resources have been downloaded to.
-    virtual_env_path: Path to the root of virtual environment.
+    download_directory_path :   Root of the directory where the resources have been downloaded to.
+    virtual_env_path        :   Path to the root of virtual environment.
     """
     # Paths.
     current_file_dir = Path(__file__).parent.resolve()
@@ -619,29 +657,32 @@ def set_up_resources(
         )
     logging.info("Using Python version: %s", sys.version_info)
 
-    json_uc_res = load_use_case_resources(current_file_dir)
+    use_case_resources = load_use_case_resources(
+        args.use_case_resources_file,
+        args.use_case_names
+    )
     setup_script_hash = get_md5sum_for_file(Path(__file__).resolve())
 
     metadata_dict, setup_script_hash_verified = initialize_resources_directory(
         download_dir,
-        check_clean_folder,
+        args.check_clean_folder,
         metadata_file_path,
         setup_script_hash
     )
 
     env_path, env_activate = set_up_python_venv(
         download_dir,
-        additional_requirements_file
+        args.additional_requirements_file
     )
 
     # 2. Download models
     logging.info("Downloading resources.")
-    for use_case in json_uc_res:
+    for use_case in use_case_resources:
         download_resources(
             use_case,
             metadata_dict,
             download_dir,
-            check_clean_folder,
+            args.check_clean_folder,
             setup_script_hash_verified
         )
 
@@ -653,14 +694,16 @@ def set_up_resources(
     #
     # Note: To avoid to run vela twice on the same model, it's supposed that
     # downloaded model names don't contain the 'vela' word.
-    if run_vela_on_models is True:
+    if args.run_vela_on_models is True:
         # Consolidate all config names while discarding duplicates:
         run_vela_on_all_models(
             current_file_dir,
             download_dir,
             env_activate,
-            arena_cache_size,
-            npu_config_names=list(set(default_npu_config_names + list(additional_npu_config_names)))
+            args.arena_cache_size,
+            npu_config_names=list(
+                set(default_npu_config_names + list(args.additional_npu_config_names))
+            )
         )
 
     # 4. Collect and write metadata
@@ -668,7 +711,7 @@ def set_up_resources(
     update_metadata(
         metadata_dict,
         setup_script_hash.strip("\n"),
-        json_uc_res,
+        use_case_resources,
         metadata_file_path
     )
 
@@ -690,6 +733,14 @@ if __name__ == "__main__":
         action="append",
     )
     parser.add_argument(
+        "--use-case",
+        help=f"""Only set up resources for the specified use case (can specify multiple times).
+        Valid values are: {get_default_use_cases_names()}
+        """,
+        default=[],
+        action="append",
+    )
+    parser.add_argument(
         "--arena-cache-size",
         help="Arena cache size in bytes (if overriding the defaults)",
         type=int,
@@ -704,24 +755,34 @@ if __name__ == "__main__":
         "--requirements-file",
         help="Path to requirements.txt file to install additional packages",
         type=str,
-        default=Path(__file__).parent.resolve() / 'scripts' / 'py' / 'requirements.txt'
+        default=default_requirements_path
+    )
+    parser.add_argument(
+        "--use-case-resources-file",
+        help="Path to the use case resources file",
+        type=str,
+        default=default_use_case_resources_path
     )
 
-    args = parser.parse_args()
+    parsed_args = parser.parse_args()
 
-    if args.arena_cache_size < 0:
+    if parsed_args.arena_cache_size < 0:
         raise ArgumentTypeError("Arena cache size cannot not be less than 0")
 
-    if not Path(args.requirements_file).is_file():
-        raise ArgumentTypeError(f"Invalid requirements file: {args.requirements_file}")
+    if not Path(parsed_args.requirements_file).is_file():
+        raise ArgumentTypeError(f"Invalid requirements file: {parsed_args.requirements_file}")
 
     logging.basicConfig(filename="log_build_default.log", level=logging.DEBUG)
     logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
 
-    set_up_resources(
-        not args.skip_vela,
-        args.additional_ethos_u_config_name,
-        args.arena_cache_size,
-        args.clean,
-        args.requirements_file,
+    setup_args = SetupArgs(
+        run_vela_on_models=not parsed_args.skip_vela,
+        additional_npu_config_names=parsed_args.additional_ethos_u_config_name,
+        use_case_names=parsed_args.use_case,
+        arena_cache_size=parsed_args.arena_cache_size,
+        check_clean_folder=parsed_args.clean,
+        additional_requirements_file=parsed_args.requirements_file,
+        use_case_resources_file=parsed_args.use_case_resources_file,
     )
+
+    set_up_resources(setup_args)