Merge branch 'feature/served-pipelines' of https://github.com/zenml-io/zenml into feature/served-pipelines

Author: safoinme

pyproject.toml

@@ -1147,7 +1147,6 @@ def deploy_snapshot(
         pipeline_name_or_id: The name or ID of the pipeline.
         deployment_name_or_id: Name or ID of the deployment to use for the
             pipeline.
-        config_path: Path to pipeline configuration file.
         update: If True, update the deployment with the same name if it
             already exists.
         overtake: If True, update the deployment with the same name if

src/zenml/cli/pipeline.py

@@ -650,17 +650,8 @@ def _compute_pipeline_spec(
             )
             for output_artifact in pipeline._output_artifacts
         ]
-        try:
-            output_schema = pipeline._compute_output_schema()
-        except Exception as e:
-            logger.warning("Failed to compute pipeline output schema: %s", e)
-            output_schema = None
-
-        parameters_model = pipeline.get_parameters_model()
-        if parameters_model:
-            input_schema = parameters_model.model_json_schema()
-        else:
-            input_schema = None
+        input_schema = pipeline._compute_input_schema()
+        output_schema = pipeline._compute_output_schema()
 
         return PipelineSpec(
             steps=step_specs,

src/zenml/config/compiler.py

@@ -42,11 +42,17 @@ class PipelineSpec(FrozenBaseModel):
     #   inputs in the step specs refer to the pipeline parameter names
     # - 0.4: New Pipeline class, the upstream steps and
     #   inputs in the step specs refer to the pipeline parameter names
-    # - 0.5: Adds outputs and output schema
+    # - 0.5: Adds input schema, outputs and output schema
     version: str = "0.5"
     source: Optional[SourceWithValidator] = None
     parameters: Dict[str, Any] = {}
-    input_schema: Dict[str, Any] = {}
+    input_schema: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="JSON schema of the pipeline inputs. This is only set "
+        "for pipeline specs with version >= 0.5. If the value is None, the "
+        "schema generation failed, which is most likely because some of the "
+        "pipeline inputs are not JSON serializable.",
+    )
     steps: List[StepSpec]
     outputs: List[OutputSpec] = []
     output_schema: Optional[Dict[str, Any]] = Field(

src/zenml/config/pipeline_spec.py

@@ -16,8 +16,13 @@
 from enum import Enum
 from typing import Literal, Optional, Union
 
-from pydantic import Field, NonNegativeInt, PositiveFloat, PositiveInt
-from pydantic_settings import SettingsConfigDict
+from pydantic import (
+    ConfigDict,
+    Field,
+    NonNegativeInt,
+    PositiveFloat,
+    PositiveInt,
+)
 
 from zenml.config.base_settings import BaseSettings
 
@@ -178,7 +183,7 @@ def get_memory(
             # Should never happen due to the regex validation
             raise ValueError(f"Unable to parse memory unit from '{memory}'.")
 
-    model_config = SettingsConfigDict(
+    model_config = ConfigDict(
         # public attributes are immutable
         frozen=True,
         # prevent extra attributes during model initialization

src/zenml/config/resource_settings.py

@@ -163,6 +163,35 @@ def _update_deployment(
             DeploymentUpdate.from_operational_state(operational_state),
         )
 
+    def _check_deployment_inputs_outputs(
+        self,
+        snapshot: PipelineSnapshotResponse,
+    ) -> None:
+        """Check if the deployment has compiled schemas for the pipeline inputs and outputs.
+
+        Args:
+            snapshot: The pipeline snapshot to check.
+
+        Raises:
+            DeploymentProvisionError: if the deployment has no compiled schemas
+            for the pipeline inputs and outputs.
+        """
+        if (
+            not snapshot.pipeline_spec
+            or not snapshot.pipeline_spec.input_schema
+            or not snapshot.pipeline_spec.output_schema
+        ):
+            raise DeploymentProvisionError(
+                f"The pipeline with name '{snapshot.pipeline.name}' referenced "
+                f"by the deployment with name or ID "
+                f"'{snapshot.name or snapshot.id}' "
+                "is missing the compiled schemas for the pipeline inputs or "
+                "outputs. This is most likely because some of the pipeline "
+                "inputs or outputs are not JSON serializable. Please check that "
+                "all the pipeline input arguments and return values have data "
+                "types that are JSON serializable."
+            )
+
     def _check_deployment_deployer(
         self,
         deployment: DeploymentResponse,
@@ -383,6 +412,8 @@ def provision_deployment(
                 "already exists"
             )
 
+        self._check_deployment_inputs_outputs(snapshot)
+
         client = Client()
 
         settings = cast(

src/zenml/deployers/base_deployer.py

@@ -414,8 +414,11 @@ def do_provision_deployment(
                 auto_remove=False,
                 ports=ports,
                 labels={
-                    "zenml-deployment-uuid": str(deployment.id),
+                    "zenml-deployment-id": str(deployment.id),
                     "zenml-deployment-name": deployment.name,
+                    "zenml-deployer-name": str(self.name),
+                    "zenml-deployer-id": str(self.id),
+                    "managed-by": "zenml",
                 },
                 extra_hosts=extra_hosts,
                 **run_args,

src/zenml/deployers/docker/docker_deployer.py

@@ -59,9 +59,5 @@ class DeploymentHTTPError(DeployerError):
     """Error raised when an HTTP request to a deployment fails."""
 
 
-class DeploymentSchemaNotFoundError(KeyError, DeployerError):
-    """Error raised when a deployment schema is not found."""
-
-
 class DeploymentInvalidParametersError(DeployerError):
     """Error raised when the parameters for a deployment are invalid."""

src/zenml/deployers/exceptions.py

@@ -63,7 +63,7 @@ def build_params_model_from_snapshot(
         logger.error(f"Failed to load pipeline class from snapshot: {e}")
         raise RuntimeError(f"Failed to load pipeline class from snapshot: {e}")
 
-    model = pipeline_class.get_parameters_model()
+    model = pipeline_class._compute_input_model()
     if not model:
         raise RuntimeError(
             f"Failed to construct parameters model from pipeline `{snapshot.pipeline_configuration.name}`."

src/zenml/deployers/server/parameters.py

@@ -137,7 +137,7 @@ def initialize(self) -> None:
 
             # Build parameter model
             self._params_model = build_params_model_from_snapshot(
-                snapshot=self.snapshot,
+                self.snapshot, strict=True
             )
 
             # Initialize orchestrator
@@ -504,29 +504,33 @@ def _build_response(
     # ----------
 
     @property
-    def input_schema(self) -> Optional[Dict[str, Any]]:
-        """Return the JSON schema for pipeline input parameters if available.
+    def input_schema(self) -> Dict[str, Any]:
+        """Return the JSON schema for pipeline input parameters.
 
         Returns:
-            The JSON schema for pipeline parameters if available.
+            The JSON schema for pipeline parameters.
         """
-        try:
-            if self.snapshot.pipeline_spec:
-                return self.snapshot.pipeline_spec.input_schema
-        except Exception:
-            return None
-        return None
+        if (
+            self.snapshot.pipeline_spec
+            and self.snapshot.pipeline_spec.input_schema
+        ):
+            return self.snapshot.pipeline_spec.input_schema
+        # This should never happen, given that we check for this in the
+        # base deployer.
+        raise RuntimeError("The pipeline input schema is not available.")
 
     @property
-    def output_schema(self) -> Optional[Dict[str, Any]]:
-        """Return the JSON schema for the deployment response if available.
+    def output_schema(self) -> Dict[str, Any]:
+        """Return the JSON schema for the pipeline outputs.
 
         Returns:
-            The JSON schema for the deployment response if available.
+            The JSON schema for the pipeline outputs.
         """
-        try:
-            if self.snapshot.pipeline_spec:
-                return self.snapshot.pipeline_spec.output_schema
-        except Exception:
-            return None
-        return None
+        if (
+            self.snapshot.pipeline_spec
+            and self.snapshot.pipeline_spec.output_schema
+        ):
+            return self.snapshot.pipeline_spec.output_schema
+        # This should never happen, given that we check for this in the
+        # base deployer.
+        raise RuntimeError("The pipeline output schema is not available.")