zenml-nightly 0.75.0.dev20250318__py3-none-any.whl → 0.80.0.dev20250321__py3-none-any.whl

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -318,30 +318,36 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     disk_tier=settings.disk_tier,
                 )
             )
-            # Set the cluster name
+            # Do not detach run if logs are being streamed
+            # Otherwise, the logs will not be streamed after the task is submitted
+            # Could also be a parameter in the settings to control this behavior
+            detach_run = not settings.stream_logs
+
+            launch_new_cluster = True
             if settings.cluster_name:
-                sky.exec(
-                    task,
-                    settings.cluster_name,
-                    down=down,
-                    stream_logs=settings.stream_logs,
-                    backend=None,
-                    detach_run=True,
+                cluster_info = sky.status(
+                    refresh=True, cluster_names=settings.cluster_name
                 )
+                if cluster_info:
+                    logger.info(
+                        f"Found existing cluster {settings.cluster_name}. Reusing..."
+                    )
+                    launch_new_cluster = False
+
+                else:
+                    logger.info(
+                        f"Cluster {settings.cluster_name} not found. Launching a new one..."
+                    )
+                    cluster_name = settings.cluster_name
             else:
-                # Find existing cluster
-                for i in sky.status(refresh=True):
-                    if isinstance(
-                        i["handle"].launched_resources.cloud, type(self.cloud)
-                    ):
-                        cluster_name = i["handle"].cluster_name
-                        logger.info(
-                            f"Found existing cluster {cluster_name}. Reusing..."
-                        )
                 cluster_name = self.sanitize_cluster_name(
                     f"{orchestrator_run_name}"
                 )
-                # Launch the cluster
+                logger.info(
+                    f"No cluster name provided. Launching a new cluster with name {cluster_name}..."
+                )
+
+            if launch_new_cluster:
                 sky.launch(
                     task,
                     cluster_name,
@@ -349,7 +355,26 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     idle_minutes_to_autostop=idle_minutes_to_autostop,
                     down=down,
                     stream_logs=settings.stream_logs,
+                    backend=None,
                     detach_setup=True,
+                    detach_run=detach_run,
+                )
+            else:
+                # Make sure the cluster is up -
+                # If the cluster is already up, this will not do anything
+                sky.start(
+                    settings.cluster_name,
+                    down=down,
+                    idle_minutes_to_autostop=idle_minutes_to_autostop,
+                    retry_until_up=settings.retry_until_up,
+                )
+                sky.exec(
+                    task,
+                    settings.cluster_name,
+                    down=down,
+                    stream_logs=settings.stream_logs,
+                    backend=None,
+                    detach_run=detach_run,
                 )
 
         except Exception as e:

zenml/integrations/skypilot_aws/__init__.py

@@ -32,7 +32,7 @@ class SkypilotAWSIntegration(Integration):
 
     NAME = SKYPILOT_AWS
     # all 0.6.x versions of skypilot[aws] are compatible
-    REQUIREMENTS = ["skypilot[aws]~=0.6.0"]
+    REQUIREMENTS = ["skypilot[aws]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -48,5 +48,3 @@ class SkypilotAWSIntegration(Integration):
 
         return [SkypilotAWSOrchestratorFlavor]
 
-
-SkypilotAWSIntegration.check_installation()

zenml/integrations/skypilot_azure/__init__.py

@@ -31,7 +31,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]>=0.6.1"]
+    REQUIREMENTS = ["skypilot[azure]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -47,4 +47,3 @@ class SkypilotAzureIntegration(Integration):
 
         return [SkypilotAzureOrchestratorFlavor]
 
-SkypilotAzureIntegration.check_installation()

zenml/integrations/skypilot_gcp/__init__.py

@@ -31,7 +31,7 @@ class SkypilotGCPIntegration(Integration):
     """Definition of Skypilot (GCP) Integration for ZenML."""
 
     NAME = SKYPILOT_GCP
-    REQUIREMENTS = ["skypilot[gcp]~=0.6.0"]
+    REQUIREMENTS = ["skypilot[gcp]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -46,5 +46,3 @@ class SkypilotGCPIntegration(Integration):
         )
 
         return [SkypilotGCPOrchestratorFlavor]
-
-SkypilotGCPIntegration.check_installation()

zenml/integrations/skypilot_kubernetes/__init__.py

@@ -32,7 +32,7 @@ class SkypilotKubernetesIntegration(Integration):
 
     NAME = SKYPILOT_KUBERNETES
     # all 0.6.x versions of skypilot[kubernetes] are compatible
-    REQUIREMENTS = ["skypilot[kubernetes]~=0.6.1"]
+    REQUIREMENTS = ["skypilot[kubernetes]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -48,5 +48,3 @@ class SkypilotKubernetesIntegration(Integration):
 
         return [SkypilotKubernetesOrchestratorFlavor]
 
-
-SkypilotKubernetesIntegration.check_installation()

zenml/integrations/skypilot_lambda/__init__.py

@@ -31,7 +31,7 @@ class SkypilotLambdaIntegration(Integration):
     """Definition of Skypilot Lambda Integration for ZenML."""
 
     NAME = SKYPILOT_LAMBDA
-    REQUIREMENTS = ["skypilot[lambda]~=0.6.0"]
+    REQUIREMENTS = ["skypilot[lambda]~=0.8.0"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
@@ -46,5 +46,3 @@ class SkypilotLambdaIntegration(Integration):
 
         return [SkypilotLambdaOrchestratorFlavor]
 
-
-SkypilotLambdaIntegration.check_installation()

zenml/integrations/slack/__init__.py

@@ -44,4 +44,3 @@ class SlackIntegration(Integration):
         return [SlackAlerterFlavor]
 
 
-SlackIntegration.check_installation()

zenml/integrations/spark/__init__.py

@@ -49,4 +49,3 @@ class SparkIntegration(Integration):
         return [KubernetesSparkStepOperatorFlavor]
 
 
-SparkIntegration.check_installation()

zenml/integrations/tekton/__init__.py

@@ -46,4 +46,3 @@ class TektonIntegration(Integration):
         return [TektonOrchestratorFlavor]
 
 
-TektonIntegration.check_installation()

zenml/integrations/tensorboard/__init__.py

@@ -42,5 +42,3 @@ class TensorBoardIntegration(Integration):
         """Activates the integration."""
         from zenml.integrations.tensorboard import services  # noqa
 
-
-TensorBoardIntegration.check_installation()

zenml/integrations/tensorboard/services/tensorboard_service.py

@@ -19,6 +19,7 @@ from typing import Any, Dict, Union
 from tensorboard import default, program  # type: ignore [import-untyped]
 
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services import (
     HTTPEndpointHealthMonitor,
     HTTPEndpointHealthMonitorConfig,
@@ -27,7 +28,6 @@ from zenml.services import (
     LocalDaemonServiceEndpoint,
     LocalDaemonServiceEndpointConfig,
     ServiceEndpointProtocol,
-    ServiceType,
 )
 
 logger = get_logger(__name__)

zenml/integrations/tensorflow/__init__.py

@@ -64,5 +64,3 @@ class TensorflowIntegration(Integration):
             requirements.append("typing-extensions>=4.6.1")
         return requirements
 
-
-TensorflowIntegration.check_installation()

zenml/integrations/vllm/__init__.py

@@ -47,4 +47,3 @@ class VLLMIntegration(Integration):
         return [VLLMModelDeployerFlavor]
 
 
-VLLMIntegration.check_installation()

zenml/integrations/vllm/services/vllm_deployment.py

@@ -19,6 +19,7 @@ from typing import Any, List, Optional, Union
 
 from zenml.constants import DEFAULT_LOCAL_SERVICE_IP_ADDRESS
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services import (
     HTTPEndpointHealthMonitor,
     HTTPEndpointHealthMonitorConfig,
@@ -27,7 +28,6 @@ from zenml.services import (
     LocalDaemonServiceEndpoint,
     LocalDaemonServiceEndpointConfig,
     ServiceEndpointProtocol,
-    ServiceType,
 )
 from zenml.services.service import BaseDeploymentService
 

zenml/integrations/wandb/__init__.py

@@ -47,4 +47,3 @@ class WandbIntegration(Integration):
         return [WandbExperimentTrackerFlavor]
 
 
-WandbIntegration.check_installation()

zenml/integrations/whylogs/__init__.py

@@ -65,4 +65,3 @@ class WhylogsIntegration(Integration):
             PandasIntegration.get_requirements(target_os=target_os)
 
 
-WhylogsIntegration.check_installation()

zenml/integrations/xgboost/__init__.py

@@ -28,5 +28,3 @@ class XgboostIntegration(Integration):
         """Activates the integration."""
         from zenml.integrations.xgboost import materializers  # noqa
 
-
-XgboostIntegration.check_installation()

zenml/login/credentials.py

@@ -20,11 +20,11 @@ from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict
 
+from zenml.enums import ServiceState
 from zenml.login.pro.constants import ZENML_PRO_API_URL, ZENML_PRO_URL
 from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
 from zenml.models import ServerModel
 from zenml.models.v2.misc.server_models import ServerDeploymentType
-from zenml.services.service_status import ServiceState
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.string_utils import get_human_readable_time
 from zenml.utils.time_utils import to_local_tz, utc_now

zenml/materializers/__init__.py

@@ -33,6 +33,7 @@ from zenml.materializers.pydantic_materializer import PydanticMaterializer
 from zenml.materializers.service_materializer import ServiceMaterializer
 from zenml.materializers.uuid_materializer import UUIDMaterializer
 
+
 __all__ = [
     "BuiltInContainerMaterializer",
     "BuiltInMaterializer",

zenml/model_deployers/base_model_deployer.py

@@ -30,9 +30,9 @@ from uuid import UUID
 from zenml.client import Client
 from zenml.enums import StackComponentType
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services import BaseService, ServiceConfig
 from zenml.services.service import BaseDeploymentService
-from zenml.services.service_type import ServiceType
 from zenml.stack import StackComponent
 from zenml.stack.flavor import Flavor
 from zenml.stack.stack_component import StackComponentConfig

zenml/models/__init__.py

@@ -29,6 +29,7 @@ from zenml.models.v2.base.base import (
 )
 from zenml.models.v2.base.scoped import (
     TaggableFilter,
+    RunMetadataFilterMixin,
     UserScopedRequest,
     UserScopedFilter,
     UserScopedResponse,
@@ -389,6 +390,7 @@ from zenml.models.v2.misc.server_models import (
     ServerDatabaseType,
     ServerDeploymentType,
 )
+from zenml.models.v2.misc.service import ServiceType
 from zenml.models.v2.core.server_settings import (
     ServerActivationRequest,
     ServerSettingsResponse,
@@ -516,6 +518,7 @@ __all__ = [
     "NumericFilter",
     "UUIDFilter",
     "TaggableFilter",
+    "RunMetadataFilterMixin",
     "Page",
     # V2 Core
     "ActionFilter",
@@ -760,6 +763,7 @@ __all__ = [
     "ServerDatabaseType",
     "ServerDeploymentType",
     "ServerStatistics",
+    "ServiceType",
     "StackDeploymentConfig",
     "StackDeploymentInfo",
     "OAuthDeviceAuthorizationRequest",

zenml/models/v2/base/filter.py

@@ -37,7 +37,7 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from sqlalchemy import Float, and_, asc, cast, desc
+from sqlalchemy import Float, and_, asc, cast, desc, or_
 from sqlmodel import SQLModel
 
 from zenml.constants import (
@@ -171,8 +171,6 @@ class BoolFilter(Filter):
 class StrFilter(Filter):
     """Filter for all string fields."""
 
-    json_encode_value: bool = False
-
     ALLOWED_OPS: ClassVar[List[str]] = [
         GenericFilterOps.EQUALS,
         GenericFilterOps.NOT_EQUALS,
@@ -201,6 +199,23 @@ class StrFilter(Filter):
                 raise ValueError(ONEOF_ERROR)
         return self
 
+    def _check_if_column_is_json_encoded(self, column: Any) -> bool:
+        """Check if the column is json encoded.
+
+        Args:
+            column: The column of an SQLModel table on which to filter.
+
+        Returns:
+            True if the column is json encoded, False otherwise.
+        """
+        from zenml.zen_stores.schemas import RunMetadataSchema
+
+        JSON_ENCODED_COLUMNS = [RunMetadataSchema.value]
+
+        if column in JSON_ENCODED_COLUMNS:
+            return True
+        return False
+
     def generate_query_conditions_from_column(self, column: Any) -> Any:
         """Generate query conditions for a string column.
 
@@ -209,70 +224,168 @@ class StrFilter(Filter):
 
         Returns:
             A list of query conditions.
-
-        Raises:
-            ValueError: the comparison of the column to a numeric value fails.
         """
+        # Handle numeric comparisons (GT, LT, GTE, LTE)
         if self.operation in {
             GenericFilterOps.GT,
             GenericFilterOps.LT,
             GenericFilterOps.GTE,
             GenericFilterOps.LTE,
         }:
-            try:
-                numeric_column = cast(column, Float)
-
-                assert self.value is not None
-
-                if self.operation == GenericFilterOps.GT:
-                    return and_(
-                        numeric_column, numeric_column > float(self.value)
-                    )
-                if self.operation == GenericFilterOps.LT:
-                    return and_(
-                        numeric_column, numeric_column < float(self.value)
-                    )
-                if self.operation == GenericFilterOps.GTE:
-                    return and_(
-                        numeric_column, numeric_column >= float(self.value)
-                    )
-                if self.operation == GenericFilterOps.LTE:
-                    return and_(
-                        numeric_column, numeric_column <= float(self.value)
-                    )
-            except Exception as e:
-                raise ValueError(
-                    f"Failed to compare the column '{column}' to the "
-                    f"value '{self.value}' (must be numeric): {e}"
-                )
+            return self._handle_numeric_comparison(column)
+
+        # Handle operations that need special treatment for JSON-encoded columns
+        is_json_encoded = self._check_if_column_is_json_encoded(column)
 
+        # Handle list operations
         if self.operation == GenericFilterOps.ONEOF:
             assert isinstance(self.value, list)
-            # Convert the list of values to a list of json strings
-            json_list = (
-                [json.dumps(v) for v in self.value]
-                if self.json_encode_value
-                else self.value
-            )
-            return column.in_(json_list)
+            return self._handle_oneof(column, is_json_encoded)
 
-        # Don't convert the value to a json string if the operation is contains
-        # because the quotes around strings will mess with the comparison
+        # Handle pattern matching operations
         if self.operation == GenericFilterOps.CONTAINS:
             return column.like(f"%{self.value}%")
 
-        json_value = (
-            json.dumps(self.value) if self.json_encode_value else self.value
-        )
-
         if self.operation == GenericFilterOps.STARTSWITH:
-            return column.startswith(f"{json_value}")
+            return self._handle_startswith(column, is_json_encoded)
+
         if self.operation == GenericFilterOps.ENDSWITH:
-            return column.endswith(f"{json_value}")
+            return self._handle_endswith(column, is_json_encoded)
+
         if self.operation == GenericFilterOps.NOT_EQUALS:
-            return column != json_value
+            return self._handle_not_equals(column, is_json_encoded)
 
-        return column == json_value
+        # Default case (EQUALS)
+        return self._handle_equals(column, is_json_encoded)
+
+    def _handle_numeric_comparison(self, column: Any) -> Any:
+        """Handle numeric comparison operations.
+
+        Args:
+            column: The column to compare.
+
+        Returns:
+            The query condition.
+
+        Raises:
+            ValueError: If the comparison fails.
+        """
+        try:
+            numeric_column = cast(column, Float)
+            assert self.value is not None
+
+            operations = {
+                GenericFilterOps.GT: lambda col, val: and_(
+                    col, col > float(val)
+                ),
+                GenericFilterOps.LT: lambda col, val: and_(
+                    col, col < float(val)
+                ),
+                GenericFilterOps.GTE: lambda col, val: and_(
+                    col, col >= float(val)
+                ),
+                GenericFilterOps.LTE: lambda col, val: and_(
+                    col, col <= float(val)
+                ),
+            }
+
+            return operations[self.operation](numeric_column, self.value)  # type: ignore[no-untyped-call]
+        except Exception as e:
+            raise ValueError(
+                f"Failed to compare the column '{column}' to the "
+                f"value '{self.value}' (must be numeric): {e}"
+            )
+
+    def _handle_oneof(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the ONEOF operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        from sqlalchemy import or_
+
+        conditions = []
+
+        assert isinstance(self.value, list)
+
+        for value in self.value:
+            if is_json_encoded:
+                # For JSON encoded columns, add conditions for both raw and JSON-quoted values
+                conditions.append(column == value)
+                conditions.append(column == f'"{value}"')
+            else:
+                conditions.append(column == value)
+
+        return or_(*conditions)
+
+    def _handle_startswith(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the STARTSWITH operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return or_(
+                column.startswith(self.value),
+                column.startswith(f'"{self.value}'),
+            )
+        else:
+            return column.startswith(self.value)
+
+    def _handle_endswith(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the ENDSWITH operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return or_(
+                column.endswith(self.value), column.endswith(f'{self.value}"')
+            )
+        else:
+            return column.endswith(self.value)
+
+    def _handle_not_equals(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the NOT_EQUALS operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return and_(column != self.value, column != f'"{self.value}"')
+        else:
+            return column != self.value
+
+    def _handle_equals(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the EQUALS operation (default).
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return or_(column == self.value, column == f'"{self.value}"')
+        else:
+            return column == self.value
 
 
 class UUIDFilter(StrFilter):
@@ -751,7 +864,6 @@ class BaseFilter(BaseModel):
         value: Any,
         table: Type[SQLModel],
         column: str,
-        json_encode_value: bool = False,
     ) -> "ColumnElement[bool]":
         """Generate custom filter conditions for a column of a table.
 
@@ -759,7 +871,6 @@ class BaseFilter(BaseModel):
             value: The filter value.
             table: The table which contains the column.
             column: The column name.
-            json_encode_value: Whether to json encode the value.
 
         Returns:
             The query conditions.
@@ -768,9 +879,6 @@ class BaseFilter(BaseModel):
         filter_ = FilterGenerator(table).define_filter(
             column=column, value=value, operator=operator
         )
-        if isinstance(filter_, StrFilter):
-            filter_.json_encode_value = json_encode_value
-
         return filter_.generate_query_conditions(table=table)
 
     @property