zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.62.0.dev20240726__py3-none-any.whl

zenml/models/v2/misc/full_stack.py

@@ -0,0 +1,129 @@
+#  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
+from zenml.models.v2.core.component import ComponentResponse
+
+
+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
+
+
+class ResourcesInfo(BaseModel):
+    """Information about the resources needed for CLI and UI."""
+
+    flavor: str
+    flavor_display_name: str
+    required_configuration: Dict[str, str] = {}
+    use_resource_value_as_fixed_config: bool = False
+
+    accessible_by_service_connector: List[str]
+    connected_through_service_connector: List[ComponentResponse]
+
+    @model_validator(mode="after")
+    def _validate_resource_info(self) -> "ResourcesInfo":
+        if (
+            self.use_resource_value_as_fixed_config
+            and len(self.required_configuration) > 1
+        ):
+            raise ValueError(
+                "Cannot use resource value as fixed config if more than one required configuration key is provided."
+            )
+        return self
+
+
+class ServiceConnectorResourcesInfo(BaseModel):
+    """Information about the service connector resources needed for CLI and UI."""
+
+    connector_type: str
+
+    components_resources_info: Dict[StackComponentType, List[ResourcesInfo]]

zenml/models/v2/misc/stack_deployment.py

@@ -0,0 +1,91 @@
+#  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.",
+    )
+    integrations: List[str] = Field(
+        title="ZenML integrations required for the stack.",
+        description="The list of ZenML integrations that need to be installed "
+        "for the stack to be usable.",
+    )
+    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.",
+    )
+    skypilot_default_regions: Dict[str, str] = Field(
+        title="The locations where the Skypilot clusters can be deployed by default.",
+        description="The locations where the Skypilot clusters can be deployed by default, as a "
+        "dictionary mapping location names to descriptions.",
+    )
+
+
+class StackDeploymentConfig(BaseModel):
+    """Configuration about a stack deployment."""
+
+    deployment_url: str = Field(
+        title="The cloud provider console URL where the stack will be deployed.",
+    )
+    deployment_url_text: str = Field(
+        title="A textual description for the cloud provider console URL.",
+    )
+    configuration: Optional[str] = Field(
+        title="Configuration for the stack deployment that the user must "
+        "manually configure into the cloud provider console.",
+    )
+
+
+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/new/pipelines/run_utils.py

@@ -250,7 +250,7 @@ def _validate_new_version_requests(
     if not is_cloud_model:
         logger.info(
             "Models can be viewed in the dashboard using ZenML Pro. Sign up "
-            "for a free trial at https://www.zenml.io/cloud/"
+            "for a free trial at https://www.zenml.io/pro/"
         )
 
 

zenml/orchestrators/__init__.py

@@ -30,6 +30,9 @@ from zenml.orchestrators.base_orchestrator import (
 from zenml.orchestrators.containerized_orchestrator import (
     ContainerizedOrchestrator,
 )
+from zenml.orchestrators.wheeled_orchestrator import (
+    WheeledOrchestrator,
+)
 from zenml.orchestrators.local.local_orchestrator import (
     LocalOrchestrator,
     LocalOrchestratorFlavor,
@@ -44,6 +47,7 @@ __all__ = [
     "BaseOrchestratorConfig",
     "BaseOrchestratorFlavor",
     "ContainerizedOrchestrator",
+    "WheeledOrchestrator",
     "LocalOrchestrator",
     "LocalOrchestratorFlavor",
     "LocalDockerOrchestrator",

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/orchestrators/step_launcher.py

@@ -291,6 +291,7 @@ class StepLauncher:
                                 logger.error(
                                     f"Failed to run step `{self._step_name}` after {max_retries} retries. Exiting."
                                 )
+                                logger.exception(e)
                                 publish_utils.publish_failed_step_run(
                                     step_run_response.id
                                 )

zenml/orchestrators/wheeled_orchestrator.py

@@ -0,0 +1,147 @@
+#  Copyright (c) ZenML GmbH 2021. 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.
+"""Wheeled orchestrator class."""
+
+import os
+import re
+import subprocess
+import tempfile
+from abc import ABC
+
+from zenml import __version__
+from zenml.io import fileio
+from zenml.logger import get_logger
+from zenml.orchestrators import BaseOrchestrator
+from zenml.utils.io_utils import copy_dir
+from zenml.utils.source_utils import get_source_root
+
+logger = get_logger(__name__)
+
+DEFAULT_PACKAGE_NAME = "zenmlproject"
+
+
+class WheeledOrchestrator(BaseOrchestrator, ABC):
+    """Base class for wheeled orchestrators."""
+
+    package_name = DEFAULT_PACKAGE_NAME
+    package_version = __version__
+
+    def copy_repository_to_temp_dir_and_add_setup_py(self) -> str:
+        """Copy the repository to a temporary directory and add a setup.py file.
+
+        Returns:
+            Path to the temporary directory containing the copied repository.
+        """
+        repo_path = get_source_root()
+
+        self.package_name = f"{DEFAULT_PACKAGE_NAME}_{self.sanitize_name(os.path.basename(repo_path))}"
+
+        # Create a temporary folder
+        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
+
+        # Create a folder within the temporary directory
+        temp_repo_path = os.path.join(temp_dir, self.package_name)
+        fileio.mkdir(temp_repo_path)
+
+        # Copy the repository to the temporary directory
+        copy_dir(repo_path, temp_repo_path)
+
+        # Create init file in the copied directory
+        init_file_path = os.path.join(temp_repo_path, "__init__.py")
+        with fileio.open(init_file_path, "w") as f:
+            f.write("")
+
+        # Create a setup.py file
+        setup_py_content = f"""
+from setuptools import setup, find_packages
+
+setup(
+    name="{self.package_name}",
+    version="{self.package_version}",
+    packages=find_packages(),
+)
+"""
+        setup_py_path = os.path.join(temp_dir, "setup.py")
+        with fileio.open(setup_py_path, "w") as f:
+            f.write(setup_py_content)
+
+        return temp_dir
+
+    def create_wheel(self, temp_dir: str) -> str:
+        """Create a wheel for the package in the given temporary directory.
+
+        Args:
+            temp_dir (str): Path to the temporary directory containing the package.
+
+        Raises:
+            RuntimeError: If the wheel file could not be created.
+
+        Returns:
+            str: Path to the created wheel file.
+        """
+        # Change to the temporary directory
+        original_dir = os.getcwd()
+        os.chdir(temp_dir)
+
+        try:
+            # Run the `pip wheel` command to create the wheel
+            result = subprocess.run(
+                ["pip", "wheel", "."], check=True, capture_output=True
+            )
+            logger.debug(f"Wheel creation stdout: {result.stdout.decode()}")
+            logger.debug(f"Wheel creation stderr: {result.stderr.decode()}")
+
+            # Find the created wheel file
+            wheel_file = next(
+                (
+                    file
+                    for file in os.listdir(temp_dir)
+                    if file.endswith(".whl")
+                ),
+                None,
+            )
+
+            if wheel_file is None:
+                raise RuntimeError("Failed to create wheel file.")
+
+            wheel_path = os.path.join(temp_dir, wheel_file)
+
+            # Verify the wheel file is a valid zip file
+            import zipfile
+
+            if not zipfile.is_zipfile(wheel_path):
+                raise RuntimeError(
+                    f"The file {wheel_path} is not a valid zip file."
+                )
+
+            return wheel_path
+        finally:
+            # Change back to the original directory
+            os.chdir(original_dir)
+
+    def sanitize_name(self, name: str) -> str:
+        """Sanitize the value to be used in a cluster name.
+
+        Args:
+            name: Arbitrary input cluster name.
+
+        Returns:
+            Sanitized cluster name.
+        """
+        name = re.sub(
+            r"[^a-z0-9-]", "-", name.lower()
+        )  # replaces any character that is not a lowercase letter, digit, or hyphen with a hyphen
+        name = re.sub(r"^[-]+", "", name)  # trim leading hyphens
+        name = re.sub(r"[-]+$", "", name)  # trim trailing hyphens
+        return name