Merge pull request #10 from RadarML/dev/upgrade

CLI tools for managing results

Author: thetianshuhuang

docs/cli.md

@@ -28,3 +28,25 @@
         separate_signature: false
         show_root_heading: false
         show_root_toc_entry: false
+
+## `nrdk upgrade-config`
+
+:::nrdk._cli.cli_upgrade
+    options:
+        heading_level: 2
+        show_symbol_type_heading: false
+        show_signature: false
+        separate_signature: false
+        show_root_heading: false
+        show_root_toc_entry: false
+
+## `nrdk validate`
+
+:::nrdk._cli.cli_validate
+    options:
+        heading_level: 2
+        show_symbol_type_heading: false
+        show_signature: false
+        separate_signature: false
+        show_root_heading: false
+        show_root_toc_entry: false

docs/index.md

@@ -12,6 +12,12 @@ Built around typed, high modular interfaces, the NRDK is designed to reduce the
 
     neural radar development kit core framework
 
+- :octicons-terminal-16: [`> nrdk ...`](cli.md)
+
+    ---
+
+    cli tools for working with NRDK output conventions
+
 - :material-book-open-page-variant-outline: [`grt`](grt/index.md)
 
     ---

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nrdk"
-version = "0.1.3"
+version = "0.1.4"
 authors = [
   { name="Tianshu Huang", email="tianshu2@andrew.cmu.edu" },
 ]

src/nrdk/_cli/__init__.py

@@ -6,6 +6,8 @@
 
 from .export import cli_export
 from .inspect import cli_inspect
+from .upgrade import cli_upgrade
+from .validate import cli_validate
 
 
 def make_annotation(name, func):
@@ -22,6 +24,8 @@ def cli_main() -> None:
     commands = {
         "inspect": cli_inspect,
         "export": cli_export,
+        "upgrade-config": cli_upgrade,
+        "validate": cli_validate,
     }
 
     return tyro.cli(Union[  # type: ignore

src/nrdk/_cli/upgrade.py

@@ -0,0 +1,139 @@
+"""Upgrade one implementation to another in hydra configs."""
+
+import os
+import re
+
+import numpy as np
+from rich import print
+from rich.columns import Columns
+from rich.panel import Panel
+
+from nrdk.framework import Result
+
+
+def _format_context(
+    context: str, line_num: int | np.integer,
+    start: int | np.integer, end: int | np.integer
+) -> str:
+    return '\n'.join([
+        f"{'>>>' if n == line_num else '   '} {line}"
+        for n, line in zip(range(start, end), context.split('\n'))
+    ])
+
+
+def _search(
+    text: str, pattern: str | re.Pattern, context_size: int = 2
+) -> list[tuple[int, str]]:
+    if isinstance(pattern, str):
+        pattern = re.compile(pattern)
+
+    newlines = np.where(np.frombuffer(
+        text.encode('utf-8'), dtype=np.uint8) == ord('\n'))[0]
+    newlines = np.concatenate([[-1], newlines, [len(text)]])
+
+    matches = []
+    search_start = 0
+    while True:
+        match = pattern.search(text, search_start)
+        if not match:
+            break
+
+        line_num = np.searchsorted(newlines, match.start(), side='right') - 1
+
+        start = max(0, line_num - context_size)
+        end = min(len(newlines) - 1, line_num + 1 + context_size)
+
+        context = text[newlines[start] + 1:newlines[end]]
+        matches.append(
+            (line_num, _format_context(context, line_num, start, end)))
+
+        search_start = match.end()
+
+    return matches
+
+
+def cli_upgrade(
+    target: str, /, to: str | None = None,
+    dry_run: bool = False, path: str = ".", follow_symlinks: bool = False
+) -> None:
+    """Upgrade implementation references in hydra configs.
+
+    !!! info "Usage"
+
+        First test with a dry run:
+        ```sh
+        nrdk upgrade-config <target> --path ./results --dry-run
+        ```
+        If you're happy with what you see, you can then run the actual upgrade:
+        ```sh
+        nrdk upgrade-config <target> <to> --path ./results
+        ```
+
+    !!! danger
+
+        This is a potentially destructive operation! Always run with
+        `--dry-run` first, and make sure that `to` does not overlap with
+        any other existing implementations in your configs.
+
+        You can also use the `upgrade-config` tool to check for this overlap
+        first:
+        ```sh
+        nrdk upgrade-config <to> --path ./results --dry-run
+        # Shouldn't return any of the config files you are planning to upgrade
+        ```
+
+    For each valid [results directory][nrdk.framework.Result] in the specified
+    `path`, search for all `_target_` fields in the hydra config, and replace
+    any occurrences of `from` with `to`.
+
+    Args:
+        target: full path name of the implementation to replace.
+        to: full path name of the implementation to replace with.
+        dry_run: if `True`, only log the changes that would be made, and do not
+            actually modify any files.
+        path: path to search for results directories.
+        follow_symlinks: whether to follow symlinks when searching for results.
+    """
+    pattern = re.compile(rf"_target_\s*:\s*{re.escape(target)}(?=\s|$)")
+    results = Result.find(path, follow_symlinks=follow_symlinks)
+
+    if dry_run:
+        all_matches = {}
+        for r in results:
+            config_path = os.path.join(r, ".hydra", "config.yaml")
+            if os.path.exists(config_path):
+                with open(config_path, "r") as f:
+                    config = f.read()
+
+                matches = _search(config, pattern, context_size=2)
+                for line_num, context in matches:
+                    if context not in all_matches:
+                        all_matches[context] = []
+                    all_matches[context].append((config_path, line_num))
+
+        for k, v in all_matches.items():
+            print(
+                f"Found {len(v)} occurrence(s) of '{target}' "
+                f"with this context:")
+            print(Panel(k))
+            print(Columns(
+                f"{os.path.relpath(config_path, path)}:{line_num}"
+                for config_path, line_num in v))
+            print()
+
+    else:
+        if to is None:
+            raise ValueError("Must specify `to` when not doing a dry run.")
+
+        for r in results:
+            config_path = os.path.join(r, ".hydra", "config.yaml")
+            if os.path.exists(config_path):
+                with open(config_path, "r") as f:
+                    config = f.read()
+
+                n = re.findall(pattern, config)
+                if n:
+                    print(f"Upgrading {len(n)} occurrence(s): {config_path}")
+                    new_config = re.sub(pattern, f"_target_: {to}", config)
+                    with open(config_path, "w") as f:
+                        f.write(new_config)

src/nrdk/_cli/validate.py

@@ -0,0 +1,81 @@
+"""Validate results directories."""
+
+import os
+
+from rich import print
+from rich.table import Table
+
+from nrdk.framework import Result
+
+
+def cli_validate(
+    path: str, /, follow_symlinks: bool = False, show_all: bool = False,
+) -> None:
+    """Validate results directories.
+
+    !!! info "Usage"
+
+        ```sh
+        nrdk validate <path> --follow_symlinks
+        ```
+
+    For each valid [results directory][nrdk.framework.Result] in the specified
+    `path`, check that all expected files are present:
+
+    | File                    | Description                                   |
+    | ----------------------- | --------------------------------------------- |
+    | `.hydra/config.yaml`    | Hydra configuration used for the run.         |
+    | `checkpoints/last.ckpt` | Last model checkpoint saved during training.  |
+    | `eval/`                 | Directory containing evaluation outputs.      |
+    | `checkpoints.yaml` | Checkpoint index; absence indicates a crashed run. |
+    | `events.out.tfevents.*` | Tensorboard log files.                        |
+
+    Args:
+        path: path to search for results directories.
+        follow_symlinks: whether to follow symlinks when searching for results.
+        show_all: show all results instead of just results with missing files.
+    """
+    results = Result.find(path, follow_symlinks=follow_symlinks, strict=False)
+
+    _check_files = [
+        ".hydra/config.yaml",
+        "checkpoints/last.ckpt",
+        "eval",
+        "checkpoints.yaml",
+    ]
+    _status = {
+        True: u'[green]\u2713[/green]',
+        False: u'[bold red]\u2718[/bold red]',
+    }
+
+    missing = 0
+
+    table = Table()
+    table.add_column("path", justify="right", style="cyan")
+    table.add_column("config.yaml", justify="left")
+    table.add_column("last.ckpt", justify="left")
+    table.add_column("eval", justify="left")
+    table.add_column("checkpoints.yaml", justify="left")
+    table.add_column("tfevents", justify="left")
+
+    for r in results:
+        row = [
+            os.path.exists(os.path.join(r, file))
+            for file in _check_files
+        ] + [any(
+            fname.startswith("events.out.tfevents.")
+            for fname in os.listdir(r)
+        )]
+        if not all(row):
+            missing += 1
+        if show_all or not all(row):
+            table.add_row(os.path.relpath(r, path), *[_status[x] for x in row])
+
+    if missing > 0:
+        print(
+            f"Found {len(results)} results directories with {missing} "
+            f"incomplete results.")
+    else:
+        print(f"All {len(results)} results directories are complete.")
+
+    print(table)

src/nrdk/framework/result.py

@@ -50,6 +50,33 @@ def __init__(self, path: str, validate: bool = True) -> None:
         if validate:
             self.validate(path)
 
+    @staticmethod
+    def find(
+        path: str, follow_symlinks: bool = False, strict: bool = True
+    ) -> list[str]:
+        """Find all results directories under the given path.
+
+        Args:
+            path: path to search under.
+            follow_symlinks: if True, follow symlinks when searching.
+            strict: if `True`, only return directories that pass validation;
+                if `False`, return directories with any `.hydra` folder or
+                `checkpoints.yaml` file.
+
+        Returns:
+            List of paths to results directories (that contain a
+                `checkpoints.yaml` and `.hydra` folder.)
+        """
+        results = []
+        for root, dirs, files in os.walk(path, followlinks=follow_symlinks):
+            if strict:
+                if ".hydra" in dirs and "checkpoints.yaml" in files:
+                    results.append(root)
+            else:
+                if ".hydra" in dirs or "checkpoints.yaml" in files:
+                    results.append(root)
+        return results
+
     @staticmethod
     def validate(path: str) -> None:
         """Validate that the given path is a valid results directory.