zenml-nightly 0.60.0.dev20240627__py3-none-any.whl → 0.61.0.dev20240711__py3-none-any.whl

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -43,7 +43,7 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
             `{'tpu_vm': True, 'runtime_version': 'tpu-vm-base'}` for TPUs.
         use_spot: whether to use spot instances. If None, defaults to
             False.
-        spot_recovery: the spot recovery strategy to use for the managed
+        job_recovery: the spot recovery strategy to use for the managed
             spot to recover the cluster from preemption. Refer to
             `recovery_strategy module <https://github.com/skypilot-org/skypilot/blob/master/sky/spot/recovery_strategy.py>`__ # pylint: disable=line-too-long
             for more details.
@@ -103,7 +103,7 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
     )
     accelerator_args: Optional[Dict[str, str]] = None
     use_spot: Optional[bool] = None
-    spot_recovery: Optional[str] = None
+    job_recovery: Optional[str] = None
     region: Optional[str] = None
     zone: Optional[str] = None
     image_id: Union[Dict[str, str], str, None] = Field(

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -303,7 +303,7 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     accelerators=settings.accelerators,
                     accelerator_args=settings.accelerator_args,
                     use_spot=settings.use_spot,
-                    spot_recovery=settings.spot_recovery,
+                    job_recovery=settings.job_recovery,
                     region=settings.region,
                     zone=settings.zone,
                     image_id=settings.image_id,

zenml/integrations/skypilot/orchestrators/skypilot_orchestrator_entrypoint.py

@@ -136,7 +136,7 @@ def main() -> None:
             settings.disk_size,  # Assuming disk_size is part of the settings
             settings.disk_tier,  # Assuming disk_tier is part of the settings
             settings.use_spot,
-            settings.spot_recovery,
+            settings.job_recovery,
             settings.region,
             settings.zone,
             accelerators_hashable,
@@ -213,7 +213,7 @@ def main() -> None:
                 accelerators=settings.accelerators,
                 accelerator_args=settings.accelerator_args,
                 use_spot=settings.use_spot,
-                spot_recovery=settings.spot_recovery,
+                job_recovery=settings.job_recovery,
                 region=settings.region,
                 zone=settings.zone,
                 image_id=settings.image_id,

zenml/integrations/skypilot_aws/__init__.py

@@ -31,7 +31,8 @@ class SkypilotAWSIntegration(Integration):
     """Definition of Skypilot AWS Integration for ZenML."""
 
     NAME = SKYPILOT_AWS
-    REQUIREMENTS = ["skypilot[aws]<=0.5.0"]
+    # all 0.6.x versions of skypilot[aws] are compatible
+    REQUIREMENTS = ["skypilot[aws]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

zenml/integrations/skypilot_azure/__init__.py

@@ -33,7 +33,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]<=0.5.0"]
+    REQUIREMENTS = ["skypilot[azure]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

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.5.0"]
+    REQUIREMENTS = ["skypilot[gcp]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

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.5.0"]
+    REQUIREMENTS = ["skypilot[lambda]<=0.6.0"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

@@ -40,7 +40,7 @@ class SkypilotLambdaOrchestratorSettings(SkypilotBaseOrchestratorSettings):
 
     _UNSUPPORTED_FEATURES = {
         "use_spot": "Spot instances not supported for Lambda orchestrator.",
-        "spot_recovery": "Spot recovery not supported for Lambda orchestrator.",
+        "job_recovery": "Job recovery not supported for Lambda orchestrator.",
         "image_id": "Custom image IDs not supported for Lambda orchestrator.",
         # Add other unsupported features as needed
     }

zenml/integrations/slack/__init__.py

@@ -31,6 +31,7 @@ class SlackIntegration(Integration):
 
     NAME = SLACK
     REQUIREMENTS = ["slack-sdk>=3.16.1", "aiohttp>=3.8.1"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["aiohttp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/tekton/__init__.py

@@ -32,6 +32,7 @@ class TektonIntegration(Integration):
 
     NAME = TEKTON
     REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kfp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/tensorboard/__init__.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Initialization for TensorBoard integration."""
 
-import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORBOARD
 from zenml.integrations.integration import Integration

zenml/integrations/tensorflow/__init__.py

@@ -18,6 +18,9 @@ import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORFLOW
 from zenml.integrations.integration import Integration
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class TensorflowIntegration(Integration):
@@ -25,20 +28,27 @@ class TensorflowIntegration(Integration):
 
     NAME = TENSORFLOW
     REQUIREMENTS = []
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["typing-extensions"]
 
     @classmethod
     def activate(cls) -> None:
         """Activates the integration."""
         # need to import this explicitly to load the Tensorflow file IO support
         # for S3 and other file systems
-        if (
-            not platform.system() == "Darwin"
-            or not platform.machine() == "arm64"
-        ):
+        if not platform.system() == "Darwin" or not platform.machine() == "arm64":
             import tensorflow_io  # type: ignore
 
         from zenml.integrations.tensorflow import materializers  # noqa
 
+        if sys.version_info.minor == 8:
+            logger.warning(
+                "Python 3.8 with TensorFlow is not fully "
+                "compatible with Pydantic 2 requirements. "
+                "Consider upgrading to a higher Python "
+                "version if you would like to use the "
+                "Tensorflow integration."
+            )
+
     @classmethod
     def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
         """Defines platform specific requirements for the integration.
@@ -52,13 +62,15 @@ class TensorflowIntegration(Integration):
         target_os = target_os or platform.system()
         if target_os == "Darwin" and platform.machine() == "arm64":
             requirements = [
-                f"tensorflow-macos>=2.12,<=2.15",
+                "tensorflow-macos>=2.12,<=2.15",
             ]
         else:
             requirements = [
-                f"tensorflow>=2.12,<=2.15",
+                "tensorflow>=2.12,<=2.15",
                 "tensorflow_io>=0.24.0",
             ]
+        if sys.version_info.minor == 8:
+            requirements.append("typing-extensions>=4.6.1")
         return requirements
 
 

zenml/integrations/wandb/__init__.py

@@ -31,6 +31,7 @@ class WandbIntegration(Integration):
 
     NAME = WANDB
     REQUIREMENTS = ["wandb>=0.12.12", "Pillow>=9.1.0"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["Pillow"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/models/__init__.py

@@ -327,6 +327,7 @@ from zenml.models.v2.misc.service_connector_type import (
     ResourceTypeModel,
 )
 from zenml.models.v2.misc.server_models import ServerDatabaseType, ServerModel
+from zenml.models.v2.misc.full_stack import FullStackRequest
 from zenml.models.v2.core.trigger import (
     TriggerRequest,
     TriggerFilter,
@@ -386,6 +387,10 @@ from zenml.models.v2.core.server_settings import (
     ServerSettingsResponseMetadata,
     ServerSettingsUpdate,
 )
+from zenml.models.v2.misc.stack_deployment import (
+    DeployedStack,
+    StackDeploymentInfo,
+)
 
 # ----------------------------- Forward References -----------------------------
 
@@ -405,6 +410,7 @@ EventSourceResponseMetadata.model_rebuild()
 EventSourceResponseResources.model_rebuild()
 FlavorResponseBody.model_rebuild()
 FlavorResponseMetadata.model_rebuild()
+FullStackRequest.model_rebuild()
 LazyArtifactVersionResponse.model_rebuild()
 LazyRunMetadataResponse.model_rebuild()
 ModelResponseBody.model_rebuild()
@@ -701,11 +707,13 @@ __all__ = [
     "WorkspaceResponseMetadata",
     # V2 Misc
     "AuthenticationMethodModel",
+    "DeployedStack",
     "ServiceConnectorResourcesModel",
     "ServiceConnectorTypeModel",
     "ServiceConnectorTypedResourcesModel",
     "ServiceConnectorRequirements",
     "ResourceTypeModel",
+    "FullStackRequest",
     "UserAuthModel",
     "ExternalUserModel",
     "BuildItem",
@@ -718,6 +726,7 @@ __all__ = [
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",
+    "StackDeploymentInfo",
     "OAuthDeviceAuthorizationRequest",
     "OAuthDeviceAuthorizationResponse",
     "OAuthDeviceTokenRequest",

zenml/models/v2/core/component.py

@@ -238,6 +238,24 @@ class ComponentResponse(
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Add the component labels to analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        if self.labels is not None:
+            metadata.update(
+                {
+                    label[6:]: value
+                    for label, value in self.labels.items()
+                    if label.startswith("zenml:")
+                }
+            )
+        return metadata
+
     def get_hydrated_version(self) -> "ComponentResponse":
         """Get the hydrated version of this component.
 

zenml/models/v2/core/model.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Models representing models."""
 
-from functools import partial
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Union
 from uuid import UUID
 
@@ -303,7 +302,7 @@ class ModelResponse(
 
         client = Client()
         model_versions = depaginate(
-            partial(client.list_model_versions, model_name_or_id=self.id)
+            client.list_model_versions, model_name_or_id=self.id
         )
         return [
             mv.to_model_class(suppress_class_validation_warnings=True)

zenml/models/v2/core/service_connector.py

@@ -493,6 +493,23 @@ class ServiceConnectorResponse(
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Add the service connector labels to analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        metadata.update(
+            {
+                label[6:]: value
+                for label, value in self.labels.items()
+                if label.startswith("zenml:")
+            }
+        )
+        return metadata
+
     def get_hydrated_version(self) -> "ServiceConnectorResponse":
         """Get the hydrated version of this service connector.
 

zenml/models/v2/core/stack.py

@@ -36,6 +36,7 @@ from zenml.models.v2.core.component import ComponentResponse
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
+
 # ------------------ Request Model ------------------
 
 
@@ -59,6 +60,10 @@ class StackRequest(WorkspaceScopedRequest):
         title="A mapping of stack component types to the actual"
         "instances of components of this type.",
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
     @property
     def is_valid(self) -> bool:
@@ -109,6 +114,10 @@ class StackUpdate(BaseUpdate):
         "instances of components of this type.",
         default=None,
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
 
 # ------------------ Response Model ------------------
@@ -134,6 +143,10 @@ class StackResponseMetadata(WorkspaceScopedResponseMetadata):
         default=None,
         title="The path to the stack spec used for mlstacks deployments.",
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
 
 class StackResponseResources(WorkspaceScopedResponseResources):
@@ -214,6 +227,15 @@ class StackResponse(
         """
         metadata = super().get_analytics_metadata()
         metadata.update({ct: c[0].flavor for ct, c in self.components.items()})
+
+        if self.labels is not None:
+            metadata.update(
+                {
+                    label[6:]: value
+                    for label, value in self.labels.items()
+                    if label.startswith("zenml:")
+                }
+            )
         return metadata
 
     @property
@@ -243,6 +265,15 @@ class StackResponse(
         """
         return self.get_metadata().components
 
+    @property
+    def labels(self) -> Optional[Dict[str, Any]]:
+        """The `labels` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().labels
+
 
 # ------------------ Filter Model ------------------
 

zenml/models/v2/misc/full_stack.py

@@ -0,0 +1,97 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Models representing full stack requests."""
+
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+
+from pydantic import BaseModel, Field, model_validator
+
+from zenml.constants import STR_FIELD_MAX_LENGTH
+from zenml.enums import StackComponentType
+from zenml.models.v2.base.base import BaseRequest
+
+
+class ServiceConnectorInfo(BaseModel):
+    """Information about the service connector when creating a full stack."""
+
+    type: str
+    auth_method: str
+    configuration: Dict[str, Any] = {}
+
+
+class ComponentInfo(BaseModel):
+    """Information about each stack components when creating a full stack."""
+
+    flavor: str
+    service_connector_index: Optional[int] = Field(
+        default=None,
+        title="The id of the service connector from the list "
+        "`service_connectors`.",
+        description="The id of the service connector from the list "
+        "`service_connectors` from `FullStackRequest`.",
+    )
+    service_connector_resource_id: Optional[str] = None
+    configuration: Dict[str, Any] = {}
+
+
+class FullStackRequest(BaseRequest):
+    """Request model for a full-stack."""
+
+    user: Optional[UUID] = None
+    workspace: Optional[UUID] = None
+
+    name: str = Field(
+        title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH
+    )
+    description: Optional[str] = Field(
+        default="",
+        title="The description of the stack",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
+    service_connectors: List[Union[UUID, ServiceConnectorInfo]] = Field(
+        default=[],
+        title="The service connectors dictionary for the full stack "
+        "registration.",
+        description="The UUID of an already existing service connector or "
+        "request information to create a service connector from "
+        "scratch.",
+    )
+    components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = Field(
+        title="The mapping for the components of the full stack registration.",
+        description="The mapping from component types to either UUIDs of "
+        "existing components or request information for brand new "
+        "components.",
+    )
+
+    @model_validator(mode="after")
+    def _validate_indexes_in_components(self) -> "FullStackRequest":
+        for component in self.components.values():
+            if isinstance(component, ComponentInfo):
+                if component.service_connector_index is not None:
+                    if (
+                        component.service_connector_index < 0
+                        or component.service_connector_index
+                        >= len(self.service_connectors)
+                    ):
+                        raise ValueError(
+                            f"Service connector index {component.service_connector_index} "
+                            "is out of range. Please provide a valid index referring to "
+                            "the position in the list of service connectors."
+                        )
+        return self

zenml/models/v2/misc/stack_deployment.py

@@ -0,0 +1,66 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Models related to cloud stack deployments."""
+
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from zenml.enums import StackDeploymentProvider
+from zenml.models.v2.core.service_connector import ServiceConnectorResponse
+from zenml.models.v2.core.stack import StackResponse
+
+
+class StackDeploymentInfo(BaseModel):
+    """Information about a stack deployment."""
+
+    provider: StackDeploymentProvider = Field(
+        title="The provider of the stack deployment."
+    )
+    description: str = Field(
+        title="The description of the stack deployment.",
+        description="The description of the stack deployment.",
+    )
+    instructions: str = Field(
+        title="The instructions for deploying the stack.",
+        description="The instructions for deploying the stack.",
+    )
+    post_deploy_instructions: str = Field(
+        title="The instructions for post-deployment.",
+        description="The instructions for post-deployment.",
+    )
+    permissions: Dict[str, List[str]] = Field(
+        title="The permissions granted to ZenML to access the cloud resources.",
+        description="The permissions granted to ZenML to access the cloud "
+        "resources, as a dictionary grouping permissions by resource.",
+    )
+    locations: Dict[str, str] = Field(
+        title="The locations where the stack can be deployed.",
+        description="The locations where the stack can be deployed, as a "
+        "dictionary mapping location names to descriptions.",
+    )
+
+
+class DeployedStack(BaseModel):
+    """Information about a deployed stack."""
+
+    stack: StackResponse = Field(
+        title="The stack that was deployed.",
+        description="The stack that was deployed.",
+    )
+    service_connector: Optional[ServiceConnectorResponse] = Field(
+        default=None,
+        title="The service connector for the deployed stack.",
+        description="The service connector for the deployed stack.",
+    )

zenml/new/pipelines/pipeline.py

@@ -281,7 +281,7 @@ class Pipeline:
         raise RuntimeError(
             f"Cannot get the model of pipeline '{self.name}' because it has "
             f"not been registered yet. Please ensure that the pipeline has "
-            f"been run and that at least one step has been executed."
+            f"been run or built and try again."
         )
 
     @contextmanager

zenml/orchestrators/input_utils.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Utilities for inputs."""
 
-import functools
 from typing import TYPE_CHECKING, Dict, List, Tuple
 from uuid import UUID
 
@@ -48,13 +47,11 @@ def resolve_step_inputs(
     """
     from zenml.models import ArtifactVersionResponse, RunMetadataResponse
 
-    list_run_steps = functools.partial(
-        Client().list_run_steps, pipeline_run_id=run_id
-    )
-
     current_run_steps = {
         run_step.name: run_step
-        for run_step in pagination_utils.depaginate(list_run_steps)
+        for run_step in pagination_utils.depaginate(
+            Client().list_run_steps, pipeline_run_id=run_id
+        )
     }
 
     input_artifacts: Dict[str, "ArtifactVersionResponse"] = {}

zenml/stack/stack.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Implementation of the ZenML Stack class."""
 
-import functools
 import itertools
 import json
 import os
@@ -155,11 +154,9 @@ class Stack:
 
         # Run a hydrated list call once to avoid one request per component
         component_models = pagination_utils.depaginate(
-            list_method=functools.partial(
-                Client().list_stack_components,
-                stack_id=stack_model.id,
-                hydrate=True,
-            )
+            Client().list_stack_components,
+            stack_id=stack_model.id,
+            hydrate=True,
         )
 
         stack_components = {

zenml/stack_deployments/__init__.py

@@ -0,0 +1,14 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""ZenML Stack Deployments."""