zenml-nightly 0.83.1.dev20250708__py3-none-any.whl → 0.83.1.dev20250710__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.83.1.dev20250708
+0.83.1.dev20250710

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/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,
             )
 

zenml/config/retry_config.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """Retry configuration for a step."""
 
+from pydantic import NonNegativeInt, PositiveInt
+
 from zenml.config.strict_base_model import StrictBaseModel
 
 
@@ -22,6 +24,6 @@ class StepRetryConfig(StrictBaseModel):
     Delay is an integer (specified in seconds).
     """
 
-    max_retries: int = 1
-    delay: int = 0  # in seconds
-    backoff: int = 0
+    max_retries: PositiveInt = 1
+    delay: NonNegativeInt = 0  # in seconds
+    backoff: NonNegativeInt = 0

zenml/config/step_configurations.py

@@ -289,6 +289,7 @@ class StepConfiguration(PartialStepConfiguration):
                 "extra",
                 "failure_hook_source",
                 "success_hook_source",
+                "retry",
                 "substitutions",
             },
             exclude_none=True,
@@ -300,12 +301,17 @@ class StepConfiguration(PartialStepConfiguration):
                     "extra",
                     "failure_hook_source",
                     "success_hook_source",
+                    "retry",
                     "substitutions",
                 },
                 exclude_none=True,
             )
 
-            updated_config = self.model_copy(update=pipeline_values, deep=True)
+            updated_config_dict = {
+                **self.model_dump(),
+                **pipeline_values,
+            }
+            updated_config = self.model_validate(updated_config_dict)
             return update_model(updated_config, original_values)
         else:
             return self.model_copy(deep=True)

zenml/console.py

@@ -18,11 +18,14 @@ from rich.style import Style
 from rich.theme import Theme
 
 zenml_style_defaults = {
-    "info": Style(color="cyan", dim=True),
+    "info": Style(color="white", dim=True),
     "warning": Style(color="yellow"),
     "danger": Style(color="red", bold=True),
     "title": Style(color="cyan", bold=True, underline=True),
     "error": Style(color="red"),
+    "success": Style(color="green"),
+    "repr.str": Style(color="white", dim=False),
+    "repr.uuid": Style(color="magenta", dim=False),
 }
 
 zenml_custom_theme = Theme(zenml_style_defaults)

zenml/constants.py

@@ -171,7 +171,6 @@ ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE = "ZENML_DISABLE_PIPELINE_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS = "ZENML_DISABLE_STEP_NAMES_IN_LOGS"
-ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"

zenml/enums.py

@@ -78,6 +78,8 @@ class ExecutionStatus(StrEnum):
     COMPLETED = "completed"
     RUNNING = "running"
     CACHED = "cached"
+    RETRYING = "retrying"
+    RETRIED = "retried"
     STOPPED = "stopped"
     STOPPING = "stopping"
 
@@ -92,9 +94,19 @@ class ExecutionStatus(StrEnum):
             ExecutionStatus.FAILED,
             ExecutionStatus.COMPLETED,
             ExecutionStatus.CACHED,
+            ExecutionStatus.RETRIED,
             ExecutionStatus.STOPPED,
         }
 
+    @property
+    def is_successful(self) -> bool:
+        """Whether the execution status refers to a successful execution.
+
+        Returns:
+            Whether the execution status refers to a successful execution.
+        """
+        return self in {ExecutionStatus.COMPLETED, ExecutionStatus.CACHED}
+
 
 class LoggingLevels(Enum):
     """Enum for logging levels."""
@@ -411,9 +423,6 @@ class OnboardingStep(StrEnum):
     DEVICE_VERIFIED = "device_verified"
     PROJECT_CREATED = "project_created"
     PIPELINE_RUN = "pipeline_run"
-    SECOND_PIPELINE_RUN = "second_pipeline_run"
-    THIRD_PIPELINE_RUN = "third_pipeline_run"
-    STARTER_SETUP_COMPLETED = "starter_setup_completed"
     STACK_WITH_REMOTE_ORCHESTRATOR_CREATED = (
         "stack_with_remote_orchestrator_created"
     )
@@ -426,7 +435,7 @@ class OnboardingStep(StrEnum):
     PIPELINE_RUN_WITH_REMOTE_ARTIFACT_STORE = (
         "pipeline_run_with_remote_artifact_store"
     )
-    PRODUCTION_SETUP_COMPLETED = "production_setup_completed"
+    OSS_ONBOARDING_COMPLETED = "oss_onboarding_completed"
     PRO_ONBOARDING_COMPLETED = "pro_onboarding_completed"