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>
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 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 @@
 # 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 @@
     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)
+@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(current_file_dir: Path) -> typing.List[UseCase]:
+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 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 @@
     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 run_vela(
-        config: NPUConfig,
+        config: NpuConfig,
         env_activate_cmd: str,
         model: Path,
         config_file: Path,
@@ -555,56 +599,50 @@
 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 @@
         )
     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 @@
     #
     # 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 @@
     update_metadata(
         metadata_dict,
         setup_script_hash.strip("\n"),
-        json_uc_res,
+        use_case_resources,
         metadata_file_path
     )
 
@@ -690,6 +733,14 @@
         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 @@
         "--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)