More code review suggestions

Author: stefannica

src/zenml/config/compiler.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/pipeline_spec.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/deployers/base_deployer.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/exceptions.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/server/parameters.py

@@ -64,15 +64,13 @@ 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:
         message = (
             f"Failed to construct parameters model from pipeline "
             f"`{snapshot.pipeline_configuration.name}`."
         )
         logger.error(message)
         raise RuntimeError(message)
-    else:
-        logger.debug(message)
 
     return model

src/zenml/deployers/server/service.py

@@ -507,29 +507,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.")

src/zenml/deployers/utils.py

@@ -25,7 +25,6 @@
     DeploymentHTTPError,
     DeploymentNotFoundError,
     DeploymentProvisionError,
-    DeploymentSchemaNotFoundError,
 )
 from zenml.enums import DeploymentStatus
 from zenml.models import DeploymentResponse
@@ -45,25 +44,18 @@ def get_deployment_input_schema(
         The schema for the deployment's input parameters.
 
     Raises:
-        DeploymentSchemaNotFoundError: If the deployment has no associated
-            snapshot, pipeline spec, or parameters schema.
+        RuntimeError: If the deployment has no associated input schema.
     """
-    if not deployment.snapshot:
-        raise DeploymentSchemaNotFoundError(
-            f"Deployment {deployment.name} has no associated snapshot."
-        )
-
-    if not deployment.snapshot.pipeline_spec:
-        raise DeploymentSchemaNotFoundError(
-            f"Deployment {deployment.name} has no associated pipeline spec."
-        )
+    if (
+        deployment.snapshot
+        and deployment.snapshot.pipeline_spec
+        and deployment.snapshot.pipeline_spec.input_schema
+    ):
+        return deployment.snapshot.pipeline_spec.input_schema
 
-    if not deployment.snapshot.pipeline_spec.input_schema:
-        raise DeploymentSchemaNotFoundError(
-            f"Deployment {deployment.name} has no associated parameters schema."
-        )
-
-    return deployment.snapshot.pipeline_spec.input_schema
+    raise RuntimeError(
+        f"Deployment {deployment.name} has no associated input schema."
+    )
 
 
 def get_deployment_output_schema(
@@ -78,25 +70,18 @@ def get_deployment_output_schema(
         The schema for the deployment's output parameters.
 
     Raises:
-        DeploymentSchemaNotFoundError: If the deployment has no associated
-            snapshot, pipeline spec, or output schema.
+        RuntimeError: If the deployment has no associated output schema.
     """
-    if not deployment.snapshot:
-        raise DeploymentSchemaNotFoundError(
-            f"Deployment {deployment.name} has no associated snapshot."
-        )
-
-    if not deployment.snapshot.pipeline_spec:
-        raise DeploymentSchemaNotFoundError(
-            f"Deployment {deployment.name} has no associated pipeline spec."
-        )
-
-    if not deployment.snapshot.pipeline_spec.output_schema:
-        raise DeploymentSchemaNotFoundError(
-            f"Deployment {deployment.name} has no associated output schema."
-        )
-
-    return deployment.snapshot.pipeline_spec.output_schema
+    if (
+        deployment.snapshot
+        and deployment.snapshot.pipeline_spec
+        and deployment.snapshot.pipeline_spec.output_schema
+    ):
+        return deployment.snapshot.pipeline_spec.output_schema
+
+    raise RuntimeError(
+        f"Deployment {deployment.name} has no associated output schema."
+    )
 
 
 def get_deployment_invocation_example(

src/zenml/integrations/aws/deployers/aws_deployer.py

@@ -1097,6 +1097,11 @@ def _select_aws_cpu_memory_combination(
         Returns:
             Tuple of (cpu, memory) that best matches requirements, in AWS App
             Runner format.
+
+        Raises:
+            ValueError: If the requested resource requirements cannot be matched
+                to any of the supported combinations for the AWS App Runner
+                service and strict_resource_matching is True.
         """
         if requested_cpu is None and requested_memory_gb is None:
             return f"{DEFAULT_CPU:g} vCPU", f"{DEFAULT_MEMORY:g} GB"

src/zenml/models/v2/core/deployment.py

@@ -100,6 +100,8 @@ class DeploymentUpdate(BaseUpdate):
     status: Optional[str] = Field(
         default=None,
         title="The new status of the deployment.",
+        description="Possible values are: "
+        f"{', '.join(DeploymentStatus.values())}",
     )
     deployment_metadata: Optional[Dict[str, Any]] = Field(
         default=None,
@@ -143,7 +145,8 @@ class DeploymentResponseBody(ProjectScopedResponseBody):
     status: Optional[str] = Field(
         default=None,
         title="The status of the deployment.",
-        description="Current operational status of the deployment.",
+        description="Current operational status of the deployment. Possible "
+        f"values are: {', '.join(DeploymentStatus.values())}",
     )