Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
ae4f71f
Try fix for click 8.2 issue
strickvl Oct 10, 2025
d28f68b
Fix typo
strickvl Oct 10, 2025
5d07137
Fixes for CI
strickvl Oct 20, 2025
5ce39d5
mypy fix
strickvl Oct 20, 2025
7476168
Darglint fixes
strickvl Oct 20, 2025
db1907a
Fix type hint for term_len in formatter.py
strickvl Oct 20, 2025
e891353
Merge branch 'feature/update-modal-step-operator' into feature/handle…
strickvl Oct 20, 2025
ad309a8
Fix CLI test failures with Click >=8.2 by using compatibility helper
strickvl Oct 20, 2025
515610f
Document CLI testing best practices for Click compatibility
strickvl Oct 20, 2025
37ba025
Add --test flag to init command tests to bypass email prompt
strickvl Oct 20, 2025
e331ce8
Improve exception handling and type hints in CLI formatter
strickvl Oct 21, 2025
92baa4d
Fix clean command test to properly use boolean flag
strickvl Oct 21, 2025
5fae563
Remove unused type: ignore comment in CLI formatter
strickvl Oct 21, 2025
65abb21
Add missing init file
strickvl Oct 21, 2025
0ba7fdf
Formatting
strickvl Oct 21, 2025
8cca137
Fix test failures by renaming utils test directory
strickvl Oct 21, 2025
be76ce5
Fix test collection failures in test_step_context
strickvl Oct 21, 2025
60f14ae
Fix failing windows DB lock in integration tests
strickvl Oct 21, 2025
6dd8121
Remove duplicated close method
strickvl Oct 21, 2025
de20aff
Define the close method on the ZenStoreInterface
strickvl Oct 22, 2025
37dd260
Improve comment clarity on Click 8.2+ help behavior
strickvl Oct 22, 2025
ff10561
Remove unnecessary formatter fallback branch
strickvl Oct 22, 2025
d9e6158
.gitignore updates
strickvl Oct 22, 2025
aac47c2
Simplify Click 8.2+ fix by reverting unnecessary formatter changes
strickvl Oct 22, 2025
9b2b3bf
Restore type ignore comment for write_dl call
strickvl Oct 22, 2025
f89bb14
Fix test_reusing_private_secret_name_succeeds on REST stores
strickvl Oct 23, 2025
85cc697
Merge branch 'feature/update-modal-step-operator' into feature/handle…
strickvl Nov 3, 2025
f81886f
Merge branch 'feature/update-modal-step-operator' into feature/handle…
strickvl Nov 3, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -205,9 +205,10 @@ examples/computer_vision/test_runs
#generated folder by zenml
zenml_tutorial/

# Claude settings
# AI code tools
**/.claude/settings.local.json
.claude/
.repoprompt/

.local/
# PLEASE KEEP THIS LINE AT THE EOF: never include here src/zenml/zen_server/dashboard, since it is affecting release flow
6 changes: 6 additions & 0 deletions CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,12 @@ Note: The MCP server indexes the latest released docs, not the develop branch. F
- For full coverage, use CI (see CI section below)
- Some tests use: `bash scripts/test-coverage-xml.sh` (but this won't run all tests)

#### CLI Testing Best Practices
- When writing new CLI tests, use the `cli_runner()` helper from `tests/integration/functional/cli/utils.py` instead of directly instantiating `CliRunner()`
- The helper handles Click version compatibility (Click 8.2+ removed the `mix_stderr` parameter)
- Always check `result.exit_code` after invoking CLI commands to catch failures early
- Existing tests with `CliRunner()` (no arguments) are fine - only new tests need the helper

## Development Workflow

### Prerequisites
Expand Down
2 changes: 2 additions & 0 deletions src/zenml/cli/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -381,6 +381,8 @@ def clean(yes: bool = False, local: bool = False) -> None:
global_zen_config = Path(get_global_config_directory())
if fileio.exists(str(global_zen_config)):
gc = GlobalConfiguration()
# Ensure any open DB handles (e.g., SQLite) are released before deletion
gc.close_store()
for dir_name in fileio.listdir(str(global_zen_config)):
if fileio.isdir(str(global_zen_config / str(dir_name))):
warning(
Expand Down
26 changes: 21 additions & 5 deletions src/zenml/cli/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,17 +137,22 @@ def format_commands(
help_ = cmd.get_short_help_str(limit=formatter.width)
rows.append((tag.value, subcommand, help_))
if rows:
colored_section_title = (
section_title = (
"[dim cyan]Available ZenML Commands (grouped)[/dim cyan]"
)
with formatter.section(colored_section_title):
with formatter.section(section_title):
formatter.write_dl(rows) # type: ignore[arg-type]


@click.group(cls=ZenMLCLI)
@click.group(cls=ZenMLCLI, invoke_without_command=True)
@click.version_option(__version__, "--version", "-v")
def cli() -> None:
"""CLI base command for ZenML."""
@click.pass_context
def cli(ctx: click.Context) -> None:
"""CLI base command for ZenML.

Args:
ctx: The click context.
"""
set_root_verbosity()
source_context.set(SourceContextTypes.CLI)
repo_root = Client.find_repository()
Expand All @@ -158,6 +163,17 @@ def cli() -> None:
# packages directory
source_utils.set_custom_source_root(source_root=os.getcwd())

# Manually show help when invoked without a subcommand to ensure our custom
# formatter is used. Without invoke_without_command=True, Click defaults to
# no_args_is_help=True, which in Click 8.2+ raises NoArgsIsHelpError before
# this callback runs. That error triggers Click's default help display,
# bypassing ZenMLCLI.get_help() and losing our custom formatting (categories,
# colors, layout). By using invoke_without_command=True and manually checking
# for no subcommand, we guarantee consistent behavior across Click versions.
if ctx.invoked_subcommand is None and not ctx.resilient_parsing:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure I really get why this is necessary. Why wouldn't click display the help anymore in newer versions with our previous setup?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The issue isn't that Click won't display help in newer versions (it will). The problem is that Click 8.2+ would display help using its default formatter instead of our custom ZenMLCLI.get_help() formatter.

(As far as I understand it, without invoke_without_command=True, Click defaults to no_args_is_help=True. In Click 8.2+, this causes Click to raise NoArgsIsHelpError before our cli() callback runs, meaning it handles the help display internally with its default formatter, completely bypassing our custom formatting (categories, rich colors, custom layout). By using invoke_without_command=True and manually checking for no subcommand in the callback, we ensure our custom formatter is always used across all Click versions.)

I will improve the comment to make it clearer maybe.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's just that I can't reproduce it. I have click 8.2+, I run zenml or zenml --help, and everything shows up correct (on develop). So it uses the correct formatter I assume?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also if the code that defined invoke_without_command=True gets executed, at that point anyway I should also know about our custom class and use that context/formatter I assume?

ctx.command.get_help(ctx)
ctx.exit(0)


if __name__ == "__main__":
cli()
34 changes: 34 additions & 0 deletions src/zenml/config/global_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,17 @@ def _reset_instance(
singleton. If None, the global GlobalConfiguration singleton is
reset to an empty value.
"""
# Proactively dispose resources held by the current singleton instance
# to avoid file-handle issues on platforms like Windows.
if cls._global_config is not None:
try:
cls._global_config.close_store()
except Exception:
logger.debug(
"Failed to close store during GlobalConfiguration reset.",
exc_info=True,
)

cls._global_config = config
if config:
config._write_config()
Expand Down Expand Up @@ -701,6 +712,29 @@ def set_store(
self._configure_store(config, skip_default_registrations, **kwargs)
logger.info("Updated the global store configuration.")

def close_store(self) -> None:
"""Close and detach the cached store if it has been initialized.

Important:
- Operates only on the cached store; does not instantiate a new one.
- Logs cleanup errors at debug level and never raises.
"""
store = self._zen_store
if store is None:
return

try:
store.close()
except Exception as e:
logger.debug(
"Error while closing zen store during cleanup: %s",
e,
exc_info=True,
)
finally:
# Clear the cached store reference irrespective of cleanup result.
self._zen_store = None

@property
def is_initialized(self) -> bool:
"""Check if the global configuration is initialized.
Expand Down
12 changes: 12 additions & 0 deletions src/zenml/zen_stores/base_zen_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,18 @@ def __init__(
else:
logger.debug("Skipping database initialization")

def close(self) -> None:
"""Release external resources held by this store instance.

Default no-op implementation. Subclasses should override this to
release any external resources (e.g., database connections, HTTP
sessions, file handles, background workers). Implementations must
remain idempotent (safe to call multiple times) and must not raise;
any cleanup failures should be handled internally and logged at
debug level instead.
"""
return

@staticmethod
def get_store_class(store_type: StoreType) -> Type["BaseZenStore"]:
"""Returns the class of the given store type.
Expand Down
34 changes: 34 additions & 0 deletions src/zenml/zen_stores/rest_zen_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -4344,6 +4344,40 @@ def batch_delete_tag_resource(
route=TAG_RESOURCES,
)

def close(self) -> None:
"""Release external resources held by this store instance.

This implementation closes the underlying HTTP session and clears transient
authentication state. It is safe to call multiple times and never raises;
unexpected cleanup issues are logged at debug level instead.
"""
try:
with self._session_lock:
if self._session is not None:
try:
self._session.close()
except Exception as e:
# Cleanup must not raise; surface at debug level only.
logger.debug(
"Error closing REST session during RestZenStore.close(): %s",
e,
exc_info=True,
)
finally:
# Clear the session reference to avoid reuse and help GC.
self._session = None

# Clear transient auth state so future use starts cleanly.
self._api_token = None
self._last_authenticated = None
except Exception as e:
# Absolute safety: no exceptions should escape close()
logger.debug(
"Unexpected error during RestZenStore.close(): %s",
e,
exc_info=True,
)

# =======================
# Internal helper methods
# =======================
Expand Down
24 changes: 24 additions & 0 deletions src/zenml/zen_stores/sql_zen_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -1248,6 +1248,30 @@ def filter_and_paginate(
# Initialization and configuration
# --------------------------------

def close(self) -> None:
"""Release DB resources held by this store.

Safe to call multiple times. Primarily ensures SQLite file handles are
released so the DB file can be removed on platforms like Windows.
"""
try:
engine = self._engine
if engine is not None:
# Disposes pooled connections and releases file handles.
engine.dispose()
except Exception as e:
# Cleanup must not raise; surface at debug level only.
logger.debug(
"Error disposing SQLAlchemy engine during SqlZenStore.close(): %s",
e,
exc_info=True,
)
finally:
# Clear references to avoid accidental reuse and help GC.
self._engine = None
self._alembic = None
self._migration_utils = None

def _initialize(self) -> None:
"""Initialize the SQL store."""
logger.debug("Initializing SqlZenStore at %s", self.config.url)
Expand Down
13 changes: 13 additions & 0 deletions src/zenml/zen_stores/zen_store_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,19 @@ def _initialize(self) -> None:
be used to set up the backend (database, connection etc.).
"""

@abstractmethod
def close(self) -> None:
"""Release external resources held by this store instance.

Implementations must:
- Release any external resources such as database connections, HTTP
sessions, open file handles, background workers, etc.
- Be safe to call multiple times (idempotent). Subsequent calls should
be no-ops and must not fail because resources were already released.
- Never raise exceptions. Any cleanup failures should be handled
internally and logged at debug level instead of being propagated.
"""

@abstractmethod
def get_store_info(self) -> ServerModel:
"""Get information about the store.
Expand Down
48 changes: 41 additions & 7 deletions tests/integration/functional/cli/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@
from pathlib import Path

import pytest
from click.testing import CliRunner

from tests.integration.functional.cli.utils import cli_runner
from zenml.cli.base import ZENML_PROJECT_TEMPLATES, clean, init
from zenml.constants import CONFIG_FILE_NAME, REPOSITORY_DIRECTORY_NAME
from zenml.utils import yaml_utils
Expand All @@ -27,8 +27,9 @@

def test_init_creates_zen_folder(tmp_path: Path) -> None:
"""Check that init command creates a .zen folder."""
runner = CliRunner()
runner.invoke(init, ["--path", str(tmp_path)])
runner = cli_runner()
result = runner.invoke(init, ["--path", str(tmp_path), "--test"])
assert result.exit_code == 0, f"Command failed: {result.output}"
assert (tmp_path / REPOSITORY_DIRECTORY_NAME).exists()


Expand All @@ -44,17 +45,19 @@ def test_init_creates_from_templates(
tmp_path: Path, template_name: str
) -> None:
"""Check that init command checks-out template."""
runner = CliRunner()
runner.invoke(
runner = cli_runner()
result = runner.invoke(
init,
[
"--path",
str(tmp_path),
"--template",
template_name,
"--template-with-defaults",
"--test",
],
)
assert result.exit_code == 0, f"Command failed: {result.output}"
assert (tmp_path / REPOSITORY_DIRECTORY_NAME).exists()
files_in_top_level = set([f.lower() for f in os.listdir(str(tmp_path))])
must_have_files = {
Expand All @@ -75,9 +78,40 @@ def test_clean_user_config(clean_client) -> None:
user_id = yaml_contents["user_id"]
analytics_opt_in = yaml_contents["analytics_opt_in"]
version = yaml_contents["version"]
runner = CliRunner()
runner.invoke(clean, ["--yes", True])
runner = cli_runner()
result = runner.invoke(clean, ["--yes"])
assert result.exit_code == 0, f"Command failed: {result.output}"
new_yaml_contents = yaml_utils.read_yaml(str(global_zen_config_yaml))
assert user_id == new_yaml_contents["user_id"]
assert analytics_opt_in == new_yaml_contents["analytics_opt_in"]
assert version == new_yaml_contents["version"]


def test_clean_calls_store_close(monkeypatch, clean_client) -> None:
# Import locally to avoid modifying module-level imports
from zenml.config.global_config import GlobalConfiguration
from zenml.zen_stores.sql_zen_store import SqlZenStore

# Track whether close was called
called = {"closed": False}

# Keep reference to the original close to preserve behavior
orig_close = SqlZenStore.close

def tracked_close(self):
called["closed"] = True
# Call through to the original to keep cleanup behavior intact
return orig_close(self)

# Patch the close method
monkeypatch.setattr(SqlZenStore, "close", tracked_close, raising=True)

# Ensure the store is initialized in this process so there is something to close
_ = GlobalConfiguration().zen_store

# Run clean
result = cli_runner().invoke(clean, ["--yes"])
assert result.exit_code == 0, f"Command failed: {result.output}"

# Verify that our close hook was executed
assert called["closed"] is True
10 changes: 6 additions & 4 deletions tests/integration/functional/cli/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,17 @@

import click
import pytest
from click.testing import CliRunner

from tests.integration.functional.cli.utils import (
cli_runner as create_cli_runner,
)
from zenml.cli.cli import ZenMLCLI, cli
from zenml.cli.formatter import ZenFormatter


@pytest.fixture(scope="function")
def runner(request):
return CliRunner()
return create_cli_runner()


def test_cli_command_defines_a_cli_group() -> None:
Expand Down Expand Up @@ -57,7 +59,7 @@ def test_cli_sets_custom_source_root_if_outside_of_repository(
mock_set_custom_source_root = mocker.patch(
"zenml.utils.source_utils.set_custom_source_root"
)
runner = CliRunner()
runner = create_cli_runner()

# Invoke a subcommand so the root CLI group gets called
runner.invoke(cli, ["version"])
Expand All @@ -73,7 +75,7 @@ def test_cli_does_not_set_custom_source_root_if_inside_repository(
mock_set_custom_source_root = mocker.patch(
"zenml.utils.source_utils.set_custom_source_root"
)
runner = CliRunner()
runner = create_cli_runner()

# Invoke a subcommand so the root CLI group gets called
runner.invoke(cli, ["version"])
Expand Down
Loading
Loading