zenml-nightly 0.54.1.dev20240118__py3-none-any.whl → 0.54.1.dev20240120__py3-none-any.whl

zenml/__init__.py

@@ -42,15 +42,19 @@ from zenml.artifacts.utils import (
     save_artifact,
     load_artifact,
 )
-from zenml.model.utils import log_model_metadata, link_artifact_to_model
+from zenml.model.utils import (
+    log_model_metadata,
+    link_artifact_to_model,
+    log_model_version_metadata,
+)
 from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.artifacts.external_artifact import ExternalArtifact
-from zenml.model.model_version import ModelVersion
-from zenml.model.utils import log_model_version_metadata
+from zenml.model.model import Model
 from zenml.new.pipelines.pipeline_context import get_pipeline_context
 from zenml.new.pipelines.pipeline_decorator import pipeline
 from zenml.new.steps.step_decorator import step
 from zenml.new.steps.step_context import get_step_context
+from zenml.steps.utils import log_step_metadata
 
 __all__ = [
     "ArtifactConfig",
@@ -60,8 +64,10 @@ __all__ = [
     "load_artifact",
     "log_artifact_metadata",
     "log_model_metadata",
+    "log_model_version_metadata",
+    "log_step_metadata",
+    "Model",
     "link_artifact_to_model",
-    "ModelVersion",
     "pipeline",
     "save_artifact",
     "show",

zenml/artifacts/artifact_config.py

@@ -22,7 +22,7 @@ from zenml.logger import get_logger
 from zenml.new.steps.step_context import get_step_context
 
 if TYPE_CHECKING:
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
 
 
 logger = get_logger(__name__)
@@ -52,7 +52,7 @@ class ArtifactConfig(BaseModel):
         version: The version of the artifact.
         tags: The tags of the artifact.
         model_name: The name of the model to link artifact to.
-        model_version: The identifier of the model version to link the artifact
+        model_version: The identifier of a version of the model to link the artifact
             to. It can be an exact version ("my_version"), exact version number
             (42), stage (ModelStages.PRODUCTION or "production"), or
             (ModelStages.LATEST or None) for the latest version (default).
@@ -88,26 +88,26 @@ class ArtifactConfig(BaseModel):
         smart_union = True
 
     @property
-    def _model_version(self) -> Optional["ModelVersion"]:
-        """The model version linked to this artifact.
+    def _model(self) -> Optional["Model"]:
+        """The model linked to this artifact.
 
         Returns:
-            The model version or None if the model version cannot be determined.
+            The model or None if the model version cannot be determined.
         """
         try:
-            model_version = get_step_context().model_version
+            model_ = get_step_context().model
         except (StepContextError, RuntimeError):
-            model_version = None
+            model_ = None
         # Check if another model name was specified
         if (self.model_name is not None) and (
-            model_version is None or model_version.name != self.model_name
+            model_ is None or model_.name != self.model_name
         ):
-            # Create a new ModelConfig instance with the provided model name and version
-            from zenml.model.model_version import ModelVersion
+            # Create a new Model instance with the provided model name and version
+            from zenml.model.model import Model
 
-            on_the_fly_config = ModelVersion(
+            on_the_fly_config = Model(
                 name=self.model_name, version=self.model_version
             )
             return on_the_fly_config
 
-        return model_version
+        return model_

zenml/artifacts/external_artifact.py

@@ -56,8 +56,8 @@ class ExternalArtifact(ExternalArtifactConfiguration):
             `version`, `pipeline_run_name`, or `pipeline_name` are set, the
             latest version of the artifact will be used.
         version: Version of the artifact to search. Only used when `name` is
-            provided. Cannot be used together with `model_version`.
-        model_version: The model version to search in. Only used when `name`
+            provided. Cannot be used together with `model`.
+        model: The model to search in. Only used when `name`
             is provided. Cannot be used together with `version`.
         materializer: The materializer to use for saving the artifact value
             to the artifact store. Only used when `value` is provided.
@@ -149,5 +149,5 @@ class ExternalArtifact(ExternalArtifactConfiguration):
             id=self.id,
             name=self.name,
             version=self.version,
-            model_version=self.model_version,
+            model=self.model,
         )

zenml/artifacts/external_artifact_config.py

@@ -18,7 +18,7 @@ from uuid import UUID
 from pydantic import BaseModel, root_validator
 
 from zenml.logger import get_logger
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
 
 logger = get_logger(__name__)
@@ -33,13 +33,13 @@ class ExternalArtifactConfiguration(BaseModel):
     id: Optional[UUID] = None
     name: Optional[str] = None
     version: Optional[str] = None
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
 
     @root_validator
     def _validate_all_eac(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get("version", None) and values.get("model_version", None):
+        if values.get("version", None) and values.get("model", None):
             raise ValueError(
-                "Cannot provide both `version` and `model_version` when "
+                "Cannot provide both `version` and `model` when "
                 "creating an external artifact."
             )
         return values
@@ -67,13 +67,13 @@ class ExternalArtifactConfiguration(BaseModel):
                 response = client.get_artifact_version(
                     self.name, version=self.version
                 )
-            elif self.model_version:
-                response_ = self.model_version.get_artifact(self.name)
+            elif self.model:
+                response_ = self.model.get_artifact(self.name)
                 if not isinstance(response_, ArtifactVersionResponse):
                     raise RuntimeError(
                         f"Failed to pull artifact `{self.name}` from the Model "
-                        f"Version (name=`{self.model_version.name}`, version="
-                        f"`{self.model_version.version}`). Please validate the "
+                        f"(name=`{self.model.name}`, version="
+                        f"`{self.model.version}`). Please validate the "
                         "input and try again."
                     )
                 response = response_

zenml/artifacts/utils.py

@@ -233,14 +233,14 @@ def save_artifact(
                     saved_artifact_versions={name: response.id}
                 ),
             )
-            error_message = "model version"
-            model_version = step_context.model_version
-            if model_version:
+            error_message = "model"
+            model = step_context.model
+            if model:
                 from zenml.model.utils import link_artifact_to_model
 
                 link_artifact_to_model(
                     artifact_version_id=response.id,
-                    model_version=model_version,
+                    model=model,
                     is_model_artifact=is_model_artifact,
                     is_deployment_artifact=is_deployment_artifact,
                 )

zenml/cli/__init__.py

@@ -796,7 +796,7 @@ Administering your Models
 ----------------------------
 
 ZenML provides several CLI commands to help you administer your models and
-model versions as part of the Model Control Plane.
+their versions as part of the Model Control Plane.
 
 To register a new model, you can use the following CLI command:
 ```bash
@@ -805,7 +805,7 @@ zenml model register --name <NAME> [--MODEL_OPTIONS]
 
 To list all registered models, use:
 ```bash
-zenml model list
+zenml model list [MODEL_FILTER_OPTIONS]
 ```
 
 To update a model, use:
@@ -827,7 +827,7 @@ zenml model delete <MODEL_NAME_OR_ID>
 The CLI interface for models also helps to navigate through artifacts linked to a specific model versions.
 ```bash
 zenml model data_artifacts <MODEL_NAME_OR_ID> [-v <VERSION>]
-zenml model endpoint_artifacts <MODEL_NAME_OR_ID> [-v <VERSION>]
+zenml model deployment_artifacts <MODEL_NAME_OR_ID> [-v <VERSION>]
 zenml model model_artifacts <MODEL_NAME_OR_ID> [-v <VERSION>]
 ```
 
@@ -838,7 +838,7 @@ zenml model runs <MODEL_NAME_OR_ID> [-v <VERSION>]
 
 To list the model versions of a specific model, use:
 ```bash
-zenml model version list <MODEL_NAME_OR_ID>
+zenml model version list [--model-name <MODEL_NAME> --name <MODEL_VERSION_NAME> OTHER_OPTIONS]
 ```
 
 To delete a model version, use:

zenml/cli/artifact.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """CLI functionality to interact with artifacts."""
 from functools import partial
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 import click
 
@@ -23,6 +23,8 @@ from zenml.client import Client
 from zenml.enums import CliCategories
 from zenml.logger import get_logger
 from zenml.models import ArtifactFilter, ArtifactVersionFilter
+from zenml.models.v2.core.artifact import ArtifactResponse
+from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
 from zenml.utils.pagination_utils import depaginate
 
 logger = get_logger(__name__)
@@ -47,10 +49,11 @@ def list_artifacts(**kwargs: Any) -> None:
         cli_utils.declare("No artifacts found.")
         return
 
-    cli_utils.print_pydantic_models(
-        artifacts,
-        exclude_columns=["created", "updated", "has_custom_name"],
-    )
+    to_print = []
+    for artifact in artifacts:
+        to_print.append(_artifact_to_print(artifact))
+
+    cli_utils.print_table(to_print)
 
 
 @artifact.command("update", help="Update an artifact.")
@@ -126,19 +129,11 @@ def list_artifact_versions(**kwargs: Any) -> None:
         cli_utils.declare("No artifact versions found.")
         return
 
-    cli_utils.print_pydantic_models(
-        artifact_versions,
-        exclude_columns=[
-            "created",
-            "updated",
-            "user",
-            "workspace",
-            "producer_step_run_id",
-            "run_metadata",
-            "artifact_store_id",
-            "visualizations",
-        ],
-    )
+    to_print = []
+    for artifact_version in artifact_versions:
+        to_print.append(_artifact_version_to_print(artifact_version))
+
+    cli_utils.print_table(to_print)
 
 
 @version.command("update", help="Update an artifact version.")
@@ -281,3 +276,28 @@ def prune_artifacts(
         except Exception as e:
             cli_utils.error(str(e))
     cli_utils.declare("All unused artifacts and artifact versions deleted.")
+
+
+def _artifact_version_to_print(
+    artifact_version: ArtifactVersionResponse,
+) -> Dict[str, Any]:
+    return {
+        "id": artifact_version.id,
+        "name": artifact_version.artifact.name,
+        "version": artifact_version.version,
+        "uri": artifact_version.uri,
+        "type": artifact_version.type,
+        "materializer": artifact_version.materializer,
+        "data_type": artifact_version.data_type,
+        "tags": [t.name for t in artifact_version.tags],
+    }
+
+
+def _artifact_to_print(
+    artifact_version: ArtifactResponse,
+) -> Dict[str, Any]:
+    return {
+        "id": artifact_version.id,
+        "name": artifact_version.name,
+        "tags": [t.name for t in artifact_version.tags],
+    }

zenml/cli/base.py

@@ -73,15 +73,15 @@ class ZenMLProjectTemplateLocation(BaseModel):
 ZENML_PROJECT_TEMPLATES = dict(
     e2e_batch=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-e2e-batch",
-        github_tag="2024.01.16",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.01.18",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     starter=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-starter",
-        github_tag="2023.12.18",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.01.12",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     nlp=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-nlp",
-        github_tag="0.45.0",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.01.12",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
 )
 

zenml/cli/model.py

@@ -23,6 +23,7 @@ from zenml.enums import CliCategories, ModelStages
 from zenml.exceptions import EntityExistsError
 from zenml.logger import get_logger
 from zenml.models import (
+    ModelFilter,
     ModelResponse,
     ModelVersionArtifactFilter,
     ModelVersionFilter,
@@ -54,13 +55,19 @@ 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,
         "name": model_version.name,
         "number": model_version.number,
         "description": model_version.description,
         "stage": model_version.stage,
-        "metadata": model_version.to_model_version().run_metadata,
+        "run_metadata": 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),
@@ -77,21 +84,17 @@ def model() -> None:
     """Interact with models and model versions in the Model Control Plane."""
 
 
-@click.option(
-    "--name",
-    "-n",
-    help="The name of the model.",
-    type=str,
-    required=False,
-)
+@cli_utils.list_options(ModelFilter)
 @model.command("list", help="List models with filter.")
-def list_models(name: str) -> None:
+def list_models(**kwargs: Any) -> None:
     """List models with filter in the Model Control Plane.
 
     Args:
-        name: The name filter for models.
+        **kwargs: Keyword arguments to filter models.
     """
-    models = Client().list_models(name=name)
+    models = Client().zen_store.list_models(
+        model_filter_model=ModelFilter(**kwargs)
+    )
 
     if not models:
         cli_utils.declare("No models found.")
@@ -369,18 +372,23 @@ def version() -> None:
 
 
 @cli_utils.list_options(ModelVersionFilter)
-@click.argument("model_name_or_id")
+@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_or_id: str, **kwargs: Any) -> None:
+def list_model_versions(model_name: str, **kwargs: Any) -> None:
     """List model versions with filter in the Model Control Plane.
 
     Args:
-        model_name_or_id: The ID or name of the model containing version.
+        model_name: The name of the parent model.
         **kwargs: Keyword arguments to filter models.
     """
-    model_id = Client().get_model(model_name_or_id=model_name_or_id).id
     model_versions = Client().zen_store.list_model_versions(
-        model_name_or_id=model_id,
+        model_name_or_id=model_name,
         model_version_filter_model=ModelVersionFilter(**kwargs),
     )
 

zenml/cli/server.py

@@ -185,7 +185,16 @@ def up(
     from zenml.zen_server.deploy.deployment import ServerDeploymentConfig
 
     server_config = ServerDeploymentConfig(**config_attrs)
+    if blocking:
+        from zenml.constants import (
+            DEFAULT_USERNAME,
+        )
 
+        cli_utils.declare(
+            "The local ZenML dashboard is about to deploy in a "
+            "blocking process. You can connect to it using the "
+            f"'{DEFAULT_USERNAME}' username and an empty password."
+        )
     server = deployer.deploy_server(server_config)
 
     assert gc.store is not None

zenml/cli/utils.py

@@ -58,8 +58,8 @@ from zenml.constants import (
 from zenml.enums import GenericFilterOps, StackComponentType
 from zenml.logger import get_logger
 from zenml.model_registries.base_model_registry import (
-    ModelVersion,
     RegisteredModel,
+    RegistryModelVersion,
 )
 from zenml.models import (
     BaseFilter,
@@ -1182,7 +1182,7 @@ def pretty_print_registered_model_table(
 
 
 def pretty_print_model_version_table(
-    model_versions: List["ModelVersion"],
+    model_versions: List["RegistryModelVersion"],
 ) -> None:
     """Given a list of model_versions, print all associated key-value pairs.
 
@@ -1206,7 +1206,7 @@ def pretty_print_model_version_table(
 
 
 def pretty_print_model_version_details(
-    model_version: "ModelVersion",
+    model_version: "RegistryModelVersion",
 ) -> None:
     """Given a model_version, print all associated key-value pairs.
 

zenml/client.py

@@ -2713,6 +2713,7 @@ class Client(metaclass=ClientMetaClass):
         name: Optional[str] = None,
         has_custom_name: Optional[bool] = None,
         hydrate: bool = False,
+        tag: Optional[str] = None,
     ) -> Page[ArtifactResponse]:
         """Get a list of artifacts.
 
@@ -2728,6 +2729,7 @@ class Client(metaclass=ClientMetaClass):
             has_custom_name: Filter artifacts with/without custom names.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
+            tag: Filter artifacts by tag.
 
         Returns:
             A list of artifacts.
@@ -2742,6 +2744,7 @@ class Client(metaclass=ClientMetaClass):
             updated=updated,
             name=name,
             has_custom_name=has_custom_name,
+            tag=tag,
         )
         return self.zen_store.list_artifacts(
             artifact_filter_model,
@@ -2865,6 +2868,7 @@ class Client(metaclass=ClientMetaClass):
         only_unused: Optional[bool] = False,
         has_custom_name: Optional[bool] = None,
         hydrate: bool = False,
+        tag: Optional[str] = None,
     ) -> Page[ArtifactVersionResponse]:
         """Get a list of artifact versions.
 
@@ -2892,6 +2896,7 @@ class Client(metaclass=ClientMetaClass):
             has_custom_name: Filter artifacts with/without custom names.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
+            tag: A tag to filter by.
 
         Returns:
             A list of artifact versions.
@@ -2917,6 +2922,7 @@ class Client(metaclass=ClientMetaClass):
             user_id=user_id,
             only_unused=only_unused,
             has_custom_name=has_custom_name,
+            tag=tag,
         )
         artifact_version_filter_model.set_scope_workspace(
             self.active_workspace.id
@@ -4788,6 +4794,7 @@ class Client(metaclass=ClientMetaClass):
         updated: Optional[Union[datetime, str]] = None,
         name: Optional[str] = None,
         hydrate: bool = False,
+        tag: Optional[str] = None,
     ) -> Page[ModelResponse]:
         """Get models by filter from Model Control Plane.
 
@@ -4801,6 +4808,7 @@ class Client(metaclass=ClientMetaClass):
             name: The name of the model to filter by.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
+            tag: The tag of the model to filter by.
 
         Returns:
             A page object with all models.
@@ -4813,11 +4821,11 @@ class Client(metaclass=ClientMetaClass):
             logical_operator=logical_operator,
             created=created,
             updated=updated,
+            tag=tag,
         )
 
         return self.zen_store.list_models(
-            model_filter_model=filter,
-            hydrate=hydrate,
+            model_filter_model=filter, hydrate=hydrate
         )
 
     #################
@@ -4976,6 +4984,7 @@ class Client(metaclass=ClientMetaClass):
         number: Optional[int] = None,
         stage: Optional[Union[str, ModelStages]] = None,
         hydrate: bool = False,
+        tag: Optional[str] = None,
     ) -> Page[ModelVersionResponse]:
         """Get model versions by filter from Model Control Plane.
 
@@ -4993,6 +5002,7 @@ class Client(metaclass=ClientMetaClass):
             stage: stage of the model version.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
+            tag: The tag to filter by.
 
         Returns:
             A page object with all model versions.
@@ -5007,6 +5017,7 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             number=number,
             stage=stage,
+            tag=tag,
         )
 
         return self.zen_store.list_model_versions(

zenml/config/compiler.py

@@ -204,7 +204,7 @@ class Compiler:
                 enable_step_logs=config.enable_step_logs,
                 settings=config.settings,
                 extra=config.extra,
-                model_version=config.model_version,
+                model=config.model,
                 parameters=config.parameters,
             )
 

zenml/config/pipeline_configurations.py

@@ -19,7 +19,7 @@ from pydantic import validator
 from zenml.config.constants import DOCKER_SETTINGS_KEY
 from zenml.config.source import Source, convert_source_validator
 from zenml.config.strict_base_model import StrictBaseModel
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings
@@ -40,7 +40,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     extra: Dict[str, Any] = {}
     failure_hook_source: Optional[Source] = None
     success_hook_source: Optional[Source] = None
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
 
     _convert_source = convert_source_validator(

zenml/config/pipeline_run_configuration.py

@@ -19,7 +19,7 @@ from zenml.config.base_settings import BaseSettings
 from zenml.config.schedule import Schedule
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.config.strict_base_model import StrictBaseModel
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 from zenml.models import PipelineBuildBase
 from zenml.utils import pydantic_utils
 
@@ -39,5 +39,5 @@ class PipelineRunConfiguration(
     steps: Dict[str, StepConfigurationUpdate] = {}
     settings: Dict[str, BaseSettings] = {}
     extra: Dict[str, Any] = {}
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None

zenml/config/step_configurations.py

@@ -35,7 +35,7 @@ from zenml.config.source import Source, convert_source_validator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.logger import get_logger
 from zenml.model.lazy_load import ModelVersionDataLazyLoader
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 from zenml.utils import deprecation_utils
 
 if TYPE_CHECKING:
@@ -135,7 +135,7 @@ class StepConfigurationUpdate(StrictBaseModel):
     extra: Dict[str, Any] = {}
     failure_hook_source: Optional[Source] = None
     success_hook_source: Optional[Source] = None
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 

zenml/integrations/__init__.py

@@ -40,7 +40,9 @@ from zenml.integrations.kserve import KServeIntegration  # noqa
 from zenml.integrations.kubeflow import KubeflowIntegration  # noqa
 from zenml.integrations.kubernetes import KubernetesIntegration  # noqa
 from zenml.integrations.label_studio import LabelStudioIntegration  # noqa
+from zenml.integrations.langchain import LangchainIntegration  # noqa
 from zenml.integrations.lightgbm import LightGBMIntegration  # noqa
+# from zenml.integrations.llama_index import LlamaIndexIntegration  # noqa
 from zenml.integrations.mlflow import MlflowIntegration  # noqa
 from zenml.integrations.neptune import NeptuneIntegration  # noqa
 from zenml.integrations.neural_prophet import NeuralProphetIntegration  # noqa
@@ -67,7 +69,3 @@ from zenml.integrations.tensorflow import TensorflowIntegration  # noqa
 from zenml.integrations.wandb import WandbIntegration  # noqa
 from zenml.integrations.whylogs import WhylogsIntegration  # noqa
 from zenml.integrations.xgboost import XgboostIntegration  # noqa
-
-if sys.version_info > (3, 7):
-    # from zenml.integrations.llama_index import LlamaIndexIntegration  # noqa
-    from zenml.integrations.langchain import LangchainIntegration  # noqa