Skip to content

perf(cli): lazy-load command groups to cut startup import cost (#801) - #805

Draft
padak wants to merge 1 commit into
mainfrom
claude/issue-801-lazy-command-groups
Draft

padak wants to merge 1 commit into
mainfrom
claude/issue-801-lazy-command-groups

Conversation

@padak

@padak padak commented Sep 27, 2026

Copy link
Copy Markdown
Member

What

Every kbagent invocation used to import all ~37 command modules, all ~36 services (and through them prompt_toolkit, the AI / Data Science / Metastore clients, ...) plus the whole SDK facade -- although one invocation only ever runs one command. This PR makes all three lazy:

  1. Command registration (cli.py): each root command/group is a LazyCommand(name, "module:attr", panel, ...) entry in LAZY_COMMANDS. The root group (LazyRootGroup, passed via typer.Typer(cls=...)) holds a commands table in which every name is a key from the start, and the module is imported only when that name is looked up. The Click object is built with Typer's own converters (get_group_from_info / get_command_from_info) from a throwaway add_typer / command registration with exactly the options the eager code passed, so the resulting commands are identical.
  2. Services (cli.py): ctx.obj is now a dict subclass that constructs a service on its first ctx.obj["x_service"] lookup. Classes are resolved through cli.<ServiceClass> (module __getattr__), so the test suite's ~1,500 mock.patch("keboola_agent_cli.cli.JobService")-style patches keep working unchanged.
  3. SDK facade (__init__.py): Client, the result models and JobIdempotencyStore resolve on first attribute access (PEP 562). __all__ is unchanged; a TYPE_CHECKING block keeps the names visible to mypy/ty consumers.

Why the commands table is a dict subclass

Plain list_commands/get_command overrides (the Click-docs lazy pattern) would leave group.commands holding only eager entries. Several consumers read .commands directly: telemetry._command_path (every invocation), Typer's typo suggestions (Did you mean 'job'?), test_permissions' registry walk, test_mcp_migration_recipe. Keeping every name as a key and resolving values on [] / get / values / items makes all of them work with no change.

Numbers (macOS, Python 3.12, warm cache, median of 5)

before after
python -X importtime -c "import keboola_agent_cli.cli" (cumulative) ~218 ms ~93 ms
len(sys.modules) after import keboola_agent_cli.cli 933 464
prompt_toolkit imported by import keboola_agent_cli.cli yes no
wall time kbagent job list --help ~0.34 s ~0.17 s

What remains is shared infrastructure nearly every real command needs: constants (imports httpx for Timeout objects, ~30 ms, and importlib.metadata, ~16 ms), models/pydantic (~23 ms, via config_store), rich, typer. telemetry still imports services.base + services.project_service (~3 ms). Those are possible follow-ups, not part of this PR.

Parity verification

  • --help for all 370 command paths (walked via list_commands/get_command), plus sl --help, mr list --help, sl model list --help, dumped before and after at COLUMNS=120: byte-identical.
  • Real-process kbagent --help and bare kbagent (non-TTY help) output: byte-identical.
  • Smoke: --json project list, typo suggestion (kbagent jbo -> "Did you mean 'job'?"), --deny-writes on a group command and on a top-level command (both exit 6 PERMISSION_DENIED), --version.
  • make check green: lint, format, ty, skill-check (SKILL.md up to date), command-sync (281 commands), version gates, error codes, sentinel guards, loc-check, full test suite (6908 passed). make gen-command-reference still sees 281 commands.
  • No version bump, no changelog.py entry.

Guardrail

tests/test_startup_imports.py runs fresh interpreters (hermetic env: auto-update off, telemetry off, temp config dir) and asserts:

  • import keboola_agent_cli.cli imports no command module, not keboola_agent_cli.lib, and none of prompt_toolkit / fastapi / uvicorn / starlette / jsonschema;
  • it loads at most KBAGENT_MODULE_BUDGET = 34 kbagent-owned modules (own modules only, so the count is identical across OS and Python versions; ratchets down only);
  • a real app(["job", "list", "--help"]) imports exactly the job command module and no job service;
  • hidden aliases sl / mr import only their group;
  • kbagent --help still resolves every registered command.

Reviewer focus

  • _LazyCommandTable overrides __getitem__, get, values, items; anything that copies the dict wholesale (dict(table), table.copy()) would see LazyCommand placeholders. Nothing in the repo does that today.
  • Services are now constructed on first use instead of in the root callback. A service whose constructor had side effects would now run them later (or not at all for commands that do not use it); the full test suite passes.
  • Adding a new root command now means one LazyCommand(...) line in LAZY_COMMANDS instead of an import + add_typer; a new service means one _SERVICES entry. The list order is the --help order within each panel.

Fixes #801

Register every root command by name and import its module only when it is
looked up, build services on first ctx.obj access, and resolve the SDK
facade in __init__.py via PEP 562 __getattr__. CLI surface, --help output
and the public __all__ are unchanged. tests/test_startup_imports.py guards
the import budget.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf: lazy-load command groups to cut ~0.4 s CLI startup, with CI guardrail

1 participant