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

zenml/cli/model.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 Model Control Plane."""
+
 from typing import Any, Dict, List, Optional
 
 import click
@@ -42,6 +43,9 @@ def _model_to_print(model: ModelResponse) -> Dict[str, Any]:
         "latest_version": model.latest_version_name,
         "description": model.description,
         "tags": [t.name for t in model.tags],
+        "save_to_registry": ":white_check_mark:"
+        if model.save_models_to_registry
+        else "",
         "use_cases": model.use_cases,
         "audience": model.audience,
         "limitations": model.limitations,
@@ -55,11 +59,6 @@ def _model_to_print(model: ModelResponse) -> Dict[str, Any]:
 def _model_version_to_print(
     model_version: ModelVersionResponse,
 ) -> Dict[str, Any]:
-    run_metadata = None
-    if model_version.run_metadata:
-        run_metadata = {
-            k: v.value for k, v in model_version.run_metadata.items()
-        }
     return {
         "id": model_version.id,
         "model": model_version.model.name,
@@ -67,7 +66,7 @@ def _model_version_to_print(
         "number": model_version.number,
         "description": model_version.description,
         "stage": model_version.stage,
-        "run_metadata": run_metadata,
+        "run_metadata": model_version.run_metadata,
         "tags": [t.name for t in model_version.tags],
         "data_artifacts_count": len(model_version.data_artifact_ids),
         "model_artifacts_count": len(model_version.model_artifact_ids),
@@ -168,6 +167,14 @@ def list_models(**kwargs: Any) -> None:
     required=False,
     multiple=True,
 )
+@click.option(
+    "--save-models-to-registry",
+    "-s",
+    help="Whether to automatically save model artifacts to the model registry.",
+    type=click.BOOL,
+    required=False,
+    default=True,
+)
 def register_model(
     name: str,
     license: Optional[str],
@@ -178,6 +185,7 @@ def register_model(
     ethical: Optional[str],
     limitations: Optional[str],
     tag: Optional[List[str]],
+    save_models_to_registry: Optional[bool],
 ) -> None:
     """Register a new model in the Model Control Plane.
 
@@ -191,18 +199,25 @@ def register_model(
         ethical: The ethical implications of the model.
         limitations: The know limitations of the model.
         tag: Tags associated with the model.
+        save_models_to_registry: Whether to save the model to the
+            registry.
     """
     try:
         model = Client().create_model(
-            name=name,
-            license=license,
-            description=description,
-            audience=audience,
-            use_cases=use_cases,
-            trade_offs=tradeoffs,
-            ethics=ethical,
-            limitations=limitations,
-            tags=tag,
+            **remove_none_values(
+                dict(
+                    name=name,
+                    license=license,
+                    description=description,
+                    audience=audience,
+                    use_cases=use_cases,
+                    trade_offs=tradeoffs,
+                    ethics=ethical,
+                    limitations=limitations,
+                    tags=tag,
+                    save_models_to_registry=save_models_to_registry,
+                )
+            )
         )
     except (EntityExistsError, ValueError) as e:
         cli_utils.error(str(e))
@@ -282,6 +297,14 @@ def register_model(
     required=False,
     multiple=True,
 )
+@click.option(
+    "--save-models-to-registry",
+    "-s",
+    help="Whether to automatically save model artifacts to the model registry.",
+    type=click.BOOL,
+    required=False,
+    default=True,
+)
 def update_model(
     model_name_or_id: str,
     name: Optional[str],
@@ -294,6 +317,7 @@ def update_model(
     limitations: Optional[str],
     tag: Optional[List[str]],
     remove_tag: Optional[List[str]],
+    save_models_to_registry: Optional[bool],
 ) -> None:
     """Register a new model in the Model Control Plane.
 
@@ -309,6 +333,8 @@ def update_model(
         limitations: The know limitations of the model.
         tag: Tags to be added to the model.
         remove_tag: Tags to be removed from the model.
+        save_models_to_registry: Whether to save the model to the
+            registry.
     """
     model_id = Client().get_model(model_name_or_id=model_name_or_id).id
     update_dict = remove_none_values(
@@ -323,6 +349,7 @@ def update_model(
             limitations=limitations,
             add_tags=tag,
             remove_tags=remove_tag,
+            save_models_to_registry=save_models_to_registry,
         )
     )
     model = Client().update_model(model_name_or_id=model_id, **update_dict)

zenml/cli/model_registry.py

@@ -183,8 +183,7 @@ def register_model_registry_subcommands() -> None:  # noqa: C901
             return
         if not yes:
             confirmation = cli_utils.confirmation(
-                f"Found Model with name {name}. Do you want to "
-                f"delete them?"
+                f"Found Model with name {name}. Do you want to delete them?"
             )
             if not confirmation:
                 cli_utils.declare("Model deletion canceled.")

zenml/cli/pipeline.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 pipelines."""
+
 import json
 import os
 from typing import Any, Dict, Optional, Union
@@ -32,13 +33,42 @@ from zenml.models import (
     PipelineRunFilter,
     ScheduleFilter,
 )
-from zenml.new.pipelines.pipeline import Pipeline
+from zenml.pipelines.pipeline_definition import Pipeline
 from zenml.utils import source_utils, uuid_utils
 from zenml.utils.yaml_utils import write_yaml
 
 logger = get_logger(__name__)
 
 
+def _import_pipeline(source: str) -> Pipeline:
+    """Import a pipeline.
+
+    Args:
+        source: The pipeline source.
+
+    Returns:
+        The pipeline.
+    """
+    try:
+        pipeline_instance = source_utils.load(source)
+    except ModuleNotFoundError as e:
+        source_root = source_utils.get_source_root()
+        cli_utils.error(
+            f"Unable to import module `{e.name}`. Make sure the source path is "
+            f"relative to your source root `{source_root}`."
+        )
+    except AttributeError as e:
+        cli_utils.error("Unable to load attribute from module: " + str(e))
+
+    if not isinstance(pipeline_instance, Pipeline):
+        cli_utils.error(
+            f"The given source path `{source}` does not resolve to a pipeline "
+            "object."
+        )
+
+    return pipeline_instance
+
+
 @cli.group(cls=TagGroup, tag=CliCategories.MANAGEMENT_TOOLS)
 def pipeline() -> None:
     """Interact with pipelines, runs and schedules."""
@@ -84,22 +114,7 @@ def register_pipeline(
             "source code root."
         )
 
-    try:
-        pipeline_instance = source_utils.load(source)
-    except ModuleNotFoundError as e:
-        source_root = source_utils.get_source_root()
-        cli_utils.error(
-            f"Unable to import module `{e.name}`. Make sure the source path is "
-            f"relative to your source root `{source_root}`."
-        )
-    except AttributeError as e:
-        cli_utils.error("Unable to load attribute from module: " + str(e))
-
-    if not isinstance(pipeline_instance, Pipeline):
-        cli_utils.error(
-            f"The given source path `{source}` does not resolve to a pipeline "
-            "object."
-        )
+    pipeline_instance = _import_pipeline(source=source)
 
     parameters: Dict[str, Any] = {}
     if parameters_path:
@@ -120,15 +135,13 @@ def register_pipeline(
     pipeline_instance.register()
 
 
-@pipeline.command("build", help="Build Docker images for a pipeline.")
-@click.argument("pipeline_name_or_id")
-@click.option(
-    "--version",
-    "-v",
-    type=str,
-    required=False,
-    help="Optional version of the pipeline.",
+@pipeline.command(
+    "build",
+    help="Build Docker images for a pipeline. The SOURCE argument needs to be "
+    " an importable source path resolving to a ZenML pipeline instance, e.g. "
+    "`my_module.my_pipeline_instance`.",
 )
+@click.argument("source")
 @click.option(
     "--config",
     "-c",
@@ -154,8 +167,7 @@ def register_pipeline(
     help="Output path for the build information.",
 )
 def build_pipeline(
-    pipeline_name_or_id: str,
-    version: Optional[str] = None,
+    source: str,
     config_path: Optional[str] = None,
     stack_name_or_id: Optional[str] = None,
     output_path: Optional[str] = None,
@@ -163,8 +175,7 @@ def build_pipeline(
     """Build Docker images for a pipeline.
 
     Args:
-        pipeline_name_or_id: Name or ID of the pipeline.
-        version: Version of the pipeline.
+        source: Importable source resolving to a pipeline instance.
         config_path: Path to pipeline configuration file.
         stack_name_or_id: Name or ID of the stack for which the images should
             be built.
@@ -179,13 +190,13 @@ def build_pipeline(
             "your source code root."
         )
 
-    pipeline_model = Client().get_pipeline(
-        name_id_or_prefix=pipeline_name_or_id, version=version
-    )
-
     with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
-        pipeline_instance = Pipeline.from_model(pipeline_model)
-        build = pipeline_instance.build(config_path=config_path)
+        pipeline_instance = _import_pipeline(source=source)
+
+        pipeline_instance = pipeline_instance.with_options(
+            config_path=config_path
+        )
+        build = pipeline_instance.build()
 
     if build:
         cli_utils.declare(f"Created pipeline build `{build.id}`.")
@@ -199,15 +210,13 @@ def build_pipeline(
         cli_utils.declare("No docker builds required.")
 
 
-@pipeline.command("run", help="Run a pipeline.")
-@click.argument("pipeline_name_or_id")
-@click.option(
-    "--version",
-    "-v",
-    type=str,
-    required=False,
-    help="Optional version of the pipeline.",
+@pipeline.command(
+    "run",
+    help="Run a pipeline. The SOURCE argument needs to be an "
+    "importable source path resolving to a ZenML pipeline instance, e.g. "
+    "`my_module.my_pipeline_instance`.",
 )
+@click.argument("source")
 @click.option(
     "--config",
     "-c",
@@ -240,8 +249,7 @@ def build_pipeline(
     help="Prevent automatic build reusing.",
 )
 def run_pipeline(
-    pipeline_name_or_id: str,
-    version: Optional[str] = None,
+    source: str,
     config_path: Optional[str] = None,
     stack_name_or_id: Optional[str] = None,
     build_path_or_id: Optional[str] = None,
@@ -250,8 +258,7 @@ def run_pipeline(
     """Run a pipeline.
 
     Args:
-        pipeline_name_or_id: Name or ID of the pipeline.
-        version: Version of the pipeline.
+        source: Importable source resolving to a pipeline instance.
         config_path: Path to pipeline configuration file.
         stack_name_or_id: Name or ID of the stack on which the pipeline should
             run.
@@ -269,24 +276,21 @@ def run_pipeline(
             "your source code root."
         )
 
-    pipeline_model = Client().get_pipeline(
-        name_id_or_prefix=pipeline_name_or_id, version=version
-    )
-
-    build: Union[str, PipelineBuildBase, None] = None
-    if build_path_or_id:
-        if uuid_utils.is_valid_uuid(build_path_or_id):
-            build = build_path_or_id
-        elif os.path.exists(build_path_or_id):
-            build = PipelineBuildBase.from_yaml(build_path_or_id)
-        else:
-            cli_utils.error(
-                f"The specified build {build_path_or_id} is not a valid UUID "
-                "or file path."
-            )
-
     with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
-        pipeline_instance = Pipeline.from_model(pipeline_model)
+        pipeline_instance = _import_pipeline(source=source)
+
+        build: Union[str, PipelineBuildBase, None] = None
+        if build_path_or_id:
+            if uuid_utils.is_valid_uuid(build_path_or_id):
+                build = build_path_or_id
+            elif os.path.exists(build_path_or_id):
+                build = PipelineBuildBase.from_yaml(build_path_or_id)
+            else:
+                cli_utils.error(
+                    f"The specified build {build_path_or_id} is not a valid UUID "
+                    "or file path."
+                )
+
         pipeline_instance = pipeline_instance.with_options(
             config_path=config_path,
             build=build,
@@ -295,6 +299,71 @@ def run_pipeline(
         pipeline_instance()
 
 
+@pipeline.command(
+    "create-run-template",
+    help="Create a run template for a pipeline. The SOURCE argument needs to "
+    "be an importable source path resolving to a ZenML pipeline instance, e.g. "
+    "`my_module.my_pipeline_instance`.",
+)
+@click.argument("source")
+@click.option(
+    "--name",
+    "-n",
+    type=str,
+    required=True,
+    help="Name for the template",
+)
+@click.option(
+    "--config",
+    "-c",
+    "config_path",
+    type=click.Path(exists=True, dir_okay=False),
+    required=False,
+    help="Path to configuration file for the build.",
+)
+@click.option(
+    "--stack",
+    "-s",
+    "stack_name_or_id",
+    type=str,
+    required=False,
+    help="Name or ID of the stack to use for the build.",
+)
+def create_run_template(
+    source: str,
+    name: str,
+    config_path: Optional[str] = None,
+    stack_name_or_id: Optional[str] = None,
+) -> None:
+    """Create a run template for a pipeline.
+
+    Args:
+        source: Importable source resolving to a pipeline instance.
+        name: Name of the run template.
+        config_path: Path to pipeline configuration file.
+        stack_name_or_id: Name or ID of the stack for which the template should
+            be created.
+    """
+    if not Client().root:
+        cli_utils.warning(
+            "You're running the `zenml pipeline create-run-template` command "
+            "without a ZenML repository. Your current working directory will "
+            "be used as the source root relative to which the registered step "
+            "classes will be resolved. To silence this warning, run `zenml "
+            "init` at your source code root."
+        )
+
+    with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+        pipeline_instance = _import_pipeline(source=source)
+
+        pipeline_instance = pipeline_instance.with_options(
+            config_path=config_path
+        )
+        template = pipeline_instance.create_run_template(name=name)
+
+    cli_utils.declare(f"Created run template `{template.id}`.")
+
+
 @pipeline.command("list", help="List all registered pipelines.")
 @list_options(PipelineFilter)
 def list_pipelines(**kwargs: Any) -> None:
@@ -319,13 +388,6 @@ def list_pipelines(**kwargs: Any) -> None:
 
 @pipeline.command("delete")
 @click.argument("pipeline_name_or_id", type=str, required=True)
-@click.option(
-    "--version",
-    "-v",
-    help="Optional pipeline version.",
-    type=str,
-    required=False,
-)
 @click.option(
     "--yes",
     "-y",
@@ -333,21 +395,19 @@ def list_pipelines(**kwargs: Any) -> None:
     help="Don't ask for confirmation.",
 )
 def delete_pipeline(
-    pipeline_name_or_id: str, version: Optional[str] = None, yes: bool = False
+    pipeline_name_or_id: str,
+    yes: bool = False,
 ) -> None:
     """Delete a pipeline.
 
     Args:
         pipeline_name_or_id: The name or ID of the pipeline to delete.
-        version: The version of the pipeline to delete.
         yes: If set, don't ask for confirmation.
     """
-    version_suffix = f" (version {version})" if version else ""
-
     if not yes:
         confirmation = cli_utils.confirmation(
             f"Are you sure you want to delete pipeline "
-            f"`{pipeline_name_or_id}{version_suffix}`? This will change all "
+            f"`{pipeline_name_or_id}`? This will change all "
             "existing runs of this pipeline to become unlisted."
         )
         if not confirmation:
@@ -356,14 +416,12 @@ def delete_pipeline(
 
     try:
         Client().delete_pipeline(
-            name_id_or_prefix=pipeline_name_or_id, version=version
+            name_id_or_prefix=pipeline_name_or_id,
         )
     except KeyError as e:
         cli_utils.error(str(e))
     else:
-        cli_utils.declare(
-            f"Deleted pipeline `{pipeline_name_or_id}{version_suffix}`."
-        )
+        cli_utils.declare(f"Deleted pipeline `{pipeline_name_or_id}`.")
 
 
 @pipeline.group()
@@ -491,6 +549,27 @@ def delete_pipeline_run(
         cli_utils.declare(f"Deleted pipeline run '{run_name_or_id}'.")
 
 
+@runs.command("refresh")
+@click.argument("run_name_or_id", type=str, required=True)
+def refresh_pipeline_run(run_name_or_id: str) -> None:
+    """Refresh the status of a pipeline run.
+
+    Args:
+        run_name_or_id: The name or ID of the pipeline run to refresh.
+    """
+    try:
+        # Fetch and update the run
+        run = Client().get_pipeline_run(name_id_or_prefix=run_name_or_id)
+        run.refresh_run_status()
+
+    except KeyError as e:
+        cli_utils.error(str(e))
+    else:
+        cli_utils.declare(
+            f"Refreshed the status of pipeline run '{run.name}'."
+        )
+
+
 @pipeline.group()
 def builds() -> None:
     """Commands for pipeline builds."""
@@ -507,7 +586,7 @@ def list_pipeline_builds(**kwargs: Any) -> None:
     client = Client()
     try:
         with console.status("Listing pipeline builds...\n"):
-            pipeline_builds = client.list_builds(**kwargs)
+            pipeline_builds = client.list_builds(hydrate=True, **kwargs)
     except KeyError as err:
         cli_utils.error(str(err))
     else:
@@ -517,7 +596,14 @@ def list_pipeline_builds(**kwargs: Any) -> None:
 
         cli_utils.print_pydantic_models(
             pipeline_builds,
-            exclude_columns=["created", "updated", "user", "workspace"],
+            exclude_columns=[
+                "created",
+                "updated",
+                "user",
+                "workspace",
+                "images",
+                "stack_checksum",
+            ],
         )
 
 

zenml/cli/secret.py

@@ -347,8 +347,7 @@ def update_secret(
             )
 
         declare(
-            "You will now have a chance to update each secret pair "
-            "one by one."
+            "You will now have a chance to update each secret pair one by one."
         )
         secret_args_add_update = {}
         for k, _ in secret.secret_values.items():