zenml-nightly 0.68.1.dev20241106__py3-none-any.whl → 0.68.1.dev20241108__py3-none-any.whl

zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py

@@ -87,41 +87,44 @@ class TensorboardVisualizer:
             *args: Additional arguments.
             **kwargs: Additional keyword arguments.
         """
-        for _, artifact_view in object.outputs.items():
-            # filter out anything but model artifacts
-            if artifact_view.type == ArtifactType.MODEL:
-                logdir = os.path.dirname(artifact_view.uri)
-
-                # first check if a TensorBoard server is already running for
-                # the same logdir location and use that one
-                running_server = self.find_running_tensorboard_server(logdir)
-                if running_server:
-                    self.visualize_tensorboard(running_server.port, height)
-                    return
-
-                if sys.platform == "win32":
-                    # Daemon service functionality is currently not supported
-                    # on Windows
-                    print(
-                        "You can run:\n"
-                        f"[italic green]    tensorboard --logdir {logdir}"
-                        "[/italic green]\n"
-                        "...to visualize the TensorBoard logs for your trained model."
+        for output in object.outputs.values():
+            for artifact_view in output:
+                # filter out anything but model artifacts
+                if artifact_view.type == ArtifactType.MODEL:
+                    logdir = os.path.dirname(artifact_view.uri)
+
+                    # first check if a TensorBoard server is already running for
+                    # the same logdir location and use that one
+                    running_server = self.find_running_tensorboard_server(
+                        logdir
                     )
-                else:
-                    # start a new TensorBoard server
-                    service = TensorboardService(
-                        TensorboardServiceConfig(
-                            logdir=logdir,
-                            name=f"zenml-tensorboard-{logdir}",
+                    if running_server:
+                        self.visualize_tensorboard(running_server.port, height)
+                        return
+
+                    if sys.platform == "win32":
+                        # Daemon service functionality is currently not supported
+                        # on Windows
+                        print(
+                            "You can run:\n"
+                            f"[italic green]    tensorboard --logdir {logdir}"
+                            "[/italic green]\n"
+                            "...to visualize the TensorBoard logs for your trained model."
                         )
-                    )
-                    service.start(timeout=60)
-                    if service.endpoint.status.port:
-                        self.visualize_tensorboard(
-                            service.endpoint.status.port, height
+                    else:
+                        # start a new TensorBoard server
+                        service = TensorboardService(
+                            TensorboardServiceConfig(
+                                logdir=logdir,
+                                name=f"zenml-tensorboard-{logdir}",
+                            )
                         )
-                return
+                        service.start(timeout=60)
+                        if service.endpoint.status.port:
+                            self.visualize_tensorboard(
+                                service.endpoint.status.port, height
+                            )
+                    return
 
     def visualize_tensorboard(
         self,
@@ -154,31 +157,34 @@ class TensorboardVisualizer:
         Args:
             object: StepRunResponseModel fetched from get_step().
         """
-        for _, artifact_view in object.outputs.items():
-            # filter out anything but model artifacts
-            if artifact_view.type == ArtifactType.MODEL:
-                logdir = os.path.dirname(artifact_view.uri)
-
-                # first check if a TensorBoard server is already running for
-                # the same logdir location and use that one
-                running_server = self.find_running_tensorboard_server(logdir)
-                if not running_server:
-                    return
+        for output in object.outputs.values():
+            for artifact_view in output:
+                # filter out anything but model artifacts
+                if artifact_view.type == ArtifactType.MODEL:
+                    logdir = os.path.dirname(artifact_view.uri)
+
+                    # first check if a TensorBoard server is already running for
+                    # the same logdir location and use that one
+                    running_server = self.find_running_tensorboard_server(
+                        logdir
+                    )
+                    if not running_server:
+                        return
 
-                logger.debug(
-                    "Stopping tensorboard server with PID '%d' ...",
-                    running_server.pid,
-                )
-                try:
-                    p = psutil.Process(running_server.pid)
-                except psutil.Error:
-                    logger.error(
-                        "Could not find process for PID '%d' ...",
+                    logger.debug(
+                        "Stopping tensorboard server with PID '%d' ...",
                         running_server.pid,
                     )
-                    continue
-                p.kill()
-                return
+                    try:
+                        p = psutil.Process(running_server.pid)
+                    except psutil.Error:
+                        logger.error(
+                            "Could not find process for PID '%d' ...",
+                            running_server.pid,
+                        )
+                        continue
+                    p.kill()
+                    return
 
 
 def get_step(pipeline_name: str, step_name: str) -> "StepRunResponse":

zenml/integrations/vllm/services/vllm_deployment.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of the vLLM Inference Server Service."""
 
+import argparse
 import os
 from typing import Any, List, Optional, Union
 
@@ -137,15 +138,23 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
         self.endpoint.prepare_for_start()
 
         import uvloop
-        from vllm.entrypoints.openai.api_server import run_server
-        from vllm.entrypoints.openai.cli_args import make_arg_parser
-        from vllm.utils import FlexibleArgumentParser
+        from vllm.entrypoints.openai.api_server import (
+            run_server,
+        )
+        from vllm.entrypoints.openai.cli_args import (
+            make_arg_parser,
+        )
+        from vllm.utils import (
+            FlexibleArgumentParser,
+        )
 
         try:
-            parser = make_arg_parser(FlexibleArgumentParser())
-            args = parser.parse_args()
+            parser: argparse.ArgumentParser = make_arg_parser(
+                FlexibleArgumentParser()
+            )
+            args: argparse.Namespace = parser.parse_args()
             # Override port with the available port
-            self.config.port = self.endpoint.status.port
+            self.config.port = self.endpoint.status.port or self.config.port
             # Update the arguments in place
             args.__dict__.update(self.config.model_dump())
             uvloop.run(run_server(args=args))
@@ -161,7 +170,7 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
         """
         if not self.is_running:
             return None
-        return self.endpoint.prediction_url_path
+        return self.endpoint.prediction_url
 
     def predict(self, data: "Any") -> "Any":
         """Make a prediction using the service.

zenml/metadata/lazy_load.py

@@ -13,10 +13,25 @@
 #  permissions and limitations under the License.
 """Run Metadata Lazy Loader definition."""
 
-from typing import TYPE_CHECKING, Optional
+from typing import Optional
 
-if TYPE_CHECKING:
-    from zenml.models import RunMetadataResponse
+from pydantic import BaseModel
+
+from zenml.metadata.metadata_types import MetadataType
+
+
+class LazyRunMetadataResponse(BaseModel):
+    """Lazy run metadata response.
+
+    Used if the run metadata is accessed from the model in
+    a pipeline context available only during pipeline compilation.
+    """
+
+    lazy_load_artifact_name: Optional[str] = None
+    lazy_load_artifact_version: Optional[str] = None
+    lazy_load_metadata_name: Optional[str] = None
+    lazy_load_model_name: str
+    lazy_load_model_version: Optional[str] = None
 
 
 class RunMetadataLazyGetter:
@@ -47,7 +62,7 @@ class RunMetadataLazyGetter:
         self._lazy_load_artifact_name = _lazy_load_artifact_name
         self._lazy_load_artifact_version = _lazy_load_artifact_version
 
-    def __getitem__(self, key: str) -> "RunMetadataResponse":
+    def __getitem__(self, key: str) -> MetadataType:
         """Get the metadata for the given key.
 
         Args:
@@ -56,9 +71,7 @@ class RunMetadataLazyGetter:
         Returns:
             The metadata lazy loader wrapper for the given key.
         """
-        from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
-
-        return LazyRunMetadataResponse(
+        return LazyRunMetadataResponse(  # type: ignore[return-value]
             lazy_load_model_name=self._lazy_load_model_name,
             lazy_load_model_version=self._lazy_load_model_version,
             lazy_load_artifact_name=self._lazy_load_artifact_name,

zenml/model/model.py

@@ -43,7 +43,6 @@ if TYPE_CHECKING:
         ModelResponse,
         ModelVersionResponse,
         PipelineRunResponse,
-        RunMetadataResponse,
         StepRunResponse,
     )
 
@@ -349,7 +348,7 @@ class Model(BaseModel):
         )
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, "MetadataType"]:
         """Get model version run metadata.
 
         Returns:

zenml/models/__init__.py

@@ -239,12 +239,7 @@ from zenml.models.v2.core.run_template import (
 )
 from zenml.models.v2.base.base_plugin_flavor import BasePluginFlavorResponse
 from zenml.models.v2.core.run_metadata import (
-    LazyRunMetadataResponse,
     RunMetadataRequest,
-    RunMetadataFilter,
-    RunMetadataResponse,
-    RunMetadataResponseBody,
-    RunMetadataResponseMetadata,
 )
 from zenml.models.v2.core.schedule import (
     ScheduleRequest,
@@ -418,7 +413,6 @@ EventSourceResponseResources.model_rebuild()
 FlavorResponseBody.model_rebuild()
 FlavorResponseMetadata.model_rebuild()
 LazyArtifactVersionResponse.model_rebuild()
-LazyRunMetadataResponse.model_rebuild()
 ModelResponseBody.model_rebuild()
 ModelResponseMetadata.model_rebuild()
 ModelVersionResponseBody.model_rebuild()
@@ -444,8 +438,6 @@ RunTemplateResponseBody.model_rebuild()
 RunTemplateResponseMetadata.model_rebuild()
 RunTemplateResponseResources.model_rebuild()
 RunTemplateResponseBody.model_rebuild()
-RunMetadataResponseBody.model_rebuild()
-RunMetadataResponseMetadata.model_rebuild()
 ScheduleResponseBody.model_rebuild()
 ScheduleResponseMetadata.model_rebuild()
 SecretResponseBody.model_rebuild()
@@ -637,10 +629,6 @@ __all__ = [
     "RunTemplateResponseResources",
     "RunTemplateFilter",
     "RunMetadataRequest",
-    "RunMetadataFilter",
-    "RunMetadataResponse",
-    "RunMetadataResponseBody",
-    "RunMetadataResponseMetadata",
     "ScheduleRequest",
     "ScheduleUpdate",
     "ScheduleFilter",

zenml/models/v2/core/artifact_version.py

@@ -34,7 +34,7 @@ from pydantic import (
 
 from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ArtifactType, GenericFilterOps
+from zenml.enums import ArtifactSaveType, ArtifactType, GenericFilterOps
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.filter import StrFilter
@@ -57,9 +57,6 @@ if TYPE_CHECKING:
         ArtifactVisualizationResponse,
     )
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.models.v2.core.step_run import StepRunResponse
 
 logger = get_logger(__name__)
@@ -107,6 +104,9 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
     visualizations: Optional[List["ArtifactVisualizationRequest"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
+    save_type: ArtifactSaveType = Field(
+        title="The save type of the artifact version.",
+    )
     metadata: Optional[Dict[str, MetadataType]] = Field(
         default=None, title="Metadata of the artifact version."
     )
@@ -193,6 +193,9 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
         title="The ID of the pipeline run that generated this artifact version.",
         default=None,
     )
+    save_type: ArtifactSaveType = Field(
+        title="The save type of the artifact version.",
+    )
     artifact_store_id: Optional[UUID] = Field(
         title="ID of the artifact store in which this artifact is stored.",
         default=None,
@@ -230,7 +233,7 @@ class ArtifactVersionResponseMetadata(WorkspaceScopedResponseMetadata):
     visualizations: Optional[List["ArtifactVisualizationResponse"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         default={}, title="Metadata of the artifact."
     )
 
@@ -313,6 +316,15 @@ class ArtifactVersionResponse(
         """
         return self.get_body().producer_pipeline_run_id
 
+    @property
+    def save_type(self) -> ArtifactSaveType:
+        """The `save_type` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().save_type
+
     @property
     def artifact_store_id(self) -> Optional[UUID]:
         """The `artifact_store_id` property.
@@ -343,7 +355,7 @@ class ArtifactVersionResponse(
         return self.get_metadata().visualizations
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `metadata` property.
 
         Returns:
@@ -671,7 +683,7 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
         )
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `metadata` property in lazy loading mode.
 
         Returns:

zenml/models/v2/core/model_version.py

@@ -29,6 +29,7 @@ from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ModelStages
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.filter import AnyQuery
 from zenml.models.v2.base.page import Page
 from zenml.models.v2.base.scoped import (
@@ -49,9 +50,6 @@ if TYPE_CHECKING:
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.model import ModelResponse
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.zen_stores.schemas import BaseSchema
 
     AnySchema = TypeVar("AnySchema", bound=BaseSchema)
@@ -193,7 +191,7 @@ class ModelVersionResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=TEXT_FIELD_MAX_LENGTH,
         default=None,
     )
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         description="Metadata linked to the model version",
         default={},
     )
@@ -304,7 +302,7 @@ class ModelVersionResponse(
         return self.get_metadata().description
 
     @property
-    def run_metadata(self) -> Optional[Dict[str, "RunMetadataResponse"]]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `run_metadata` property.
 
         Returns:

zenml/models/v2/core/pipeline_run.py

@@ -30,6 +30,7 @@ from pydantic import BaseModel, ConfigDict, Field
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
@@ -51,9 +52,6 @@ if TYPE_CHECKING:
     from zenml.models.v2.core.pipeline_build import (
         PipelineBuildResponse,
     )
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.models.v2.core.schedule import ScheduleResponse
     from zenml.models.v2.core.stack import StackResponse
     from zenml.models.v2.core.step_run import StepRunResponse
@@ -190,7 +188,7 @@ class PipelineRunResponseBody(WorkspaceScopedResponseBody):
 class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for pipeline runs."""
 
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         default={},
         title="Metadata associated with this pipeline run.",
     )
@@ -450,7 +448,7 @@ class PipelineRunResponse(
         return self.get_body().model_version_id
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `run_metadata` property.
 
         Returns:

zenml/models/v2/core/run_metadata.py

@@ -13,21 +13,15 @@
 #  permissions and limitations under the License.
 """Models representing run metadata."""
 
-from typing import Any, Dict, Optional, Union
+from typing import Dict, Optional
 from uuid import UUID
 
-from pydantic import Field, field_validator
+from pydantic import Field
 
-from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import MetadataResourceTypes
 from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
-    WorkspaceScopedResponse,
-    WorkspaceScopedResponseBody,
-    WorkspaceScopedResponseMetadata,
-    WorkspaceScopedResponseResources,
 )
 
 # ------------------ Request Model ------------------
@@ -51,212 +45,3 @@ class RunMetadataRequest(WorkspaceScopedRequest):
     types: Dict[str, "MetadataTypeEnum"] = Field(
         title="The types of the metadata to be created.",
     )
-
-
-# ------------------ Update Model ------------------
-
-# There is no update model for run metadata.
-
-# ------------------ Response Model ------------------
-
-
-class RunMetadataResponseBody(WorkspaceScopedResponseBody):
-    """Response body for run metadata."""
-
-    key: str = Field(title="The key of the metadata.")
-    value: MetadataType = Field(
-        title="The value of the metadata.", union_mode="smart"
-    )
-    type: MetadataTypeEnum = Field(title="The type of the metadata.")
-
-    @field_validator("key", "type")
-    @classmethod
-    def str_field_max_length_check(cls, value: Any) -> Any:
-        """Checks if the length of the value exceeds the maximum str length.
-
-        Args:
-            value: the value set in the field
-
-        Returns:
-            the value itself.
-
-        Raises:
-            AssertionError: if the length of the field is longer than the
-                maximum threshold.
-        """
-        assert len(str(value)) < STR_FIELD_MAX_LENGTH, (
-            "The length of the value for this field can not "
-            f"exceed {STR_FIELD_MAX_LENGTH}"
-        )
-        return value
-
-    @field_validator("value")
-    @classmethod
-    def text_field_max_length_check(cls, value: Any) -> Any:
-        """Checks if the length of the value exceeds the maximum text length.
-
-        Args:
-            value: the value set in the field
-
-        Returns:
-            the value itself.
-
-        Raises:
-            AssertionError: if the length of the field is longer than the
-                maximum threshold.
-        """
-        assert len(str(value)) < TEXT_FIELD_MAX_LENGTH, (
-            "The length of the value for this field can not "
-            f"exceed {TEXT_FIELD_MAX_LENGTH}"
-        )
-        return value
-
-
-class RunMetadataResponseMetadata(WorkspaceScopedResponseMetadata):
-    """Response metadata for run metadata."""
-
-    resource_id: UUID = Field(
-        title="The ID of the resource that this metadata belongs to.",
-    )
-    resource_type: MetadataResourceTypes = Field(
-        title="The type of the resource that this metadata belongs to.",
-    )
-    stack_component_id: Optional[UUID] = Field(
-        title="The ID of the stack component that this metadata belongs to."
-    )
-
-
-class RunMetadataResponseResources(WorkspaceScopedResponseResources):
-    """Class for all resource models associated with the run metadata entity."""
-
-
-class RunMetadataResponse(
-    WorkspaceScopedResponse[
-        RunMetadataResponseBody,
-        RunMetadataResponseMetadata,
-        RunMetadataResponseResources,
-    ]
-):
-    """Response model for run metadata."""
-
-    def get_hydrated_version(self) -> "RunMetadataResponse":
-        """Get the hydrated version of this run metadata.
-
-        Returns:
-            an instance of the same entity with the metadata field attached.
-        """
-        from zenml.client import Client
-
-        return Client().zen_store.get_run_metadata(self.id)
-
-    # Body and metadata properties
-    @property
-    def key(self) -> str:
-        """The `key` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().key
-
-    @property
-    def value(self) -> MetadataType:
-        """The `value` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().value
-
-    @property
-    def type(self) -> MetadataTypeEnum:
-        """The `type` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().type
-
-    @property
-    def resource_id(self) -> UUID:
-        """The `resource_id` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().resource_id
-
-    @property
-    def resource_type(self) -> MetadataResourceTypes:
-        """The `resource_type` property.
-
-        Returns:
-            the value of the property.
-        """
-        return MetadataResourceTypes(self.get_metadata().resource_type)
-
-    @property
-    def stack_component_id(self) -> Optional[UUID]:
-        """The `stack_component_id` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().stack_component_id
-
-
-# ------------------ Filter Model ------------------
-
-
-class RunMetadataFilter(WorkspaceScopedFilter):
-    """Model to enable advanced filtering of run metadata."""
-
-    resource_id: Optional[Union[str, UUID]] = Field(
-        default=None, union_mode="left_to_right"
-    )
-    resource_type: Optional[MetadataResourceTypes] = None
-    stack_component_id: Optional[Union[str, UUID]] = Field(
-        default=None, union_mode="left_to_right"
-    )
-    key: Optional[str] = None
-    type: Optional[Union[str, MetadataTypeEnum]] = Field(
-        default=None, union_mode="left_to_right"
-    )
-
-
-# -------------------- Lazy Loader --------------------
-
-
-class LazyRunMetadataResponse(RunMetadataResponse):
-    """Lazy run metadata response.
-
-    Used if the run metadata is accessed from the model in
-    a pipeline context available only during pipeline compilation.
-    """
-
-    id: Optional[UUID] = None  # type: ignore[assignment]
-    lazy_load_artifact_name: Optional[str] = None
-    lazy_load_artifact_version: Optional[str] = None
-    lazy_load_metadata_name: Optional[str] = None
-    lazy_load_model_name: str
-    lazy_load_model_version: Optional[str] = None
-
-    def get_body(self) -> None:  # type: ignore[override]
-        """Protects from misuse of the lazy loader.
-
-        Raises:
-            RuntimeError: always
-        """
-        raise RuntimeError(
-            "Cannot access run metadata body before pipeline runs."
-        )
-
-    def get_metadata(self) -> None:  # type: ignore[override]
-        """Protects from misuse of the lazy loader.
-
-        Raises:
-            RuntimeError: always
-        """
-        raise RuntimeError(
-            "Cannot access run metadata metadata before pipeline runs."
-        )