zenml-nightly 0.83.1.dev20250709__py3-none-any.whl → 0.83.1.dev20250711__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.83.1.dev20250709
+0.83.1.dev20250711

zenml/artifact_stores/base_artifact_store.py

@@ -33,11 +33,15 @@ from typing import (
     cast,
 )
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
-from zenml.constants import ENV_ZENML_SERVER
+from zenml.constants import (
+    ENV_ZENML_SERVER,
+    ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS,
+    handle_bool_env_var,
+)
 from zenml.enums import StackComponentType
-from zenml.exceptions import ArtifactStoreInterfaceError
+from zenml.exceptions import ArtifactStoreInterfaceError, IllegalOperationError
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.stack import Flavor, StackComponent, StackComponentConfig
@@ -73,6 +77,12 @@ class _sanitize_paths:
         """
         self.func = func
         self.fixed_root_path = fixed_root_path
+        if ENV_ZENML_SERVER in os.environ:
+            self.allow_local_file_access = handle_bool_env_var(
+                ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS, False
+            )
+        else:
+            self.allow_local_file_access = True
 
         self.path_args: List[int] = []
         self.path_kwargs: List[str] = []
@@ -93,7 +103,15 @@ class _sanitize_paths:
         Raises:
             FileNotFoundError: If the path is outside of the artifact store
                 bounds.
+            IllegalOperationError: If the path is a local file and the server
+                is not configured to allow local file access.
         """
+        if not self.allow_local_file_access and not io_utils.is_remote(path):
+            raise IllegalOperationError(
+                "Files in a local artifact store cannot be accessed from the "
+                "server."
+            )
+
         if not path.startswith(self.fixed_root_path):
             raise FileNotFoundError(
                 f"File `{path}` is outside of "
@@ -169,9 +187,18 @@ class _sanitize_paths:
 
 
 class BaseArtifactStoreConfig(StackComponentConfig):
-    """Config class for `BaseArtifactStore`."""
+    """Config class for `BaseArtifactStore`.
 
-    path: str
+    Base configuration for artifact storage backends.
+    Field descriptions are defined inline using Field() descriptors.
+    """
+
+    path: str = Field(
+        description="Root path for artifact storage. Must be a valid URI supported by the "
+        "specific artifact store implementation. Examples: 's3://my-bucket/artifacts', "
+        "'/local/storage/path', 'gs://bucket-name/zenml-artifacts', 'azure://container/path'. "
+        "Path must be accessible with the configured credentials and permissions"
+    )
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]]
     IS_IMMUTABLE_FILESYSTEM: ClassVar[bool] = False
@@ -435,40 +462,41 @@ class BaseArtifactStore(StackComponent):
             **kwargs: The keyword arguments to pass to the Pydantic object.
         """
         super(BaseArtifactStore, self).__init__(*args, **kwargs)
+        self._add_path_sanitization()
 
         # If running in a ZenML server environment, we don't register
         # the filesystems. We always use the artifact stores directly.
         if ENV_ZENML_SERVER not in os.environ:
             self._register()
 
+    def _add_path_sanitization(self) -> None:
+        """Add path sanitization to the artifact store."""
+        for method_name, method in inspect.getmembers(BaseArtifactStore):
+            if getattr(method, "__isabstractmethod__", False):
+                method_implementation = getattr(self, method_name)
+                sanitized_method = _sanitize_paths(
+                    method_implementation, self.path
+                )
+                setattr(self, method_name, sanitized_method)
+
     def _register(self) -> None:
         """Create and register a filesystem within the filesystem registry."""
         from zenml.io.filesystem import BaseFilesystem
         from zenml.io.filesystem_registry import default_filesystem_registry
         from zenml.io.local_filesystem import LocalFilesystem
 
-        overloads: Dict[str, Any] = {
-            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
-        }
-        for abc_method in inspect.getmembers(BaseArtifactStore):
-            if getattr(abc_method[1], "__isabstractmethod__", False):
-                sanitized_method = _sanitize_paths(
-                    getattr(self, abc_method[0]), self.path
-                )
-                # prepare overloads for filesystem methods
-                overloads[abc_method[0]] = staticmethod(sanitized_method)
-
-                # decorate artifact store methods
-                setattr(
-                    self,
-                    abc_method[0],
-                    sanitized_method,
-                )
-
         # Local filesystem is always registered, no point in doing it again.
         if isinstance(self, LocalFilesystem):
             return
 
+        overloads: Dict[str, Any] = {
+            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
+        }
+        for method_name, method in inspect.getmembers(BaseArtifactStore):
+            if getattr(method, "__isabstractmethod__", False):
+                method_implementation = getattr(self, method_name)
+                overloads[method_name] = staticmethod(method_implementation)
+
         filesystem_class = type(
             self.__class__.__name__, (BaseFilesystem,), overloads
         )

zenml/artifacts/utils.py

@@ -49,6 +49,7 @@ from zenml.enums import (
 )
 from zenml.exceptions import (
     DoesNotExistException,
+    IllegalOperationError,
     StepContextError,
 )
 from zenml.io import fileio
@@ -925,6 +926,7 @@ def _load_file_from_artifact_store(
         DoesNotExistException: If the file does not exist in the artifact store.
         NotImplementedError: If the artifact store cannot open the file.
         IOError: If the artifact store rejects the request.
+        IllegalOperationError: If the artifact store rejects the request.
     """
     try:
         with artifact_store.open(uri, mode) as text_file:
@@ -946,7 +948,7 @@ def _load_file_from_artifact_store(
             f"File '{uri}' does not exist in artifact store "
             f"'{artifact_store.name}'."
         )
-    except IOError as e:
+    except (IOError, IllegalOperationError) as e:
         raise e
     except Exception as e:
         logger.exception(e)

zenml/cli/login.py

@@ -22,6 +22,10 @@ from typing import Any, Dict, Optional, Tuple, Union
 from uuid import UUID
 
 import click
+from pydantic import BaseModel
+from rich.panel import Panel
+from rich.prompt import Prompt
+from rich.text import Text
 
 from zenml.cli import utils as cli_utils
 from zenml.cli.cli import cli
@@ -46,6 +50,77 @@ from zenml.utils.server_utils import (
 logger = get_logger(__name__)
 
 
+class LoginMethod(BaseModel):
+    """Login method class."""
+
+    name: str
+    description: str
+    help: str
+
+
+possible_login_methods = [
+    LoginMethod(
+        name="local",
+        description="Login to a local server",
+        help="(zenml login --local)",
+    ),
+    LoginMethod(
+        name="pro",
+        description="Login to ZenML Pro",
+        help="(https://cloud.zenml.io)",
+    ),
+    LoginMethod(
+        name="cloud",
+        description="Login to a cloud server",
+        help="(custom URL)",
+    ),
+]
+
+
+def _display_login_menu() -> LoginMethod:
+    """Display an interactive login menu and return the user's choice.
+
+    Returns:
+        The selected login method enum value.
+    """
+    title_text = Text("ZenML Login", style="bold cyan")
+
+    options_text = Text()
+    options_text.append("Choose your login method:\n\n", style="dim")
+
+    for i, login_method in enumerate(possible_login_methods):
+        options_text.append(f"{i + 1}. ", style="bold purple")
+        options_text.append(
+            f"{login_method.description}"
+            + " " * (25 - len(login_method.description)),
+            style="white",
+        )
+        options_text.append(f"{login_method.help}", style="dim")
+        if i < len(possible_login_methods) - 1:
+            options_text.append("\n")
+
+    panel = Panel(
+        options_text,
+        title=title_text,
+        border_style="white dim",
+        padding=(1, 2),
+        width=60,
+    )
+
+    console.print(panel)
+
+    # Get user choice with validation
+    while True:
+        choice = Prompt.ask(
+            "[bold]Enter your choice[/bold]",
+            choices=list(
+                str(i) for i in range(1, len(possible_login_methods) + 1)
+            ),
+            default="2",
+        )
+        return possible_login_methods[int(choice) - 1]
+
+
 def start_local_server(
     docker: bool = False,
     ip_address: Union[
@@ -424,7 +499,7 @@ def connect_to_pro_server(
         )
 
     cli_utils.declare(
-        f"Connecting to ZenML Pro server: {server.name} [{str(server.id)}] "
+        f"Connecting to ZenML Pro server: '{server.name}' [{str(server.id)}] "
     )
 
     connect_to_server(
@@ -435,7 +510,7 @@ def connect_to_pro_server(
     # ZenML Pro workspace object.
     credentials_store.update_server_info(server.url, server)
 
-    cli_utils.declare(f"Connected to ZenML Pro server: {server.name}.")
+    cli_utils.success(f"✔ Connected to ZenML Pro server: {server.name}.")
 
 
 def is_pro_server(
@@ -523,8 +598,10 @@ def _fail_if_authentication_environment_variables_set() -> None:
     on the current client state:
 
       * if the client is not connected to a non-local ZenML server, the command
-        will take the user to the ZenML Pro login / signup page to authenticate
-        and connect to a ZenML Pro server.
+        will display an interactive menu with three login options:
+        1. Login to a local server (equivalent to `zenml login --local`)
+        2. Login to ZenML Pro (https://cloud.zenml.io)
+        3. Login to a custom cloud server (you'll be prompted for the URL)
 
       * if the client is already connected to a non-local ZenML server, the
         command triggers a new web login flow with the same server. This allows
@@ -895,20 +972,66 @@ def login(
             )
     else:
         # If no server argument is provided, and the client is not currently
-        # connected to any non-local server, we default to logging in to ZenML
-        # Pro.
-        cli_utils.declare(
-            "No server argument was provided. Logging to ZenML Pro...\n"
-            "Hint: You can run 'zenml login --local' to start a local ZenML "
-            "server and connect to it or 'zenml login <server-url>' to connect "
-            "to a specific ZenML server. If you wish to login to a ZenML Pro "
-            "server, you can run 'zenml login --pro'."
-        )
-        connect_to_pro_server(
-            api_key=api_key_value,
-            pro_api_url=pro_api_url,
-            verify_ssl=verify_ssl,
-        )
+        # connected to any non-local server, show the interactive login menu.
+        login_method = _display_login_menu()
+
+        if login_method.name == "local":
+            # Start a local ZenML server and connect to it
+            start_local_server(
+                docker=docker,
+                ip_address=ip_address,
+                port=port,
+                blocking=blocking,
+                image=image,
+                ngrok_token=ngrok_token,
+                restart=restart,
+            )
+        elif login_method.name == "pro":
+            # Connect to ZenML Pro
+            connect_to_pro_server(
+                api_key=api_key_value,
+                pro_api_url=pro_api_url,
+                verify_ssl=verify_ssl,
+            )
+        elif login_method.name == "cloud":
+            # Get custom server URL from user
+            console.print()
+            server_url = Prompt.ask(
+                "[bold]Enter the ZenML server URL[/bold]", default="http/https"
+            )
+
+            if not server_url.strip():
+                cli_utils.error("Server URL cannot be empty.")
+                return
+
+            # Validate URL format
+            if not re.match(r"^https?://", server_url):
+                cli_utils.error(
+                    "Invalid URL format. Please provide a URL starting with "
+                    "http:// or https://"
+                )
+                return
+
+            # Connect to the custom server
+            # First, try to discover if the server is a ZenML Pro server or not
+            server_is_pro, server_pro_api_url = is_pro_server(server_url)
+            if server_is_pro:
+                connect_to_pro_server(
+                    pro_server=server_url,
+                    api_key=api_key_value,
+                    refresh=True,  # Force refresh for manually entered URLs
+                    # Prefer the pro API URL extracted from the server info if
+                    # available
+                    pro_api_url=server_pro_api_url or pro_api_url,
+                    verify_ssl=verify_ssl,
+                )
+            else:
+                connect_to_server(
+                    url=server_url,
+                    api_key=api_key_value,
+                    verify_ssl=verify_ssl,
+                    refresh=True,  # Force refresh for manually entered URLs
+                )
 
 
 @cli.command(

zenml/cli/pipeline.py

@@ -604,16 +604,27 @@ def delete_pipeline_run(
 
 @runs.command("refresh")
 @click.argument("run_name_or_id", type=str, required=True)
-def refresh_pipeline_run(run_name_or_id: str) -> None:
+@click.option(
+    "--include-steps",
+    is_flag=True,
+    default=False,
+    help="Also refresh the status of individual steps.",
+)
+def refresh_pipeline_run(
+    run_name_or_id: str, include_steps: bool = False
+) -> None:
     """Refresh the status of a pipeline run.
 
     Args:
         run_name_or_id: The name or ID of the pipeline run to refresh.
+        include_steps: If True, also refresh the status of individual steps.
     """
     try:
         # Fetch and update the run
         run = Client().get_pipeline_run(name_id_or_prefix=run_name_or_id)
-        run.refresh_run_status()
+        run_utils.refresh_run_status(
+            run=run, include_step_updates=include_steps
+        )
 
     except KeyError as e:
         cli_utils.error(str(e))

zenml/cli/project.py

@@ -109,21 +109,23 @@ def register_project(
                 description="",
                 display_name=display_name,
             )
-            cli_utils.declare("Project created successfully.")
+            cli_utils.success("✔ Project created successfully.")
         except Exception as e:
             cli_utils.error(str(e))
 
     if set_project:
         client.set_active_project(project_name)
-        cli_utils.declare(f"The active project has been set to {project_name}")
+        cli_utils.success(
+            f"✔ The active project has been set to {project_name}"
+        )
 
     if set_default:
         client.update_user(
             name_id_or_prefix=client.active_user.id,
             updated_default_project_id=project.id,
         )
-        cli_utils.declare(
-            f"The default project has been set to {project.name}"
+        cli_utils.success(
+            f"✔ The default project has been set to {project.name}"
         )
 
 
@@ -147,8 +149,8 @@ def set_project(project_name_or_id: str, default: bool = False) -> None:
     with console.status("Setting project...\n"):
         try:
             project = client.set_active_project(project_name_or_id)
-            cli_utils.declare(
-                f"The active project has been set to {project_name_or_id}"
+            cli_utils.success(
+                f"✔ The active project has been set to {project_name_or_id}"
             )
         except Exception as e:
             cli_utils.error(str(e))

zenml/cli/utils.py

@@ -22,6 +22,7 @@ import shutil
 import subprocess
 import sys
 from typing import (
+    IO,
     TYPE_CHECKING,
     AbstractSet,
     Any,
@@ -162,9 +163,20 @@ def error(text: str) -> NoReturn:
         text: Input text string.
 
     Raises:
-        ClickException: when called.
+        StyledClickException: when called.
     """
-    raise click.ClickException(message=click.style(text, fg="red", bold=True))
+    error_prefix = click.style("Error: ", fg="red", bold=True)
+    error_message = click.style(text, fg="red", bold=False)
+
+    # Create a custom ClickException that bypasses Click's default "Error: " prefix
+    class StyledClickException(click.ClickException):
+        def show(self, file: Optional[IO[Any]] = None) -> None:
+            if file is None:
+                file = click.get_text_stream("stderr")
+            # Print our custom styled message directly without Click's prefix
+            click.echo(self.message, file=file)
+
+    raise StyledClickException(message=error_prefix + error_message)
 
 
 def warning(
@@ -186,6 +198,25 @@ def warning(
     console.print(text, style=style, **kwargs)
 
 
+def success(
+    text: str,
+    bold: Optional[bool] = None,
+    italic: Optional[bool] = None,
+    **kwargs: Any,
+) -> None:
+    """Echo a success string on the CLI.
+
+    Args:
+        text: Input text string.
+        bold: Optional boolean to bold the text.
+        italic: Optional boolean to italicize the text.
+        **kwargs: Optional kwargs to be passed to console.print().
+    """
+    base_style = zenml_style_defaults["success"]
+    style = Style.chain(base_style, Style(bold=bold, italic=italic))
+    console.print(text, style=style, **kwargs)
+
+
 def print_markdown(text: str) -> None:
     """Prints a string as markdown.
 
@@ -230,7 +261,11 @@ def print_table(
     column_keys = {key: None for dict_ in obj for key in dict_}
     column_names = [columns.get(key, key.upper()) for key in column_keys]
     rich_table = table.Table(
-        box=box.HEAVY_EDGE, show_lines=True, title=title, caption=caption
+        box=box.ROUNDED,
+        show_lines=True,
+        title=title,
+        caption=caption,
+        border_style="dim",
     )
     for col_name in column_names:
         if isinstance(col_name, str):
@@ -434,9 +469,10 @@ def print_pydantic_model(
         columns: Optionally specify subset and order of columns to display.
     """
     rich_table = table.Table(
-        box=box.HEAVY_EDGE,
+        box=box.ROUNDED,
         title=title,
         show_lines=True,
+        border_style="dim",
     )
     rich_table.add_column("PROPERTY", overflow="fold")
     rich_table.add_column("VALUE", overflow="fold")
@@ -552,10 +588,11 @@ def print_stack_configuration(stack: "StackResponse", active: bool) -> None:
     if active:
         stack_caption += " (ACTIVE)"
     rich_table = table.Table(
-        box=box.HEAVY_EDGE,
+        box=box.ROUNDED,
         title="Stack Configuration",
         caption=stack_caption,
         show_lines=True,
+        border_style="dim",
     )
     rich_table.add_column("COMPONENT_TYPE", overflow="fold")
     rich_table.add_column("COMPONENT_NAME", overflow="fold")
@@ -573,9 +610,10 @@ def print_stack_configuration(stack: "StackResponse", active: bool) -> None:
         declare("No labels are set for this stack.")
     else:
         rich_table = table.Table(
-            box=box.HEAVY_EDGE,
+            box=box.ROUNDED,
             title="Labels",
             show_lines=True,
+            border_style="dim",
         )
         rich_table.add_column("LABEL")
         rich_table.add_column("VALUE", overflow="fold")
@@ -649,9 +687,10 @@ def print_stack_component_configuration(
         if active_status:
             title_ += " (ACTIVE)"
         rich_table = table.Table(
-            box=box.HEAVY_EDGE,
+            box=box.ROUNDED,
             title=title_,
             show_lines=True,
+            border_style="dim",
         )
         rich_table.add_column("COMPONENT_PROPERTY")
         rich_table.add_column("VALUE", overflow="fold")
@@ -672,9 +711,10 @@ def print_stack_component_configuration(
         declare("No labels are set for this component.")
     else:
         rich_table = table.Table(
-            box=box.HEAVY_EDGE,
+            box=box.ROUNDED,
             title="Labels",
             show_lines=True,
+            border_style="dim",
         )
         rich_table.add_column("LABEL")
         rich_table.add_column("VALUE", overflow="fold")
@@ -688,9 +728,10 @@ def print_stack_component_configuration(
         declare("No connector is set for this component.")
     else:
         rich_table = table.Table(
-            box=box.HEAVY_EDGE,
+            box=box.ROUNDED,
             title="Service Connector",
             show_lines=True,
+            border_style="dim",
         )
         rich_table.add_column("PROPERTY")
         rich_table.add_column("VALUE", overflow="fold")
@@ -1165,8 +1206,9 @@ def print_list_items(list_items: List[str], column_title: str) -> None:
         column_title: Title of the column
     """
     rich_table = table.Table(
-        box=box.HEAVY_EDGE,
+        box=box.ROUNDED,
         show_lines=True,
+        border_style="dim",
     )
     rich_table.add_column(column_title.upper(), overflow="fold")
     list_items.sort()
@@ -1291,9 +1333,10 @@ def pretty_print_model_version_details(
     title_ = f"Properties of model `{model_version.registered_model.name}` version `{model_version.version}`"
 
     rich_table = table.Table(
-        box=box.HEAVY_EDGE,
+        box=box.ROUNDED,
         title=title_,
         show_lines=True,
+        border_style="dim",
     )
     rich_table.add_column("MODEL VERSION PROPERTY", overflow="fold")
     rich_table.add_column("VALUE", overflow="fold")
@@ -1343,9 +1386,10 @@ def print_served_model_configuration(
     title_ = f"Properties of Served Model {model_service.uuid}"
 
     rich_table = table.Table(
-        box=box.HEAVY_EDGE,
+        box=box.ROUNDED,
         title=title_,
         show_lines=True,
+        border_style="dim",
     )
     rich_table.add_column("MODEL SERVICE PROPERTY", overflow="fold")
     rich_table.add_column("VALUE", overflow="fold")
@@ -1745,9 +1789,10 @@ def print_service_connector_configuration(
     if active_status:
         title_ += " (ACTIVE)"
     rich_table = table.Table(
-        box=box.HEAVY_EDGE,
+        box=box.ROUNDED,
         title=title_,
         show_lines=True,
+        border_style="dim",
     )
     rich_table.add_column("PROPERTY")
     rich_table.add_column("VALUE", overflow="fold")
@@ -1819,9 +1864,10 @@ def print_service_connector_configuration(
 
     else:
         rich_table = table.Table(
-            box=box.HEAVY_EDGE,
+            box=box.ROUNDED,
             title="Configuration",
             show_lines=True,
+            border_style="dim",
         )
         rich_table.add_column("PROPERTY")
         rich_table.add_column("VALUE", overflow="fold")
@@ -1847,9 +1893,10 @@ def print_service_connector_configuration(
         return
 
     rich_table = table.Table(
-        box=box.HEAVY_EDGE,
+        box=box.ROUNDED,
         title="Labels",
         show_lines=True,
+        border_style="dim",
     )
     rich_table.add_column("LABEL")
     rich_table.add_column("VALUE", overflow="fold")
@@ -2609,7 +2656,7 @@ def multi_choice_prompt(
     table = Table(
         title=f"Available {object_type}",
         show_header=True,
-        border_style=None,
+        border_style="dim",
         expand=True,
         show_lines=True,
     )

zenml/client.py

@@ -4051,6 +4051,7 @@ class Client(metaclass=ClientMetaClass):
         model_version_id: Optional[Union[str, UUID]] = None,
         model: Optional[Union[UUID, str]] = None,
         run_metadata: Optional[List[str]] = None,
+        exclude_retried: Optional[bool] = None,
         hydrate: bool = False,
     ) -> Page[StepRunResponse]:
         """List all pipelines.
@@ -4077,6 +4078,7 @@ class Client(metaclass=ClientMetaClass):
             code_hash: The code hash of the step run to filter by.
             status: The name of the run to filter by.
             run_metadata: Filter by run metadata.
+            exclude_retried: Whether to exclude retried step runs.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -4105,6 +4107,7 @@ class Client(metaclass=ClientMetaClass):
             model_version_id=model_version_id,
             model=model,
             run_metadata=run_metadata,
+            exclude_retried=exclude_retried,
         )
         return self.zen_store.list_run_steps(
             step_run_filter_model=step_run_filter_model,
@@ -4352,7 +4355,7 @@ class Client(metaclass=ClientMetaClass):
         version: Optional[Union[str, int]] = None,
         version_number: Optional[int] = None,
         artifact_store_id: Optional[Union[str, UUID]] = None,
-        type: Optional[ArtifactType] = None,
+        type: Optional[Union[ArtifactType, str]] = None,
         data_type: Optional[str] = None,
         uri: Optional[str] = None,
         materializer: Optional[str] = None,

zenml/config/compiler.py

@@ -206,6 +206,7 @@ class Compiler:
                 tags=config.tags,
                 extra=config.extra,
                 model=config.model,
+                retry=config.retry,
                 parameters=config.parameters,
             )