Speed up shell and CLI invocation (#367)

* Experiment with speeding up CLI

* Speed up CLI invocation further by late imports in cli itself

* Remove __future__ annotations

* Fix lint

* Change completion script to not slow down shell start

* CHANGELOG and bump to 0.14.24

Author: gatli

CHANGELOG.md

@@ -5,6 +5,13 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.14.24](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.14.24) - 2022-10-19
+
+### Fixed
+- Late imports for seldomly used heavy libraries. Sped up CLI invocation and autocomplation.
+  If you had shell completions installed before we recommend removeing them from your .(bash|zsh)rc
+  file and reinstalling with nu install-completions
+
 ## [0.14.23](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.14.23) - 2022-10-17
 
 ### Added

cli/client.py

@@ -1,14 +1,19 @@
 import functools
-import os
 
-import nucleus
+from rich.console import Console
 
 
 @functools.lru_cache()
 def init_client():
-    api_key = os.environ.get("NUCLEUS_API_KEY", None)
-    if api_key:
-        client = nucleus.NucleusClient(api_key)
-    else:
-        raise RuntimeError("No NUCLEUS_API_KEY set")
-    return client
+    console = Console()
+    with console.status("Initializing client"):
+        import os
+
+        import nucleus
+
+        api_key = os.environ.get("NUCLEUS_API_KEY", None)
+        if api_key:
+            client = nucleus.NucleusClient(api_key)
+        else:
+            raise RuntimeError("No NUCLEUS_API_KEY set")
+        return client

cli/install_completion.py

@@ -1,33 +1,48 @@
 import os
 import shutil
+import subprocess
 
 import click
 from shellingham import detect_shell
 
 
+def generate_completions(shell, completion_path):
+    command = f"_NU_COMPLETE={shell}_source nu > {completion_path}"
+    subprocess.run(command, shell=True)
+    click.echo(f"Generated completions for {shell}: {completion_path}")
+
+
 @click.command("install-completion")
 def install_completion():
     """Install shell completion script to your rc file"""
     shell, _ = detect_shell()
+    os.makedirs("~/.config", exist_ok=True)
+    append_to_file = None
     if shell == "zsh":
         rc_path = "~/.zshrc"
-        append_to_file = 'eval "$(_NU_COMPLETE=zsh_source nu)"'
+        completion_path = "~/.config/nu-cli-completions.zsh"
+        append_to_file = f". {completion_path}"
     elif shell == "bash":
         rc_path = "~/.bashrc"
-        append_to_file = 'eval "$(_NU_COMPLETE=bash_source nu)"'
+        completion_path = "~/.config/nu-cli-completions.bash"
+        append_to_file = f". {completion_path}"
     elif shell == "fish":
         rc_path = "~/.config/fish/completions/foo-bar.fish"
-        append_to_file = "eval (env _NU_COMPLETE=fish_source nu)"
+        completion_path = "~/.config/fish/completions/nu.fish"
     else:
         raise RuntimeError(f"Unsupported shell {shell} for completions")
 
-    rc_path_expanded = os.path.expanduser(rc_path)
-    rc_bak = f"{rc_path_expanded}.bak"
-    shutil.copy(rc_path_expanded, rc_bak)
-    click.echo(f"Backed up {rc_path} to {rc_bak}")
-    with open(rc_path_expanded, mode="a") as rc_file:
-        rc_file.write("\n")
-        rc_file.write("# Shell completion for nu\n")
-        rc_file.write(append_to_file)
+    generate_completions(shell, completion_path)
+
+    if append_to_file is not None:
+        rc_path_expanded = os.path.expanduser(rc_path)
+        rc_bak = f"{rc_path_expanded}.bak"
+        shutil.copy(rc_path_expanded, rc_bak)
+        click.echo(f"Backed up {rc_path} to {rc_bak}")
+        with open(rc_path_expanded, mode="a") as rc_file:
+            rc_file.write("\n")
+            rc_file.write("# Shell completion for nu\n")
+            rc_file.write(append_to_file)
+
     click.echo(f"Completion script added to {rc_path}")
     click.echo(f"Don't forget to `source {rc_path}")

cli/models.py

@@ -1,5 +1,4 @@
 import click
-import questionary
 from rich.console import Console
 from rich.pretty import pretty_repr
 from rich.table import Column, Table
@@ -35,6 +34,8 @@ def json_string_to_string(s: str) -> str:
 
 @models.command("calculate-metrics")
 def metrics():
+    import questionary
+
     client = init_client()
     models = client.models
     prompt_to_id = {f"{m.id}: {m.name}": m.id for m in models}

cli/slices.py

@@ -1,14 +1,10 @@
 import click
 from rich.console import Console
-from rich.live import Live
-from rich.spinner import Spinner
 from rich.table import Column, Table
-from rich.tree import Tree
 
 from cli.client import init_client
 from cli.helpers.nucleus_url import nucleus_url
 from cli.helpers.web_helper import launch_web_or_invoke
-from nucleus import NucleusAPIError
 
 
 @click.group("slices", invoke_without_command=True)

cli/tests.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 import click
 from rich.console import Console
 from rich.live import Live
@@ -7,12 +9,9 @@
 from cli.client import init_client
 from cli.helpers.nucleus_url import nucleus_url
 from cli.helpers.web_helper import launch_web_or_invoke
-from nucleus import NucleusAPIError
-from nucleus.validate import (
-    AvailableEvalFunctions,
-    ScenarioTestMetric,
-    ThresholdComparison,
-)
+
+if TYPE_CHECKING:
+    from nucleus.validate import AvailableEvalFunctions, ScenarioTestMetric
 
 
 @click.group("tests", invoke_without_command=True)
@@ -47,8 +46,10 @@ def list_tests():
 
 
 def format_criterion(
-    criterion: ScenarioTestMetric, eval_functions: AvailableEvalFunctions
+    criterion: "ScenarioTestMetric", eval_functions: "AvailableEvalFunctions"
 ):
+    from nucleus.validate import ThresholdComparison
+
     op_map = {
         ThresholdComparison.GREATER_THAN: ">",
         ThresholdComparison.GREATER_THAN_EQUAL_TO: ">=",
@@ -95,6 +96,8 @@ def describe_test(scenario_test_id, all):
 
 
 def build_scenario_test_info_tree(client, scenario_test, tree):
+    from nucleus import NucleusAPIError
+
     try:
         slc = client.get_slice(scenario_test.slice_id)
         info_branch = tree.add(":mag: Details")

docs/conf.py

@@ -13,6 +13,8 @@
 import os
 import sys
 
+import pkg_resources
+
 sys.path.insert(0, os.path.abspath("../../"))
 
 
@@ -22,10 +24,9 @@
 copyright = "2021, Scale"
 author = "Scale"
 
-# The full version, including alpha/beta/rc tags
-from nucleus import __version__  # noqa: E402
 
-release = "v" + str(__version__)
+# The full version, including alpha/beta/rc tags
+release = "v" + str(pkg_resources.get_distribution("scale-nucleus").version)
 
 
 # -- General configuration ---------------------------------------------------

nucleus/__init__.py

@@ -44,11 +44,9 @@
 import warnings
 from typing import Any, Dict, List, Optional, Sequence, Union
 
-import pkg_resources
 import pydantic
 import requests
 import tqdm
-import tqdm.notebook as tqdm_notebook
 
 from nucleus.url_utils import sanitize_string_args
 
@@ -150,9 +148,6 @@
 # pylint: disable=C0302
 
 
-__version__ = pkg_resources.get_distribution("scale-nucleus").version
-
-
 class NucleusClient:
     """Client to interact with the Nucleus API via Python SDK.
 
@@ -180,6 +175,8 @@ def __init__(
             self.endpoint = endpoint
         self._use_notebook = use_notebook
         if use_notebook:
+            import tqdm.notebook as tqdm_notebook
+
             self.tqdm_bar = tqdm_notebook.tqdm
         self._connection = Connection(self.api_key, self.endpoint)
         self.validate = Validate(self.api_key, self.endpoint)

nucleus/errors.py

@@ -1,9 +1,3 @@
-import pkg_resources
-
-nucleus_client_version = pkg_resources.get_distribution(
-    "scale-nucleus"
-).version
-
 INFRA_FLAKE_MESSAGES = [
     "downstream duration timeout",
     "upstream connect error or disconnect/reset before headers. reset reason: local reset",
@@ -40,6 +34,12 @@ class NucleusAPIError(Exception):
     def __init__(
         self, endpoint, command, requests_response=None, aiohttp_response=None
     ):
+        import pkg_resources
+
+        nucleus_client_version = pkg_resources.get_distribution(
+            "scale-nucleus"
+        ).version
+
         self.message = f"Your client is on version {nucleus_client_version}. If you have not recently done so, please make sure you have updated to the latest version of the client by running pip install --upgrade scale-nucleus\n"
         if requests_response is not None:
             self.status_code = requests_response.status_code

nucleus/metrics/cuboid_utils.py

@@ -99,15 +99,15 @@ def process_dataitem(dataitem):
 
 
 def compute_outer_iou(
-    xyz_0: np.ndarray,
-    wlh_0: np.ndarray,
-    yaw_0: np.ndarray,
-    xyz_1: np.ndarray,
-    wlh_1: np.ndarray,
-    yaw_1: np.ndarray,
+    xyz_0: "np.ndarray",
+    wlh_0: "np.ndarray",
+    yaw_0: "np.ndarray",
+    xyz_1: "np.ndarray",
+    wlh_1: "np.ndarray",
+    yaw_1: "np.ndarray",
     scale_convention: bool = True,
     distance_threshold=25,
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> Tuple["np.ndarray", "np.ndarray"]:
     """
     Computes outer 3D and 2D IoU
     :param xyz_0: (n, 3)
@@ -173,13 +173,13 @@ def compute_outer_iou(
 
 
 def get_batch_cuboid_corners(
-    xyz: np.ndarray,
-    wlh: np.ndarray,
-    yaw: np.ndarray,
-    pitch: np.ndarray = None,
-    roll: np.ndarray = None,
+    xyz: "np.ndarray",
+    wlh: "np.ndarray",
+    yaw: "np.ndarray",
+    pitch: "np.ndarray" = None,
+    roll: "np.ndarray" = None,
     scale_convention: bool = True,
-) -> np.ndarray:
+) -> "np.ndarray":
     """
     Vectorized batch version of get_cuboid_corners
     :param xyz: (n, 3)
@@ -211,8 +211,8 @@ def get_batch_cuboid_corners(
 
 
 def get_batch_rotation_matrices(
-    yaw: np.ndarray, pitch: np.ndarray = None, roll: np.ndarray = None
-) -> np.ndarray:
+    yaw: "np.ndarray", pitch: "np.ndarray" = None, roll: "np.ndarray" = None
+) -> "np.ndarray":
     if pitch is None:
         pitch = np.zeros_like(yaw)
     if roll is None:
@@ -238,12 +238,12 @@ def get_batch_rotation_matrices(
 
 
 def associate_cuboids_on_iou(
-    xyz_0: np.ndarray,
-    wlh_0: np.ndarray,
-    yaw_0: np.ndarray,
-    xyz_1: np.ndarray,
-    wlh_1: np.ndarray,
-    yaw_1: np.ndarray,
+    xyz_0: "np.ndarray",
+    wlh_0: "np.ndarray",
+    yaw_0: "np.ndarray",
+    xyz_1: "np.ndarray",
+    wlh_1: "np.ndarray",
+    yaw_1: "np.ndarray",
     threshold_in_overlap_ratio: float = 0.1,
 ) -> List[Tuple[int, int]]:
     if xyz_0.shape[0] < 1 or xyz_1.shape[0] < 1:
@@ -322,7 +322,7 @@ def detection_iou(
     prediction: List[CuboidPrediction],
     groundtruth: List[CuboidAnnotation],
     threshold_in_overlap_ratio: float,
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> Tuple["np.ndarray", "np.ndarray"]:
     """
     Calculates the 2D IOU and 3D IOU overlap between predictions and groundtruth.
     Uses linear sum assignment to associate cuboids.