diff --git a/architecture/cli.md b/architecture/cli.md index 1c9e9cf4a..ecdcece1e 100644 --- a/architecture/cli.md +++ b/architecture/cli.md @@ -20,6 +20,33 @@ The CLI is built on Typer with lazy command loading to keep startup fast. Config `create_lazy_typer_group` and `_LazyCommand` stubs defer importing command modules until a command is actually invoked. This keeps `data-designer --help` fast — only the command names and descriptions are loaded eagerly; the full module (and its dependencies) loads on first use. +The root group also discovers optional command groups from the `data_designer.cli` +entry-point group. The entry-point name is the top-level command name. Its target +must be a zero-argument callable that returns a `click.Command`, normally a Click +group produced from a Typer application. Root help uses the providing +distribution's `Summary` metadata and does not load the target. + +```toml +[project.entry-points."data_designer.cli"] +slurm = "data_designer.slurm.cli:create_cli" +``` + +Extension distributions must declare a compatible dependency on +`data-designer`. Compatibility and the loaded object are validated only when the +extension command is selected. The entry-point name is authoritative and is +assigned to the returned command. Built-in command names are reserved. A +built-in collision makes that command unavailable with a targeted error rather +than allowing the extension to replace it. A duplicate extension name remains +visible as unavailable in root help and fails deterministically without loading +any conflicting target. Invalid extension names and discovery failures are +reported as warnings without disabling built-in commands. + +The installed-wheel smoke test measures nine warm root-help subprocesses for a +base-only environment and an environment with the Slurm extension. The accepted +base median is at most one second, and the median overhead from installing the +extension is at most 100 ms. It also verifies that root help does not import the +extension or compatibility-checking modules. + ### Layering Pattern (Setup Workflows) Config management commands (models, providers, MCP providers, tools) follow a consistent four-layer pattern: diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 17d2906df..70866a129 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -19,6 +19,9 @@ classifiers = [ "Programming Language :: Python :: 3.14", ] +[project.entry-points."data_designer.cli"] +slurm = "data_designer.slurm.cli:create_cli" + [build-system] requires = ["hatchling", "uv-dynamic-versioning>=0.7.0"] build-backend = "hatchling.build" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py new file mode 100644 index 000000000..a77cc462b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click +import typer + +app = typer.Typer( + name="slurm", + help="Run Data Designer workloads on Slurm", + no_args_is_help=True, +) + + +@app.callback() +def slurm_callback() -> None: + pass + + +def create_cli() -> click.Command: + """Create the Slurm CLI group.""" + return typer.main.get_command(app) diff --git a/packages/data-designer-slurm/tests/test_package.py b/packages/data-designer-slurm/tests/test_package.py index b9ed63874..2544dd973 100644 --- a/packages/data-designer-slurm/tests/test_package.py +++ b/packages/data-designer-slurm/tests/test_package.py @@ -3,6 +3,7 @@ from __future__ import annotations +import importlib.metadata from pathlib import Path import data_designer @@ -20,3 +21,14 @@ def test_slurm_is_published_before_base_extra() -> None: publish_script = (REPO_ROOT / "scripts" / "publish.sh").read_text() assert publish_script.index('"packages/data-designer-slurm"') < publish_script.index('"packages/data-designer"') + + +def test_slurm_registers_lazy_cli_extension() -> None: + entry_points = importlib.metadata.entry_points(group="data_designer.cli") + slurm_entry_point = next( + entry_point + for entry_point in entry_points + if entry_point.name == "slurm" and entry_point.dist.name == "data-designer-slurm" + ) + + assert slurm_entry_point.value == "data_designer.slurm.cli:create_cli" diff --git a/packages/data-designer/src/data_designer/cli/lazy_group.py b/packages/data-designer/src/data_designer/cli/lazy_group.py index 6b6f055aa..f0f0a49bf 100644 --- a/packages/data-designer/src/data_designer/cli/lazy_group.py +++ b/packages/data-designer/src/data_designer/cli/lazy_group.py @@ -4,11 +4,205 @@ from __future__ import annotations import importlib +import importlib.metadata +from collections import defaultdict from typing import Any import click import typer -from typer.core import TyperGroup +from typer.core import TyperCommand, TyperGroup + +CLI_EXTENSION_ENTRY_POINT_GROUP = "data_designer.cli" + + +def create_lazy_typer_group( + lazy_subcommands: dict[str, dict[str, str]], + *, + entry_point_group: str | None = None, +) -> type[TyperGroup]: + """Factory that returns a ``TyperGroup`` subclass with lazy-loaded commands. + + ``list_commands`` includes lazy command names so that ``--help`` works + without importing any command module. ``get_command`` returns a lightweight + ``_LazyCommand`` stub for lazy entries; the real Typer/Click command is only + built when the stub is invoked. + + Args: + lazy_subcommands: Mapping of command names to metadata dicts with keys: + - ``module``: Dotted module path (e.g. ``data_designer.cli.commands.preview``) + - ``attr``: Function attribute name in the module (e.g. ``preview_command``) + - ``help``: (optional) Short help text for group listing + - ``rich_help_panel``: (optional) Rich help panel name + entry_point_group: Optional entry-point group for lazy top-level command + extensions. Entry-point targets must be zero-argument callables that + return a ``click.Command``. + + Returns: + A ``TyperGroup`` subclass. + """ + + class LazyTyperGroup(TyperGroup): + _extension_entry_points: dict[str, list[importlib.metadata.EntryPoint]] | None = None + + def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]: + if not args and self.no_args_is_help and not ctx.resilient_parsing: + click.echo(ctx.get_help(), color=ctx.color) + ctx.exit(0) + return super().parse_args(ctx, args) + + def list_commands(self, ctx: click.Context) -> list[str]: + eager = super().list_commands(ctx) + lazy_names = [name for name in lazy_subcommands if name not in eager] + built_in_names = set(eager) | set(lazy_names) + extension_names = [name for name in self._discover_extension_entry_points() if name not in built_in_names] + return eager + sorted(lazy_names) + sorted(extension_names) + + def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None: + cmd = super().get_command(ctx, cmd_name) + extension_entry_points = self._discover_extension_entry_points().get(cmd_name, []) + if extension_entry_points and (cmd is not None or cmd_name in lazy_subcommands): + providers = ", ".join(_entry_point_label(entry_point) for entry_point in extension_entry_points) + return _UnavailableCommand( + cmd_name, + f"CLI extension command {cmd_name!r} from {providers} conflicts with a built-in command.", + ) + if len(extension_entry_points) > 1: + providers = ", ".join(_entry_point_label(entry_point) for entry_point in extension_entry_points) + return _UnavailableCommand( + cmd_name, + f"CLI command {cmd_name!r} is provided by multiple extensions: {providers}.", + ) + if cmd is not None: + return cmd + if cmd_name in lazy_subcommands: + info = lazy_subcommands[cmd_name] + return _LazyCommand( + name=cmd_name, + module_path=info["module"], + attr_name=info["attr"], + help=info.get("help"), + rich_help_panel=info.get("rich_help_panel"), + ) + if extension_entry_points: + return _LazyEntryPointCommand(extension_entry_points[0]) + return None + + def _discover_extension_entry_points(self) -> dict[str, list[importlib.metadata.EntryPoint]]: + if self._extension_entry_points is not None: + return self._extension_entry_points + if entry_point_group is None: + self._extension_entry_points = {} + return self._extension_entry_points + + try: + entry_points = sorted( + importlib.metadata.entry_points(group=entry_point_group), + key=_entry_point_sort_key, + ) + except Exception as e: + click.echo(f"Warning: Failed to discover CLI extensions from {entry_point_group!r}: {e}.", err=True) + self._extension_entry_points = {} + return self._extension_entry_points + + discovered: defaultdict[str, list[importlib.metadata.EntryPoint]] = defaultdict(list) + for entry_point in entry_points: + if ( + not entry_point.name + or entry_point.name.startswith("-") + or any(character.isspace() for character in entry_point.name) + ): + click.echo( + f"Warning: Ignoring invalid CLI extension command name {entry_point.name!r} " + f"from {_entry_point_label(entry_point)}.", + err=True, + ) + continue + discovered[entry_point.name].append(entry_point) + self._extension_entry_points = dict(discovered) + return self._extension_entry_points + + return LazyTyperGroup + + +_DATA_DESIGNER_DISTRIBUTION = "data-designer" + + +def _distribution_name(entry_point: importlib.metadata.EntryPoint) -> str: + distribution = getattr(entry_point, "dist", None) + if distribution is None: + return "unknown-distribution" + return distribution.metadata.get("Name") or "unknown-distribution" + + +def _distribution_version(entry_point: importlib.metadata.EntryPoint) -> str: + distribution = getattr(entry_point, "dist", None) + if distribution is None: + return "unknown-version" + return distribution.version + + +def _entry_point_label(entry_point: importlib.metadata.EntryPoint) -> str: + return f"{_distribution_name(entry_point)} {_distribution_version(entry_point)} ({entry_point.value})" + + +def _entry_point_help(entry_point: importlib.metadata.EntryPoint) -> str: + distribution = getattr(entry_point, "dist", None) + if distribution is not None: + summary = distribution.metadata.get("Summary") + if summary: + return summary + return f"CLI extension provided by {_distribution_name(entry_point)}" + + +def _entry_point_sort_key(entry_point: importlib.metadata.EntryPoint) -> tuple[str, str, str, str, str]: + distribution_name = _distribution_name(entry_point) + return ( + entry_point.name, + distribution_name.casefold(), + distribution_name, + _distribution_version(entry_point), + entry_point.value, + ) + + +def _validate_entry_point_compatibility(entry_point: importlib.metadata.EntryPoint) -> None: + packaging_requirements = importlib.import_module("packaging.requirements") + packaging_utils = importlib.import_module("packaging.utils") + + distribution = getattr(entry_point, "dist", None) + label = _entry_point_label(entry_point) + if distribution is None: + raise click.ClickException(f"CLI extension {entry_point.name!r} has no owning distribution: {label}.") + + try: + requirements = [packaging_requirements.Requirement(value) for value in distribution.requires or []] + except packaging_requirements.InvalidRequirement as e: + raise click.ClickException( + f"CLI extension {entry_point.name!r} from {label} has invalid dependency metadata: {e}." + ) from None + + data_designer_requirements = [ + requirement + for requirement in requirements + if packaging_utils.canonicalize_name(requirement.name) == _DATA_DESIGNER_DISTRIBUTION + and (requirement.marker is None or requirement.marker.evaluate()) + ] + if not data_designer_requirements: + raise click.ClickException( + f"CLI extension {entry_point.name!r} from {label} must declare a dependency on data-designer." + ) + + try: + installed_version = importlib.metadata.version(_DATA_DESIGNER_DISTRIBUTION) + except importlib.metadata.PackageNotFoundError: + raise click.ClickException("Unable to resolve the installed data-designer version.") from None + + if not all(installed_version in requirement.specifier for requirement in data_designer_requirements): + expected = " and ".join(str(requirement) for requirement in data_designer_requirements) + raise click.ClickException( + f"CLI extension {entry_point.name!r} from {label} is incompatible with " + f"data-designer {installed_version}; requires {expected}." + ) class _LazyCommand(click.Command): @@ -35,6 +229,15 @@ def __init__( self._resolved: click.Command | None = None self.rich_help_panel = rich_help_panel + def make_context( + self, + info_name: str, + args: list[str], + parent: click.Context | None = None, + **extra: Any, + ) -> click.Context: + return self._resolve().make_context(info_name, args, parent, **extra) + def _resolve(self) -> click.Command: if self._resolved is not None: return self._resolved @@ -51,6 +254,14 @@ def _resolve(self) -> click.Command: self._resolved = click_cmd return self._resolved + +class _LazyEntryPointCommand(click.Command): + def __init__(self, entry_point: importlib.metadata.EntryPoint) -> None: + super().__init__(name=entry_point.name, help=_entry_point_help(entry_point)) + self._entry_point = entry_point + self._resolved: click.Command | None = None + self.rich_help_panel = "Extensions" + def make_context( self, info_name: str, @@ -60,53 +271,44 @@ def make_context( ) -> click.Context: return self._resolve().make_context(info_name, args, parent, **extra) + def _resolve(self) -> click.Command: + if self._resolved is not None: + return self._resolved -def create_lazy_typer_group( - lazy_subcommands: dict[str, dict[str, str]], -) -> type[TyperGroup]: - """Factory that returns a ``TyperGroup`` subclass with lazy-loaded commands. + _validate_entry_point_compatibility(self._entry_point) + label = _entry_point_label(self._entry_point) + try: + factory = self._entry_point.load() + except Exception as e: + raise click.ClickException(f"Failed to load CLI extension {self.name!r} from {label}: {e}.") from None + if not callable(factory): + raise click.ClickException(f"CLI extension {self.name!r} from {label} must load a zero-argument callable.") - ``list_commands`` includes lazy command names so that ``--help`` works - without importing any command module. ``get_command`` returns a lightweight - ``_LazyCommand`` stub for lazy entries; the real Typer/Click command is only - built when the stub is invoked. + try: + command = factory() + except Exception as e: + raise click.ClickException(f"Failed to create CLI extension {self.name!r} from {label}: {e}.") from None + if not isinstance(command, (click.Command, TyperCommand, TyperGroup)): + raise click.ClickException( + f"CLI extension {self.name!r} from {label} returned {type(command).__name__}, expected click.Command." + ) - Args: - lazy_subcommands: Mapping of command names to metadata dicts with keys: - - ``module``: Dotted module path (e.g. ``data_designer.cli.commands.preview``) - - ``attr``: Function attribute name in the module (e.g. ``preview_command``) - - ``help``: (optional) Short help text for group listing - - ``rich_help_panel``: (optional) Rich help panel name + command.name = self.name + self._resolved = command + return command - Returns: - A ``TyperGroup`` subclass. - """ - - class LazyTyperGroup(TyperGroup): - def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]: - if not args and self.no_args_is_help and not ctx.resilient_parsing: - click.echo(ctx.get_help(), color=ctx.color) - ctx.exit(0) - return super().parse_args(ctx, args) - def list_commands(self, ctx: click.Context) -> list[str]: - eager = super().list_commands(ctx) - lazy_names = [name for name in lazy_subcommands if name not in eager] - return eager + sorted(lazy_names) +class _UnavailableCommand(click.Command): + def __init__(self, name: str, message: str) -> None: + super().__init__(name=name, help=f"Unavailable: {message}") + self._message = message + self.rich_help_panel = "Extensions" - def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None: - cmd = super().get_command(ctx, cmd_name) - if cmd is not None: - return cmd - if cmd_name in lazy_subcommands: - info = lazy_subcommands[cmd_name] - return _LazyCommand( - name=cmd_name, - module_path=info["module"], - attr_name=info["attr"], - help=info.get("help"), - rich_help_panel=info.get("rich_help_panel"), - ) - return None - - return LazyTyperGroup + def make_context( + self, + info_name: str, + args: list[str], + parent: click.Context | None = None, + **extra: Any, + ) -> click.Context: + raise click.ClickException(self._message) diff --git a/packages/data-designer/src/data_designer/cli/main.py b/packages/data-designer/src/data_designer/cli/main.py index 502631048..26bf335f3 100644 --- a/packages/data-designer/src/data_designer/cli/main.py +++ b/packages/data-designer/src/data_designer/cli/main.py @@ -10,7 +10,7 @@ import typer from data_designer.cli.agent_command_defs import AGENT_COMMANDS -from data_designer.cli.lazy_group import create_lazy_typer_group +from data_designer.cli.lazy_group import CLI_EXTENSION_ENTRY_POINT_GROUP, create_lazy_typer_group from data_designer.cli.runtime import ensure_cli_default_model_settings from data_designer.config.utils.constants import DATA_DESIGNER_PACKAGE_NAME @@ -84,7 +84,8 @@ def _is_version_request(args: list[str]) -> bool: "help": "Check that every referenced model and MCP tool is reachable", "rich_help_panel": "Generation", }, - } + }, + entry_point_group=CLI_EXTENSION_ENTRY_POINT_GROUP, ), add_completion=False, no_args_is_help=True, diff --git a/packages/data-designer/tests/cli/test_lazy_group.py b/packages/data-designer/tests/cli/test_lazy_group.py new file mode 100644 index 000000000..26203f014 --- /dev/null +++ b/packages/data-designer/tests/cli/test_lazy_group.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib.metadata +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import click +import pytest +import typer +from typer.testing import CliRunner + +from data_designer.cli.lazy_group import create_lazy_typer_group + +ENTRY_POINT_GROUP = "test.data_designer.cli" +runner = CliRunner() + + +def _normalized_output(output: str) -> str: + return " ".join(line.strip(" │") for line in click.unstyle(output).splitlines()) + + +def _app() -> typer.Typer: + app = typer.Typer(cls=create_lazy_typer_group({}, entry_point_group=ENTRY_POINT_GROUP)) + + @app.callback() + def callback() -> None: + pass + + @app.command() + def base() -> None: + click.echo("base") + + return app + + +def _command(output: str = "extension") -> click.Command: + @click.command() + def run() -> None: + click.echo(output) + + return click.Group(name="extension", commands={"run": run}) + + +def _nested_command() -> click.Command: + @click.command() + @click.option("--value", required=True) + @click.pass_context + def init(ctx: click.Context, value: str) -> None: + click.echo(value) + ctx.exit(3) + + profile = click.Group(name="profile", commands={"init": init}) + return click.Group(name="extension", commands={"profile": profile}) + + +def _entry_point( + name: str = "slurm", + *, + distribution_name: str = "data-designer-slurm", + requirements: list[str] | None = None, + loaded: object | None = None, + load_error: Exception | None = None, +) -> SimpleNamespace: + distribution = SimpleNamespace( + metadata={"Name": distribution_name, "Summary": f"{distribution_name} summary"}, + requires=requirements if requirements is not None else ["data-designer>=0"], + version="1.0.0", + ) + load = ( + Mock(side_effect=load_error) + if load_error is not None + else Mock(return_value=loaded or Mock(return_value=_command())) + ) + return SimpleNamespace( + name=name, + value=f"{distribution_name}.cli:create_cli", + group=ENTRY_POINT_GROUP, + dist=distribution, + load=load, + ) + + +def test_no_extensions_preserves_help_and_built_in_commands() -> None: + app = _app() + with patch.object(importlib.metadata, "entry_points", return_value=[]): + help_result = runner.invoke(app, ["--help"]) + command_result = runner.invoke(app, ["base"]) + + assert help_result.exit_code == 0 + assert "base" in help_result.output + assert "slurm" not in help_result.output + assert command_result.exit_code == 0 + assert command_result.output == "base\n" + + +def test_root_help_lists_extension_without_loading_it() -> None: + entry_point = _entry_point() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["--help"]) + + assert result.exit_code == 0 + assert "slurm" in result.output + assert "data-designer-slurm summary" in result.output + entry_point.load.assert_not_called() + + +def test_root_help_uses_fallback_when_distribution_summary_is_missing() -> None: + entry_point = _entry_point() + entry_point.dist.metadata.pop("Summary") + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["--help"]) + + assert result.exit_code == 0 + assert "CLI extension provided by data-designer-slurm" in result.output + entry_point.load.assert_not_called() + + +def test_built_in_command_does_not_load_extension() -> None: + entry_point = _entry_point() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["base"]) + + assert result.exit_code == 0 + assert result.output == "base\n" + entry_point.load.assert_not_called() + + +def test_selected_extension_loads_and_dispatches() -> None: + factory = Mock(return_value=_command("selected")) + entry_point = _entry_point(loaded=factory) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm", "run"]) + + assert result.exit_code == 0 + assert result.output == "selected\n" + entry_point.load.assert_called_once_with() + factory.assert_called_once_with() + + +def test_selected_extension_dispatches_nested_command_and_preserves_exit_code() -> None: + factory = Mock(return_value=_nested_command()) + entry_point = _entry_point(loaded=factory) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm", "profile", "init", "--value", "created"]) + + assert result.exit_code == 3 + assert result.output == "created\n" + entry_point.load.assert_called_once_with() + factory.assert_called_once_with() + + +def test_selecting_one_extension_does_not_load_another() -> None: + alpha = _entry_point("alpha", distribution_name="alpha-extension") + beta = _entry_point("beta", distribution_name="beta-extension") + with patch.object(importlib.metadata, "entry_points", return_value=[beta, alpha]): + result = runner.invoke(_app(), ["alpha", "run"]) + + assert result.exit_code == 0 + alpha.load.assert_called_once_with() + beta.load.assert_not_called() + + +def test_broken_extension_fails_only_when_selected() -> None: + entry_point = _entry_point(load_error=ImportError("missing dependency")) + app = _app() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + help_result = runner.invoke(app, ["--help"]) + built_in_result = runner.invoke(app, ["base"]) + command_result = runner.invoke(app, ["slurm"]) + + assert help_result.exit_code == 0 + assert built_in_result.exit_code == 0 + assert built_in_result.output == "base\n" + assert command_result.exit_code == 1 + assert "Failed to load CLI extension 'slurm'" in command_result.output + assert "data-designer-slurm 1.0.0" in command_result.output + assert "missing dependency" in command_result.output + + +def test_extension_target_must_be_callable() -> None: + entry_point = _entry_point(loaded=object()) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "must load a zero-argument callable" in result.output + + +def test_extension_factory_must_return_click_command() -> None: + entry_point = _entry_point(loaded=Mock(return_value=object())) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "returned object" in result.output + assert "click.Command" in result.output + + +def test_extension_factory_error_has_distribution_context() -> None: + entry_point = _entry_point(loaded=Mock(side_effect=RuntimeError("factory failed"))) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "Failed to create CLI extension 'slurm'" in result.output + assert "data-designer-slurm 1.0.0" in result.output + assert "factory failed" in result.output + + +def test_extension_without_owning_distribution_fails_before_import() -> None: + entry_point = _entry_point() + entry_point.dist = None + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "has no owning distribution" in result.output + entry_point.load.assert_not_called() + + +def test_invalid_dependency_metadata_fails_before_import() -> None: + entry_point = _entry_point(requirements=["data-designer=>1"]) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "has invalid dependency metadata" in result.output + entry_point.load.assert_not_called() + + +def test_missing_base_distribution_version_fails_before_import() -> None: + entry_point = _entry_point() + with ( + patch.object(importlib.metadata, "entry_points", return_value=[entry_point]), + patch.object( + importlib.metadata, + "version", + side_effect=importlib.metadata.PackageNotFoundError("data-designer"), + ), + ): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "Unable to resolve the installed data-designer version" in result.output + entry_point.load.assert_not_called() + + +def test_duplicate_extensions_fail_deterministically_without_loading() -> None: + alpha = _entry_point(distribution_name="alpha-extension") + beta = _entry_point(distribution_name="beta-extension") + outputs = [] + + for entry_points in ([beta, alpha], [alpha, beta]): + with patch.object(importlib.metadata, "entry_points", return_value=entry_points): + result = runner.invoke(_app(), ["slurm"]) + assert result.exit_code == 1 + outputs.append(result.output) + + assert outputs[0] == outputs[1] + assert "CLI command 'slurm' is provided by multiple extensions" in outputs[0] + assert outputs[0].index("alpha-extension") < outputs[0].index("beta-extension") + alpha.load.assert_not_called() + beta.load.assert_not_called() + + +def test_extension_cannot_replace_built_in_command() -> None: + entry_point = _entry_point("base") + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + help_result = runner.invoke(_app(), ["--help"]) + command_result = runner.invoke(_app(), ["base"]) + + assert help_result.exit_code == 0 + assert "Unavailable" in help_result.output + assert command_result.exit_code == 1 + assert "conflicts with a built-in command" in command_result.output + entry_point.load.assert_not_called() + + +@pytest.mark.parametrize( + ("requirements", "message"), + [ + ([], "must declare a dependency on data-designer"), + (["data-designer<0"], "is incompatible with data-designer"), + (["data-designer>=0", "data-designer<0"], "is incompatible with data-designer"), + ], +) +def test_incompatible_extension_fails_before_import(requirements: list[str], message: str) -> None: + entry_point = _entry_point(requirements=requirements) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert message in _normalized_output(result.output) + entry_point.load.assert_not_called() + + +def test_invalid_extension_name_is_ignored_without_breaking_built_in_commands() -> None: + entry_point = _entry_point("bad name") + app = _app() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + help_result = runner.invoke(app, ["--help"]) + command_result = runner.invoke(app, ["base"]) + + assert help_result.exit_code == 0 + assert "Ignoring invalid CLI extension command name 'bad name'" in help_result.output + assert "data-designer-slurm 1.0.0" in help_result.output + assert command_result.exit_code == 0 + assert command_result.output.endswith("base\n") + entry_point.load.assert_not_called() + + +def test_discovery_error_warns_without_breaking_built_in_commands() -> None: + app = _app() + with patch.object(importlib.metadata, "entry_points", side_effect=RuntimeError("corrupt metadata")): + help_result = runner.invoke(app, ["--help"]) + command_result = runner.invoke(app, ["base"]) + + assert help_result.exit_code == 0 + assert f"Failed to discover CLI extensions from {ENTRY_POINT_GROUP!r}: corrupt metadata" in help_result.output + assert command_result.exit_code == 0 + assert command_result.output.endswith("base\n") diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index c42bdf477..ff6fd6b6a 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -5,9 +5,11 @@ import os import shutil +import statistics import subprocess import sys import tempfile +import time from email.message import Message from email.parser import BytesParser from pathlib import Path @@ -23,6 +25,9 @@ "packages/data-designer", "packages/data-designer-slurm", ) +CLI_HELP_SAMPLES = 9 +MAX_BASE_CLI_HELP_SECONDS = 1.0 +MAX_EXTENSION_CLI_HELP_OVERHEAD_SECONDS = 0.1 def run(command: list[str], *, cwd: Path, check: bool = True) -> subprocess.CompletedProcess[str]: @@ -90,6 +95,7 @@ def install(uv: str, python: Path, wheel_directory: Path, package: str, *, cwd: def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> None: statement = f""" +import sys from importlib.metadata import version from importlib.util import find_spec @@ -97,16 +103,53 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non import data_designer.config import data_designer.engine import data_designer.interface +from typer.testing import CliRunner + +from data_designer.cli.main import app assert data_designer.__file__ is None assert version("data-designer") == {version!r} assert (find_spec("data_designer.slurm") is not None) is {slurm!r} +assert "data_designer.slurm" not in sys.modules +assert "packaging.requirements" not in sys.modules + +help_result = CliRunner().invoke(app, ["--help"]) +assert help_result.exit_code == 0, help_result.output +assert ("slurm" in help_result.output) is {slurm!r} +assert "data_designer.slurm" not in sys.modules +assert "packaging.requirements" not in sys.modules """ if slurm: - statement += f'\nimport data_designer.slurm\nassert version("data-designer-slurm") == {version!r}\n' + statement += f""" +slurm_help_result = CliRunner().invoke(app, ["slurm", "--help"]) +assert slurm_help_result.exit_code == 0, (slurm_help_result.output, repr(slurm_help_result.exception)) +assert "data_designer.slurm.cli" in sys.modules +assert version("data-designer-slurm") == {version!r} +""" run([str(python), "-c", statement], cwd=cwd) +def cli_help_medians(base_python: Path, extension_python: Path, *, cwd: Path) -> tuple[float, float]: + statement = """ +from typer.testing import CliRunner +from data_designer.cli.main import app + +result = CliRunner().invoke(app, ["--help"]) +assert result.exit_code == 0, result.output +""" + for python in (base_python, extension_python): + run([str(python), "-c", statement], cwd=cwd) + + samples: dict[Path, list[float]] = {base_python: [], extension_python: []} + for index in range(CLI_HELP_SAMPLES): + environments = (base_python, extension_python) if index % 2 == 0 else (extension_python, base_python) + for python in environments: + start = time.perf_counter() + run([str(python), "-c", statement], cwd=cwd) + samples[python].append(time.perf_counter() - start) + return statistics.median(samples[base_python]), statistics.median(samples[extension_python]) + + def main() -> None: uv = shutil.which("uv") if uv is None: @@ -143,6 +186,19 @@ def main() -> None: extra_python = create_environment(uv, root / "extra", cwd=root) install(uv, extra_python, wheel_directory, f"data-designer[slurm]=={version}", cwd=root) verify_install(extra_python, version, slurm=True, cwd=root) + base_cli_help, extension_cli_help = cli_help_medians(base_python, extra_python, cwd=root) + extension_overhead = extension_cli_help - base_cli_help + assert base_cli_help <= MAX_BASE_CLI_HELP_SECONDS, ( + f"Base CLI root help took {base_cli_help:.3f}s; budget is {MAX_BASE_CLI_HELP_SECONDS:.3f}s" + ) + assert extension_overhead <= MAX_EXTENSION_CLI_HELP_OVERHEAD_SECONDS, ( + f"CLI extension added {extension_overhead:.3f}s to root help " + f"(base={base_cli_help:.3f}s, extension={extension_cli_help:.3f}s)" + ) + print( + f"CLI help median: base={base_cli_help:.3f}s, extension={extension_cli_help:.3f}s, " + f"overhead={extension_overhead:.3f}s" + ) leaf_python = create_environment(uv, root / "leaf", cwd=root) install(uv, leaf_python, wheel_directory, f"data-designer-slurm=={version}", cwd=root)