zenml-nightly 0.62.0.dev20240729__py3-none-any.whl → 0.63.0.dev20240731__py3-none-any.whl

zenml/zen_stores/template_utils.py

@@ -0,0 +1,261 @@
+#  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.
+"""Utilities for run templates."""
+
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from pydantic import create_model
+from pydantic.fields import FieldInfo
+
+from zenml.config import ResourceSettings
+from zenml.config.base_settings import BaseSettings
+from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
+from zenml.config.source import SourceWithValidator
+from zenml.config.step_configurations import StepConfigurationUpdate
+from zenml.enums import StackComponentType
+from zenml.logger import get_logger
+from zenml.stack import Flavor
+from zenml.zen_stores.schemas import PipelineDeploymentSchema
+
+logger = get_logger(__name__)
+
+
+def validate_deployment_is_templatable(
+    deployment: PipelineDeploymentSchema,
+) -> None:
+    """Validate that a deployment is templatable.
+
+    Args:
+        deployment: The deployment to validate.
+
+    Raises:
+        ValueError: If the deployment is not templatable.
+    """
+    if not deployment.build:
+        raise ValueError(
+            "Unable to create run template as there is no associated build."
+        )
+
+    if not deployment.build.stack:
+        raise ValueError(
+            "Unable to create run template as the associated build has no "
+            "stack reference."
+        )
+
+    for component in deployment.build.stack.components:
+        if not component.flavor_schema:
+            raise ValueError(
+                "Unable to create run template as a component of the "
+                "associated stack has no flavor."
+            )
+
+        if component.flavor_schema.workspace_id:
+            raise ValueError(
+                "Unable to create run template as a component of the "
+                "associated stack has a custom flavor."
+            )
+
+        flavor_model = component.flavor_schema.to_model()
+        flavor = Flavor.from_model(flavor_model)
+        component_config = flavor.config_class(
+            **component.to_model(include_metadata=True).configuration
+        )
+
+        if component_config.is_local:
+            raise ValueError(
+                "Unable to create run template as the associated stack "
+                "contains local components."
+            )
+
+
+def generate_config_template(
+    deployment: PipelineDeploymentSchema,
+) -> Dict[str, Any]:
+    """Generate a run configuration template for a deployment.
+
+    Args:
+        deployment: The deployment.
+
+    Returns:
+        The run configuration template.
+    """
+    deployment_model = deployment.to_model(include_metadata=True)
+
+    steps_configs = {
+        name: step.config.model_dump(
+            include=set(StepConfigurationUpdate.model_fields),
+            exclude={"name", "outputs"},
+        )
+        for name, step in deployment_model.step_configurations.items()
+    }
+
+    for config in steps_configs.values():
+        config["settings"].pop("docker", None)
+
+    pipeline_config = deployment_model.pipeline_configuration.model_dump(
+        include=set(PipelineRunConfiguration.model_fields),
+        exclude={"schedule", "build", "parameters"},
+    )
+
+    pipeline_config["settings"].pop("docker", None)
+
+    config_template = {
+        "run_name": deployment_model.run_name_template,
+        "steps": steps_configs,
+        **pipeline_config,
+    }
+    return config_template
+
+
+def generate_config_schema(
+    deployment: PipelineDeploymentSchema,
+) -> Dict[str, Any]:
+    """Generate a run configuration schema for the deployment and stack.
+
+    Args:
+        deployment: The deployment schema.
+
+    Returns:
+        The generated schema dictionary.
+    """
+    # Config schema can only be generated for a runnable template, so this is
+    # guaranteed by checks in the run template schema
+    assert deployment.build
+    assert deployment.build.stack
+
+    stack = deployment.build.stack
+    experiment_trackers = []
+    step_operators = []
+
+    settings_fields: Dict[str, Any] = {"resources": (ResourceSettings, None)}
+    for component in stack.components:
+        if not component.flavor_schema:
+            continue
+
+        flavor_model = component.flavor_schema.to_model()
+        flavor = Flavor.from_model(flavor_model)
+
+        for class_ in flavor.config_class.__mro__[1:]:
+            # Ugly hack to get the settings class of a flavor without having
+            # the integration installed. This is based on the convention that
+            # the static config of a stack component should always inherit
+            # from the dynamic settings.
+            if issubclass(class_, BaseSettings):
+                if len(class_.model_fields) > 0:
+                    settings_key = f"{component.type}.{component.flavor}"
+                    settings_fields[settings_key] = (
+                        Optional[class_],
+                        None,
+                    )
+
+                break
+
+        if component.type == StackComponentType.EXPERIMENT_TRACKER:
+            experiment_trackers.append(component.name)
+        if component.type == StackComponentType.STEP_OPERATOR:
+            step_operators.append(component.name)
+
+    settings_model = create_model("Settings", **settings_fields)
+
+    generic_step_fields: Dict[str, Any] = {}
+
+    for key, field_info in StepConfigurationUpdate.model_fields.items():
+        if key in [
+            "name",
+            "outputs",
+            "step_operator",
+            "experiment_tracker",
+            "parameters",
+        ]:
+            continue
+
+        if field_info.annotation == Optional[SourceWithValidator]:
+            generic_step_fields[key] = (Optional[str], None)
+        else:
+            generic_step_fields[key] = (field_info.annotation, field_info)
+
+    if experiment_trackers:
+        experiment_tracker_enum = Enum(  # type: ignore[misc]
+            "ExperimentTrackers", {e: e for e in experiment_trackers}
+        )
+        generic_step_fields["experiment_tracker"] = (
+            Optional[experiment_tracker_enum],
+            None,
+        )
+    if step_operators:
+        step_operator_enum = Enum(  # type: ignore[misc]
+            "StepOperators", {s: s for s in step_operators}
+        )
+        generic_step_fields["step_operator"] = (
+            Optional[step_operator_enum],
+            None,
+        )
+
+    generic_step_fields["settings"] = (Optional[settings_model], None)
+
+    all_steps: Dict[str, Any] = {}
+    all_steps_required = False
+    for name, step in deployment.to_model(
+        include_metadata=True
+    ).step_configurations.items():
+        step_fields = generic_step_fields.copy()
+        if step.config.parameters:
+            parameter_fields: Dict[str, Any] = {
+                name: (Any, FieldInfo(default=...))
+                for name in step.config.parameters
+            }
+            parameters_class = create_model(
+                f"{name}_parameters", **parameter_fields
+            )
+            step_fields["parameters"] = (
+                parameters_class,
+                FieldInfo(default=...),
+            )
+
+        step_model = create_model(name, **step_fields)
+
+        if step.config.parameters:
+            # This step has required parameters -> we make this attribute
+            # required and also the parent attribute so these parameters must
+            # always be included
+            all_steps_required = True
+            all_steps[name] = (step_model, FieldInfo(default=...))
+        else:
+            all_steps[name] = (Optional[step_model], FieldInfo(default=None))
+
+    all_steps_model = create_model("Steps", **all_steps)
+
+    top_level_fields: Dict[str, Any] = {}
+
+    for key, field_info in PipelineRunConfiguration.model_fields.items():
+        if key in ["schedule", "build", "steps", "settings", "parameters"]:
+            continue
+
+        if field_info.annotation == Optional[SourceWithValidator]:
+            top_level_fields[key] = (Optional[str], None)
+        else:
+            top_level_fields[key] = (field_info.annotation, field_info)
+
+    top_level_fields["settings"] = (Optional[settings_model], None)
+
+    if all_steps_required:
+        top_level_fields["steps"] = (all_steps_model, FieldInfo(default=...))
+    else:
+        top_level_fields["steps"] = (
+            Optional[all_steps_model],
+            FieldInfo(default=None),
+        )
+
+    return create_model("Result", **top_level_fields).model_json_schema()  # type: ignore[no-any-return]

zenml/zen_stores/zen_store_interface.py

@@ -94,6 +94,10 @@ from zenml.models import (
     RunMetadataFilter,
     RunMetadataRequest,
     RunMetadataResponse,
+    RunTemplateFilter,
+    RunTemplateRequest,
+    RunTemplateResponse,
+    RunTemplateUpdate,
     ScheduleFilter,
     ScheduleRequest,
     ScheduleResponse,
@@ -1263,22 +1267,6 @@ class ZenStoreInterface(ABC):
             KeyError: if the build doesn't exist.
         """
 
-    @abstractmethod
-    def run_build(
-        self,
-        build_id: UUID,
-        run_configuration: Optional[PipelineRunConfiguration] = None,
-    ) -> PipelineRunResponse:
-        """Run a pipeline from a build.
-
-        Args:
-            build_id: The ID of the build to run.
-            run_configuration: Configuration for the run.
-
-        Returns:
-            Model of the pipeline run.
-        """
-
     # -------------------- Pipeline deployments --------------------
 
     @abstractmethod
@@ -1346,16 +1334,101 @@ class ZenStoreInterface(ABC):
             KeyError: If the deployment doesn't exist.
         """
 
+    # -------------------- Run templates --------------------
+
+    @abstractmethod
+    def create_run_template(
+        self,
+        template: RunTemplateRequest,
+    ) -> RunTemplateResponse:
+        """Create a new run template.
+
+        Args:
+            template: The template to create.
+
+        Returns:
+            The newly created template.
+
+        Raises:
+            EntityExistsError: If a template with the same name already exists.
+        """
+
+    @abstractmethod
+    def get_run_template(
+        self, template_id: UUID, hydrate: bool = True
+    ) -> RunTemplateResponse:
+        """Get a run template with a given ID.
+
+        Args:
+            template_id: ID of the template.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The template.
+
+        Raises:
+            KeyError: If the template does not exist.
+        """
+
+    @abstractmethod
+    def list_run_templates(
+        self,
+        template_filter_model: RunTemplateFilter,
+        hydrate: bool = False,
+    ) -> Page[RunTemplateResponse]:
+        """List all run templates matching the given filter criteria.
+
+        Args:
+            template_filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all templates matching the filter criteria.
+        """
+
+    @abstractmethod
+    def update_run_template(
+        self,
+        template_id: UUID,
+        template_update: RunTemplateUpdate,
+    ) -> RunTemplateResponse:
+        """Updates a run template.
+
+        Args:
+            template_id: The ID of the template to update.
+            template_update: The update to apply.
+
+        Returns:
+            The updated template.
+
+        Raises:
+            KeyError: If the template does not exist.
+        """
+
+    @abstractmethod
+    def delete_run_template(self, template_id: UUID) -> None:
+        """Delete a run template.
+
+        Args:
+            template_id: The ID of the template to delete.
+
+        Raises:
+            KeyError: If the template does not exist.
+        """
+
     @abstractmethod
-    def run_deployment(
+    def run_template(
         self,
-        deployment_id: UUID,
+        template_id: UUID,
         run_configuration: Optional[PipelineRunConfiguration] = None,
     ) -> PipelineRunResponse:
-        """Run a pipeline from a deployment.
+        """Run a template.
 
         Args:
-            deployment_id: The ID of the deployment to run.
+            template_id: The ID of the template to run.
             run_configuration: Configuration for the run.
 
         Returns:

{zenml_nightly-0.62.0.dev20240729.dist-info → zenml_nightly-0.63.0.dev20240731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.62.0.dev20240729
+Version: 0.63.0.dev20240731
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0
@@ -467,7 +467,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.62.0 is out. Check out the release notes
+    🎉 Version 0.63.0 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.