From 7f38490ec086558552fe28e68a6552105b993fbd Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Tue, 27 Sep 2022 22:23:22 -0700 Subject: [PATCH] Install completions for entrypoints, misc config tweaks (#569) * Install completions for entrypoints, misc config tweaks * Bump dcargs --- README.md | 4 +- docs/tutorials/quickstart.md | 4 +- nerfactory/configs/base.py | 6 +- nerfactory/configs/base_configs.py | 27 +------- pyproject.toml | 4 +- .../completions/{configure.py => install.py} | 62 +++++++++++++------ scripts/completions/setup.zsh | 9 ++- scripts/train.py | 2 + 8 files changed, 65 insertions(+), 53 deletions(-) rename scripts/completions/{configure.py => install.py} (82%) diff --git a/README.md b/README.md index a76190fe9..a9e6c5cdf 100644 --- a/README.md +++ b/README.md @@ -94,10 +94,10 @@ pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.py pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` -Experimental: install tab completion for nerfactory (bash & zsh) +Install tab completion for nerfactory (bash & zsh) ``` -python scripts/completions/configure.py install +nf-install-cli ``` Run the test cases diff --git a/docs/tutorials/quickstart.md b/docs/tutorials/quickstart.md index 8fbcafc77..04dd615a3 100644 --- a/docs/tutorials/quickstart.md +++ b/docs/tutorials/quickstart.md @@ -32,11 +32,11 @@ pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.py pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` -Experimental: install tab completion for nerfactory (bash and zsh) +Install tab completion for nerfactory (bash and zsh) ```bash # If the CLI changes, we can also re-install to synchronize completions -python scripts/completions/configure.py install +nf-install-cli ``` # Downloading data diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 76b8bc79d..27ab47cdb 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -240,7 +240,7 @@ class Config(PrintableConfig): """Method name. Required to set in python or via cli""" experiment_name: Optional[str] = None """Experiment name. If None, will automatically be set to dataset name""" - timestamp: str = datetime.now().strftime("%Y-%m-%d_%H%M%S") + timestamp: str = "{timestamp}" """Experiment timestamp.""" machine: MachineConfig = MachineConfig() logging: LoggingConfig = LoggingConfig() @@ -256,6 +256,10 @@ class Config(PrintableConfig): } ) + def set_timestamp(self) -> None: + """Dynamically set the experiment timestamp""" + self.timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S") + def set_experiment_name(self) -> None: """Dynamically set the experiment name""" if self.experiment_name is None: diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index 7355f0656..60aea1a13 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -18,11 +18,9 @@ Put all the method implementations in one location. from __future__ import annotations -import copy -from typing import Dict, Type +from typing import Dict import dcargs -from typeguard import typeguard_ignore from nerfactory.configs.base import ( AdamOptimizerConfig, @@ -232,28 +230,7 @@ base_configs["tensorf"] = Config( ) -@typeguard_ignore # TypeGuard doesn't understand the generic alias that's returned here. -def _make_base_config_subcommand_type() -> Type[Config]: - """Generate a Union[] type over the possible base config types, with runtime - annotations containing default values. Used to generate CLIs. - - Returns: - An annotated Union type, which can be used to pick a base configuration. - """ - # When a base config is used to generate a CLI: replace the auto-generated timestamp - # with {timestamp}. This makes the CLI helptext (and, for zsh, autocomplete - # generation) consistent everytime you run a script with --help. - # - # Note that when a config is instantiated with dcargs.cli(), the __post_init__ - # called when the config is instantiated will set the correct timestamp. - base_configs_placeholder_timestamp = {} - for name, config in base_configs.items(): - base_configs_placeholder_timestamp[name] = copy.deepcopy(config) - - return dcargs.extras.subcommand_type_from_defaults(base_configs_placeholder_timestamp) - - -AnnotatedBaseConfigUnion = _make_base_config_subcommand_type() +AnnotatedBaseConfigUnion = dcargs.extras.subcommand_type_from_defaults(base_configs) """Union[] type over config types, annotated with default instances for use with dcargs.cli(). Allows the user to pick between one of several base configurations, and then override values in it.""" diff --git a/pyproject.toml b/pyproject.toml index c22be853e..a07451b48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ dependencies = [ "aiortc==1.3.2", "appdirs>=1.4", "av==9.2.0", - "dcargs>=0.3.9", + "dcargs>=0.3.11", "gdown==4.5.1", "ninja==1.10.2.3", "functorch==0.2.1", @@ -74,7 +74,7 @@ docs = [ ] [project.scripts] -nf-install-cli = "scripts.completions.configure:entrypoint" +nf-install-cli = "scripts.completions.install:entrypoint" nf-process-data = "scripts.process_data:entrypoint" nf-train = "scripts.train:entrypoint" nf-eval = "scripts.eval:entrypoint" diff --git a/scripts/completions/configure.py b/scripts/completions/install.py similarity index 82% rename from scripts/completions/configure.py rename to scripts/completions/install.py index a8aad955f..6065f077b 100755 --- a/scripts/completions/configure.py +++ b/scripts/completions/install.py @@ -10,7 +10,7 @@ import shutil import stat import subprocess import sys -from typing import List, Literal +from typing import List, Literal, Union import dcargs from rich.console import Console @@ -58,7 +58,9 @@ def _check_dcargs_cli(script_path: pathlib.Path) -> bool: return False -def _generate_completion(script_path: pathlib.Path, shell: ShellType, completions_dir: pathlib.Path) -> pathlib.Path: +def _generate_completion( + path_or_entrypoint: Union[pathlib.Path, str], shell: ShellType, completions_dir: pathlib.Path +) -> pathlib.Path: """Given a path to a dcargs CLI, write a completion script to a target directory. Args: @@ -69,12 +71,22 @@ def _generate_completion(script_path: pathlib.Path, shell: ShellType, completion Returns: Success flag. """ - # Use zsh standard for naming completion scripts. - target_path = completions_dir / shell / ("_" + script_path.name.replace(".", "_")) + if isinstance(path_or_entrypoint, pathlib.Path): + # Scripts. + target_name = "_" + path_or_entrypoint.name.replace(".", "_") + args = [sys.executable, str(path_or_entrypoint), "--dcargs-print-completion", shell] + elif isinstance(path_or_entrypoint, str): + # Entry points. + target_name = "_" + path_or_entrypoint + args = [path_or_entrypoint, "--dcargs-print-completion", shell] + else: + assert_never(path_or_entrypoint) + + target_path = completions_dir / shell / target_name # Generate and write the new completion. new = subprocess.run( - args=[sys.executable, str(script_path), "--dcargs-print-completion", shell], + args=args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf8", @@ -151,7 +163,7 @@ def _update_rc( def main( - mode: ConfigureMode, + mode: ConfigureMode = "install", /, ) -> None: """Main script. @@ -194,8 +206,28 @@ def main( else: CONSOLE.log(f":heavy_check_mark: No existing completions at: {target_dir}.") elif mode == "install": + # Set to True to install completions for scripts as well. + include_scripts = False + # Find dcargs CLIs. - script_paths = list(filter(_check_dcargs_cli, scripts_dir.glob("**/*.py"))) + script_paths = list(filter(_check_dcargs_cli, scripts_dir.glob("**/*.py"))) if include_scripts else [] + entry_points = ( + # Read pyproject.toml. + (scripts_dir.parent / "pyproject.toml") + .read_text() + # Get entrypoint lines. + .partition("[project.scripts]")[2] + .partition("[")[0] + # Split into list of strings, which should each be of the format: + # `entrypoint = "scripts.something:entrypoint"` + .strip() + .split("\n") + ) + entry_points = list( + # nf-train = "..." ==> nf-train + e.partition("=")[0].strip() + for e in entry_points + ) script_names = tuple(p.name for p in script_paths) assert len(set(script_names)) == len(script_names) @@ -211,8 +243,10 @@ def main( with CONSOLE.status("[bold]:writing_hand: Generating completions...", spinner="bouncingBall"): completion_paths = list( concurrent_executor.map( - lambda path_and_shell: _generate_completion(path_and_shell[0], path_and_shell[1], completions_dir), - itertools.product(script_paths, shells_found), + lambda path_or_entrypoint_and_shell: _generate_completion( + path_or_entrypoint_and_shell[0], path_or_entrypoint_and_shell[1], completions_dir + ), + itertools.product(script_paths + entry_points, shells_found), ) ) @@ -233,16 +267,6 @@ def main( _update_rc(completions_dir, mode, shell) CONSOLE.print("[bold]All done![/bold]") - if mode == "install": - CONSOLE.print() - CONSOLE.print("Notes:") - CONSOLE.print( - ":warning: In bash, completions will trigger via [dim white]./scripts/train.py [/dim white], but" - " not [dim white]python ./scripts/train.py [/dim white]." - ) - CONSOLE.print( - ":warning: Completions are in an experimental state. If you run into any issues, please file an issue!" - ) def entrypoint(): diff --git a/scripts/completions/setup.zsh b/scripts/completions/setup.zsh index 0a629e6b9..26ca1cef2 100644 --- a/scripts/completions/setup.zsh +++ b/scripts/completions/setup.zsh @@ -21,8 +21,13 @@ for completion_path in ${completions_dir}/* do # /some/path/to/_our_completion_py => _our_completion_py completion_name=${completion_path##*/} - # _our_completion_py => our_completion.py - script_name="${completion_name:1:-3}.py" + if [[ $name == *_py ]]; then + # _our_completion_py => our_completion.py + script_name="${completion_name:1:-3}.py" + else + # _entry-point => entry-point + script_name="${completion_name:1}" + fi autoload -Uz $completion_path compdef $completion_name $script_name diff --git a/scripts/train.py b/scripts/train.py index 167712c2e..95edfcc22 100755 --- a/scripts/train.py +++ b/scripts/train.py @@ -186,6 +186,8 @@ def launch( def main(config: cfg.Config) -> None: """Main function.""" + config.set_timestamp() + if config.trainer.load_config: logging.info(f"Loading pre-set config from: {config.trainer.load_config}") config = yaml.load(config.trainer.load_config.read_text(), Loader=yaml.Loader)