zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/cli/annotator.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Functionality for annotator CLI subcommands."""
 
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Tuple, cast
 
 import click
 
@@ -104,9 +104,11 @@ def register_annotator_subcommands() -> None:
         )
         cli_utils.declare(f"Total annotation tasks: {total_task_count}")
         cli_utils.declare(f"Labeled annotation tasks: {labeled_task_count}")
-        cli_utils.declare(
-            f"Unlabeled annotation tasks: {unlabeled_task_count}"
-        )
+        if annotator.flavor != "prodigy":
+            # Prodigy doesn't allow you to get the unlabeled task count
+            cli_utils.declare(
+                f"Unlabeled annotation tasks: {unlabeled_task_count}"
+            )
 
     @dataset.command("delete")
     @click.argument("dataset_name", type=click.STRING)
@@ -131,29 +133,43 @@ def register_annotator_subcommands() -> None:
             dataset_name: Name of the dataset to delete.
             all_: Whether to delete all datasets.
         """
+        if not cli_utils.confirmation(
+            f"Are you sure you want to delete dataset '{dataset_name}'?"
+        ):
+            return
         cli_utils.declare(f"Deleting your dataset '{dataset_name}'")
         dataset_names = (
             annotator.get_dataset_names() if all_ else [dataset_name]
         )
         for dataset_name in dataset_names:
-            annotator.delete_dataset(dataset_name=dataset_name)
-            cli_utils.declare(
-                f"Dataset '{dataset_name}' has now been deleted."
-            )
+            try:
+                annotator.delete_dataset(dataset_name=dataset_name)
+                cli_utils.declare(
+                    f"Dataset '{dataset_name}' has now been deleted."
+                )
+            except ValueError as e:
+                cli_utils.error(
+                    f"Failed to delete dataset '{dataset_name}': {e}"
+                )
 
     @dataset.command(
         "annotate", context_settings={"ignore_unknown_options": True}
     )
     @click.argument("dataset_name", type=click.STRING)
+    @click.argument("kwargs", nargs=-1, type=click.UNPROCESSED)
     @click.pass_obj
     def dataset_annotate(
-        annotator: "BaseAnnotator", dataset_name: str
+        annotator: "BaseAnnotator",
+        dataset_name: str,
+        kwargs: Tuple[str, ...],
     ) -> None:
         """Command to launch the annotation interface for a dataset.
 
         Args:
             annotator: The annotator stack component.
             dataset_name: Name of the dataset
+            kwargs: Additional keyword arguments to pass to the
+                annotation client.
 
         Raises:
             ValueError: If the dataset does not exist.
@@ -161,8 +177,26 @@ def register_annotator_subcommands() -> None:
         cli_utils.declare(
             f"Launching the annotation interface for dataset '{dataset_name}'."
         )
-        try:
-            annotator.get_dataset(dataset_name=dataset_name)
-            annotator.launch(url=annotator.get_url_for_dataset(dataset_name))
-        except ValueError as e:
-            raise ValueError("Dataset does not exist.") from e
+
+        # Process the arbitrary keyword arguments
+        kwargs_dict = {}
+        for arg in kwargs:
+            if arg.startswith("--"):
+                key, value = arg.lstrip("--").split("=", 1)
+                kwargs_dict[key] = value
+
+        if annotator.flavor == "prodigy":
+            command = kwargs_dict.get("command")
+            if not command:
+                raise ValueError(
+                    "The 'command' keyword argument is required for launching the Prodigy interface."
+                )
+            annotator.launch(**kwargs_dict)
+        else:
+            try:
+                annotator.get_dataset(dataset_name=dataset_name)
+                annotator.launch(
+                    url=annotator.get_url_for_dataset(dataset_name)
+                )
+            except ValueError as e:
+                raise ValueError("Dataset does not exist.") from e

zenml/cli/artifact.py

@@ -12,7 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """CLI functionality to interact with artifacts."""
-from functools import partial
+
 from typing import Any, Dict, List, Optional
 
 import click
@@ -228,8 +228,17 @@ def update_artifact_version(
     is_flag=True,
     help="Don't ask for confirmation.",
 )
+@click.option(
+    "--ignore-errors",
+    "-i",
+    is_flag=True,
+    help="Ignore errors and continue with the next artifact version.",
+)
 def prune_artifacts(
-    only_artifact: bool = False, only_metadata: bool = False, yes: bool = False
+    only_artifact: bool = False,
+    only_metadata: bool = False,
+    yes: bool = False,
+    ignore_errors: bool = False,
 ) -> None:
     """Delete all unused artifacts and artifact versions.
 
@@ -243,10 +252,12 @@ def prune_artifacts(
         only_metadata: If set, only delete metadata and not the actual artifact
             objects stored in the artifact store.
         yes: If set, don't ask for confirmation.
+        ignore_errors: If set, ignore errors and continue with the next
+            artifact version.
     """
     client = Client()
     unused_artifact_versions = depaginate(
-        partial(client.list_artifact_versions, only_unused=True)
+        client.list_artifact_versions, only_unused=True
     )
 
     if not unused_artifact_versions:
@@ -274,7 +285,14 @@ def prune_artifacts(
                 Client().delete_artifact(unused_artifact.id)
 
         except Exception as e:
-            cli_utils.error(str(e))
+            if ignore_errors:
+                cli_utils.warning(
+                    f"Failed to delete artifact version {unused_artifact_version.id}: {str(e)}"
+                )
+            else:
+                cli_utils.error(
+                    f"Failed to delete artifact version {unused_artifact_version.id}: {str(e)}"
+                )
     cli_utils.declare("All unused artifacts and artifact versions deleted.")
 
 

zenml/cli/authorized_device.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """CLI functionality to interact with authorized devices."""
+
 from typing import Any
 
 import click

zenml/cli/base.py

@@ -28,7 +28,13 @@ from zenml.analytics.enums import AnalyticsEvent
 from zenml.analytics.utils import email_opt_int, track_handler
 from zenml.cli import utils as cli_utils
 from zenml.cli.cli import cli
-from zenml.cli.utils import confirmation, declare, error, warning
+from zenml.cli.utils import (
+    confirmation,
+    declare,
+    error,
+    is_jupyter_installed,
+    warning,
+)
 from zenml.client import Client
 from zenml.config.global_config import GlobalConfiguration
 from zenml.console import console
@@ -36,7 +42,7 @@ from zenml.constants import (
     ENV_ZENML_ENABLE_REPO_INIT_WARNINGS,
     REPOSITORY_DIRECTORY_NAME,
 )
-from zenml.enums import AnalyticsEventSource, StoreType
+from zenml.enums import AnalyticsEventSource, DatabaseBackupStrategy, StoreType
 from zenml.environment import Environment, get_environment
 from zenml.exceptions import GitNotFoundError, InitializationException
 from zenml.integrations.registry import integration_registry
@@ -44,7 +50,7 @@ from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils.io_utils import copy_dir, get_global_config_directory
 from zenml.utils.yaml_utils import write_yaml
-from zenml.zen_server.utils import get_active_deployment
+from zenml.zen_server.utils import get_local_server
 
 logger = get_logger(__name__)
 # WT_SESSION is a Windows Terminal specific environment variable. If it
@@ -73,15 +79,19 @@ class ZenMLProjectTemplateLocation(BaseModel):
 ZENML_PROJECT_TEMPLATES = dict(
     e2e_batch=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-e2e-batch",
-        github_tag="2024.01.18",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.11.28",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     starter=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-starter",
-        github_tag="2024.01.12",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.11.28",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     nlp=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-nlp",
-        github_tag="2024.01.12",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2025.01.08",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+    ),
+    llm_finetuning=ZenMLProjectTemplateLocation(
+        github_url="zenml-io/template-llm-finetuning",
+        github_tag="2024.11.28",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
 )
 
@@ -98,9 +108,9 @@ ZENML_PROJECT_TEMPLATES = dict(
     type=str,
     required=False,
     help="Name or URL of the ZenML project template to use to initialize the "
-    "repository, Can be a string like `e2e_batch`, `nlp`, `starter` etc. or a "
-    "copier URL like gh:owner/repo_name. If not specified, no template is "
-    "used.",
+    "repository, Can be a string like `e2e_batch`, `nlp`, `llm_finetuning`, "
+    "`starter` etc. or a copier URL like gh:owner/repo_name. If not specified, "
+    "no template is used.",
 )
 @click.option(
     "--template-tag",
@@ -254,7 +264,7 @@ def init(
         f"will only take effect when you're running ZenML from the initialized "
         f"repository root, or from a subdirectory. For more information on "
         f"repositories and configurations, please visit "
-        f"https://docs.zenml.io/user-guide/starter-guide/understand-stacks."
+        f"https://docs.zenml.io/user-guide/production-guide/understand-stacks."
     )
 
 
@@ -348,13 +358,13 @@ def clean(yes: bool = False, local: bool = False) -> None:
         )
 
     if yes or confirm:
-        server = get_active_deployment(local=True)
+        server = get_local_server()
 
         if server:
-            from zenml.zen_server.deploy.deployer import ServerDeployer
+            from zenml.zen_server.deploy.deployer import LocalServerDeployer
 
-            deployer = ServerDeployer()
-            deployer.remove_server(server.config.name)
+            deployer = LocalServerDeployer()
+            deployer.remove_server()
             cli_utils.declare("The local ZenML dashboard has been shut down.")
 
         # delete the .zen folder
@@ -378,6 +388,7 @@ def clean(yes: bool = False, local: bool = False) -> None:
                     )
             fileio.rmtree(str(global_zen_config))
             declare(f"Deleted global ZenML config from {global_zen_config}.")
+            GlobalConfiguration._reset_instance()
             fresh_gc = GlobalConfiguration(
                 user_id=gc.user_id,
                 analytics_opt_in=gc.analytics_opt_in,
@@ -396,6 +407,7 @@ def go() -> None:
 
     Raises:
         GitNotFoundError: If git is not installed.
+        e: when Jupyter Notebook fails to launch.
     """
     from zenml.cli.text_utils import (
         zenml_cli_privacy_message,
@@ -420,6 +432,13 @@ def go() -> None:
 
     zenml_tutorial_path = os.path.join(os.getcwd(), "zenml_tutorial")
 
+    if not is_jupyter_installed():
+        cli_utils.error(
+            "Jupyter Notebook or JupyterLab is not installed. "
+            "Please install the 'notebook' package with `pip` "
+            "first so you can run the tutorial notebooks."
+        )
+
     with track_handler(event=AnalyticsEvent.RUN_ZENML_GO, metadata=metadata):
         console.print(zenml_cli_privacy_message, width=80)
 
@@ -447,6 +466,7 @@ def go() -> None:
                         TUTORIAL_REPO,
                         tmp_cloned_dir,
                         branch=f"release/{zenml_version}",
+                        depth=1,  # to prevent timeouts when downloading
                     )
                 example_dir = os.path.join(
                     tmp_cloned_dir, "examples/quickstart"
@@ -470,7 +490,18 @@ def go() -> None:
             zenml_go_notebook_tutorial_message(ipynb_files), width=80
         )
         input("Press ENTER to continue...")
-    subprocess.check_call(["jupyter", "notebook"], cwd=zenml_tutorial_path)
+
+    try:
+        subprocess.check_call(
+            ["jupyter", "notebook", "--ContentsManager.allow_hidden=True"],
+            cwd=zenml_tutorial_path,
+        )
+    except subprocess.CalledProcessError as e:
+        cli_utils.error(
+            "An error occurred while launching Jupyter Notebook. "
+            "Please make sure Jupyter is properly installed and try again."
+        )
+        raise e
 
 
 def _prompt_email(event_source: AnalyticsEventSource) -> bool:
@@ -502,6 +533,7 @@ def _prompt_email(event_source: AnalyticsEventSource) -> bool:
             email_opt_int(opted_in=True, email=email, source=event_source)
 
             GlobalConfiguration().user_email_opt_in = True
+            GlobalConfiguration().user_email = email
 
             # Add consent and email to user model
             client.update_user(
@@ -576,7 +608,7 @@ def info(
     client = Client()
     store_info = client.zen_store.get_store_info()
 
-    store_cfg = gc.store
+    store_cfg = gc.store_configuration
 
     user_info = {
         "zenml_local_version": zenml_version,
@@ -585,7 +617,7 @@ def info(
         "zenml_server_deployment_type": str(store_info.deployment_type),
         "zenml_config_dir": gc.config_directory,
         "zenml_local_store_dir": gc.local_stores_path,
-        "zenml_server_url": "" if store_cfg is None else store_cfg.url,
+        "zenml_server_url": store_cfg.url,
         "zenml_active_repository_root": str(client.root),
         "python_version": environment.python_version(),
         "environment": get_environment(),
@@ -624,7 +656,15 @@ def info(
         cli_utils.print_user_info(user_info)
 
     if stack:
-        cli_utils.print_debug_stack()
+        try:
+            cli_utils.print_debug_stack()
+        except ModuleNotFoundError as e:
+            cli_utils.warning(
+                "Could not print debug stack information. Please make sure "
+                "you have the necessary dependencies and integrations "
+                "installed for all your stack components."
+            )
+            cli_utils.warning(f"The missing package is: '{e.name}'")
 
 
 @cli.command(
@@ -646,10 +686,7 @@ def migrate_database(skip_default_registrations: bool = False) -> None:
     """
     from zenml.zen_stores.base_zen_store import BaseZenStore
 
-    store_config = (
-        GlobalConfiguration().store
-        or GlobalConfiguration().get_default_store()
-    )
+    store_config = GlobalConfiguration().store_configuration
     if store_config.type == StoreType.SQL:
         BaseZenStore.create_store(
             store_config, skip_default_registrations=skip_default_registrations
@@ -659,3 +696,129 @@ def migrate_database(skip_default_registrations: bool = False) -> None:
         cli_utils.warning(
             "Unable to migrate database while connected to a ZenML server."
         )
+
+
+@cli.command("backup-database", help="Create a database backup.", hidden=True)
+@click.option(
+    "--strategy",
+    "-s",
+    help="Custom backup strategy to use. Defaults to whatever is configured "
+    "in the store config.",
+    type=click.Choice(choices=DatabaseBackupStrategy.values()),
+    required=False,
+    default=None,
+)
+@click.option(
+    "--location",
+    default=None,
+    help="Custom location to store the backup. Defaults to whatever is "
+    "configured in the store config. Depending on the strategy, this can be "
+    "a local path or a database name.",
+    type=str,
+)
+@click.option(
+    "--overwrite",
+    "-o",
+    is_flag=True,
+    default=False,
+    help="Overwrite the existing backup.",
+    type=bool,
+)
+def backup_database(
+    strategy: Optional[str] = None,
+    location: Optional[str] = None,
+    overwrite: bool = False,
+) -> None:
+    """Backup the ZenML database.
+
+    Args:
+        strategy: Custom backup strategy to use. Defaults to whatever is
+            configured in the store config.
+        location: Custom location to store the backup. Defaults to whatever is
+            configured in the store config. Depending on the strategy, this can
+            be a local path or a database name.
+        overwrite: Whether to overwrite the existing backup.
+    """
+    from zenml.zen_stores.base_zen_store import BaseZenStore
+    from zenml.zen_stores.sql_zen_store import SqlZenStore
+
+    store_config = GlobalConfiguration().store_configuration
+    if store_config.type == StoreType.SQL:
+        store = BaseZenStore.create_store(
+            store_config, skip_default_registrations=True, skip_migrations=True
+        )
+        assert isinstance(store, SqlZenStore)
+        msg, location = store.backup_database(
+            strategy=DatabaseBackupStrategy(strategy) if strategy else None,
+            location=location,
+            overwrite=overwrite,
+        )
+        cli_utils.declare(f"Database was backed up to {msg}.")
+    else:
+        cli_utils.warning(
+            "Cannot backup database while connected to a ZenML server."
+        )
+
+
+@cli.command(
+    "restore-database", help="Restore the database from a backup.", hidden=True
+)
+@click.option(
+    "--strategy",
+    "-s",
+    help="Custom backup strategy to use. Defaults to whatever is configured "
+    "in the store config.",
+    type=click.Choice(choices=DatabaseBackupStrategy.values()),
+    required=False,
+    default=None,
+)
+@click.option(
+    "--location",
+    default=None,
+    help="Custom location where the backup is stored. Defaults to whatever is "
+    "configured in the store config. Depending on the strategy, this can be "
+    "a local path or a database name.",
+    type=str,
+)
+@click.option(
+    "--cleanup",
+    "-c",
+    is_flag=True,
+    default=False,
+    help="Cleanup the backup after restoring.",
+    type=bool,
+)
+def restore_database(
+    strategy: Optional[str] = None,
+    location: Optional[str] = None,
+    cleanup: bool = False,
+) -> None:
+    """Restore the ZenML database.
+
+    Args:
+        strategy: Custom backup strategy to use. Defaults to whatever is
+            configured in the store config.
+        location: Custom location where the backup is stored. Defaults to
+            whatever is configured in the store config. Depending on the
+            strategy, this can be a local path or a database name.
+        cleanup: Whether to cleanup the backup after restoring.
+    """
+    from zenml.zen_stores.base_zen_store import BaseZenStore
+    from zenml.zen_stores.sql_zen_store import SqlZenStore
+
+    store_config = GlobalConfiguration().store_configuration
+    if store_config.type == StoreType.SQL:
+        store = BaseZenStore.create_store(
+            store_config, skip_default_registrations=True, skip_migrations=True
+        )
+        assert isinstance(store, SqlZenStore)
+        store.restore_database(
+            strategy=DatabaseBackupStrategy(strategy) if strategy else None,
+            location=location,
+            cleanup=cleanup,
+        )
+        cli_utils.declare("Database restore finished.")
+    else:
+        cli_utils.warning(
+            "Cannot restore database while connected to a ZenML server."
+        )

zenml/cli/code_repository.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """CLI functionality to interact with code repositories."""
+
 from typing import Any, List, Optional
 
 import click

zenml/cli/formatter.py

@@ -173,7 +173,7 @@ class ZenFormatter(formatting.HelpFormatter):
 
                     for line in lines[1:]:
                         self.write(
-                            f"{'':>{second_col + self.current_indent * 4 }}{line}\n"
+                            f"{'':>{second_col + self.current_indent * 4}}{line}\n"
                         )
                 else:
                     self.write("\n")