zenml-nightly 0.75.0.dev20250311__py3-none-any.whl → 0.75.0.dev20250313__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.75.0.dev20250311
+0.75.0.dev20250313

zenml/__init__.py

@@ -49,9 +49,11 @@ from zenml.pipelines import get_pipeline_context, pipeline
 from zenml.steps import step, get_step_context
 from zenml.steps.utils import log_step_metadata
 from zenml.utils.metadata_utils import log_metadata
+from zenml.utils.tag_utils import Tag, add_tags, remove_tags
 from zenml.entrypoints import entrypoint
 
 __all__ = [
+    "add_tags",
     "ArtifactConfig",
     "ExternalArtifact",
     "get_pipeline_context",

zenml/analytics/context.py

@@ -118,7 +118,10 @@ class AnalyticsContext:
                     self.external_user_id = auth_context.user.external_user_id
 
                 self.external_server_id = server_config().external_server_id
+                self.workspace_id = None
             else:
+                from zenml.client import Client
+
                 # If the code is running on the client, use the default user.
                 active_user = gc.zen_store.get_user()
                 self.user_id = active_user.id
@@ -126,6 +129,7 @@ class AnalyticsContext:
                     active_user.is_service_account
                 )
                 self.external_user_id = active_user.external_user_id
+                self.workspace_id = Client().active_workspace.id
 
             # Fetch the `client_id`
             if self.in_server:
@@ -304,6 +308,9 @@ class AnalyticsContext:
         if self.server_metadata:
             properties.update(self.server_metadata)
 
+        if self.workspace_id:
+            properties.setdefault("workspace_id", str(self.workspace_id))
+
         for k, v in properties.items():
             if isinstance(v, UUID):
                 properties[k] = str(v)

zenml/artifacts/utils.py

@@ -182,7 +182,6 @@ def _store_artifact_data_and_prepare_request(
         uri=materializer.uri,
         materializer=source_utils.resolve(materializer.__class__),
         data_type=source_utils.resolve(data_type),
-        user=Client().active_user.id,
         workspace=Client().active_workspace.id,
         artifact_store_id=artifact_store.id,
         visualizations=visualizations,
@@ -350,7 +349,6 @@ def register_artifact(
         uri=folder_or_file_uri,
         materializer=source_utils.resolve(PreexistingDataMaterializer),
         data_type=source_utils.resolve(Path),
-        user=Client().active_user.id,
         workspace=Client().active_workspace.id,
         artifact_store_id=artifact_store.id,
         has_custom_name=has_custom_name,

zenml/cli/login.py

@@ -291,6 +291,11 @@ def connect_to_pro_server(
             login = True
 
     if login or refresh:
+        # If we reached this point, then we need to start a new login flow.
+        # We also need to remove all existing API tokens associated with the
+        # target ZenML Pro API, otherwise they will continue to be used after
+        # the re-login flow.
+        credentials_store.clear_all_pro_tokens(pro_api_url)
         try:
             token = web_login(
                 pro_api_url=pro_api_url,
@@ -976,6 +981,7 @@ def logout(
             )
 
         pro_api_url = pro_api_url or ZENML_PRO_API_URL
+        pro_api_url = pro_api_url.rstrip("/")
         if credentials_store.has_valid_pro_authentication(pro_api_url):
             credentials_store.clear_pro_credentials(pro_api_url)
             cli_utils.declare("Logged out from ZenML Pro.")

zenml/cli/model.py

@@ -91,9 +91,7 @@ def list_models(**kwargs: Any) -> None:
     Args:
         **kwargs: Keyword arguments to filter models.
     """
-    models = Client().zen_store.list_models(
-        model_filter_model=ModelFilter(**kwargs)
-    )
+    models = Client().list_models(**kwargs)
 
     if not models:
         cli_utils.declare("No models found.")
@@ -399,24 +397,18 @@ def version() -> None:
 
 
 @cli_utils.list_options(ModelVersionFilter)
-@click.option(
-    "--model-name",
-    "-n",
-    help="The name of the parent model.",
-    type=str,
-    required=False,
-)
 @version.command("list", help="List model versions with filter.")
-def list_model_versions(model_name: str, **kwargs: Any) -> None:
+@click.argument("model_name_or_id")
+def list_model_versions(model_name_or_id: str, **kwargs: Any) -> None:
     """List model versions with filter in the Model Control Plane.
 
     Args:
-        model_name: The name of the parent model.
+        model_name_or_id: The name or ID of the parent model.
         **kwargs: Keyword arguments to filter models.
     """
-    model_versions = Client().zen_store.list_model_versions(
-        model_name_or_id=model_name,
-        model_version_filter_model=ModelVersionFilter(**kwargs),
+    model_versions = Client().list_model_versions(
+        model_name_or_id=model_name_or_id,
+        **kwargs,
     )
 
     if not model_versions:

zenml/cli/secret.py

@@ -38,7 +38,6 @@ from zenml.console import console
 from zenml.constants import SECRET_VALUES
 from zenml.enums import (
     CliCategories,
-    SecretScope,
 )
 from zenml.exceptions import EntityExistsError, ZenKeyError
 from zenml.logger import get_logger
@@ -59,11 +58,12 @@ def secret() -> None:
 )
 @click.argument("name", type=click.STRING)
 @click.option(
-    "--scope",
-    "-s",
-    "scope",
-    type=click.Choice([scope.value for scope in list(SecretScope)]),
-    default=SecretScope.WORKSPACE.value,
+    "--private",
+    "-p",
+    "private",
+    is_flag=True,
+    help="Whether the secret is private. A private secret is only accessible "
+    "to the user who creates it.",
 )
 @click.option(
     "--interactive",
@@ -84,13 +84,13 @@ def secret() -> None:
 )
 @click.argument("args", nargs=-1, type=click.UNPROCESSED)
 def create_secret(
-    name: str, scope: str, interactive: bool, values: str, args: List[str]
+    name: str, private: bool, interactive: bool, values: str, args: List[str]
 ) -> None:
     """Create a secret.
 
     Args:
         name: The name of the secret to create.
-        scope: The scope of the secret to create.
+        private: Whether the secret is private.
         interactive: Whether to use interactive mode to enter the secret values.
         values: Secret key-value pairs to be passed as JSON or YAML.
         args: The arguments to pass to the secret.
@@ -152,7 +152,7 @@ def create_secret(
         with console.status(f"Saving secret `{name}`..."):
             try:
                 client.create_secret(
-                    name=name, values=parsed_args, scope=SecretScope(scope)
+                    name=name, values=parsed_args, private=private
                 )
                 declare(f"Secret '{name}' successfully created.")
             except EntityExistsError as e:
@@ -186,7 +186,7 @@ def list_secrets(**kwargs: Any) -> None:
             dict(
                 name=secret.name,
                 id=str(secret.id),
-                scope=secret.scope.value,
+                private=secret.private,
             )
             for secret in secrets.items
         ]
@@ -200,22 +200,26 @@ def list_secrets(**kwargs: Any) -> None:
     type=click.STRING,
 )
 @click.option(
-    "--scope",
-    "-s",
-    type=click.Choice([scope.value for scope in list(SecretScope)]),
-    default=None,
+    "--private",
+    "-p",
+    "private",
+    type=click.BOOL,
+    required=False,
+    help="Use this flag to explicitly fetch a private secret or a public secret.",
 )
-def get_secret(name_id_or_prefix: str, scope: Optional[str] = None) -> None:
+def get_secret(name_id_or_prefix: str, private: Optional[bool] = None) -> None:
     """Get a secret and print it to the console.
 
     Args:
         name_id_or_prefix: The name of the secret to get.
-        scope: The scope of the secret to get.
+        private: Private status of the secret to filter for.
     """
-    secret = _get_secret(name_id_or_prefix, scope)
+    secret = _get_secret(name_id_or_prefix, private)
+    scope = ""
+    if private is not None:
+        scope = "private " if private else "public "
     declare(
-        f"Fetched secret with name `{secret.name}` and ID `{secret.id}` in "
-        f"scope `{secret.scope.value}`:"
+        f"Fetched {scope}secret with name `{secret.name}` and ID `{secret.id}`:"
     )
     if not secret.secret_values:
         warning(f"Secret with name `{name_id_or_prefix}` is empty.")
@@ -224,25 +228,22 @@ def get_secret(name_id_or_prefix: str, scope: Optional[str] = None) -> None:
 
 
 def _get_secret(
-    name_id_or_prefix: str, scope: Optional[str] = None
+    name_id_or_prefix: str, private: Optional[bool] = None
 ) -> SecretResponse:
     """Get a secret with a given name, prefix or id.
 
     Args:
         name_id_or_prefix: The name of the secret to get.
-        scope: The scope of the secret to get.
+        private: Private status of the secret to filter for.
 
     Returns:
         The secret response model.
     """
     client = Client()
     try:
-        if scope:
-            return client.get_secret(
-                name_id_or_prefix=name_id_or_prefix, scope=SecretScope(scope)
-            )
-        else:
-            return client.get_secret(name_id_or_prefix=name_id_or_prefix)
+        return client.get_secret(
+            name_id_or_prefix=name_id_or_prefix, private=private
+        )
     except ZenKeyError as e:
         error(
             f"Error fetching secret with name id or prefix "
@@ -267,9 +268,12 @@ def _get_secret(
     type=click.STRING,
 )
 @click.option(
-    "--new-scope",
-    "-s",
-    type=click.Choice([scope.value for scope in list(SecretScope)]),
+    "--private",
+    "-p",
+    "private",
+    type=click.BOOL,
+    required=False,
+    help="Update the private status of the secret.",
 )
 @click.option(
     "--interactive",
@@ -293,7 +297,7 @@ def _get_secret(
 def update_secret(
     name_or_id: str,
     extra_args: List[str],
-    new_scope: Optional[str] = None,
+    private: Optional[bool] = None,
     remove_keys: List[str] = [],
     interactive: bool = False,
     values: str = "",
@@ -302,7 +306,7 @@ def update_secret(
 
     Args:
         name_or_id: The name or id of the secret to update.
-        new_scope: The new scope of the secret.
+        private: Private status of the secret to update.
         extra_args: The arguments to pass to the secret.
         interactive: Whether to use interactive mode to update the secret.
         remove_keys: The keys to remove from the secret.
@@ -331,10 +335,7 @@ def update_secret(
         except NotImplementedError as e:
             error(f"Centralized secrets management is disabled: {str(e)}")
 
-    declare(
-        f"Updating secret with name '{secret.name}' and ID '{secret.id}' in "
-        f"scope '{secret.scope.value}:"
-    )
+    declare(f"Updating secret with name '{secret.name}' and ID '{secret.id}'")
 
     if "name" in parsed_args:
         error("The word 'name' cannot be used as a key for a secret.")
@@ -388,7 +389,7 @@ def update_secret(
 
     client.update_secret(
         name_id_or_prefix=secret.id,
-        new_scope=SecretScope(new_scope) if new_scope else None,
+        update_private=private,
         add_or_update_values=secret_args_add_update,
         remove_values=remove_keys,
     )
@@ -492,10 +493,12 @@ def delete_secret(name_or_id: str, yes: bool = False) -> None:
     type=click.STRING,
 )
 @click.option(
-    "--scope",
-    "-s",
-    type=click.Choice([scope.value for scope in list(SecretScope)]),
-    default=None,
+    "--private",
+    "-p",
+    "private",
+    type=click.BOOL,
+    required=False,
+    help="Use this flag to explicitly fetch a private secret or a public secret.",
 )
 @click.option(
     "--filename",
@@ -509,7 +512,7 @@ def delete_secret(name_or_id: str, yes: bool = False) -> None:
 )
 def export_secret(
     name_id_or_prefix: str,
-    scope: Optional[str] = None,
+    private: Optional[bool] = None,
     filename: Optional[str] = None,
 ) -> None:
     """Export a secret as a YAML file.
@@ -519,12 +522,12 @@ def export_secret(
 
     Args:
         name_id_or_prefix: The name of the secret to export.
-        scope: The scope of the secret to export.
+        private: Private status of the secret to export.
         filename: The name of the file to export the secret to.
     """
     from zenml.utils.yaml_utils import write_yaml
 
-    secret = _get_secret(name_id_or_prefix=name_id_or_prefix, scope=scope)
+    secret = _get_secret(name_id_or_prefix=name_id_or_prefix, private=private)
     if not secret.secret_values:
         warning(f"Secret with name `{name_id_or_prefix}` is empty.")
         return

zenml/cli/service_connectors.py

@@ -1131,7 +1131,6 @@ def describe_service_connector(
             )
 
         connector = connector_client.to_response_model(
-            workspace=client.active_workspace,
             user=client.active_user,
         )
     else:

zenml/cli/stack.py

@@ -503,7 +503,6 @@ def register_stack(
         try:
             created_stack = client.zen_store.create_stack(
                 stack=StackRequest(
-                    user=client.active_user.id,
                     workspace=client.active_workspace.id,
                     name=stack_name,
                     components=components,

zenml/cli/tag.py

@@ -26,8 +26,6 @@ from zenml.exceptions import EntityExistsError
 from zenml.logger import get_logger
 from zenml.models import (
     TagFilter,
-    TagRequest,
-    TagUpdate,
 )
 from zenml.utils.dict_utils import remove_none_values
 
@@ -47,7 +45,7 @@ def list_tags(**kwargs: Any) -> None:
     Args:
         **kwargs: Keyword arguments to filter models.
     """
-    tags = Client().list_tags(TagFilter(**kwargs))
+    tags = Client().list_tags(**kwargs)
 
     if not tags:
         cli_utils.declare("No tags found.")
@@ -83,7 +81,7 @@ def register_tag(name: str, color: Optional[ColorVariants]) -> None:
     """
     request_dict = remove_none_values(dict(name=name, color=color))
     try:
-        tag = Client().create_tag(TagRequest(**request_dict))
+        tag = Client().create_tag(**request_dict)
     except (EntityExistsError, ValueError) as e:
         cli_utils.error(str(e))
 
@@ -126,7 +124,7 @@ def update_tag(
 
     tag = Client().update_tag(
         tag_name_or_id=tag_name_or_id,
-        tag_update_model=TagUpdate(**update_dict),
+        **update_dict,
     )
 
     cli_utils.print_pydantic_models(

zenml/cli/utils.py

@@ -81,6 +81,7 @@ from zenml.stack.flavor import Flavor
 from zenml.stack.stack_component import StackComponentConfig
 from zenml.utils import secret_utils
 from zenml.utils.time_utils import expires_in
+from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
     from uuid import UUID
@@ -1778,7 +1779,6 @@ def print_service_connector_configuration(
             "AUTH METHOD": connector.auth_method,
             "RESOURCE TYPES": ", ".join(connector.emojified_resource_types),
             "RESOURCE NAME": connector.resource_id or "<multiple>",
-            "SECRET ID": connector.secret_id or "",
             "SESSION DURATION": expiration,
             "EXPIRES IN": (
                 expires_in(
@@ -1795,7 +1795,6 @@ def print_service_connector_configuration(
                 else "N/A"
             ),
             "OWNER": user_name,
-            "WORKSPACE": connector.workspace.name,
             "CREATED_AT": connector.created,
             "UPDATED_AT": connector.updated,
         }
@@ -2161,9 +2160,6 @@ def print_debug_stack() -> None:
     console.print(f"ID: {str(stack.id)}")
     if stack.user and stack.user.name and stack.user.id:  # mypy check
         console.print(f"User: {stack.user.name} / {str(stack.user.id)}")
-    console.print(
-        f"Workspace: {stack.workspace.name} / {str(stack.workspace.id)}"
-    )
 
     for component_type, components in stack.components.items():
         component = components[0]
@@ -2190,9 +2186,6 @@ def print_debug_stack() -> None:
             console.print(
                 f"User: {component_response.user.name} / {str(component_response.user.id)}"
             )
-        console.print(
-            f"Workspace: {component_response.workspace.name} / {str(component_response.workspace.id)}"
-        )
 
 
 def _component_display_name(
@@ -2275,22 +2268,13 @@ def print_pipeline_runs_table(
     print_table(runs_dicts)
 
 
-def warn_unsupported_non_default_workspace() -> None:
-    """Warning for unsupported non-default workspace."""
-    from zenml.constants import (
-        ENV_ZENML_DISABLE_WORKSPACE_WARNINGS,
-        handle_bool_env_var,
-    )
-
-    disable_warnings = handle_bool_env_var(
-        ENV_ZENML_DISABLE_WORKSPACE_WARNINGS, False
-    )
-    if not disable_warnings:
+def check_zenml_pro_workspace_availability() -> None:
+    """Check if the ZenML Pro workspace feature is available."""
+    client = Client()
+    if not client.zen_store.get_store_info().is_pro_server():
         warning(
-            "Currently the concept of `workspace` is not supported "
-            "within the Dashboard. The Project functionality will be "
-            "completed in the coming weeks. For the time being it "
-            "is recommended to stay within the `default` workspace."
+            "The ZenML workspace feature is available only on ZenML Pro. "
+            "Please visit https://zenml.io/pro to learn more."
         )
 
 
@@ -2402,6 +2386,23 @@ def create_data_type_help_text(
         return f"{field}"
 
 
+def _is_list_field(field_info: Any) -> bool:
+    """Check if a field is a list field.
+
+    Args:
+        field_info: The field info to check.
+
+    Returns:
+        True if the field is a list field, False otherwise.
+    """
+    field_type = field_info.annotation
+    origin = get_origin(field_type)
+    return origin is list or (
+        is_union(origin)
+        and any(get_origin(arg) is list for arg in field_type.__args__)
+    )
+
+
 def list_options(filter_model: Type[BaseFilter]) -> Callable[[F], F]:
     """Create a decorator to generate the correct list of filter parameters.
 
@@ -2430,6 +2431,7 @@ def list_options(filter_model: Type[BaseFilter]) -> Callable[[F], F]:
                         type=str,
                         default=v.default,
                         required=False,
+                        multiple=_is_list_field(v),
                         help=create_filter_help_text(filter_model, k),
                     )
                 )

zenml/cli/workspace.py

@@ -20,9 +20,9 @@ import click
 from zenml.cli import utils as cli_utils
 from zenml.cli.cli import TagGroup, cli
 from zenml.cli.utils import (
+    check_zenml_pro_workspace_availability,
     is_sorted_or_filtered,
     list_options,
-    warn_unsupported_non_default_workspace,
 )
 from zenml.client import Client
 from zenml.console import console
@@ -35,7 +35,7 @@ def workspace() -> None:
     """Commands for workspace management."""
 
 
-@workspace.command("list", hidden=True)
+@workspace.command("list")
 @list_options(WorkspaceFilter)
 @click.pass_context
 def list_workspaces(ctx: click.Context, **kwargs: Any) -> None:
@@ -45,7 +45,7 @@ def list_workspaces(ctx: click.Context, **kwargs: Any) -> None:
         ctx: The click context object
         **kwargs: Keyword arguments to filter the list of workspaces.
     """
-    warn_unsupported_non_default_workspace()
+    check_zenml_pro_workspace_availability()
     client = Client()
     with console.status("Listing workspaces...\n"):
         workspaces = client.list_workspaces(**kwargs)
@@ -60,7 +60,61 @@ def list_workspaces(ctx: click.Context, **kwargs: Any) -> None:
             cli_utils.declare("No workspaces found for the given filter.")
 
 
-@workspace.command("describe", hidden=True)
+@workspace.command("register")
+@click.option(
+    "--set",
+    "set_workspace",
+    is_flag=True,
+    help="Immediately set this workspace as active.",
+    type=click.BOOL,
+)
+@click.argument("workspace_name", type=str, required=True)
+def register_workspace(
+    workspace_name: str, set_workspace: bool = False
+) -> None:
+    """Register a new workspace.
+
+    Args:
+        workspace_name: The name of the workspace to register.
+        set_workspace: Whether to set the workspace as active.
+    """
+    check_zenml_pro_workspace_availability()
+    client = Client()
+    with console.status("Creating workspace...\n"):
+        try:
+            client.create_workspace(workspace_name, description="")
+            cli_utils.declare("Workspace created successfully.")
+        except Exception as e:
+            cli_utils.error(str(e))
+
+    if set_workspace:
+        client.set_active_workspace(workspace_name)
+        cli_utils.declare(
+            f"The active workspace has been set to {workspace_name}"
+        )
+
+
+@workspace.command("set")
+@click.argument("workspace_name_or_id", type=str, required=True)
+def set_workspace(workspace_name_or_id: str) -> None:
+    """Set the active workspace.
+
+    Args:
+        workspace_name_or_id: The name or ID of the workspace to set as active.
+    """
+    check_zenml_pro_workspace_availability()
+    client = Client()
+    with console.status("Setting workspace...\n"):
+        try:
+            client.set_active_workspace(workspace_name_or_id)
+            cli_utils.declare(
+                f"The active workspace has been set to {workspace_name_or_id}"
+            )
+        except Exception as e:
+            cli_utils.error(str(e))
+
+
+@workspace.command("describe")
 @click.argument("workspace_name_or_id", type=str, required=False)
 def describe_workspace(workspace_name_or_id: Optional[str] = None) -> None:
     """Get the workspace.
@@ -68,7 +122,7 @@ def describe_workspace(workspace_name_or_id: Optional[str] = None) -> None:
     Args:
         workspace_name_or_id: The name or ID of the workspace to set as active.
     """
-    warn_unsupported_non_default_workspace()
+    check_zenml_pro_workspace_availability()
     client = Client()
     if not workspace_name_or_id:
         active_workspace = client.active_workspace
@@ -84,3 +138,23 @@ def describe_workspace(workspace_name_or_id: Optional[str] = None) -> None:
             cli_utils.print_pydantic_models(
                 [workspace_], exclude_columns=["created", "updated"]
             )
+
+
+@workspace.command("delete")
+@click.argument("workspace_name_or_id", type=str, required=True)
+def delete_workspace(workspace_name_or_id: str) -> None:
+    """Delete a workspace.
+
+    Args:
+        workspace_name_or_id: The name or ID of the workspace to delete.
+    """
+    check_zenml_pro_workspace_availability()
+    client = Client()
+    with console.status("Deleting workspace...\n"):
+        try:
+            client.delete_workspace(workspace_name_or_id)
+            cli_utils.declare(
+                f"Workspace '{workspace_name_or_id}' deleted successfully."
+            )
+        except Exception as e:
+            cli_utils.error(str(e))