Improve CLI

Author: cbalioglu

src/fairseq2/recipes/__init__.py

@@ -61,10 +61,12 @@ def main() -> None:
 
 
 def _run() -> int:
-    from fairseq2 import __version__
+    from fairseq2 import __version__, setup_fairseq2
 
     setup_basic_logging()
 
+    setup_fairseq2()
+
     cli = Cli(
         name="fairseq2",
         origin_module="fairseq2",
@@ -86,33 +88,4 @@ def _setup_cli(cli: Cli) -> None:
     _setup_wav2vec2_asr_cli(cli)
     _setup_hg_cli(cli)
 
-    run_extensions("setup_fairseq2_cli", cli)
-
-    _setup_legacy_extensions(cli)
-
-
-def _setup_legacy_extensions(cli: Cli) -> None:
-    should_trace = "FAIRSEQ2_EXTENSION_TRACE" in os.environ
-
-    for entry_point in entry_points(group="fairseq2.cli"):
-        try:
-            extension = entry_point.load()
-
-            extension(cli)
-        except TypeError:
-            if should_trace:
-                raise ExtensionError(
-                    entry_point.value, f"The '{entry_point.value}' entry point is not a valid CLI extension function."  # fmt: skip
-                ) from None
-
-            log.warning("The '{}' entry point is not a valid CLI extension function. Set `FAIRSEQ2_EXTENSION_TRACE` environment variable to print the stack trace.", entry_point.value)  # fmt: skip
-        except Exception as ex:
-            if should_trace:
-                raise ExtensionError(
-                    entry_point.value, f"The '{entry_point.value}' CLI extension function has failed. See the nested exception for details."  # fmt: skip
-                ) from ex
-
-            log.warning("The '{}' CLI extension function has failed. Set `FAIRSEQ2_EXTENSION_TRACE` environment variable to print the stack trace.", entry_point.value)  # fmt: skip
-
-        if should_trace:
-            log.info("The `{}` CLI extension function run successfully.", entry_point.value)  # fmt: skip
+    run_extensions("fairseq2.cli", cli)

src/fairseq2/recipes/assets.py

@@ -6,27 +6,19 @@
 
 from __future__ import annotations
 
-import sys
 from argparse import ArgumentParser, Namespace
 from collections import defaultdict
-from typing import Any, cast, final
+from typing import final
 
 from rich.console import Console
 from rich.pretty import pretty_repr
 from typing_extensions import override
 
-from fairseq2.assets import (
-    AssetCard,
-    AssetNotFoundError,
-    AssetStore,
-    default_asset_store,
-)
-from fairseq2.datasets import is_dataset_card
+from fairseq2.assets import AssetCard, AssetCardNotFoundError, AssetStore
+from fairseq2.context import get_runtime_context
 from fairseq2.logging import get_log_writer
-from fairseq2.models import is_model_card
 from fairseq2.recipes.cli import Cli, CliCommandHandler
 from fairseq2.recipes.console import get_console
-from fairseq2.setup import setup_fairseq2
 
 log = get_log_writer(__name__)
 
@@ -38,31 +30,21 @@ def _setup_asset_cli(cli: Cli) -> None:
 
     group.add_command(
         "list",
-        ListAssetsCommand(),
+        ListAssetsHandler(),
         help="list assets",
     )
 
     group.add_command(
         "show",
-        ShowAssetCommand(),
+        ShowAssetHandler(),
         help="show asset",
     )
 
 
 @final
-class ListAssetsCommand(CliCommandHandler):
+class ListAssetsHandler(CliCommandHandler):
     """Lists assets available in the current Python environment."""
 
-    _asset_store: AssetStore
-
-    def __init__(self, asset_store: AssetStore | None = None) -> None:
-        """
-        :param asset_store:
-            The asset store from which to retrieve the asset cards. If ``None``,
-            the default asset store will be used.
-        """
-        self._asset_store = asset_store or default_asset_store
-
     @override
     def init_parser(self, parser: ArgumentParser) -> None:
         parser.add_argument(
@@ -73,89 +55,102 @@ def init_parser(self, parser: ArgumentParser) -> None:
         )
 
     @override
-    def run(self, args: Namespace) -> int:
-        setup_fairseq2()
-
-        usr_assets = self._retrieve_assets(args, user=True)
-        glb_assets = self._retrieve_assets(args, user=False)
+    def run(self, parser: ArgumentParser, args: Namespace) -> int:
+        context = get_runtime_context()
 
         console = get_console()
 
         console.print("[green bold]user:")
 
-        self._dump_assets(console, usr_assets)
+        assets = self._retrieve_assets(context.asset_store, args.type, user=True)
+
+        self._dump_assets(console, assets)
 
         console.print("[green bold]global:")
 
-        self._dump_assets(console, glb_assets)
+        assets = self._retrieve_assets(context.asset_store, args.type, user=False)
+
+        self._dump_assets(console, assets)
 
         return 0
 
+    @classmethod
     def _retrieve_assets(
-        self, args: Namespace, user: bool
+        cls, asset_store: AssetStore, asset_type: str, user: bool
     ) -> list[tuple[str, list[str]]]:
         assets: dict[str, list[str]] = defaultdict(list)
 
-        names = self._asset_store.retrieve_names(scope="user" if user else "global")
+        asset_names = asset_store.retrieve_names(scope="user" if user else "global")
 
-        for name in names:
+        for asset_name in asset_names:
             try:
-                card = self._asset_store.retrieve_card(
-                    name, scope="all" if user else "global"
+                card = asset_store.retrieve_card(
+                    asset_name, scope="all" if user else "global"
                 )
-            except AssetNotFoundError:
-                log.warning("The asset '{}' has an invalid card. Skipping.", name)
+            except AssetCardNotFoundError:
+                log.warning("The '{}' asset card is not valid. Skipping.", asset_name)
 
                 continue
 
-            if name[-1] == "@":
-                name = name[:-1]
+            if asset_name[-1] == "@":
+                asset_name = asset_name[:-1]
 
-            try:
-                source = cast(str, card.metadata["__source__"])
-            except KeyError:
+            source = card.metadata.get("__source__", "unknown source")
+            if not isinstance(source, str):
                 source = "unknown source"
 
-            types = []
+            asset_types = []
 
-            if args.type == "all" or args.type == "model":
-                if is_model_card(card):
-                    types.append("model")
+            if asset_type == "all" or asset_type == "model":
+                if cls._is_model_card(card):
+                    asset_types.append("model")
 
-            if args.type == "all" or args.type == "dataset":
-                if is_dataset_card(card):
-                    types.append("dataset")
+            if asset_type == "all" or asset_type == "dataset":
+                if cls._is_dataset_card(card):
+                    asset_types.append("dataset")
 
-            if args.type == "all" or args.type == "tokenizer":
-                if self._is_tokenizer_card(card):
-                    types.append("tokenizer")
+            if asset_type == "all" or asset_type == "tokenizer":
+                if cls._is_tokenizer_card(card):
+                    asset_types.append("tokenizer")
 
-            if args.type == "all" and not types:
-                types.append("other")
+            if asset_type == "all" and not asset_types:
+                asset_types.append("other")
 
-            if not types:
+            if not asset_types:
                 continue
 
             source_assets = assets[source]
 
-            for t in types:
-                source_assets.append(f"{t}:{name}")
+            for t in asset_types:
+                source_assets.append(f"{t}:{asset_name}")
 
-        return [(source, names) for source, names in assets.items()]
+        output = []
+
+        for source, asset_names in assets.items():
+            asset_names.sort()
+
+            output.append((source, asset_names))
+
+        output.sort(key=lambda e: e[0])  # sort by source
+
+        return output
+
+    @staticmethod
+    def _is_model_card(card: AssetCard) -> bool:
+        return card.field("model_family").exists()
 
     @staticmethod
     def _is_tokenizer_card(card: AssetCard) -> bool:
         return card.field("tokenizer_family").exists()
 
-    def _dump_assets(
-        self, console: Console, assets: list[tuple[str, list[str]]]
-    ) -> None:
-        if assets:
-            assets.sort(key=lambda a: a[0])  # sort by source.
+    @staticmethod
+    def _is_dataset_card(card: AssetCard) -> bool:
+        return card.field("dataset_family").exists()
 
+    @staticmethod
+    def _dump_assets(console: Console, assets: list[tuple[str, list[str]]]) -> None:
+        if assets:
             for source, names in assets:
-                names.sort(key=lambda n: n[0])  # sort by name.
-
                 console.print(f"  [blue bold]{source}")
 
                 for idx, name in enumerate(names):
@@ -167,21 +162,9 @@ def _dump_assets(
             console.print()
 
 
-class ShowAssetCommand(CliCommandHandler):
+class ShowAssetHandler(CliCommandHandler):
     """Shows the metadata of an asset."""
 
-    _asset_store: AssetStore
-
-    def __init__(self, asset_store: AssetStore | None = None) -> None:
-        """
-        :param asset_store:
-            The asset store from which to retrieve the asset cards. If ``None``,
-            the default asset store will be used.
-        """
-        setup_fairseq2()
-
-        self._asset_store = asset_store or default_asset_store
-
     @override
     def init_parser(self, parser: ArgumentParser) -> None:
         parser.add_argument(
@@ -202,15 +185,12 @@ def init_parser(self, parser: ArgumentParser) -> None:
         parser.add_argument("name", help="name of the asset")
 
     @override
-    def run(self, args: Namespace) -> int:
-        try:
-            card: AssetCard | None = self._asset_store.retrieve_card(
-                args.name, envs=args.envs, scope=args.scope
-            )
-        except AssetNotFoundError:
-            log.error("An asset with the name '{}' cannot be found.", args.name)
+    def run(self, parser: ArgumentParser, args: Namespace) -> int:
+        context = get_runtime_context()
 
-            sys.exit(1)
+        card: AssetCard | None = context.asset_store.retrieve_card(
+            args.name, envs=args.envs, scope=args.scope
+        )
 
         while card is not None:
             self._print_metadata(dict(card.metadata))
@@ -219,17 +199,14 @@ def run(self, args: Namespace) -> int:
 
         return 0
 
-    def _print_metadata(self, metadata: dict[str, Any]) -> None:
+    def _print_metadata(self, metadata: dict[str, object]) -> None:
         console = get_console()
 
         name = metadata.pop("name")
 
         console.print(f"[green bold]{name}")
 
-        try:
-            source = metadata.pop("__source__")
-        except KeyError:
-            source = "unknown"
+        source = metadata.pop("__source__", "unknown")
 
         items = list(metadata.items())