scale-nucleus 0.1.22__py3-none-any.whl → 0.6.4__py3-none-any.whl

cli/client.py

@@ -0,0 +1,14 @@
+import functools
+import os
+
+import nucleus
+
+
+@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

cli/datasets.py

@@ -0,0 +1,77 @@
+import click
+from rich.console import Console
+from rich.table import Column, Table
+
+from cli.client import init_client
+from cli.helpers.nucleus_url import nucleus_url
+from cli.helpers.web_helper import launch_web_or_invoke
+
+
+@click.group("datasets", invoke_without_command=True)
+@click.option("--web", is_flag=True, help="Launch browser")
+@click.pass_context
+def datasets(ctx, web):
+    """Datasets are the base collections of items in Nucleus
+
+    https://dashboard.scale.com/nucleus/datasets
+    """
+    launch_web_or_invoke(
+        sub_url="datasets", ctx=ctx, launch_browser=web, command=list_datasets
+    )
+
+
+@datasets.command("list")
+@click.option(
+    "-m", "--machine-readable", is_flag=True, help="Removes pretty printing"
+)
+def list_datasets(machine_readable):
+    """List all available Datasets"""
+    console = Console()
+    with console.status("Finding your Datasets!", spinner="dots4"):
+        client = init_client()
+        all_datasets = client.datasets
+        if machine_readable:
+            table_params = {"box": None, "pad_edge": False}
+        else:
+            table_params = {
+                "title": ":fire: Datasets",
+                "title_justify": "left",
+            }
+
+        table = Table(
+            "id", "Name", Column("url", overflow="fold"), **table_params
+        )
+        for ds in all_datasets:
+            table.add_row(ds.id, ds.name, nucleus_url(ds.id))
+    console.print(table)
+
+
+@datasets.command("delete")
+@click.option("--id", prompt=True)
+@click.option(
+    "--no-confirm-deletion",
+    is_flag=True,
+    help="WARNING: No confirmation for deletion",
+)
+@click.pass_context
+def delete_dataset(ctx, id, no_confirm_deletion):
+    """Delete a Dataset"""
+    console = Console()
+    client = init_client()
+    id = id.strip()
+    dataset = client.get_dataset(id)
+    delete_string = ""
+    if not no_confirm_deletion:
+        delete_string = click.prompt(
+            click.style(
+                f"Type 'DELETE' to delete dataset: {dataset}", fg="red"
+            )
+        )
+    if no_confirm_deletion or delete_string == "DELETE":
+        client.delete_dataset(dataset.id)
+        console.print(f":fire: :anguished: Deleted {id}")
+    else:
+        console.print(
+            f":rotating_light: Refusing to delete {id}. Received '{delete_string}' instead of 'DELETE'"
+        )
+        ctx.abort()

cli/helpers/nucleus_url.py

@@ -0,0 +1,10 @@
+import os
+from urllib.parse import urljoin
+
+
+def nucleus_url(sub_path: str):
+    nucleus_base = os.environ.get(
+        "NUCLEUS_DASHBOARD", "https://dashboard.scale.com/nucleus/"
+    )
+    extra_params = os.environ.get("NUCLEUS_DASH_PARAMS", "")
+    return urljoin(nucleus_base, sub_path.lstrip("/") + extra_params)

cli/helpers/web_helper.py

@@ -0,0 +1,40 @@
+import click
+
+from cli.helpers.nucleus_url import nucleus_url
+
+
+def launch_web_or_show_help(
+    sub_url: str, ctx: click.Context, launch_browser: bool
+):
+    """ Launches the sub_url (composed with nuclues_url(sub_url)) in the browser if requested"""
+    if not ctx.invoked_subcommand:
+        if launch_browser:
+            url = nucleus_url(sub_url)
+            click.launch(url)
+        else:
+            click.echo(ctx.get_help())
+    else:
+        if launch_browser:
+            click.echo(click.style("--web does not work with sub-commands"))
+            ctx.abort()
+
+
+def launch_web_or_invoke(
+    sub_url: str,
+    ctx: click.Context,
+    launch_browser: bool,
+    command: click.BaseCommand,
+):
+    """Launches the sub_url (composed with nuclues_url(sub_url)) in the browser if requested, otherwise invokes
+    the passed command
+    """
+    if not ctx.invoked_subcommand:
+        if launch_browser:
+            url = nucleus_url(sub_url)
+            click.launch(url)
+        else:
+            ctx.invoke(command)
+    else:
+        if launch_browser:
+            click.echo(click.style("--web does not work with sub-commands"))
+            ctx.abort()

cli/install_completion.py

@@ -0,0 +1,33 @@
+import os
+import shutil
+
+import click
+from shellingham import detect_shell
+
+
+@click.command("install-completion")
+def install_completion():
+    """Install shell completion script to your rc file"""
+    shell, _ = detect_shell()
+    if shell == "zsh":
+        rc_path = "~/.zshrc"
+        append_to_file = 'eval "$(_NU_COMPLETE=zsh_source nu)"'
+    elif shell == "bash":
+        rc_path = "~/.bashrc"
+        append_to_file = 'eval "$(_NU_COMPLETE=bash_source nu)"'
+    elif shell == "fish":
+        rc_path = "~/.config/fish/completions/foo-bar.fish"
+        append_to_file = "eval (env _NU_COMPLETE=fish_source nu)"
+    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)
+    click.echo(f"Completion script added to {rc_path}")
+    click.echo(f"Don't forget to `source {rc_path}")

cli/jobs.py

@@ -0,0 +1,42 @@
+import click
+from rich.live import Live
+from rich.spinner import Spinner
+from rich.table import Column, Table
+
+from cli.client import init_client
+from cli.helpers.web_helper import launch_web_or_invoke
+
+
+@click.group("jobs", invoke_without_command=True)
+@click.option("--web", is_flag=True, help="Launch browser")
+@click.pass_context
+def jobs(ctx, web):
+    """Jobs are a wrapper around various long-running tasks withing Nucleus
+
+    https://dashboard.scale.com/nucleus/jobs
+    """
+    launch_web_or_invoke("jobs", ctx, web, list_jobs)
+
+
+@jobs.command("list")
+def list_jobs():
+    """List all of your Jobs"""
+    client = init_client()
+    table = Table(
+        Column("id", overflow="fold", min_width=24),
+        "status",
+        "type",
+        "created at",
+        title=":satellite: Jobs",
+        title_justify="left",
+    )
+    with Live(Spinner("dots4", text="Finding your Jobs!")) as live:
+        all_jobs = client.jobs
+        for job in all_jobs:
+            table.add_row(
+                job.job_id,
+                job.job_last_known_status,
+                job.job_type,
+                job.job_creation_time,
+            )
+            live.update(table)

cli/models.py

@@ -0,0 +1,35 @@
+import click
+from rich.console import Console
+from rich.table import Column, Table
+
+from cli.client import init_client
+from cli.helpers.nucleus_url import nucleus_url
+from cli.helpers.web_helper import launch_web_or_invoke
+
+
+@click.group("models", invoke_without_command=True)
+@click.option("--web", is_flag=True, help="Launch browser")
+@click.pass_context
+def models(ctx, web):
+    """Models help you store and access your ML model data
+
+    https://dashboard.scale.com/nucleus/models
+    """
+    launch_web_or_invoke("models", ctx, web, list_models)
+
+
+@models.command("list")
+def list_models():
+    """List your Models"""
+    console = Console()
+    with console.status("Finding your Models!", spinner="dots4"):
+        client = init_client()
+        table = Table(
+            Column("id", overflow="fold", min_width=24),
+            "name",
+            Column("url", overflow="fold"),
+        )
+        models = client.models
+        for m in models:
+            table.add_row(m.id, m.name, nucleus_url(m.id))
+    console.print(table)

cli/nu.py

@@ -0,0 +1,42 @@
+import click
+
+from cli.datasets import datasets
+from cli.helpers.web_helper import launch_web_or_show_help
+from cli.install_completion import install_completion
+from cli.jobs import jobs
+from cli.models import models
+from cli.reference import reference
+from cli.slices import slices
+from cli.tests import tests
+
+
+@click.group("cli", invoke_without_command=True)
+@click.option("--web", is_flag=True, help="Launch browser")
+@click.pass_context
+def nu(ctx, web):
+    """Nucleus CLI (nu)
+
+        \b
+    ███╗   ██╗██╗   ██╗ ██████╗██╗     ███████╗██╗   ██╗███████╗
+    ████╗  ██║██║   ██║██╔════╝██║     ██╔════╝██║   ██║██╔════╝
+    ██╔██╗ ██║██║   ██║██║     ██║     █████╗  ██║   ██║███████╗
+    ██║╚██╗██║██║   ██║██║     ██║     ██╔══╝  ██║   ██║╚════██║
+    ██║ ╚████║╚██████╔╝╚██████╗███████╗███████╗╚██████╔╝███████║
+    ╚═╝  ╚═══╝ ╚═════╝  ╚═════╝╚══════╝╚══════╝ ╚═════╝ ╚══════╝
+
+        `nu` is a command line interface to interact with Scale Nucleus (https://dashboard.scale.com/nucleus)
+    """
+    launch_web_or_show_help(sub_url="", ctx=ctx, launch_browser=web)
+
+
+nu.add_command(datasets)  # type: ignore
+nu.add_command(install_completion)  # type: ignore
+nu.add_command(jobs)  # type: ignore
+nu.add_command(models)  # type: ignore
+nu.add_command(reference)  # type: ignore
+nu.add_command(slices)  # type: ignore
+nu.add_command(tests)  # type: ignore
+
+if __name__ == "__main__":
+    """To debug, run this script followed by request command tree e.g. `cli/nu.py datasets list`"""
+    nu()

cli/reference.py

@@ -0,0 +1,8 @@
+import click
+
+
+# NOTE: Named reference instead of docs to not clash with dataset autocomplete
+@click.command("reference")
+def reference():
+    """ View the Nucleus reference documentation in the browser"""
+    click.launch("https://nucleus.scale.com/docs")

cli/slices.py

@@ -0,0 +1,62 @@
+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)
+@click.option("--web", is_flag=True, help="Launch browser")
+@click.pass_context
+def slices(ctx, web):
+    """Slices are named subsets of Datasets"""
+    # TODO(gunnar): We don't have a natural landing for slices overview, until then we land on "/"
+    launch_web_or_invoke("", ctx, web, list_slices)
+
+
+@slices.command("list")
+def list_slices():
+    """List all available Slices"""
+    with Live(
+        Spinner("dots4", text="Finding your Slices!"),
+        vertical_overflow="visible",
+    ) as live:
+        client = init_client()
+        datasets = client.datasets
+        table = Table(
+            Column("id", overflow="fold", min_width=24),
+            "name",
+            "dataset_name",
+            Column("url", overflow="fold"),
+            title=":cake: Slices",
+            title_justify="left",
+        )
+        errors = {}
+        for ds in datasets:
+            try:
+                ds_slices = ds.slices
+                if ds_slices:
+                    for slc_id in ds_slices:
+                        slice_url = nucleus_url(f"{ds.id}/{slc_id}")
+                        slice_info = client.get_slice(slc_id).info()
+                        table.add_row(
+                            slc_id, slice_info["name"], ds.name, slice_url
+                        )
+                        live.update(table)
+            except NucleusAPIError as e:
+                errors[ds.id] = e
+
+        error_tree = Tree(
+            ":x: Encountered the following errors while fetching information"
+        )
+        for ds_id, error in errors.items():
+            dataset_branch = error_tree.add(f"Dataset: {ds_id}")
+            dataset_branch.add(f"Error: {error}")
+
+        Console().print(error_tree)

cli/tests.py

@@ -0,0 +1,121 @@
+import click
+from rich.console import Console
+from rich.live import Live
+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
+from nucleus.validate import (
+    AvailableEvalFunctions,
+    ScenarioTestMetric,
+    ThresholdComparison,
+)
+
+
+@click.group("tests", invoke_without_command=True)
+@click.option("--web", is_flag=True, help="Launch browser")
+@click.pass_context
+def tests(ctx, web):
+    """Scenario Tests allow you to test your Models
+
+    https://dashboard.scale.com/nucleus/scenario-tests
+    """
+    launch_web_or_invoke("scenario-tests", ctx, web, list_tests)
+
+
+@tests.command("list")
+def list_tests():
+    """List all your Scenario Tests"""
+    console = Console()
+    with console.status("Finding your Scenario Tests", spinner="dots4"):
+        client = init_client()
+        scenario_tests = client.validate.list_scenario_tests()
+        table = Table(
+            Column("id", overflow="fold", min_width=24),
+            "Name",
+            "slice_id",
+            Column("url", overflow="fold"),
+            title=":chart_with_upwards_trend: Scenario Tests",
+            title_justify="left",
+        )
+        for ut in scenario_tests:
+            table.add_row(ut.id, ut.name, ut.slice_id, nucleus_url(ut.id))
+    console.print(table)
+
+
+def format_criterion(
+    criterion: ScenarioTestMetric, eval_functions: AvailableEvalFunctions
+):
+    op_map = {
+        ThresholdComparison.GREATER_THAN: ">",
+        ThresholdComparison.GREATER_THAN_EQUAL_TO: ">=",
+        ThresholdComparison.LESS_THAN: "<",
+        ThresholdComparison.LESS_THAN_EQUAL_TO: "<=",
+    }
+    eval_function_name = eval_functions.from_id(
+        criterion.eval_function_id
+    ).name
+    return f"{eval_function_name} {op_map[criterion.threshold_comparison]} {criterion.threshold}"
+    pass
+
+
+@tests.command("describe")
+@click.argument("scenario-test-id", default=None, required=False)
+@click.option(
+    "--all", "-a", is_flag=True, help="View details about all Scenario Tests"
+)
+def describe_test(scenario_test_id, all):
+    """View detailed information about a test or all tests"""
+    console = Console()
+    # scenario_test = client.validate.get_scenario_test(scenario_test_id)
+    assert scenario_test_id or all, "Must pass a scenario_test_id or --all"
+    client = init_client()
+    scenario_tests = client.validate.list_scenario_tests()
+    if all:
+        tree = Tree(":chart_with_upwards_trend: All Scenario Tests")
+        with Live(
+            "Fetching description of all Scenario Tests",
+            vertical_overflow="visible",
+        ) as live:
+            for idx, ut in enumerate(scenario_tests):
+                test_branch = tree.add(f"{idx}: Scenario Test")
+                build_scenario_test_info_tree(client, ut, test_branch)
+                live.update(tree)
+    else:
+        with console.status("Fetching Scenario Test information"):
+            scenario_test = [
+                ut for ut in scenario_tests if ut.id == scenario_test_id
+            ][0]
+            tree = Tree(":chart_with_upwards_trend: Scenario Test")
+            build_scenario_test_info_tree(client, scenario_test, tree)
+            console.print(tree)
+
+
+def build_scenario_test_info_tree(client, scenario_test, tree):
+    try:
+        slc = client.get_slice(scenario_test.slice_id)
+        info_branch = tree.add(":mag: Details")
+        info_branch.add(f"id: '{scenario_test.id}'")
+        info_branch.add(f"name: '{scenario_test.name}'")
+        scenario_test_url = nucleus_url(scenario_test.id)
+        info_branch.add(f"url: {scenario_test_url}")
+        slice_url = nucleus_url(f"{slc.dataset_id}/{slc.slice_id}")
+        slice_branch = tree.add(":cake: Slice")
+        slice_branch.add(f"id: '{slc.id}'")
+        slice_info = slc.info()
+        slice_branch.add(f"name: '{slice_info['name']}'")
+        slice_branch.add(f"len: {len(slc.items)}")
+        slice_branch.add(f"url: {slice_url}")
+        criteria = scenario_test.get_criteria()
+        criteria_branch = tree.add(":crossed_flags: Criteria")
+        for criterion in criteria:
+            pretty_criterion = format_criterion(
+                criterion, client.validate.eval_functions
+            )
+            criteria_branch.add(pretty_criterion)
+    except NucleusAPIError as e:
+        error_branch = tree.add(":x: Error")
+        error_branch.add(f"detail: {str(e)}")